Discrete Langevin Samplers via Wasserstein Gradient Flow

Langevin dynamics, the particle version of the gradient flow that minimizes the KL divergence on Wasserstein manifold, induces gradient based Markov chain Monte Carlo (MCMC) samplers like Langevin Monte Carlo (LMC) in continuous space. The superior efficiency of gradient based MCMC samplers stimulates the recent attempts to generalize LMC to discrete space. However, the principled extension of the Langevin dynamics for discrete space is still missing due to the lack of well-defined gradients. In this work, we generalize the Wasserstein gradient flow to discrete spaces and derive the corresponding discrete counterpart of Langevin dynamics. With this new understanding, the recent ``gradient''-based samplers in discrete space can be obtained by choosing proper discretizations. This new framework also enables us to derive a new algorithm named \textit{Discrete Langevin Monte Carlo} (DLMC) by simulating the Wasserstein gradient flow with respect to simulation time. As a result, DLMC has a convenient parallel implementation and location-dependent velocities that allow larger average jump distance. We demonstrate the advantages of DLMC for sampling and learning in various binary and categorical distributions.