DPVI: A Dynamic-Weight Particle-Based Variational Inference Framework

1 Introduction

Recently, Particle-based Variational Inference (ParVI) methods have drawn much attention in the Bayesian inference literature, due to their success in efficiently approximating the target posterior distribution π [Liu and Wang, 2016; Liu and Zhu, 2018; Liu and Wang, 2018; Pu et al., 2017; Zhu et al., 2020]. The core of ParVIs lies at evolving the empirical distribution of M fixed-weight particles by simulating a continuity equation through its easy-to-calculate finiteparticle position transport approximation [Liu et al., 2019a]. Typically, the continuity equation is constructed according to the Wasserstein gradient flow of certain dissimilarity functional F(µ) := D(µ|π) vanishing at µ = π [Liu et al., 2019a]. By using different dissimilarity functionals F and position transport approximations, several ParVIs have been proposed, e.g., Stein Variational Gradient Descent (SVGD)

∗Chao Zhang and Zhijian Li contribute equally.

†Contact Author.

method [Liu and Wang, 2016], Blob[Chen et al., 2018a] GFSD [Liu et al., 2019a], the Kernel Stein Discrepancy Descent (KSDD) [Korba et al., 2021].

Fixed weight restriction. Existing ParVIs have a common fixed weight restriction, i.e., they all keep the particles' weights fixed during the whole procedure, and only update the positions of particles according to the position transport approximation derived from a continuity equation. This restriction severely confines the empirical distribution's approximation ability, especially when the particle number M is limited (depicted in Figure 1). To mitigate the influence of this restriction, existing ParVIs require plenty of particles to obtain satisfying approximation accuracy [Liu and Wang, 2018; Korba et al., 2020]. As a result, a large amount of computation is usually needed since the per-iteration computational cost in ParVIs is typically in the square order of M. Actually, a huge deviation between the empirical distribution and the target π is often observed when the particle number is insufficient due to a limited computational budget [Zhang et al., 2020a; Zhang et al., 2020b].

Thus, it is in great need to find an effective weight adjustment strategy and design dynamic-weight ParVI algorithms which could achieve a high approximation accuracy with fewer particles and hence less computation. Note that, though the continuity equation underlying existing fixedweight ParVIs can be directly transformed into an easy-tocalculate position transport approximation, adjusting weights according to the continuity equation is generally infeasible as the divergence operator in it would introduce great computational challenge. Constructing effective algorithms to evolve a set of dynamic-weight particles towards the target π is still an open problem in the ParVI field.

To tackle this problem, we propose to dynamically adjust particles' weights according to the Fisher-Rao reaction flow of the underlying dissimilarity F, and design a continuoustime composite flow, which evolves the positions and weights of M particles simultaneously. Specifically, the composite flow is a combination of a finite-particle approximation of the reaction flow and a finite-particle position transport approximation of the continuity equation. Different dynamic-weight ParVI algorithms can be obtained by discretizing the continuous flow with different discretization schemes and dissimilarities F. The contribution of our paper are listed as follows:

• We show that the mean-field limit of our proposed composite flow is actually the gradient flow of F in the Wasserstein-Fisher-Rao space, which leads to an extra decrease of F compared with both the Wasserstein gradient flow underlying existing fixed-weight ParVI algorithms and the Fisher-Rao reaction flow.
• We propose a general Dynamic-weight Particle-based Variational Inference (DPVI) framework, which utilizes an Euler discretization of the composite flow and adopts a Gauss-Siedel-type strategy to update the positions and the weights. Note that the weight adjustment step can be implemented without bringing much extra computation compared to the position update step. By adopting different dissimilarities F in DPVI, we can obtain different efficient dynamic-weight ParVI algorithms.
• We propose three efficient DPVI algorithms by using different dissimilarities F and their associated finiteparticle approximations in the general framework. Besides, we also derive three duplicate/kill variants of our proposed algorithms, where a probabilistic discretization to the weight adjustment part in the composite flow is used to dynamically duplicate/kill particles, instead of adjusting the particles' weights continuously.

