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: In this paper the authors generalise the Stein Variational Gradient Descent (SVGD) algorithm by considering an alternative flow, essentially equipping the Stein Geometry with a Riemmanian Poisson structure.   This enables a particle evolution scheme which is no longer a gradient flow of the KL divergence in the Stein geometry.    The resulting flow differs from the original SVGD flow through the introduction of an "irreversible term" characterised by a Skew-Symmetric matrix $C$.    Adopting this approach, the authors are able to consider augmented variable strategies to improve the performance of SVGD by promoting superior exploration, using analogues of underdamped Langevin and Nose-Hoover type dynamics.   This is demonstrated on some toy examples and on a Bayesian Neural Network model.