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: The paper suggests a method for approximating the 2-Wasserstein gradient flow for the relative entropy. The proposed particle-based method uses a neural network function approximation-based approach to estimating the necessary density ratios. Experiments verify reasonable performance compared to MALA and ULA.