We evaluate our algorithms on various synthetic and realworld tasks. The empirical results demonstrate the superiority of our dynamic weight strategy over the fixed weight strategy, and our DPVI algorithms constantly outperform their fixedweight counterparts in all the tasks.

Notation. Given a probability measure µ on Rd , we denote µ ∈ P2(Rd ) if its second moment is finite. For a given functional F(µ) : P2(Rd ) → R, δF(ρ) δρ : Rd → R denote its first variation at µ = ρ. Besides, we use ∇ and ∇ · () to denote the gradient and the divergence operator, respectively.

2 Preliminaries

2.1 Particle-Based Variational Inference Methods

When dealing with Bayesian inference tasks, classical Variational Inference methods approximate the target posterior π with an easy-to-sample distribution µ, and recast the inference task as an optimization problem over P2(Rd ) (or its subspace) [Ranganath et al., 2014]:

To solve this optimization problem, one can consider a descent flow of F(µ) in the Wasserstein space, which transports any initial distribution µ0 towards the target π [Wibisono, 2018]. Specifically, the descent flow of F(µ) is described by the following continuity equation [Ambrosio et al., 2008]:

where vµt : Rd →Rd is a vector field that defines the direction of position transportation. To ensure a descent of F(µt) over time t, the vector field vµt should satisfy the following inequality [Ambrosio et al., 2008]:

A straightforward choice of vµt is vµt = -∇ δF(µt) δµt , which is actually the steepest descent direction of F(µt). From now on, we fix this choice of vµt in (2). We note that the continuity equation (2) with this vµt is also known as the Wasserstein gradient flow of F.

To simulate the Wasserstein gradient flow of F, existing ParVIs evolve a set of M fixed-weight particles and use the empirical distribution ˜µt = ∑Mi=1 ait δxit to approximate µt in (2), where xit and ait (usually set to 1/M) denote the position and the weight of the i-th particle at time t, respectively. Specifically, ParVIs update the position of each particle xit according to the following finite-particle position transport approximation of the continuity equation (2) [Chen et al., 2018a; Liu, 2017; Craig and Bertozzi, 2016]:

where v˜µt is an approximation of vµt through the empirical distribution ˜µt. It can be verified that the empirical distribution ˜µt weakly converges to µt defined in (2) when M → ∞ under mild conditions [Korba et al., 2020; Liu, 2017; Liu and Wang, 2018]. Therefore, one can obtain different ParVIs by choosing proper v˜µt and discretizing (4) with certain scheme (the first-order explicit Euler discretization is set as a default). Note that existing ParVIs only update the position xit and keep the weight ait fixed during the whole procedure, since adjusting weight ait according to (2) involves the calculation of the second-order derivative of the first variation due to the divergence operator, which is usually expensive.

To develop a ParVI method, it remains to select a proper dissimilarity functional F and construct an approximation v˜µt of the vector field vµt . From the seminal SVGD to the subsequent Blob and GFSD, KL-divergence is widely adopted as the underlying dissimilarity functional [Liu and Wang, 2016; Chen and Zhang, 2017; Chen et al., 2018a; Liu et al., 2019a; Liu et al., 2019b; Zhang et al., 2020a]. The associated vector field is defined as follows [Jordan et al., 1998; Liu et al., 2019a]:

As ∇logµt is undefined with the empirical distribution ˜µt, the KL-divergence based ParVIs use different approaches to construct suitable approximations to the vector field. In SVGD, Liu and Wang [2016] restrict the vector field vµt within the unit ball of a Reproducing Kernel Hilbert Space (RKHS), and propose to approximate vµt by

