Relative Entropy Gradient Sampler for Unnormalized Distributions

1 INTRODUCTION . Sampling from unnormalized distributions plays a fundamental role in statistical inference and machine learning . This problem is frequently encountered in Bayesian statistics . Conducting Bayesian analysis requires evaluation of multi-dimensional integrals where analytical expressions for unnormalized posterior distributions are usually not available . Consequently , sampling is necessary for Monte Carlo approximation of these integrals . In this work , we propose a general purpose sampling algorithm for unnormalized distributions . Markov chain Monte Carlo ( MCMC ) methods ( Andrieu et al. , 2003 ; Brooks et al. , 2011 ) are widely used to sample from unnormalized distributions . Sampling with MCMC relies on defining an appropriate transition kernel to construct a Markov chain whose equilibrium distribution is precisely the target distribution . Based on rejection sampling , the Metropolis–Hastings algorithm ( Metropolis et al. , 1953 ; Hastings , 1970 ; Tierney , 1994 ; Dunson & Johndrow , 2019 ) provides a flexible framework for general MCMC sampling . To implement a Metropolis–Hastings algorithm , one needs to specify a proposal density and an acceptance policy . However , without a careful design of these two aspects , the Metropolis–Hastings algorithm can be inefficient due to strong correlations , slow mixing , or low acceptance rates , especially in the large-scale and high-dimensional settings . Moreover , proposals through discretizing some continuous processes like Langevin diffusion and Hamiltonian dynamics are introduced ( Roberts & Tweedie , 1996 ; Roberts & Stramer , 2002 ; Duane et al. , 1987 ; Neal , 2011 ; Hoffman & Gelman , 2014 ) and further enhanced by stochastic gradient estimation ( Welling & Teh , 2011 ; Chen et al. , 2014 ) . Variational Bayesian inference ( Beal , 2003 ) , often simply referred to as variational inference ( VI ) ( Wainwright & Jordan , 2008 ; Blei et al. , 2017 ) , is another prominent approach to sampling from unnormalized distributions . VI approximates the unnormalized posterior distribution with a restricted parametric variational posterior distribution by minimizing the Kullback-Leibler ( KL ) divergence between them . Since the true posterior distribution is intractable , VI turns to maximize a surrogate variational objective called the evidence lower bound ( ELBO ) . However , one is required to trade off the parameterization flexibility of variational posteriors against the optimization complexity of ELBO in practice . In the spirit of VI , particle-based variational inference ( ParVI ) ( Liu & Wang , 2016 ; Chen et al. , 2018 ; Zhu et al. , 2020 ) iteratively optimizes a set of particles to mimic a functional gradient descent for minimizing the KL divergence . ParVI seeks to move a variational distribution towards the unnormalized target distribution , along a steepest descent direction of the KL divergence . In a continuous view , these movements of variational distributions can be understood as a gradient flow in probability measure spaces ( Liu et al. , 2019a ; b ) . A key part of ParVI is how to estimate the desired steepest descent direction ( i.e. , functional gradient ) from the evolving random particles . An elegant approach is the Stein variational gradient descent ( SVGD ) Liu & Wang ( 2016 ) . In SVGD , the functional gradient descent is embedded in a reproducing kernel Hilbert space ( RKHS ) , which is further recognized as a gradient flow under the Stein geometry ( Liu , 2017 ; Lu et al. , 2019 ; Duncan et al. , 2019 ) . A drawback of SVGD is that it tends to collapse at part of the modes of the target , due to a negative correlation between the data dimensionality and the repulsive force in the RKHS ( Zhuo et al. , 2018 ) . In this work , we propose a relative entropy gradient sampler ( REGS ) for sampling from unnormalized target distributions . To approximate a target distribution , we consider the Wasserstein gradient flow of relative entropy ( or KL divergence ) , named relative entropy gradient flow . The relative entropy gradient flow represents a path of probability distributions that follows the functional gradient descent direction of relative entropy . There exists an ODE system of random particles that uniquely determines the spatial and temporal dynamics of the relative entropy gradient flow . Therefore , to sample with REGS , we only need to simulate the ODE system with particle evolution . Evaluating the velocity fields of this ODE system can be transformed into estimating the logarithmic density ratio between the density of evolving particles and the unnormalized target density . Based on this observation , we propose a novel logarithmic density ratio estimation method for unnormalized distributions . By alternating between particle evolution and velocity field estimation , we can collect a set of stable particles which are approximately distributed as the target distribution . Our contributions can be summarized as follows : ( 1 ) Building upon the relative entropy gradient flow , we propose the relative entropy gradient sampler ( REGS ) for unnormalized target distributions . REGS preserves high efficiency and strong stability with respect to increasing singularity in mixtures of Gaussians , when the number of components increases ( as shown in Figure 1 ) , the variance of each component decreases , and the distance between any two components increases . ( 2 ) We propose to directly estimate velocity fields of the relative entropy gradient flow as gradients of logarithmic density ratios , that is computationally stable and efficient . ( 3 ) We develop a nonparametric approach to estimating the density ratio between an unnormalized density and an underlying density represented by samples , which is of independent interest . ( 4 ) We present experimental comparisons on varieties of multi-mode synthetic data and benchmark data and demonstrate that REGS is a more accurate sampler than the popular samplers including ULA , MALA and SVGD . Related work The proposed REGS is most related to sampling methods based on the relative entropy gradient flow , in particular , the recently proposed SVGD ( Liu & Wang , 2016 ; Liu , 2017 ) , which estimates the velocity fields of the relative entropy gradient flow in a reproducing kernel Hilbert space . See also Korba et al . ( 2020 ) ; Salim et al . ( 2021 ; 2020 ) for theoretical analysis of SVGD . In contrast , REGS approximates the velocity fields based on a novel logarithmic density ratio estimation approach with deep neural networks . The undesirable mode collapse feature of SVGD is not inevitable for REGS since the approximation and expressive powers of deep neural networks is known to surpass those of kernel methods . MCMC algorithms constructed from overdamped Langevin diffusion can be studied as discretization of the relative entropy gradient flow ( Jordan et al. , 1998 ) . Based on the Euler-Maruyama discretization of overdamped Langevin diffusion , unadjusted Langevin algorithm ( ULA ) ( Roberts & Tweedie , 1996 ) aims at generating samples from an approximation of the unnormalized target , but is biased for fixed step size . When a Metropolis-Hastings step is included , Metropolis-adjusted Langevin algorithm ( MALA ) ( Roberts & Tweedie , 1996 ) is capable of correcting the bias , but leaves a large number of intermediate samples rejected . In REGS , one only needs to estimate the deterministic velocity fields of the relative entropy gradient flow , which differs from running ULA and MALA with randomness from diffusion processes . All particles produced by REGS are generated from an approximation of the target distribution . Another line of work ( Gao et al. , 2019 ; 2021 ) uses Wasserstein gradient flows of f -divergences for generative learning with samples from the underlying target distribution . In their work , evaluating velocity fields of gradient flows also boils down to estimating density ratios . However , our current problem is to sample from an unnormalized target density . Furthermore , we propose a novel density ratio estimation procedure when the target distribution is only known up to a normalizing constant . Notation Let P2 ( X ) be the space of Borel probability measures on a support space X ⊂ Rd with a finite second moment , and let Pa2 ( X ) be a subspace of P2 ( X ) whose measures are absolutely continuous w.r.t . the Lebesgue measure . All probability measures we considered thereinafter are assumed to belong to Pa2 ( X ) . To ease the notation , we use probability density functions such as q ( x ) , p ( x ) , x ∈ X to express probability distributions in Pa2 ( X ) . Let ( Pa2 ( X ) , W2 ) denote the metric space Pa2 ( X ) endowed with the 2-Wasserstein distance W2 , which is referred to as the quadratic Wasserstein space . We use ∇ and Div to denote the gradient operator and the divergence operator , respectively . 2 PROBLEM FORMULATION . Consider an unnormalized probability density function u : X → [ 0 , ∞ ) , where X ⊆ Rd is the support of u . Suppose u has an intractable normalizing constant Z = ∫ X u ( x ) dx < ∞ . Our goal is to generate random samples from the underlying distribution p ∈ Pa2 ( X ) , whose probability density function is only known up to proportionality , i.e. , p ( x ) = u ( x ) /Z , x ∈ X . The basic idea is to gradually optimize samples from a given distribution q ∈ Pa2 ( X ) to approximate samples from p , where it is easy to sample from q. Optimizing samples leads to functional optimization of distributions . We then introduce the classical relative entropy as the functional optimization objective . The relative entropy , a.k.a. , the Kullback–Leibler divergence , for q , p ∈ Pa2 ( X ) is the average logarithmic density ratio , which is defined as Dre ( q‖p ) = ∫ X q ( x ) log ( q ( x ) p ( x ) ) dx . ( 1 ) It holds that Dre ( q‖p ) ≥ 0 and Dre ( q‖p ) = 0 iff q ( x ) = p ( x ) a.e . x ∈ X . Moreover , we denote the relative entropy functional as F [ · ] : = Dre ( ·‖p ) : Pa2 ( X ) → [ 0 , ∞ ] . ( 2 ) To sample from the unnormalized density u = pZ , we consider the functional minimization problem min q∈Pa2 ( X ) F [ q ] , ( 3 ) where F [ q ] is always minimized at the underlying target distribution p , i.e. , q ( x ) = p ( x ) a.e . x ∈ X . In a nutshell , problem ( 3 ) is an energy functional minimization problem in a metric space . To minimize the energy functional F , it suffices to move along the corresponding gradient flow in a metric space until the flow converges . For example , a gradient flow in the Euclidean space refers to a curve whose tangent space contains the steepest descent direction of a given function . Analogously , a gradient flow in the space of probability measures means a curve that points in the steepest descent direction of a given energy functional . When equipped with the 2-Wasserstein distance , minimization of the energy functional F naturally corresponds to a continuous path on the quadratic Wasserstein space of distributions , which is commonly known as a Wasserstein gradient flow of the relative entropy . We call this flow a relative entropy gradient flow for briefness .

Summary: This paper considers the problem of sampling from an unnormalized distribution. The unnormalized target distribution can be regarded as a stationary point of the Wasserstein gradient flow of the corresponding relative entropy functional, which can be equivalently identified from a microscopic perspective by defining a time-varying velocity field of the particles. While the exact time-varying velocity field is not exactly available, the authors propose to estimate such a quantity by approximating the corresponding logarithmic density ratio through minimizing the Bregman score. Such an approximation requires only samples from the variable distribution which can obtain by simulating particles following the estimated velocity field.