De-randomizing MCMC dynamics with the diffusion Stein operator
Year: 2021

Abstract:
Approximate Bayesian inference estimates descriptors of an intractable target distribution – in essence, an optimization problem within a family of distributions. For example, Langevin dynamics (LD) extracts asymptotically exact samples from a diffusion process because the time evolution of its marginal distributions constitutes a curve that minimizes the KL-divergence via steepest descent in the Wasserstein space. Parallel to LD, Stein variational gradient descent (SVGD) similarly minimizes the KL , albeit endowed with a novel Stein-Wasserstein distance, by deterministically transporting a set of particle samples, thus de-randomizes the stochastic diffusion process. We propose de-randomized kernel-based particle samplers to all diffusion-based samplers known as MCMC dynamics. Following previous work in interpreting MCMC dynamics, we equip the Stein-Wasserstein metric with a fiber-Riemannian Poisson structure, with the capacity of characterizing a fiber-gradient Hamiltonian flow that simulates MCMC dynamics. Such dynamics discretize into generalized SVGD (GSVGD), a Stein-type deterministic particle sampler, with particle updates coinciding with applying the diffusion Stein operator to a kernel function. We demonstrate empirically that GSVGD can de-randomize complicated MCMC dynamics, which combine the advantages of auxiliary momentum variables and Riemannian structure, while maintaining the high sample quality from an interacting particle system.

heading: 1 Introduction
text: 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.

heading: 2 MCMC dynamics – the relevant bits
text: 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.

heading: 2.1 Langevin dynamics and its Fokker-Planck equation
text: 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.

heading: 2.2 Gradient flow on (P(Ω),W2)
text: 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.

heading: 2.3 SVGD as gradient flow
text: 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.

heading: 2.4 MCMC dynamics as fiber-gradient Hamiltonian flow
text: 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).

heading: 2.5 Stein’s method and other relevant works
text: Our work enriches the practical application of Stein’s method (Stein, 1972), which studies the class of operators that maps functions to ones with expectation zero w.r.t. a distribution π. Gorham et al. (2019) draw from the findings of the generator method (Barbour, 1988, 1990; Gotze, 1991) and note the same property with infinitesimal generators of Feller processes and their stationary measures, which further extends into a mapping between diffusion-based MCMC sampling and Stein operators, noted as diffusion Stein operators. In this work, we extend beyond sample quality measurement (Gorham et al., 2019) and parameter estimation (Barp et al., 2019) and establish an application of the diffusion Stein operator as a deterministic alternative of MCMC dynamics.
Myriad works (Chen et al., 2018; Liu et al., 2019a; Zhang et al., 2020) seek alternatives to approximating the gradient flow (2) through kernelization. Notably, Chen et al. (2018); Liu et al. (2019a) construct PARVI algorithms by approximating the term ∇ log ρt(θt) in the gradient flow, namely with Blob method (Carrillo et al., 2019), kernel density estimation (Liu et al., 2019a) and Stein gradient estimator (Li and Turner, 2017). The de-randomization of underdamped LD connects to the viewpoint of accelerating PARVI methods. Ma et al. (2019) demonstrate that underdamped LD (Chen et al., 2014) accelerates the steepest descent steps taken by the overdamped LD, forming an analog of Nesterov acceleration for MCMC methods. Wang and Li (2019) present a framework for Nesterov’s accelerated gradient method in the Wasserstein space, which consists of augmenting the energy functional with the kinetic energy of an additional momentum variable.

heading: 3 PARVI for MCMC dynamics – a general recipe
text: In this section, we outline the main contribution of this work, which traces the roadmap delineated by previous works to explore the flow interpretation, as well as the approximation by interacting particle systems of general form MCMC dynamics (Ma et al., 2015) in the form of Itō diffusion:
f(θ) = 1
π(θ) ∇ · [π(θ) (A(θ) + C(θ))] , (9)
Σ(θ) = A(θ), (10)
where A and C are positive-semidefinite and skew-symmetric matrix-valued functions, respectively. This general framework covers all continuous Itō diffusion processes with stationary distribution π, most notably LD, stochastic gradient HMC (Chen et al., 2014), stochastic gradient Nosé-Hoover thermostat (SGNHT) (Ding et al., 2014) and stochastic gradient Riemannian Hamiltonian Monte Carlo (SGRHMC) (Ma et al., 2015).
We demonstrate in this section that applying the fiber-gradient Hamiltonian flow structure of MCMC dynamics to the Stein-Wasserstein metric yields a flow on (P(Ω),WH) that discretizes into GSVGD, a PARVI that updates particles with the diffusion Stein operator (Gorham et al., 2019), suggesting that the infinitesimal generator of MCMC diffusion processes offers a “kernel smoothing” derandomization. Apart from the Stein-Wasserstein metric, the analog of GSVGD generalizing SVGD extends to other interpretations of SVGD, namely that GSVGD takes steepest descent minimizing the KL-divergence in incremental transformation of particles, and that that GSVGD projects the fiber-gradient Hamiltonian flow onto an RKHS.

