SHPOS: A Theoretical Guaranteed Accelerated Particle Optimization Sampling Method

1 Introduction

Sampling from a Bayesian posterior distribution lies at the core of many modern machine learning tasks, such as topic modelling [Gan et al., 2015], reinforcement learning [Liu et al., 2017], and Bayesian neural networks [Hern´andezLobato and Adams, 2015]. Particle based Variational Inference (ParVI) methods have recently drawn great attention due to their empirical success in approximating the target posterior distribution [Liu and Wang, 2016; Liu et al., 2017; Feng et al., 2017; Chen et al., 2018; Liu and Zhu, 2018]. Typically, these methods update a finite set of interacting particles deterministically to approximately simulate infinite-particle gradient flows on the Wasserstein space P2(X).

One representative method of this type is the Stein Variational Gradient Descent (SVGD) method [Liu and Wang, 2016], which updates the particles according to a gradient flow described by the Vlasov equation [Liu, 2017; Braun and Hepp, 1977]. Subsequently, by exploiting the Riemannian structure of the Wasserstein space P2(X), [Liu et al., 2019] proposed a Nesterov's-acceleration variant of SVGD called SVGD Wasserstein Nesterov's method (SVGD-WNes).

∗Zhijian Li and Chao Zhang contribute equally.

†Contact Author.

However, ParVI methods have an intractable pitfall that particles tend to collapse under certain condition due to the deterministic-update fashion with a limited number of particles [Zhang et al., 2020a; Zhuo et al., 2018], which indicates a large deviation of the particles' empirical distribution to target distribution. To understand the influence of finiteparticle approximation to the infinite-particle gradient flows, Liu et al. [2019] provided a unified theory on the approximation property of different ParVI methods. They show that existing finite-particle ParVIs can be regarded as essentially smoothing operations on gradient flows, in the form of either smoothing the density or smoothing functions. Though the underlying gradient flows usually evolve towards the target distribution, the smoothing operations deteriorate the convergence of the particles' empirical distributions in ParVIs to the target distribution.

Inspired by the random exploration in the dynamics-based Markov Chain Monte Carlo methods, recent researches relieve the particle-collapsing phenomenon by introducing additional stochasticity into ParVI methods [Zhang et al., 2020a; Zhang et al., 2020b]. Zhang et al. [2020a] integrated the gradient flow of the first-order Overdamped Langevin Dynamics (OLD) into the Vlasov equation used in SVGD, and proposed a new method called Stochastic Particle Optimization Sampling (SPOS). Specifically, SPOS updates a set of particles following the same mechanism as in SVGD with an additional drift term and an extra Gaussian random noise induced by OLD. Following the analysis framework of OLD based MCMC methods, they conduct the non-asymptotic convergence analysis of SPOS, and establish a bound on the 2-Wasserstein distance between the empirical distribution of particles and the target distribution. However, observations in the dynamics-based MCMC indicate that OLD is not the most efficient dynamics although it is simple and easy-toanalyze. Actually, the error bound in SPOS increases with the iterations, therefore the step size should be samll in order to restrict the error in a desirable level.

In this work, we propose an accelerated particle optimization sampling method called Stochastic Hamiltonian Particle Optimization Sampling (SHPOS) by introducing the secondorder Underdamped Langevin Dynamics (ULD) [Jaakkola and Jordan, 1997] into the flow of SVGD. Compared with the first-order OLD, ULD possesses an auxiliary momentum term, and can be regarded as an accelerated secondorder dynamics of OLD. Typically, ULD based methods have both better theoretical guarantees and more competitive practical performance than their OLD based counterparts [Chen et al., 2014; Cheng et al., 2018; Zou et al., 2018; Zou et al., 2019]. We utilize a variant of the first order exponential integrator scheme to discretize the corresponding stochastic differential equation of the augmented flow, which is obtained by integrating the gradient flow of ULD into the Vlasov flow. The contributions of our paper are listed as follows:

• As far as we know, SHPOS firstly introduce secondorder ULD into the Vlasov flow to accelerate particle optimization sampling. The particles in SHPOS can be regarded as drifted by the (stochastic-)gradient of the log-posterior of the target distribution, the repulsive force between particles, and an extra random standard Gaussian force. By balancing these forces, particles explore the high probability space of the target distribution rapidly and avoid particle-collapsing (refer to the experiment for more details). We also propose a variancereduced stochastic gradient variant of SHPOS by using the historical gradient information to build a recursively updated gradient estimator. This variant is shown to be more suitable for Bayesian learning tasks with largescale datasets and complex models.
• We establish the non-asymptotic convergence guarantee of SHPOS. Specifically, we show that the 2-Wasserstein distance between the particles' distribution in the k-th iteration of SHPOS and the target distribution is in the order of ˜O(M -1/2 + exp(-kη) + ηd1/2 ), where M is the number of particles, d denotes the dimension of variable, and η is the step size. The result shows that the approximation error decreases as the particle number M and the iteration k increase, while the error bound ˜O(M -1/2 + exp(-kη) + Mηd3/2k1/2 ) of SPOS increases with particle numbers and iterations. Thus, SHPOS could use a much larger step size compared with SPOS to achieve the same sampling accuracy.

We evaluate our method on a list of tasks, including both synthetic and real datasets. The empirical results demonstrate the superiority of our method over the state-of-the-art.

2 Preliminaries

In this paper, we focus on Bayesian learning tasks, i.e., sampling from the target posterior distribution

where x ∈ X ⊂ Rd is the model parameter, D = {Di}Ni=1 denotes the dataset and U(x) = - log p(x) - log p(D|x) denotes the potential energy function. SVGD and ULD based MCMC methods are two recently proposed inference methods for solving Bayesian learning tasks.

2.1 Stein Variational Gradient Descent

SVGD is firstly proposed by [Liu and Wang, 2016] as an efficient approximate inference method for sampling from the target distribution. Basically, it simulates the following density gradient flow

where φt denotes the evolutionary density at time t and ϕφt is the evolution function of the following form

with F(x) = ∇U(x) denoting the gradient of the potential function and κ(x, x′ ) denoting a kernel function, such as RBF kernel. Actually, ϕφt is a direction of steepest descent that minimize the KL divergence between the evolutionary density φt and the target distribution p∗ over the unit ball in Reproducing Kernel Hilbert Space induced by the kernel function κ [Liu and Wang, 2016]. The density gradient flow (1) is a special type of Vlasov equation and its stationary distribution is p∗ under certain regularity conditions [Liu, 2017].

SVGD draws a set of particles {xi,0}Mi=1 from an initial distribution q0, and updates the particles iteratively according to the following rule:

where η denotes the step size. On the RHS of the update rule, the -F (xj,k)κ(xi,k, xj,k) term pushes particles towards high probability regions, while the second term ∇xj,k κ(xi,k, xj,k) forces the particles to move away from each other. The stochastic gradient technique, i.e., constructing a stochastic approximation from a mini-batch of data to replace the full gradient, has been adopted by SVGD to reduce the per-iteration computational cost in large-scale and complex Bayesian learning tasks. Moreover, SVGD uses empirical distribution of the particles to approximate target distribution. However, Zhang et al.[2020a] pointed out that SVGD has an unintended particle-collapsing pitfall, i.e., particles tend to collapse to a local mode under particular conditions.

2.2 Dynamics-BasedMCMC

Recently, different dynamics-based MCMC methods have been proposed by discretizing certain dynamics whose stationary distribution (or its marginal) is the target distribution. Two mostly adopted dynamics are the Overdamped Langevin Dynamics (OLD) and the Underdamped Langevin Dynamics (ULD). OLD is a first-order dynamics with only one variable, while ULD involves an extra momentum term, thus can be viewed as a second-order dynamics. Basically, ULD based methods usually outperform their OLD based counterparts in both practice and theory, due to the extra momentum term. However, their analysis are much more complicated and elaborate as their structures are more complex than their OLD based counterparts.

The Underdamped Langevin Dynamics is described by the following stochastic differential equation (SDE):

where γ > 0 is the friction parameter, u is the inverse mass, xt, vt ∈ Rd denote the position and velocity variables of the dynamics respectively, and Bt represents standard Brownian motion. Let zTt = [xTt , vT t ], according to the Fokker-Planck (FP) equation [Risken, 1996], the density gradient flow of z is as follow:

where Ψt denotes the density of zt. The stationary distribution of variable zt is p(z) ∝ exp (-U(x) - ∥v∥22 /(2u)), and the marginal distribution of xt converges to the target distribution p∗ .

Directly using the Euler-Maruyama discretization scheme of ULD (3) gives rise to the Underdamped Langevin MCMC (UL-MCMC) method [Kloeden and Platen, 1992]:

where η > 0 is the step size, ϵk ∼ N(0, Id×d) is a standard Gaussian random vector, and Id×d is a d × d identity matrix. To make the algorithm scalable in large-scale tasks, SGHMC method [Chen et al., 2014] is proposed by replacing the full gradient with a mini-batch estimator

where Bk denotes a mini-batch of size |B|, and F(xk, ξk) is an unbiased estimate of F(xk). Ever since, different variance-reduced stochastic gradient ULD MCMC methods has been proposed to reduce the influence of stochastic variance, e.g. SVR-HMC [Zou et al., 2018], SAGA-HMC [Li et al., 2019] and HSG-HMC [Zhang et al., 2021].

3 Methodology

In this section, we introduce our Stochastic Hamiltonian Particle Optimization Sampling method as in Algorithm 1. We first construct the following mixed density gradient flow by integrating the gradient flow of ULD into the Vlasov flow:

where zTt = [xTt , vT t ], Φt denotes the density of zt, φt denotes the marginal distribution of variable xt, and ϕφt is defined as in (2).

In the following lemma, we give the corresponding SDE of flow (7) and show that the marginal distribution φt of variable xt converges to the target distribution p∗ .

Lemma 1. The density flow of the following SDE is exact (7).

dx =vdt,

Algorithm 1 Stochastic Hamiltonian Particle Optimization Sampling

Input: Initial particles {xi,0}Mi=1 and the corresponding momentum variable {vi,0}Mi=1 . Weight parameter β, friction parameter γ, inverse mass u, and step size η.

• 1: for k = 0, 1, ..., T -1 do
• 2: Uniformly sample a subset of index Bk ⊂ [n]
• 3: for i = 1, 2, ...,M do
• 4: Update xi,k+1 and vi,k+1 as

with gi,k as the full gradient F(xi,k) (or its stochastic estimate Gmini i,k (14) or Gvr i,k (15)).

• 5: end for
• 6: end for
• 7: Output:{xi,T}Mi=1 .

where Y ∈ Rd is a random sample from evolutionary density φt but independent of x and v, and Bt represents standard Brownian motion. The stationary distribution of this SDE is p(z) ∝ exp (-U(x) - ∥v∥22 /(2u)), and its marginal distribution φt converges to the target distribution p∗ .

According to the system (8), x and v can be regarded as the position of a particle and its associated velocity. Besides, it can be verified that the Vlasov flow endows an repulsive force on the particle and β controls the magnitude of particle interaction. As this system contains an extra Hamiltonian velocity part compared to the system used in SPOS, we call our method a Hamiltonian method. We list a feature-by-feature comparison of different sampling methods in Table 1 to give a more intuitive understanding of SHPOS method.

We construct our SHPOS method by simulating system (8). Note that (8) contains the marginal distribution φt of xt, which is intractable to calculate. Following the idea in ParVI, we maintains a set of M particles and use the empirical distribution to approximate φt in SDE (8). Consequently, we obtain an multi-particle approximate dynamics of (8), where the dynamics of each particle is as follow:

Here, we focus on the RBF kernel κ(x, x′ ) = exp(-∥x - x′ ∥22 /h), and replace κ(x, x′ ) with K(x-x′ ) = exp(-∥xx′ ∥22 /h). Note that this kernel is also adopted in SPOS and many ParVIs [Liu and Wang, 2016; Liu et al., 2019].

In order to sample from the target distribution, we can apply numerical integrators to discretize the multi-particle dynamics (11). Specifically, we utilize a variant of the first order exponential integrator scheme [Zou et al., 2018], which leads to the following updates for variables xi,k and vi,k:

where γ, u, η, β are tunable parameters, ϵxi,k and ϵvi,k are Gaussian random vectors with zero mean and covariance matrices satisfy

