De-randomizing MCMC dynamics with the diffusion Stein operator

1 Introduction . Evaluating an un-normalized target distribution π is a centerpiece of Bayesian inference , due to its ubiquitous presence in posterior distributions . Markov chain Monte Carlo ( MCMC ) methods fulfill this objective by generating asymptotically exact random samples from the distribution , a significant subset of which involves discretization of continuous-time diffusion processes , most notably Langevin diffusion , stochastic gradient Hamiltonian Monte Carlo ( HMC ) ( Chen et al. , 2014 ) and their further generalizations , which we collectively call MCMC dynamics ( Ma et al. , 2015 ) . Despite its simplicity and theoretical soundness , this diffusion-based sampling often suffers from slow convergence and small effective sample sizes , largely due to the auto-correlation of the samples . As an alternative to the simulation of stochastic systems , particle-based variational inference ( Liu and Wang , 2016 ; Chen et al. , 2018 ; Liu et al. , 2019a ) partially addresses the shortcomings of MCMC by replacing the Langevin diffusion , the simplest MCMC dynamics , with a deterministic interacting particle system that transports a set of interacting particles towards the target distribution . Theoretically speaking , Langevin diffusion encodes an evolution of density that minimizes the KL-divergence through steepest descent in the 2-Wasserstein space ( Jordan et al. , 1998 ) , and particle- 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . based variational inference ( PARVI ) approximates such evolution using reproducing kernel Hilbert space ( RKHS ) ( Liu , 2017 ; Liu et al. , 2019a ) , thus de-randomizing the Langevin diffusion process . Given the elegant theoretical link between Langevin diffusion and PARVI , one crucial question arises : can we leverage the advantages of general MCMC dynamics onto de-randomized particle systems ? While it is difficult to formulate them as a direct Wasserstein gradient flow , Liu et al . ( 2019b ) propose an interpretation for “ regular ” MCMC dynamics that augments the 2-Wasserstein space with a fiber bundle , forming a fiber-Riemannian Poisson manifold , under which MCMC dynamics follows a Hamiltonian flow on the fiber bundle , and a gradient flow on each fiber . We show that adapting the fiber-Riemannian Hamiltonian flow to the Stein-Wasserstein metric ( Liu , 2017 ; Duncan et al. , 2019 ) yields a vector field in the form of applying the diffusion Stein operator to the kernel function k ( · , θ ) . A discretization of this vector field gives generalized SVGD ( GSVGD ) , an interacting particle system that generalizes SVGD , with particle updates balancing an attractive force maximizing the log-likelihood and a repulsive force preventing a “ mode collapse ” of particles . We further demonstrate that the connection drawn between LD and SVGD ( Liu and Wang , 2016 ; Liu , 2017 ; Liu et al. , 2019a ) is retraced by GSVGD , reaffirming our claim that GSVGD mirrors SVGD in approximating a larger class of MCMC diffusion processes . Within the generality of our framework , we can develop PARVI algorithms that exploit two key types of possible acceleration in MCMC dynamics : auxiliary momentum variables and an adaptive Riemannian parametrization that allows for fast and efficient exploration of the probability space ( Girolami and Calderhead , 2011 ) , as shown in Table 1 . 2 MCMC dynamics – the relevant bits . In this paper , we consider the problem of extracting samples θ ∈ RD = Ω from an un-normalized target distribution π ( θ ) ∝ e−H ( θ ) . We begin by reviewing the key characteristics of MCMC dynamics , with particular emphasis on Langevin dynamics and its Fokker-Planck equation ( FPE ) , the partial differential equation ( PDE ) that depicts the time evolution of a diffusion process . As the FPE of LD conforms to a gradient flow structure in the 2-Wasserstein metric space of probability measures ( P ( Ω ) , W2 ) ( Jordan et al. , 1998 ) , we then briefly cover gradient flow on P ( Ω ) as an infinite-dimensional Riemannian manifold . Through the lens of gradient flow , we can see SVGD as a deterministic interacting particle system that approximates the gradient flow of LD through ( i ) gradient flow on a novel Stein-Wasserstein metric ( Liu , 2017 ) or ( ii ) projecting the gradient flow direction onto an RKHS ( Liu et al. , 2019a ) . With Wasserstein gradient flow ( WGF ) neatly linking to both LD and SVGD , we see that a diffusion process LD and an interacting particle system SVGD are , respectively , a stochastic instance and a deterministic approximation of the same Wasserstein gradient flow . Notations : We use∇f to denote the gradient of a scalar-valued function , and∇ · f the divergence of a vector-valued function , ∇ ·A applies the divergence operator to each row of a matrix-valued function , ρ̇t the “ partial derivative ” with respect to t . 2.1 Langevin dynamics and its Fokker-Planck equation . In this paper , we consider MCMC dynamics in the form of Itō diffusion processes following the stochastic differential equation ( SDE ) formula dθt = f ( θt ) dt+ √ 2Σ ( θt ) dWt , ( 1 ) consisting of drift coefficient f : RD 7→ RD , diffusion coefficient √ 2Σ and a D-dimensional Brownian motion Wt . Ma et al . ( 2015 ) provide a complete recipe of all Itō diffusion processes converging to the target measure π . The simplest MCMC dynamics takes the form of Langevin diffusion ( Langevin , 1908 ) , which moves towards higher densities with f ( θ ) = ∇ log π ( θ ) perturbed by white noise Σ = I . Given an initial distribution θ0 ∼ ρ0 , the Fokker-Planck equation describes the time evolution of the density of θt ( Risken , 1996 ) ρ̇t +∇ · ( ρt∇ log π ) − ( ∇∇ ) : ( ρtI ) = 0 , ( 2 ) where X : Y = tr ( X > Y ) . Given that ( ∇∇ ) : M = ∇ · ( ∇ ·M ) , we can rewrite ( 2 ) as ρ̇t = ∇ · ( ∇ρt − ρt∇ log π ) = ∇ · ( ρt∇ δKL [ ρt ‖ π ] δρt︸ ︷︷ ︸ first variation of KL [ ρt ‖ π ] ) , ( 3 ) a crucial step in developing the Wasserstein gradient flow perspective of LD , as ∇δE/δρ coincides with the differential of functionals induced by the Wasserstein metric . 2.2 Gradient flow on ( P ( Ω ) , W2 ) . The conventional gradient flow ( or the steepest descent curve ) that minimizes a smooth function F : RD 7→ R follows the PDE : ẋt + ∇F ( xt ) = 0 . Defining gradient flow on P ( Ω ) , informally written as ρ̇t +∇WE ( ρt ) = 0 , requires a definition of the differentiation ∇W , which requires an understanding of the Wasserstein metric , i.e. , the inner product gρ defined on its tangent space . To simplify the discussion , we restrict the discussion on measures with a density function and with finite second-order moments . For each ρ , the tangent space constitutes smooth functions integrating to zero , TρP ( Ω ) = { f |f ∈ C∞ ( Ω ) , ∫ f ( θ ) dθ = 0 } , in that a curve on P ( Ω ) preserves volume . The cotangent space constitutes an equivalence class of smooth functions in differing constant , noted as the quotient space T ∗ρP ( Ω ) = C∞ ( Ω ) /R . We characterize the inner product space using the metric tensor G ( ρ ) : TρP 7→ T ∗ρP , a one-to-one mapping between elements of the tangent space and those of the cotangent space . The inner product gρ is defined by the inverse of the metric tensor , often denoted as Onsager operators ( Onsager , 1931 ) gρ ( f1 , f2 ) = ∫ f1G ( ρ ) f2dθ = ∫ φ1G ( ρ ) −1φ2dx , φ1 = G ( ρ ) f1 , φ2 = G ( ρ ) f2 . ( 4 ) The Wasserstein Onsager operator takes the form G ( ρ ) −1 : φ 7→ −∇ · ( ρ∇φ ) . In the context of Wasserstein gradient flow , we define∇WE ( ρ ) as G ( ρ ) −1 δE ( ρ ) δρ , leading to the formulation : 0 = ρ̇t −G ( ρ ) −1 δE ( ρt ) δρt = ρ̇t +∇ · ( ρt∇ δE ( ρt ) δρt ) . ( 5 ) The first variation of KL [ ρ ‖ π ] takes the form of δKL [ ρ ‖ π ] δρ = log ρ/π + 1 , leading to the conclusion by Jordan et al . ( 1998 ) that the time evolution of the Langevin diffusion follows a curve of steepest descent in the Wasserstein space , laying the foundation for approximation using particle interaction . 2.3 SVGD as gradient flow . PARVI transports a set N of interacting particles { θ ( i ) t } 1≤i≤N towards the target distribution over time . Take SVGD for example , denoting the empirical measure at time t as ρ̂t = 1N ∑N i=1 δθ ( i ) t , SVGD ( Liu and Wang , 2016 ) follows the update rule θ̇t = Eθ′∼ρ̂t [ k ( θt , θ ′ ) ∇ log π ( θ′ ) +∇2k ( θt , θ′ ) ] = vH ( θt|ρ̂t ) , ( 6 ) where k ( · , · ) defines a positive-definite kernel , and∇2 denotes the gradient of the second argument . The FPE of SVGD can be written as ρ̇t −∇ · ( ρtKρ∇ δKL [ ρt ‖ π ] δρt ) = 0 , ( 7 ) where Kρ is an integral operator : Kρf ( θ ) = ∫ k ( θ , θ′ ) f ( θ′ ) dρ ( θ ) . We shall denote the RKHS defined by k as H. While significantly different from LD at first glance , SVGD approximates the WGF of LD by • kernelizing the Wasserstein Onsager operator to give GH ( ρ ) −1 : φ 7→ −∇ · ( ρKρ∇φ ) , thus forming the Stein-Wasserstein metric WH ( Liu , 2017 ; Duncan et al. , 2019 ) ; • projecting the gradient flow vector field v ( ρ ) = −∇ δKL [ ρ ‖ π ] δρ = ∇ log π −∇ log ρ ontoH ( Liu et al. , 2019a ) . The interpretation of SVGD yields valuable insights . While gradient flow ( 7 ) on ( P ( Ω ) , W2 ) yields no closed-form energy functional ( Chen et al. , 2018 ) , it behooves to absorb the operator Kρ into the definition of the Stein-Wasserstein metric to guarantee a tractable energy functional . In the meantime , the kernelization trick transforms a gradient flow ( 2 ) simulated by a diffusion process into an approximate deterministic transportation of particles – in other words , de-randomizes it . 2.4 MCMC dynamics as fiber-gradient Hamiltonian flow . The concept of flow on ( P ( Ω ) , W2 ) is a generalization of the gradient flow : given a vector field v : P ( Ω ) 7→ ⋃ ρ∈P Tρ , the flow of v is defined as ρ̇t = v ( ρt ) . Liu et al . ( 2019b ) interpret “ regular ” MCMC dynamics on ( P ( Ω ) , W2 ) as a flow , by combining gradient flow and Hamiltonian flow on the Wasserstein space , yielding a fiber-Riemannian manifold structure of P ( Ω ) , a fiber bundle consisting of Riemannian manifolds . In the context of a “ regular ” MCMC dynamics taking the form of underdamped Langevin dynamics ( Chen et al. , 2014 ) , the diffusion matrix A determines the gradient flow on each fiber , and the curl matrix C determines the Hamiltonian flow on the fiber bundle . The Hamiltonian flow keeps KL [ ρ ‖ π ] constant while encouraging fast exploration of the probability space ; the fiber gradient flow minimizes the KL . The fiber-gradient Hamiltonian flow determines a PARVI with the particle update θ̇t = ( A ( θt ) + C ( θt ) ) ( ∇ log π ( θt ) − ∇̂ log ρt ( θt ) ) = vA , C ( ρt ) , ( 8 ) where the intractable ∇̂ log ρt is approximated via the Blob method ( Carrillo et al. , 2019 ) .

Summary: This paper proposes a generalization of the Stein variational gradient descent (SVGD) algorithm via a discretization of the fiber-gradient Hamiltonian flow. The approach that the authors take to extend the framework of particle-based variational inference is similar to the way in which MCMC dynamics are extended through the "complete recipe" framework in Ma et al. (2015). The authors provide background material to establish the connection between MCMC dynamics, the Fokker-Planck equation and gradient flow methods. The GSVGD then follows by applying the fiber-gradient Hamiltonian flow ideas from MCMC dynamics to the Stein-Wasserstein metric, which is then discretized to produce the GSVGD algorithm. The authors test the efficacy of their algorithm on a toy data model and a Bayesian neural network.