heading: 3.1 Fiber-gradient Hamiltonian flow on (P(Ω),WH)
text: Similar to (2), we derive the FPE of MCMC dynamics (9)-(10)
ρ̇t = −∇ · (ρtf) + (∇∇) : (ρtA) (11) = −∇ · (ρtf) + (∇∇) : (ρtA) + (∇∇) : (ρtC)︸ ︷︷ ︸
=0
(12)
= ∇ · ( ρt (A + C)∇ log
ρt π
) = ∇ · ( ρt (A + C)∇
δKL [ρt ‖ π] δρt
) , (13)
which is a curve in P(Ω) following the continuity equation ρ̇t + ∇ · (ρtvt) = 0,vt = (A + C)∇ log ρt/π.
We can derive GSVGD from the fiber-Riemannian manifold perspective of (P(Ω),WH) (Liu et al., 2019b). The Riemannian structure of the Stein-Wasserstein metric induces the tangent space TρP(Ω) = {Kρ∇f | f ∈ C∞(Ω)} HD
. An orthogonal projection pρ from L2ρ(Ω) to TρP is uniquely defined as such vector in TρP that satisfies∇ · (ρKρφ) = ∇ · (ρpρ(φ)). Using Theorem 5 from Liu et al. (2019b), we arrive at a fiber-gradient Hamiltonian flow
WKLπ (ρ) = pρ ((A + C)∇ log π/ρ) . (14)
It is straightforward to verify that vA,CH (ρ) = Kρ(A + C)∇ log π/ρ induces the same evolution of distribution asWKLπ (ρ).

heading: 3.2 fRH flow induces the diffusion Stein operator
text: The Wasserstein flow on (P(Ω),WH) following ρ̇t = ∇ · ( ρtv A,C H (ρt) ) requires spatial and
temporal discretization suitable for a particle update. While vA,CH constitutes a vector field on P(Ω), its calculation does not involve an explicit term of∇ log ρ inside the expectation.
vA,CH (ρ) = Kρ [(A + C)∇ log π/ρ] (15) = Eθ′∼ρk(·,θ′)(A(θ′) + C(θ′))∇ log π(θ′)− ∫ k(·,θ′)(A(θ′) + C(θ′))∇ρ(θ′)dθ′
= Eθ′∼ρ(A(θ′) + C(θ′))∇ log π(θ′)k(·,θ′) + ∫ ∇ · ( (A(θ′) + C(θ′))k(·,θ′) ) dρ︸ ︷︷ ︸
integration by parts = Eθ′∼ρ [
f(θ′)k(·,θ′)︸ ︷︷ ︸ weighted drift coefficient (9) + (A + C)(θ′)∇2k(·,θ′)︸ ︷︷ ︸ repulsive force
] . (16)
= Eθ′∼ρ 1 π ∇ · ( π(A + C)k(·,θ′) ) ︸ ︷︷ ︸
diffusion Stein operator T A,Cπ k(·,θ′)
, (17)
where∇2 refers to the gradient w.r.t. the second argument of the kernel function. Interestingly, the kernelized flow vA,CH coincides with calculating the expectation of the diffusion Stein operator T A,Cπ (Gorham et al., 2019) applied to the kernel function k(θ, ·), a family of operators mapping functions to zero-expectation functions under an (un-normalized) distribution π. When ρt is approximated by the empirical measure of its particles ρ̂t = 1/N ∑N i=1 δθ(i)t , we can derive the particle update of
θ̇t = v A,C H (ρ̂t) = Eθ′∼ρ̂tT A,C π k(θt,θ
′) = 1
N N∑ i=1 T A,Cπ k(θt,θ (i) t ),θ (i) 0 ∼ ρ0. (18)
Similar to SVGD, the particle update of GSVGD leverages between a weighted average of the drift vector and a repulsive force, which drives particles towards the target distribution while preventing a mode collapse. The fiber-gradient Hamiltonian flow on (P(Ω),WH) informs a duality between diffusion Stein operator and MCMC dynamics, extending the use of Stein operators beyond a generalization of kernel Stein discrepancy (Gorham et al., 2019).