where K is a kernel function, such as the Radial Basis Function (RBF) kernel. Subsequently, Blob [Chen et al., 2018a] reformulates the intractable term ∇logµt by partly smoothing the density with a kernel function K, while GFSD [Liu et al., 2019a] directly approximates µt by ˜µt ∗ K, where ∗ denotes the convolution operator. Recently, Kernel Stein Discrepancy Descent (KSDD) [Korba et al., 2021] method considers the Kernel Stein Discrepancy (KSD) [Liu et al., 2016] as the dissimilarity F, whose first variation is compatible with empirical distribution and can be calculated directly.

2.2 Fisher-Rao Distance and Reaction Flow

The Fisher-Rao distance [Rao, 1945; Kakutani, 1948] is a metric defined for general positive Radon measures and allows comparing measures with mass variations. The positive Radon measure space equipped with the Fisher-Rao distance is known as the Fisher-Rao space. For a given dissimilarity functional F(µ), its gradient flow in the Fisher-Rao space, also named as the Fisher-Rao reaction flow [Gallou¨et and Monsaingeon, 2017], is described by the following equation [Wang and Li, 2019; Liero et al., 2016]:

where αµt : Rd → R represents a construction/destruction function of mass. Since the average variation of µt equals zero due to the second integral term, the total mass of µt is conserved during the whole procedure [Lu et al., 2019; Rotskoff et al., 2019]. The target distribution π is actually an invariant distribution of this flow, as the first variation of the dissimilarity functional F vanishes at π, i.e., δF(π) δπ = 0. It can be verified that with a proper F, the process (6) starting from a given µ0 evolves towards the target distribution π [Kondratyev et al., 2016].

For a fixed position x, the process (6) provides an effective way to adjust its density (weight) µt(x) at each time t according to the function αµt . Thus, given a set of weighted particles and its empirical distribution ˜µt, one can adjust the weight by discretizing the reaction flow with an empirical approximate construction/deconstruction function α˜µt . Though the reaction flow has been adopted in particle-based methods in other literature, such as MCMC [Lu et al., 2019], global minimization [Rotskoff et al., 2019] and generative models [Mroueh and Rigotti, 2020], to the best of our knowledge, it has never been adopted in the ParVI literature to adjust the weights of particles.

3 Methodology

In this section, we first construct a composite flow that evolves positions and weights of particles simultaneously and investigate its mean-field property. Then, we develop our DPVI framework by discretizing this composite flow. We finally provide three effective DPVI algorithms by using different dissimilarity functionals F (KL-divergence and KSD) and finite-particle approximations. Besides, we also derive the duplicate/kill variants of our proposed algorithms.

3.1 Continuous-Time Composite Flow

Based on the position transport approximation (4) for displacing the position and the Fisher-Rao reaction flow (6) for adjusting weight in previous sections, we consider the following composite flow that evolves the positions xi 's and the weights ai 's of M particles simultaneously.

For ease of notation, we use Uµ and vµ to denote the first variation of F(µ) at µ and the vector field associated with it, respectively, i.e., Uµ = δF(µ) δµ and vµ = -∇Uµ. Although the mean-field limit of the empirical distribution with either the position update part dxit or the weight adjustment part dait alone has the target distribution π as its stationary distribution, the behaviour of the empirical distribution ˜µt in (7) remains unknown. Thus, we first investigate the mean-field limit of ˜µt and show that it actually follows the gradient flow of F in the Wasserstein-Fisher-Rao space when M → ∞.

Proposition 1. Suppose the empirical distribution ˜µM0 of M weighted particles weakly converges to a distribution µ0 when M → ∞. Then, the path of (7) starting from ˜µM0 weakly converges to a solution of the following partial differential equation starting from µ0 as M → ∞:

which is actually the gradient flow of F in the WassersteinFisher-Rao space.

Compared to the Wasserstein gradient flow (2) of F, the Wasserstein-Fisher-Rao gradient flow (8) has an additional density adjustment part (second term in the r.h.s. of equation (8)). Due to the limit of space, we defer the discussion of the descending property of (8) in our long version. Actually, it can be verified that (8) results in an extra decreasing term than the Wasserstein gradient flow used in the fixed-weight ParVIs due to the additional density adjustment part.