In many real-world Bayesian learning tasks with largescale datasets and complex models, obtaining the exact gradient F(x) is computationally expensive or even prohibitive, and only unbiased estimates F(x, ξ) of F(x) are available. To deal with this situation, one typical approach is to replace the full gradient with a mini-batch estimator

However, this estimator usually suffers from high stochastic variance, which will degrade the performance of algorithms. To relieve the influence caused by stochastic variance, we propose to use the following estimator

where ˜ Fk(x) = 1 |B| ∑ξk∈Bk F(x, ξk). This estimator has been widely adopted in the optimization literature [Cutkosky and Orabona, 2019] and stochastic gradient MCMC methods [Zhang et al., 2021] to relieve the influence of variance. Specifically, Gvr i,k is a combination of two parts, gu,k is an unbiased stochastic gradient estimator with high variance, while gb,k is a biased estimator with low variance. This merit lies in using a proper weight parameter ρk to strike a balance between bias and variance.

With different gradient estimators, we obtain different variants of SHPOS. We denote the variant with Gmini and Gvr as SHPOS-MINI and SHPOS-VR, respectively.

4 Theoretical Analysis

In this section, we provide the non-asymptotic convergence analysis for the proposed method under 2-Wasserstein distance W2. Given two probability measures ρ1 and ρ2, the W2 distance is defined as

where Γ(ρ1, ρ2) denotes a set of joint distributions on Rd×Rd with marginal distribution ρ1 and ρ2. In our analysis, we first provide the approximate error between the stationary distribution ν∞ of ˜xi,t in the continuous-time dynamics (11) and target distribution in terms of W2 distance, then bound the W2 distance between ν∞ and the density µkη of the discretetime iterative xi,k (9), and draw our conclusion by combining these results together as

If we initialize all the particles with the same distribution, they would endow the same distribution during the evolution due to the exchangeability of the particles according to [Zhang et al., 2020a], and we denote the density of all particles in the k-th iteration as µkη. We only list the main theorems here due to the limit of space. We refer readers to Appendix for more detailed proofs.

4.1 Assumptions

Following [Zhang et al., 2020a; Zhang et al., 2020b], we make the following standard assumptions on F and K.

Assumption 1. F satisfies the following conditions:

• F is LF -Lipschitz, i.e., ∥F(x)-F(y)∥2 ≤ LF∥x-y∥2.
• There exists positive mF such that ⟨F(x) - F(y), x - y⟩ ≥ mF ∥x - y∥2.

Assumption 1 means that the potential function is smooth and strongly convex. This assumption is quite standard in the analysis of existing dynamics-based sampling methods, including SPOS and many underdamped Langevin MCMC methods [Cheng et al., 2018; Zou et al., 2018]. It is satisfied in many Bayesian learning tasks, such as Bayesian Logistic Regression and Bayesian Ridge Regression. It is techinically possible to extend our theoretical analysis to the non-convex setting, and we leave this as future work.

Assumption 2. F(x) is bounded, i.e. there exists an positive constant HF such that ∥F(x)∥2 ≤ HF , and F(0) = 0.

It is easy to be varified that the first part of this assumption holds if F is Lipschitz and the domain X of x is a compact set. Besides, we can make F(0) = 0 by transforming coordinate system.

Assumption 3. The kernel function K is LK-Lipschitz and L∇K smooth, i.e. ∥∇K(x1 - x2) - ∇K(y1 - y2)∥2 ≤ L∇K∥x1 - x2 - y1 + y2∥2.

Since we focus on RBF kernel in this paper, this assumption can be satisfied by setting the bandwidth h large enough.

Besides, we also assume that the stochastic estimator F(x, ξ) is unbiased and of bounded variance.

Assumption 4 (Unbiasedness and Bounded Variance). For all x ∈ Rd , the unbiased estimator F(x, ξ) of F(x) has a bounded variance E∥F (x, ξ)-F (x)∥22 ≤ σ2 , where the sample variable ξ is drawn from certain fixed distribution (e.g. the uniform distribution on the indices set or the hyperparameter prior).

6

4.2 Main Theorems

In the following theorem, we establish an upper bound for the 2-Wasserstein distance W2(ν∞, p∗ ) between the stationary distribution ν∞ of ˜xi,t in the continuous-time dynamics (11) and the target distribution p∗ .