heading: 3.3 How does GSVGD approximate MCMC dynamics?
text: While the derivation of GSVGD originates from the flow interpretation on WH, the analogy between SVGD and GSVGD expresses in equivalent alternative forms discussed by previous literature, further shedding light on the interpretation of GSVGD.
Originally, the SVGD update is expressed in the form of the functional gradient (Liu and Wang, 2016). For MCMC dynamics, we similarly derive the functional gradient with respect to the push-forward measure ρ = ( id + (A + C)>v ) # ρ,v ∈ HD,
∇v KL [ρ ‖ π]|v=0 = −Eθ′∼ρT A,C π k(·,θ ′) = −vA,CH (ρ). (19)
This presents a generalization of the discussion in Liu and Wang (2016) that GSVGD takes incremental transformation of particles in the direction that minimizes KL-divergence, with such transformation warped by (A,C) and constrained by RKHS.
Alternatively, Liu et al. (2019a) explains SVGD as the projection of the original gradient flow direction v(ρt) = ∇ log π/ρt onto an RKHS. We similarly generalize this explanation for SVGD, as it projects vA,C(ρt) = (A + C)∇ log π/ρt ∈ L2ρ ontoHD, as
vA,CH (ρ) = max · arg max v∈HD,‖v‖HD=1 〈vA,C(ρ),v〉L2ρ . (20)
The analogy of GSVGD goes further when we consider the transportation of N particles as inferring a joint distribution π⊗N , which will be covered in the supplements.

heading: 3.4 GSVGD in practice
text: We can apply GSVGD updates (16) to de-randomize myriad MCMC dynamics (e.g., Table 1). Notably, GSVGD establishes the previously unexplored duality between Riemannian Langevin diffusion (Girolami and Calderhead, 2011) and Riemannian SVGD (Liu and Zhu, 2017).
To fully harness the capacity of our framework, we can introduce auxiliary (momentum) variables r to help with the exploration of probability space, namely to augment the target distribution as π(θ, r) = π(θ)N (r|0,Σ), therefore achieving de-randomized PARVI variant of SGHMC (Chen et al., 2014). Further leveraging a positive definite G(θ) to efficiently explore the target distribution yields a PARVI variant of SGRHMC (Ma et al., 2015).
When GSVGD is used in conjunction with auxiliary momentum variables, we can employ symmetric splitting for leapfrog-like steps for GSVGDs de-randomizing with momentum,
r (i) k+1/2 = r (i) k + 2 vA,CH
( r (i) k ∣∣∣ ρ̂(θk, rk)) , (21) θ (i) k+1 = θ (i) k + v A,C H ( θ (i) k
∣∣∣ ρ̂(θk, rk+1/2)) , (22) r (i) k+1 = r (i) k+1/2 +
2 vA,CH
( r (i) k+1/2 ∣∣∣ ρ̂(θk+1, rk+1/2)) . (23)

heading: 4 Experiments

heading: 4.1 Stein PARVI on toy data
text: To demonstrate the efficacy of our methods, we use various PARVI to infer a 3-component mixture of crescent-shaped target measure inspired by Ma et al. (2015). The parametrizations of Riemannian SVGD and SGRHMC-Stein coincide with those in Ma et al. (2015). While all methods explore the nearest mode well (See figure 2), only the SGRHMC-Stein manage to effectively use the Riemannian formulation of (A,C) to explore the 2 other modes (see figure 3), showcasing the gain in efficiency with application-specific dynamics.