3.2 Dynamic-Weight ParVI Framework

Generally, it is impossible to obtain an analytic solution of the continuous composite flow (7), thus a numerical integration method is required to derive an approximate solution. Note that any numerical solver, such as the implicit Euler method [Platen and Bruti-Liberati, 2010] and higher-order Runge-Kutta method [Butcher, 1964] can be used. Here, we adopt the first-order explicit Euler discretization [S¨uli and Mayers, 2003] since it is simple and easy-to-implement, and propose our Dynamic-weight Particle-based Variational Inference (DPVI) framework, as listed in Algorithm 1.

Starting from M weighted particles located at {xi0 }Mi=1 with weights {ai0 }Mi=0 , DPVI first updates the positions of particles according to the following rule:

where ˜µk = ∑Mi=1 aik δxik . Then, it adjusts the particles' weights following

˜µk+1/2 = ∑Mi=1 aik δxik+1 represents the empirical distribution after the position update (9). Here, we assume that a suitable empirical approximation U˜µk of the first variation is already constructed, and we will discuss this comprehensively in the following subsection. It can be verified that the total mass of ˜µk is conserved as the following equality holds:

Thus, ˜µk remains a valid probability distribution during the whole procedure of DPVI, i.e. ∑i aik = 1 for all k.

In developing our DPVI framework, we adopt a GaussSiedel-type strategy to update the position xik and the weight aik , i.e., we adjust the weight based on the newly obtained position xik+1 in the k-th iteration. As the weight adjustment step (10) only involves calculating the first variation approximation U˜µk+1/2 , it would bring little extra computational cost compared with the position update step (9), which involves calculating the gradient of U˜µk . It can be verified that, the overall computational complexity of DPVI is in the same order as their fixed weight counterpart, as the cost of gradient calculation in (9) is typically d times of the cost needed to compute function values in (10). Besides, one can further reduce the computational cost by adopting the Jacobi-type strategy, i.e., update weight aik in (10) with U˜µk at position xik . In this case, the weight adjustment step boils down to M scalar additions since the term U˜µk can be directly obtained when calculating v˜µk = ∇U˜µk in the position update step (9). The empirical result in the experiment section also shows that the extra weight-adjustment step would not bring much extra computational cost.

3.3 DPVI Algorithms and Their Duplicate/Kill Variants

To derive an efficient DPVI algorithm, it remains to decide the underlying functional F, and construct proper empirical approximations to its first variation Uµ and the associated vector field vµ = -∇Uµ (denoted as U˜µ and v˜µ, respectively). Unfortunately, there exists no systematic approach to design proper approximations U˜µ and v˜µ for arbitrary dissimilarity functional F. Here, we propose three efficient DPVI algorithms based on different approximations utilized in existing fixed-weight ParVIs, two with the KL-divergence as the underlying functional F (D-GFSD-CA and D-Blob-CA) and one with the KSD (D-KSDD-CA). Note that these three algorithms are actually the dynamic weight-adjustment counterparts of the fixed weight ParVI algorithms, GFSD, Blob and Algorithm 1 Dynamic-weight Particle-based Variational Inference (DPVI) Framework Input: Initial distribution ˜µ0 = ∑Mi=1 ai0 δxi0 , step-size η, weight parameter λ.

• 1: for k = 0, 1, ..., T -1 do
• 2: for i = 1, 2, ...,M do
• 3: Update positions xik+1 's according to (9) .
• 4: end for
• 5: for i = 1, 2, ...,M do
• 6: Adjust weights aik+1 's according to (10).
• 7: end for
• 8: end for
• 9: Output:˜µT = ∑Mi=1 aiT δxiT .

KSDD, respectively. Due to the limit of space, we only list the D-GFSD-CA algorithm here, and refer the readers to our long version for the details of the other two algorithms. Moreover, we emphasize that other divergences, such as MMD and W2 (which are more related with the generative task), can also be chosen as the underlying functional.