Theorem 1 (Finite Particle Error). Under assumption 1, 2, and 3, let ν0 = φ0, there exists a positive constant c1 such that

Recall that p∗ is the stationary distribution of the continuous dynamics (8), and (11) is an finite-particle approximation of (8). This theorem shows that the approximation accuracy increases with the growth of particle number .

The following theorem spells out the W2 distance between ν∞ and the density µkη of the discrete-time iterative xi,k (9).

Theorem 2 (Time Discretization and Stochastic Gradient Error). Under assumption 1, 2, 3, and 4, if we run Algorithm 1 with stochastic gradient estimator (14) and small step size η such that exp (-λη) + c2η2 < 1, the 2-Wasserstein distance W2(µkη, ν∞) is bounded as follow:

where λ = mF 2LF - βL∇K - β 1 2 LF - βHFLK, (c2, c3) are positive constants independent of (M, η, k), and Ω is in the order of ˜O(√d).

Note that the third term in the RHS of (19) belongs to the stochastic gradient error, and this term diminishes if we use the full gradient F in SHPOS.

Based on Theorem 1 and 2, we present our main theorem about the 2-Wasserstein distance between the density µkη of particles generated in SHPOS and the target distribution p∗ .

Theorem 3 (The Overall Bound). Under the assumptions of theorem 1 and 2, the 2-Wasserstein distance between µkη and target distribution p∗ is in the order of

If we use the full gradient F, the 2-Wasserstein distance between µkη and target distribution p∗ is in the order of

6

(b) SVGD-WNes

6

(d) UL-MCMC

Figure 1: Results on sampling from a Gaussian mixture distribution.

The error bound of SHPOS in Theroem 3 is better than that of SPOS under the same condition. In SPOS [Zhang et al., 2020a], the authors give a bound in terms of 1Wasserstein distance between the empirical distribution of particles and the target distribution p∗ , which is in the order of ˜O(M -1/2 + exp (-λ′kη) +Mηd3/2k1/2 ), where λ′ is a constant in the same order as λ. This result demonstrates that the sample accuracy of SPOS decreases with particle numbers and iterarions, and a small step size η is needed to restrict the 1-Wasserstein distance in a desirable level. Note that 1Wasserstein distance is a weaker metric than 2-Wasserstein distance (W1(ρ1, ρ2) ≤ W2(ρ1, ρ2)), thus our results in Theorem 3 also hold in terms of the 1-Wasserstein distance. Consequently, it can be verified that the bound (21) of SHPOS is strictly better than the result of SPOS. Futhermore, the sampling accuracy of SHPOS increases as the particle number M and the iteraion k increase, and a much larger step size can be used compared with SPOS to achieve the same accuracy.

5 Experiments

We follow the conventions in dynamics-based sampling literature [Liu and Wang, 2016; Zhang et al., 2020b] and conduct empirical studies on one synthetic experiment and two real-world applications. Three types of methods are selected as baselines: ParVI methods (SVGD, SVGD-WNes), dynamics-based MCMC method (UL-MCMC), and Stochastic Particle Optimization Sampling method (SPOS). For a fair comparison, we use a parallel version of the sequential UL-MCMC method, i.e., maintain M independent chains simultaneously. We focus on the stochastic gradient methods in the two real-world applications since it is difficult to calculate the full gradient. We compare our SHPOS-MINI and SHPOS-VR methods with the mini-batch estimator based methods, i.e. SVGD-MINI, SPOS-MINI and SVGD-WNesMINI, and the variance-reduced stochastic gradient methods SVRG-SPOS [Zhang et al., 2020b] and UL-MCMCVR [Zhang et al., 2021]. We list the information of datasets and parameter tuning of different methods in our Appendix.

5.1 Synthetic Experiment

We consider sampling from a 2-D Gaussian mixture distribution with three modes, whose density is defined as follow:

where a = [2, 2]T , and Σ11 = Σ22 = 6 with correlation Σ21 = Σ12 = -0.98 ∗ 6 = -5.88. We use 1000 particles

-2

-4

-6

6

(e) SHPOS

that initialized by drawing from a Gaussian distribution with mean [-4, 2]T and variance 0.252 .