heading: 4.2 GSVGD on Bayesian neural networks
text: We apply GSVGD of advanced MCMC dynamics to the inference of Bayesian neural networks, taking a simple structure of one hidden layer, and an output with Gaussian likelihood. We opt for the fully Bayesian specification of BNN, where the precision parameter of the weight prior and the Gaussian likelihood follows Gamma(1, 0.1). The GSVGD variant used in the experiments include SGHMC (Chen et al., 2014) and SGNHT (Ding et al., 2014), with tar-
get measure π(θ, r, ξ) = π(θ)N (r|0, σ2I)N (ξ|A1, µ−1I), and A = ( 0 0 0 0 AI 0 0 0 0 ) ,C =0 −I 0I 0 (µσ2)−1diag(r)
0 −(µσ2)−1diag(r) 0 . As a MCMC form suitable when used in conjunction with stochastic gradients, SGNHT takes additional temperature parameter ξ. The log-likelihood results in Table 2 demonstrates that the de-randomized MCMC dynamics achieve superior predictive performance than their LD counterparts, largely thanks to the exploration of probability space from the momentum variables (See figure 5), with a notable gain from symmetric splitting. SGNHT-Stein methods show robustness to hyperparameter selection and stochastic gradient noise.

heading: 5 Discussion
text: In this section, we discuss open questions pertinent to the “optimality” of the diffusion Stein operator as a PARVI update, and a direction for future work regarding the fiber-gradient Hamiltonian flow of the chi-squared divergence, inspired by Chewi et al. (2020).

heading: 5.1 Alternatives to the diffusion Stein operator
text: The vector field vA,C(ρ) = (A + C)∇ log πρ is not the unique vector field that induces the time evolution of MCMC dynamics: indeed, two vector fields v1,v2 induce the same curve in P(Ω) when they differ by a divergence-free field: ∇ · (ρ(v1(ρ)− v2(ρ))) = 0. Notably, Liu et al. (2019b) propose an alternative vector field ṽA,C(ρ) = vA,C(ρ)−C∇ log ρ+∇ ·C 1. A projection onto RKHS yields a Stein PARVI with update
θ̇t = Eρt [ 1
π ∇ · [π (A + C)] k(θt, ·) + A∇2k(θt, ·)
] , (24)
consistent with the standard formulation infinitesimal generator. However, the expectation in (24) notably does not converge to zero when ρt = π (Gorham et al., 2019), suggesting that the particles will keep rotating along the trajectories of a divergence-free vector field, as opposed to stopping, in the equilibrium state. Interested readers can find experiments with respect to this alternative form in the supplementary materials.

heading: 5.2 Generalizing LAWGD with the diffusion Stein operator
text: The application of the diffusion Stein operator generalizes Laplacian Adjusted Wasserstein Gradient Descent (LAWGD) (Chewi et al., 2020), which views SVGD as a kernelized Wasserstein gradient flow driven by the chi-squared divergence. Given the first variation of the chi-squared divergence δχ2[ρ ‖ π]
δρ = 2ρ π , we can verify that the vector field v A,C H = Kρ [ (A + C)∇ log πρ ] is rewritten as
2vA,CH = Kπ [ (A + C)∇2π
ρ
] = Kπ [ (A + C)∇δχ
2 [ρt ‖ π] δρt
] . (25)
The chi-squared divergence perspectives enriches GSVGD as a kernelized fiber-gradient Hamiltonian flow minimizing the chi-squared divergence. Replacing the vector field Kπ [ (A + C)∇ dρdπ ] with
u = ∇ · ( Kπ [ (A + C) dρdπ ]) , we get the dissipation formula for the KL-divergence
dKL [ρt ‖ π] dt = −Eπ 〈 ∇ ρ π ,∇ · ( Kπ [ (A + C) ρ π ])〉 = −Eπ [ ρ π AA,Cπ Kπ ρ π ] , (26)
1In fact, a divergence-free vector field C0∇ log ρ − ∇ ·C0 can be constructed out of an arbitrary skewsymmetric matrix-valued function C0
whereAA,Cπ represents the infinitesimal generator that induces the diffusion Stein operator: AA,Cπ g = 1 2π∇ · (π(A + C)∇g). Following the construction of LAWGD, we choose the kernel so that Kπ =( AA,Cπ )−1 . While it remains to be seen what kernel induces this equality, the diffusion Stein operator yields a more flexible formulation, while retaining the superior convergence of LAWGD.