D-GFSD-CA algorithm. As we have discussed in the preliminaries, a large portion of existing fixed-weight ParVIs adopt the KL-divergence as underlying functional F, whose associated vector field and first variation are defined as follows:

In order to deal with the intractable logµ(x) which is undefined with empirical distribution ˜µk, Our proposed D-GFSDCA algorithm adopts the approximation techniques used in the fixed-weight ParVI methods GFSD Specifically, we directly approximate µ by smoothing the empirical distribution ˜µ with a kernel function K: ˆµ = ˜µ ∗ K = ∑Mi=1 aiK(·, xi ), which leads to the following approximations:

The duplicate/kill variants. In fact, there is a probabilistic discretization strategy of the approximate reaction flow in (7). This strategy duplicates/kills particle xik+1 according to an exponential clock with instantaneous rate:

Specifically, if Rik+1 > 0, duplicate the particle xik+1 with probability 1 - exp (-Rik+1 ), and kill another one with uniform probability to conserve the total mass; if Rik+1 < 0, kill the particle xik+1 with probability 1-exp (Rik+1 ), and duplicate another one with uniform probability. By replacing the CA strategy (10) in the DPVI framework, we could obtain the DK variants of DPVIs.

4 Experiments

In this section, we conduct empirical studies with our DPVI algorithms (D-GFSD-CA, D-Blob-CA, and D-KSDDCA), their duplicate/kill variants (D-GFSD-DK, D-Blob-DK, and D-KSDD-DK), and the fixed-weight ParVI algorithms (SVGD, GFSD, Blob and KSDD). Besides, we also include three MCMC algorithms as our baseline: the Hamiltonian Monte Carlo (HMC) method [Neal and others, 2011], the Unadjusted Langevin Dynamics (ULD) method [Ma et al., 2015], and the Birth Death Langevin Sampling (BDLS) method [Lu et al., 2019] (ULD with DK weight-adjustment). We compare the performance of these algorithms on two simulations, i.e., a 10-D Single-mode Gaussian model (SG) and a Gaussian mixture model (GMM), and two real-world applications, i.e. a Gaussian Process (GP) regression and a Bayesian neural network. Due to limited space, only part of the results are reported in this section. We refer readers to our long version for the results on SG and additional results for GMM, GP and BNN.

4.1 GaussianMixture Model

We consider approximating a 10-D Gaussian mixture model with two components, weighted by 1/3 and 2/3 respectively. To investigate the influence of particle number M in fixedweight ParVIs and the dynamic-weight algorithms, we run all the algorithms with M ∈ {32, 64, 128, 256, 512}.

In Table 1, we report the 2-Wasserstein (W2) distance between the empirical distribution generated by each algorithm and the target distribution. We generate 5,000 samples from the target distribution π as reference to evaluate the W2 distance by using the POT library 1 . It can be observed that both the CA and DK weight strategies contribute to obtain a higher approximation accuracy. The DPVI algorithms constantly outperform their fixed-weight counterparts with the same or even less number of particles. For instance, D-GFSD-CA with M = 32 achieves a lower W2 than the fixed-weight GFSD algorithm with M = 512. Note that the KSDD-type algorithms (KSDD, D-KSDD-CA, and D-KSDD-DK) perform poorly in this task. Actually, this phenomenon has already been reported in KSDD [Korba et al., 2021], and the authors claim that particles in KSDD may get stuck in spurious local minima when dealing with multi-mode models.

1 http://jmlr.org/papers/v22/20-451.html

In Figure 2, we plot the W2 w.r.t iterations and time of each algorithm. Due to the limit of space, we only report the results when M = 32 (the smallest) and M = 512 (the largest). From this figure, we can observe that, compared with the fixed-weight ParVI algorithms, both the CA and DK strategies result in a better performance w.r.t. both iterations and time. Actually, as we have discussed in the end of Section 3.2, while the weight-adjustment step in DPVI greatly enhances the expressiveness of particles' empirical distribution, it would not bring much extra computational costs compared with fixed-weight ParVI methods. Note that the performance of certain dynamic weight algorithms with DK (e.g., BDLS, D-GFSD-DK, and D-Blob-DK in the first subfigure) oscillates around its best value, which results from the fact that the step-size are tuned via grid search for the fixed-weight ParVI algorithm and then fixed for their dynamic-weight counterpart to show a direct influence of the weight adjustment strategy.