In Figure 1, we report the sampling results generated by each algorithm. From the results, the particle-collapsing phenomenon can be observed in the deterministic ParVI methods SVGD and the accelerated SVGD-WNes method. The particles in these two methods only aggregate near two of the three modes and fail to explore region near the third mode. For the methods with stochastic exploration, i.e. SPOS, UL-MCMC and SHPOS, though SPOS and UL-MCMC could explore all the three high probability space, our SHPOS method has the best performance as it contains both an extra momentum term and a repulsive force to explore the whole space more efficiently. In Table 2, we report the W2 distance between density of random samples generated by each algorithm and the true distribution. We use Sinkhorn method with penalty 10-2 to calculate the W2 distance. The results show that our SHPOS method achieves the best sampling accuracy.

5.2 Bayesian Logistic Regression

Bayesian Logistic Regression (BLR) is a robust binary classification task. Given a dataset D = {xi, yi}Ni=1 , the likelihood of the BLR task is Sigmoid(yiwT xi), where xi ∈ Rd denotes the sample covariate vector, yi ∈ {-1, 1} denotes the label, and w is the model parameter with a standard multivariate Gaussian prior N(0, I). Four publicly available benchmark datasets from LIBSVM 1 , a3a, w8a, a8a, and ijcnn1 are used for evaluation. We report negative log-likelihood versus the number of data passes with 50 particles on datasets a3a and w8a in Figure 2, and we refer readers to Appendix for more results. The experimental results demonstrate that our variance-reduced SHPOS method SHPOS-VR achieves the best performance among all the comparisons. Moreover, it can be verified that the variance-reduced methods outperform their mini-batch counterparts, and both the particle interaction and the momentum contribute to a better performance.

5.3 Bayesian Neural Network

In this experiment, we study a regression Bayesian posterior learning task with Bayesian Neural Network (BNN) on 6 datasets from UCI 2 and LIBSVM. Following the settings from [Liu and Wang, 2016], we use a Gamma(1, 0.1) prior for the inverse covariance and adopt a one-hidden-layer neural network with 50 hidden units. In Figure 3, we report the root mean squared error (RMSE) versus the number of data passes with 20 particles, and we refer readers to Appendix for more results. The proposed SHPOS-VR method achieves the best performance, which shows its superiority over other methods. Besides, the results demonstrate that all the stochastic noise, the repulsive force, the extra momentum term, and

1 https://www.csie.ntu.edu.tw/ cjlin/libsvmtools/datasets/

2 http://archive.ics.uci.edu/ml/datasets.php

1e-2

1e4

the variance reduction technique improve the performance. Note that the SVRG estimator based method SVRG-SPOS performs badly in this neural network task. Actually, similar phenomenon has been widely observed in the deep learning literature, and [Defazio and Bottou, 2019] shows that the SVRG estimator is not suitable for nonconvex models. Conversely, the recursive estimator Gvr used in SHPOS-VR and UL-MCMC-VR has been shown to be a better estimator for nonconvex problem [Cutkosky and Orabona, 2019; Zhang et al., 2021], which agrees with our results.

6 Conclusion

In this paper, we propose a second-order particle optimization sampling method called Stochastic Hamiltonian Particle Optimization Sampling (SHPOS), which involves an extra momentum term to achieve acceleration. In SHPOS, the particles can be regarded as drifted by the (stochastic-)gradient of the target distribution, the repulsive force between particles, and an extra random standard Gaussian force. By balancing these forces, particles in SHPOS explore the high probability space of the target distribution rapidly and avoid particle-collapsing. Besides, we also propose a variancereduced stochastic gradient variant of SHPOS for large-scale tasks. We establish the non-asymptotic convergence analysis for SHPOS, and show that it enjoys a faster convergence rate than the first-order SPOS. The empirical results also demonstrate the superiority of our methods over the state-of-the-art.

Acknowledgements

This work is supported by National Key Research and Development Program of China under Grant 2020AAA0107400, Zhejiang Provincial Natural Science Foundation of China under Grant No. LZ18F020002, Alibaba-Zhejiang University Joint Research Institute of Frontier Technologies, and National Natural Science Foundation of China (Grant No: 61672376, 61751209, 61472347).