heading: 6 Conclusion
text: We further the application of the Stein operator and its generalizations by proposing GSVGD, an interacting particle system transporting a particle set towards a target distribution π in an emulation of the corresponding MCMC dynamics. We showcase theoretically and empirically that GSVGD helps augment the well-researched field of efficient MCMC dynamics with deterministic interacting particle systems with high-quality samples.
Acknowledgments and Disclosure of Funding
We wish to acknowledge the computational resources provided by Aalto Science-IT Project from Computer Science IT and CSC–IT Center for Science, Finland. This work has been supported by the Academy of Finland (Flagship programme: Finnish Center for Artificial Intelligence FCAI and grants no. 292334, 294238, 319264, 334600) and UKRI Turing AI World-Leading Researcher Fellowship, EP/W002973/1.

Review summaries:
The paper combines the results of two previous works: (Liu, 2017) and (Liu et al., 2019). I'll first describe the contribution of the previous works and then will move to the contribution of the current paper.   (Liu, 2017) provides the theoretical analysis of the Stein Variational Gradient Descent (SVGD). It shows that the evolution of the density under SVGD is described by the Vlasov equation (non-linear analog of the Fokker-Planck equation); moreover, in the space of distributions, SVGD defines the gradient descent minimizing the KL divergence between the current point and the target density (similarly, Langevin Dynamics {LD} defines the gradient descent in the Wasserstein space). Importantly, to derive this result (Liu, 2017) introduces a novel metric in the distribution space (H-Wasserstein distance) based on the Stein operator.   In a similar fashion, (Liu et al., 2019) provides the analysis of different MCMC dynamics (Langevin, HMC, RMHMC, SGHMC) by considering the time evolution of the densities in the space of distributions. They first show how any MCMC dynamics can be reformulated as a completely deterministic procedure by deriving the corresponding vector field in the state space, which induces a vector field in the Wasserstein space. Further, they derive a kind of a Hodge decomposition for the vector field in the Wasserstein space, which provides many insights into existing MCMC dynamics. To describe SGHMC and SG Nose-Hoover thermostat algorithms, the authors introduce the Fiber-Riemannian manifold, which is a Fiber Bundle, but with the Riemannian structure on each fiber. Unlike (Liu, 2017), this paper works in the Wasserstein space (distributions with a finite second momentum equipped with the 2-Wasserstein distance).  The current paper goes the same way as (Liu et al., 2019), but uses the H-Wasserstein distance (introduced in {Liu, 2017}) for the space of distributions. To be more precise, it starts with the same family of MCMC dynamics (introduced in (Ma et al., 2015)) and then uses Theorem 5 from (Liu et al., 2019) changing the projection operator onto the tangent space, which changes because of the different metric. Substituting the different metric into Theorem 5 from (Liu et al., 2019) the authors extend the previous results by additionally describing the SGRHMC algorithm.   - (Liu, 2017) Qiang Liu. Stein variational gradient descent as gradient flow. In Advances in Neural Information Processing Systems, volume 30, pages 3115–3123, 2017. - (Liu et al., 2019) Chang Liu, Jingwei Zhuo, and Jun Zhu. Understanding MCMC dynamics as flows on the Wasserstein space. In International Conference on Machine Learning, pages 4093–4103, 2019. - (Ma et al., 2015) Yi-An Ma, Tianqi Chen, and Emily B. Fox. A complete recipe for stochastic gradient MCMC. In NIPS’15 Proceedings of the 28th International Conference on Neural Information Processing Systems - Volume 2, volume 28, pages 2917–2925, 2015.

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. 

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.  

The authors propose a generalization of Stein Variational Gradient Descent (SVGD) that hopes to produce more reliable surrogate measures of an invariant measure pi. SVGD can be viewed as a procedure that is minimizing the KL divergence by taking (functional) gradient steps in an RKHS; it can similarly be viewed as an approximation to the gradient flow induced by Langevin dynamics by accommodating H^d as its function space. This paper studies the generalization of SVGD that occurs by replacing Langevin dynamics by the dynamics induced from an Ito diffusion with invariant measure pi. Using Ito diffusions in lieu of the overdamped Langevin diffusion produces deterministic analogs of SGRLD, SGHMC and SGRHMC. This also permits the formulation of GSVGD, a variant of SVGD that incorporates auxiliary variables in a flavor similar to HMC. The authors demonstrate the efficacy of this PARVI method on a toy 3-component GMM and on Bayesian neural nets for a suite of datasets.