4.2 Gaussian Process Regression

The Gaussian Process (GP) model is widely adopted for the uncertainty quantification in regression problems [Rasmussen, 2003]. We follow the experiment setting in [Chen et al., 2018b], and use the dataset LIDAR , which consists of 221 observations. In this task, we set the particle number to M = 128 for all the algorithms.

Figure 3 gives the contour line of the log posterior, and particles generated by each algorithm. It is shown that DPVIs with CA achieve better approximation results compared to other algorithms. We can also observe that D-Blob-CA has the best performance and covers a wider range of area due to both the dynamic weight adjustment strategy and extra repulsive term.

To evaluate how well the particles approximate the posterior p(ϕ|D), we also report the W2 distance between the empirical distribution and the target distribution w.r.t iterations and time in Figure 4, given 10000 reference particles generated by the HMC method after it achieves its equilibrium[Brooks et al., 2011]. The result demonstrates the effectiveness of our proposed weight strategy: DPVIs with CA constantly outperform their fixed-weight counterparts, and D-Blob-CA achieves the highest approximation accuracy w.r.t. both iterations and time among all the algorithms. Note that the oscillation phenomena of algorithms with DK is more obvious than that in GMM . This is mainly due to the fact that the single-mode nature of GP greatly weakens the advantage of DK, i.e., transferring particles from low-probability region to distant high-probability area (e.g. among different local modes), and the gain of DK is not enough to overweight the influence of a sub-optimal stepsize.

4.3 Bayesian Neural Network

In this experiment, we study a Bayesian regression task with Bayesian neural network on 4 datasets from UCI and LIBSVM. We follow the experiment setting from [Liu and Wang, 2016], which models the output as a Gaussian distribution and uses a Gamma(1, 0.1) prior for the inverse covariance. We use a one-hidden-layer neural network with 50 hidden units and maintain 128 particles. For all the datasets, we set the batchsize as 128. Since KSDD-type algorithms require evaluating the Hessian matrix of the objective function and the HMC method need calculate the full gradient of the logposterior, which will induce enormous computational burden in neural network based tasks, we exclude them in this task.

We report the Root Mean Squared Error (RMSE) of each algorithm in Table 2, and we refer readers to Appendix for the negative log-likelihood results . The results show that both the CA and DK weight strategies contribute to a lower RMSE, and DPVI algorithms with CA achieve the best performance. It can be observed that GFSD-type algorithms obtain similar results as Blob-type algorithms, which may be ascribed to the negative correlation between the magnitude of the repulsive force and the dimensionality in ParVIs [Zhuo et al., 2018].

5 Conclusion

In this paper, we propose a general Dynamic-weight Particlebased Variational Inference (DPVI) framework, which maintains a set of weighted particles to approximate a given target distribution π and updates both the particles' positions and weights iteratively. Our DPVI framework is developed by discretizing a novel composite flow, which is a combination of a finite-particle approximation of a reaction flow and the finite-particle position transport approximation adopted in existing fixed-weight ParVIs. We show that the mean-field limit of the proposed composite flow is actually the gradient flow of the associated dissimilarity functional F in the Wasserstein-Fisher-Rao space, which leads to an extra decrease of F than the Wasserstein gradient flow underlying existing fixed-weight ParVIs. We provide three effective DPVI algorithms with different finite-particle approximations, and derive three variants of them by using the duplicate/kill strategy. The empirical results show that the proposed DPVI algorithms constantly outperform their fixed-weight counterparts and D-Blob-CA usually obtains the best performance.

Acknowledgements

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