A Unified Particle-Optimization Framework for Scalable Bayesian Sampling

missingum@section Introduction

Bayesian methods have been playing an important role in modern machine learning, especially in unsupervised learning (Kingma and Welling,, 2014; Li et al.,, 2017), and recently in deep reinforcement learning (Houthooft et al.,, 2016; Liu et al.,, 2017). When dealing with big data, two lines of research directions have been developed to scale up Bayesian methods, e.g., variational-Bayes-based and sampling-based methods. Stochastic gradient Markov chain Monte Carlo (SG-MCMC) is a family of scalable Bayesian learning algorithms designed to efficiently sample from a target distribution such as a posterior distribution (Welling and Teh,, 2011; Chen et al.,, 2014; Ding et al.,, 2014; Chen et al.,, 2015). In principle, SG-MCMC generates samples from a Markov chain, which are used to approximate a target distribution. Under a standard setting, samples from SG-MCMC are able to match a target distribution exactly with an infinite number of samples (Teh et al.,, 2016; Chen et al.,, 2015). However, this is practically infeasible, as only a finite number of samples are obtained. Although nonasymptotic approximation bounds w.r.t.​ the number of samples have been investigated (Teh et al.,, 2016; Vollmer et al.,, 2016; Chen et al.,, 2015), there are no theory/algorithms to guide learning an optimal set of fixed-size samples/particles. This is an undesirable property of SG-MCMC, because in practice one often seeks to learn the optimal samples of a finite size that best approximate a target distribution.

A remedy for this issue is to adopt the idea of particle-based sampling methods, where a set of particles (or samples) are initialized from some simple distribution, followed by iterative updates to better approximate a target distribution. The updating procedure is usually done by optimizing some metrics such as a distance measure between the target distribution and the current approximation. There is not much work in this direction for large-scale Bayesian sampling, with an outstanding representative being the Stein variational gradient descent (SVGD) (Liu and Wang, 2016a, ). In SVGD, the update of particles are done by optimizing the KL-divergence between the empirical particle distribution and a target distribution, thus the samples are designed to be updated optimally to reduce the KL-divergence in each iteration. Because of this property, SVGD is found to perform better than SG-MCMC when the number of samples used to approximate a target distribution is limited, and has been applied to other problems such as deep generative models (Feng et al.,, 2017) and deep reinforcement learning (Liu et al.,, 2017; Haarnoja et al.,, 2017; Zhang et al., 2018b, ).

Though often achieving comparable performance in practice, little work has been done on investigating connections between SG-MCMC and SVGD, and on developing particle-optimization schemes for SG-MCMC. In this paper, adopting ideas from Waserstein-gradient-flow literature, we propose a unified particle-optimization framework for scalable Bayesian sampling. The idea of our framework is to work directly on the evolution of a density functions on the space of probability measures, e.g., the Fokker-Planck equation in SG-MCMC. To make the evolution solution computationally feasible, particle approximations are adopted for densities, where particles can be optimized during the evolution process. Both SG-MCMC and SVGD are special cases of our framework, and are shown to be highly related. Notably, sampling with SG-MCMC becomes a deterministic particle-optimization problem as SVGD on the space of probability measures, overcoming the aforementioned correlated-sample issue. Furthermore, we are able to develop new unified particle-optimization algorithms by combing SG-MCMC and SVGD, which is less prone to high-dimension space and thus obtains better performance for large-scale Bayesian sampling. We conduct extensive experiments on both synthetic data and Bayesian learning of deep neural networks, verifying the effectiveness and efficiency of our proposed framework.

missingum@section Preliminaries

In this section, we review related concepts and algorithms for SG-MCMC, SVGD, and Wasserstein gradient flows (WGF) on the space of probability measures.

Generating random samples from a distribution (e.g., a posterior distribution) is one of the fundamental problems in Bayesian statistics, which has many important applications in machine learning. Traditional Markov Chain Monte Carlo methods (MCMC), such as the Metropolis–Hastings algorithm (Metropolis et al.,, 1953) produces unbiased samples from a desired distribution when the density function is known up to a normalizing constant. However, most of these methods are based on random walk proposals which suffer from high dimensionality and often lead to highly correlated samples. On the other hand, dynamics-based sampling methods such as the Metropolis adjusted Langevin algorithm (MALA) (Xifara et al.,, 2014) avoid this high degree of correlation by combining dynamical systems with the Metropolis step. In fact, these dynamical systems are derived from a more general mathematical technique called diffusion process, or more specifically, Itó diffusion (Øksendal,, 1985).

is referred to as the potential energy based on an i.i.d.​ assumption of the model, and $Z$ is the normalizing constant. In Bayesian sampling, the posterior distribution corresponds to the (marginal) stationary distribution of a (continuous-time) Itó diffusion, defined as a stochastic differential equation of the form:

Let the density of ${\bm{\mathchar 28930\relax}}_{t}$ be $\mu_{t}$, it is known $\mu_{t}$ is characterized by the Fokker-Planck (FP) equation (Risken,, 1989):

where $\Sigma({\bm{\mathchar 28930\relax}}_{t})\triangleq g({\bm{\mathchar 28930\relax}}_{t})g^{\top}({\bm{\mathchar 28930\relax}}_{t})$, $\operatorname{\mathbf{a}}\cdot\operatorname{\mathbf{b}}\triangleq\operatorname{\mathbf{a}}^{\top}\operatorname{\mathbf{b}}$ for vectors $\operatorname{\mathbf{a}}$ and $\operatorname{\mathbf{b}}$, $\operatorname{\mathbf{A}}\!:\!\operatorname{\mathbf{B}}\triangleq\mbox{trace}(\operatorname{\mathbf{A}}^{\top}\operatorname{\mathbf{B}})$ for matrices $\operatorname{\mathbf{A}}$ and $\operatorname{\mathbf{B}}$. The FP equation is the key to develop our particle-optimization framework for SG-MCMC. In the following, we focus on the simplest case of 1st-order Langevin dynamics if not stated explicitly, though the derivations apply to other variants.

Stochastic gradient MCMC

SG-MCMC algorithms are discretized numerical approximations of the Itó diffusion (1). They mitigate the slow mixing and non-scalability issues encountered in traditional MCMC algorithms by $\textup{\it i})$ adopting gradient information of the posterior distribution, $\textup{\it ii})$ using minibatches of the data in each iteration of the algorithm to generate samples, and $\textup{\it iii})$ ignoring the rejection step as in standard MCMC. To make the algorithms scalable in a big-data setting, three developments will be implemented based on the Itó diffusion: $\textup{\it i})$ define appropriate functions $F$ and $g$ in the Itó-diffusion formula so that the (marginal) stationary distributions coincide with the target posterior distribution $p({\bm{\mathchar 28946\relax}}|\operatorname{\mathbf{X}})$; $\textup{\it ii})$ replace $F$ or $g$ with unbiased stochastic approximations to reduce the computational complexity, e.g., approximating $F$ with a random subset of the data instead of using the full data. For example, in the 1st-order Langevin dynamics, $\nabla_{{\bm{\mathchar 28946\relax}}}U({\bm{\mathchar 28946\relax}})$ could be approximated by an unbiased estimator with a subset of data:

where $\pi$ is a size-$n$ random subset of $\{1,2,\cdots,N\}$, leading to the first SG-MCMC algorithm in machine learning – stochastic gradient Langevin dynamics (SGLD) (Welling and Teh,, 2011); and $\textup{\it iii})$ solve the generally intractable continuous-time Itô diffusions with a numerical method, e.g., the Euler method (Chen et al.,, 2015). For example, this leads to the following update in SGLD:

2 Stein variational gradient descent

Different from SG-MCMC, SVGD initializes a set of particles which are iteratively updated so that the empirical particle distribution approximates the posterior distribution. Specifically, we consider a set of particles $\{{\bm{\mathchar 28946\relax}}^{(i)}\}_{i=1}^{M}$ drawn from some distribution $q$. SVGD tries to update these particles by doing gradient descent on the interactive particle system via

where $\phi$ is a function perturbation direction chosen to minimize the KL divergence between the updated density $q_{[h\phi]}$ estimated by the particles and the posterior $p({\bm{\mathchar 28946\relax}}|\operatorname{\mathbf{X}})$ ($p$ for short). Since $\textsf{KL}(q\|p)$ is convex in $q$, global optimum of $q=p$ can be guaranteed. SVGD considers $\mathcal{F}$ as the unit ball of a vector-valued reproducing kernel Hilbert space (RKHS) $\mathcal{H}$ associated with a kernel $\kappa({\bm{\mathchar 28946\relax}},{\bm{\mathchar 28946\relax}}^{\prime})$. In such as setting, Liu and Wang, 2016b shown:

where $\Gamma_{p}$ is called the Stein operator. Assuming that the update function $\phi({\bm{\mathchar 28946\relax}})$ is in a RKHS with kernel $\kappa(\cdot,\cdot)$, it was shown in (Liu and Wang, 2016b, ) that (4) is maximized with:

SVGD applies updates (6) repeatedly, moving the samples to a target distribution $p$.

3 Wasserstein Gradient Flows

For a better motivation of WGF, we start from gradient flows defined on the Euclidean space.

Wasserstein gradient flows

Let $\{\mu_{t}\}_{t\in}$ be an absolutely continuous curve in $\mathcal{P}(\Omega)$ with finite second-order moments. We consider to define the change of $\mu_{t}$’s by investigating $W_{2}^{2}(\mu_{t},\mu_{t+h})$. Motivated by the Euclidean-space case, this is reflected by a vector field, $\operatorname{\mathbf{v}}_{t}({\bm{\mathchar 28946\relax}})\triangleq\lim_{h\rightarrow 0}\frac{\mathcal{T}({\bm{\mathchar 28946\relax}}_{t})-{\bm{\mathchar 28946\relax}}_{t}}{h}$ called the velocity of the particle. A gradient flow can be defined on $\mathcal{P}(\Omega)$ correspondingly (Ambrosio et al.,, 2005).

Let $\{\mu_{t}\}_{t\in}$ be an absolutely-continuous curve in $\mathcal{P}(\Omega)$ with finite second-order moments. Then for a.e.​ $t\in$, the above vector field $\operatorname{\mathbf{v}}_{t}$ defines a gradient flow on $\mathcal{P}(\Omega)$ as: $\partial_{t}\mu_{t}+\nabla\cdot(\operatorname{\mathbf{v}}_{t}\mu_{t})=0$.

Intuitively, an energy functional $E$ characterizes the landscape structure (appearance) of the corresponding manifold in $\mathcal{P}(\Omega)$, and the gradient flow (7) defines a geodesic path on this manifold. Usually, by choosing appropriate $E$, the landscape is convex, e.g., for the cases of both SG-MCMC and SVGD described below. This provides a theoretical guarantee on the optimal convergence of a gradient flow.

missingum@section Particle-Optimization-based Sampling

In this section, we interpret the continuous versions of both SG-MCMC and SVGD as WGFs, followed by several techniques for particle optimization in the next section. In the following, $\mu_{t}$ denotes the distribution of ${\bm{\mathchar 28946\relax}}_{t}$.

The continuous-time and infinite-particle limit of SVGD with full gradients, denoted as SVGD∞, is known to be a special instance of the Vlasov equation in nonlinear partial-differential-equation literature (Liu,, 2017):

Under this setting, we can specify the function $\operatorname{\mathbf{W}}(\cdot,\cdot)$ for SVGD∞ as

As will be shown in Section 4, $\operatorname{\mathbf{W}}$ in (3.1) naturally leads to the SVGD algorithm, without the need to derive from an RKHS perspective.

The stationary distribution of (8) is $\lim_{t\rightarrow\infty}\mu_{t}\triangleq\mu=p({\bm{\mathchar 28946\relax}}|\operatorname{\mathbf{X}})$.

To interpret SVGD∞ as a WGF, we need to specify two quantities, the energy functional and an underlying metric to measure distances between density functions.

There are two ways to derive energy functionals for SVGD∞, depending on the underlying metrics for probability distributions. When adopting the WGF framework where $W_{2}$ is used as the underlying metric, according to (7), the energy functional $E_{s}$ must satisfy

In general, there is no close-form solution for the above equation. Alternatively, Liu, (2017) proved another form of the energy functional by defining a different distance metric on the space of probability measures, called $\mathcal{H}$-Wasserstein distance:

where $\phi_{t}\triangleq\operatorname{\mathbf{W}}*\mu_{t}$, and $\|\cdot\|_{\mathcal{H}}$ is the norm in the Hilbert space induced by $\kappa(\cdot,\cdot)$. Under this metric, the underlying energy functional is proved to be the standard KL-divergence between $\mu_{t}$ and $p$, e.g.,

As can be seen in Section 4, this interpretation allows one to derive SVGD, a particle-optimization-based algorithm to approximate the continuous-time equation (8).

2 SG-MCMC as WGF

The continuous-time limit of SG-MCMC, when considering gradients to be exact, corresponds to standard Itó diffusions. We consider the Itó diffusion of SGLD for simplicity, e.g.,

The energy functional for SG-MCMC is easily seen by noting that the corresponding FP equation (2) is in the gradient-flow form of (7). Specifically, the energy functional $E$ is defined as:

Note $E_{2}$ is the energy functional of a pure Brownian motion (e.g., $U({\bm{\mathchar 28946\relax}})=0$ in (12)). We can verify (13) by showing that it satisfies that FP equation. According to (7), the first variation of $E_{1}$ and $E_{2}$ is calculated as

Substituting (14) into (7) recovers the FP equation (2) for the Itó diffusion (12).

missingum@section Particle Optimization

Lemma 3 suggests the discrete gradient flow can approximate the original WGF arbitrarily well if a small enough stepsize $h$ is adopted. Consequently, one solves (16) through a sequence of optimization procedures to update the particles. We will derive a particle-approximation method for the $W_{2}$ term in (15), which allows us to solve SG-MCMC efficiently. However, this technique is not applicable to SVGD, as we neither have an explicit form of the energy functional in (10) when adopting the $W_{2}$ metric, nor have an explicit form for the metric $W_{\mathcal{H}}$ in (3.1) when adopting the KL-divergence as the energy functional. Fortunately, this can be solved by the second approximation method called blob methods.

Particle approximation by blob methods

The name of blob methods comes from the classical fluids literature, where instead of evolving the density in (7), one evolves all particles on a grid with time-spacing $h$ (Carrillo et al.,, 2017). Specifically, note the function $\operatorname{\mathbf{v}}_{t}$ in (7) represents velocity of particles via transportation map $\mathcal{T}$, thus solving a WGF is equivalent to evolving the particles along their velocity in each iteration. Formally, one can prove

Let $\mu_{0}\approx\frac{1}{M}\sum_{i=1}^{M}\delta({\bm{\mathchar 28946\relax}}_{0}^{(i)})$. Assume $\operatorname{\mathbf{v}}_{t}$ in (7) is well-defined and continuous w.r.t.​ each ${\bm{\mathchar 28946\relax}}_{t}^{(i)}$ at time $t$. Then solving the PDE (7) reduces to solving a system of ordinary differential equations for the locations of the Dirac masses:

Proposition 4 suggests evolving each particle along the directions defined by $\operatorname{\mathbf{v}}_{t}$, eliminating the requirement to know an explicit form of the energy functional. In the following, we apply the above particle-optimization techniques to derive algorithms for SVGD and SG-MCMC.

1 A particle-optimization algorithm for SVGD

As mentioned above, discrete-gradient-flow approximation does not apply to SVGD. We thus rely on the blob method. From Section 3.1, $\operatorname{\mathbf{v}}_{t}$ in SVGD is defined as $\operatorname{\mathbf{v}}_{t}({\bm{\mathchar 28946\relax}})=(\operatorname{\mathbf{W}}*\mu_{t})({\bm{\mathchar 28946\relax}})$. When $\mu_{t}({\bm{\mathchar 28946\relax}})$ is approximated by particles, $\operatorname{\mathbf{v}}_{t}({\bm{\mathchar 28946\relax}}_{t}^{(i)})$ is simplified as:

As a result, with the definition of $\operatorname{\mathbf{W}}$ in (3.1), updating $\{{\bm{\mathchar 28946\relax}}_{t}^{(i)}\}$ by time discretizing (17) recovers the update equations for standard SVGD in (6).

2 Particle-optimization algorithms for SG-MCMC

Both the discrete-gradient-flow and the blob methods can be applied for SG-MCMC, which are detailed below.

We first specify Lemma 3 in the case of SG-MCMC in Lemma 5, which is known as the Jordan-Kinderlehrer-Otto scheme (Jordan et al.,, 1998).

According to Lemma 5, it is apparent that SG-MCMC can be implemented by iteratively solving the optimization problem in (18). However, particle approximations for both terms in (18) are challenging. In the following, we develop efficient techniques to solve the problem.

First, rewrite the optimization problem in (18) as

We aim at deriving gradient formulas for both the $F_{1}$ and $F_{2}$ terms under a particle approximation in order to perform gradient descent for the particles. Let $\mu\approx\frac{1}{M}\sum_{i=1}^{M}\delta({\bm{\mathchar 28946\relax}}^{(i)})$. The gradient of $F_{1}$ is easily approximated as

where $\lambda$ is the weight for the regularizer. The optimal $p_{ij}$’s can be obtained by applying KKT conditions to set the derivative w.r.t.​ $p_{ij}$ to be zero, ending up with the following form:

where $u_{i}\triangleq e^{-\frac{1}{2}-\frac{\alpha_{i}}{\lambda}}$, $v_{j}=e^{-\frac{1}{2}-\frac{\beta_{j}}{\lambda}}$. As a result, the particle gradients on $F_{2}$ can be approximated as

Theoretically, we need to adaptively update $\{u_{i},v_{j}\}$ as well to ensure the constraints in (20). In practice, however, we use a fixed scaling factor $\gamma$ to approximate $u_{i}v_{j}$ for the sake of simplicity.

Particle gradients are obtained by combining (19) and (21), which are then used to update the particles $\{{\bm{\mathchar 28946\relax}}^{(i)}\}$ by standard gradient descent. Intuitively, (19) encourages particles move to local modes while (21) regularizes particle interactions. Different from SVGD, our scheme imposes both attractive and repulsive forces for the particles. Specifically, by inspecting (21), we can conclude that: $\textup{\it i})$ When ${\bm{\mathchar 28946\relax}}^{(i)}$ is far from a previous particle ${\bm{\mathchar 28946\relax}}_{k}^{(j)}$, i.e., $\frac{d_{ij}}{\lambda}>1$, ${\bm{\mathchar 28946\relax}}^{(i)}$ is pulled close to $\{{\bm{\mathchar 28946\relax}}_{k}^{(j)}\}$ with force proportional to $(\frac{d_{ij}}{\lambda}-1)e^{-d_{ij}/\lambda}$; $\textup{\it ii})$ when ${\bm{\mathchar 28946\relax}}^{(i)}$ is close enough to a previous particle ${\bm{\mathchar 28946\relax}}_{k}^{(j)}$, i.e., $\frac{d_{ij}}{\lambda}<1$, ${\bm{\mathchar 28946\relax}}^{(i)}$ is pushed away, preventing it from collapsing to ${\bm{\mathchar 28946\relax}}_{k}^{(j)}$.

Particle optimization with blob methods

Given $\operatorname{\mathbf{v}}_{t}$, particle updates can be obtained by solving (17) numerically as in SVGD. By inspecting the formula of $\operatorname{\mathbf{v}}_{t}$ in (4.2), the last two terms both act as repulsive forces. Interestingly, the mechanism is similar to SVGD, but with adaptive force between different particle pairs.

missingum@section The General Recipe

Based on the above development, a more general particle-optimization framework is proposed by combining the PDEs of both SG-MCMC and SVGD. As a result, we propose the following PDE to drive evolution of densities

where $\lambda_{1}$ and $\lambda_{2}$ are two constants. It is easily seen that to ensure the stationary distribution of (5) to be equal to $p({\bm{\mathchar 28946\relax}}|\operatorname{\mathbf{X}})$, the following condition must be satisfied:

There are many feasible choices for the functions and parameters $\{F({\bm{\mathchar 28946\relax}}),\operatorname{\mathbf{W}},g({\bm{\mathchar 28946\relax}}),\lambda_{1},\lambda_{2}\}$ to satisfy (5). However, the verification procedure might be complicated given the present of a convolutional term in (5). We recommend the following choices for simplicity:

$F({\bm{\mathchar 28946\relax}})=\frac{1}{2}U({\bm{\mathchar 28946\relax}})$, $\operatorname{\mathbf{W}}=0$, $g({\bm{\mathchar 28946\relax}})=\operatorname{\mathbf{I}}$ and $\lambda_{2}=1$: this reduces to the Wasserstein-based SGLD with particle optimization. Specifically, when the discrete-gradient-flow approximation is adopted, the algorithm is denoted as $w$-SGLD; whereas when the blob method is adopted, it is denoted as $w$-SGLD-B.

$F({\bm{\mathchar 28946\relax}})=0$, $g({\bm{\mathchar 28946\relax}})=0$, $\operatorname{\mathbf{W}}$ is defined as (3.1): this reduces to standard SVGD.

$F({\bm{\mathchar 28946\relax}})=\frac{1}{2}U({\bm{\mathchar 28946\relax}})$, $g({\bm{\mathchar 28946\relax}})=\operatorname{\mathbf{I}}$, $\operatorname{\mathbf{W}}$ is defined as (3.1), and $\lambda_{2}=1$: this is the combination of SGLD and SVGD, and is called particle interactive SGLD, denoted as PI-SGLD or $\pi$-SGLD.

It is easy to verify that condition (5) is satisfied for all the above three particle-optimization algorithms. Furthermore, particle updates are readily developed by applying either the discrete-gradient-flow or blob-based methods.

missingum@section Related Particle-Based MCMC Methods

There have been related particle-based MCMC algorithms. Representative methods are sequential Monte Carlo (SMC) (Moral et al.,, 2006), particle MCMC (PMCMC) (Andrieu et al.,, 2010) and many variants. In SMC, particles are sample from a proposal distribution, and the corresponding weights are updated by a resampling step. PMCMC extends SMC by sampling from an extended distribution interacted with a MH-rejection step. Compared to our framework, their proposal distributions are typically hard to choose; furthermore, optimality of the particles from both methods can not be guaranteed. Furthermore, the methods are typically much more computationally expensive. Recently, Dai et al., (2016) proposed a particle-based MCMC algorithm by approximating a target distribution with weighted kernel density estimator, which updates particle weights based on likelihoods of the corresponding particles. This approach is theoretically sound but lacks an underlying geometry interpretation. Finally, we note that $w$-SGLD has been successfully applied to reinforcement learning recently for improved policy optimization (Zhang et al., 2018a, ).

missingum@section Experiments

We verify our framework on a set of experiments, including a number of toy experiments and applications to Bayesian sampling of deep neural networks (DNNs).

We compare various sampling methods on multi-mode toy examples, i.e., SGLD, SVGD, $w$-SGLD, $w$-SGLD-B and $\pi$-SGLD. We aim to sample from four unnormalized 2D densities $p(z)\propto\exp\{U(z)\}$, with detailed functional form provided in the SM. We optimize/sample 2000 particles to approximate target distributions. The results are shown in Figure 1. It can be seen from Figure 1 that though SGLD maintains good asymptotic properties, it is inaccurate to approximate distributions with only a few samples; in some case, the samples cannot even cover all the modes. Interestingly, all other particle-optimization-based algorithms successfully find all the modes and fit the distributions well. $w$-SGLD is good at finding modes, but worse at modeling the correct variance due to difficulty of controlling the balance between attractive and repulsive forces between particles. $w$-SGLD-B is better than $w$-SGLD at modeling the distribution variance, performing similarly to SVGD and $\pi$-SGLD. Even though, we note that $w$-SGLD is very useful when the number of particles is small, which fits a distribution better, as shown in Section E of the SM.

Bayesian Logistic regression

We next compare the three variants of our framework (i.e.SVGD, $w$-SGLD and $w$-SGLD-B) on a simple logistic-regression task with quantitative evaluations. We use the same model, data and experimental settings as Liu and Wang, 2016a . The Covertype dataset contains 581,012 data points and 54 features. We perform 5 runs for each setting and report the mean of testing accuracies/log-likelihoods. Figure 2 plots both test accuracies and test log-likelihoods w.r.t.​ the number of training iterations. It is clearly that while all methods converge to the same accuracy/likelihood level, both $w$-SGLD and $w$-SGLD-B converge slightly faster than SVGD. In addition, $w$-SGLD and $w$-SGLD-B have similar convergence behaviors, thus we only use $w$-SGLD in the DNN experiments below.

Parameter Sensitivity

Now we study the role of hyperparameters in $\pi$-SGLD: the number of particles $M$ and the scaling factor $\gamma$ to replace the $u_{i}v_{j}$-term in (21). We use the same dataset and model as the above experiment. Figure 3 plots test accuracies along with different parameter settings. As expected, the best performance is achieved with appropriate scale of $W^{2}_{2}$. The performance keep improving with increasing particles. Interestingly, the Wasserstein regularization is more important when the number of particles is small, demonstrating the superiority when approximate distributions with very few particles.

2 Applications on deep neural networks

We conduct experiments for Bayesian learning of DNNs. Different from traditional optimization for DNNs, we are interested in modeling weight uncertainty of neural networks, an important topic that has been well explored (Hernández-Lobato and Adams,, 2015; Blundell et al., 2015a, ; Li et al.,, 2016; Louizos and Welling,, 2016). We assign priors to the weights, which are simple isotropic Gaussian priors in our case, and perform posterior sampling with the proposed particle-optimization-based algorithms, as well as other standard algorithms such as SGLD and SGD. We use the RMSprop optimizer for feed-forward networks (FNN), and Adam for for convolutional neural networks (CNNs) and recurrent neural networks (RNNs). For all methods, we use a RBF kernel $K({\bm{\mathchar 28946\relax}},{\bm{\mathchar 28946\relax}}^{\prime})=\exp(-\|{\bm{\mathchar 28946\relax}}-{\bm{\mathchar 28946\relax}}^{\prime}\|_{2}^{2}/h)$, with the bandwidth set to $h=\mathtt{med}^{2}/\log M$. Here $\mathtt{med}$ is the median of the pairwise distance between particles. All experiments are conducted on a single TITAN X GPU.

We perform the classification tasks on the standard MNIST dataset. A two-layer model 784-X-X-10 with ReLU activation function is used, with X being the number of hidden units for each layer. The training epoch is set to 100. The test errors are reported in Table LABEL:Table:FNN. Not surprisingly, Bayesian methods generally perform better than their optimization counterparts. The new $\pi$-SGLD which combines $w$-SGLD and SVGD improves both methods with little computational overhead. In additional, $w$-SGLD seems to perform better than SVGD in this case, partially due to a better asymptotic property mentioned in (Liu,, 2017). Furthermore, standard SGLD which is based on MCMC obtains higher test errors compared to particle-optimization-based algorithms, partially due to the correlated-sample issue discussed in the introduction. See (Blundell et al., 2015b, ) for details on the other methods in Table LABEL:Table:FNN.

We follow the standard set up as Wang et al., (2017). Specifically, we lower case all the word tokens and filter out word tokens that occur less than 10 times. All the datasets are divided into training, development and testing sets. For the language model set up, we consider a 1-layer LSTM model with 600 hidden units. The sequence length is fixed to be 30. In order to alleviate overfitting, dropout with a rate of 0.4 is used in each LSTM layer. Results in terms of test perplexities are presented in Table 3. Again, we see that $\pi$-SGLD performs best among all algorithms, and $w$-SGLD is slightly better than SVGD, both of which are better than other algorithms.

We propose a unified particle-optimization framework for large-scale Bayesian sampling. Our framework defines gradient flows on the space of probability measures, and uses particles to approximate the corresponding densities. Consequently, solving gradient flows reduces to optimizing particles on the parameter space. Our framework includes the standard SVGD as a special case, and also allows us to develop efficient particle-optimization algorithms for SG-MCMC, which is highly related to SVGD. Extensive experiments are conducted, demonstrating the effectiveness and efficiency of our proposed framework. Interesting future work includes designing more practically efficient variants of the proposed particle-optimization framework, and developing theory to study general convergence behaviors of the algorithms, in addition to the asymptotic results presented in (Liu,, 2017).

Appendix A Proof of Proposition 2

Proof A stationary distribution $\mu$ of (8) means $\partial_{t}\mu=0$. Assuming $\mu=p({\bm{\mathchar 28946\relax}}|\operatorname{\mathbf{X}})\triangleq p$, then we need to prove that

By the definition of $\operatorname{\mathbf{W}}$ in (3.1), and applying Stein’s identity Liu and Wang, 2016a , we have $\operatorname{\mathbf{W}}*p=0$. Consequently, we have $\nabla\cdot\left((\operatorname{\mathbf{W}}*p)p\right)=0$.

The above argument indicates $p({\bm{\mathchar 28946\relax}}|\operatorname{\mathbf{X}})$ is a stationary distribution of (8). This completes the proof.

Appendix B More Details on Lemma 3

We first specify the conditions the energy functional $E$ needs to satisfy in Assumption 1.

proper: $D(E)\triangleq\{{\bm{\mathchar 28946\relax}}\in\Omega:E({\bm{\mathchar 28946\relax}})<+\infty\}\neq\emptyset$.

coercive: There exists $\tau_{0}>0$, ${\bm{\mathchar 28946\relax}}_{0}\in\Omega$ such that $\inf\left\{\frac{1}{2\tau_{0}}W_{2}^{2}({\bm{\mathchar 28946\relax}}_{0},\operatorname{\mathbf{v}})+E(\operatorname{\mathbf{v}}):\operatorname{\mathbf{v}}\in\Omega\right\}>-\infty$.

lower semicontinuous: For all ${\bm{\mathchar 28946\relax}}_{n},{\bm{\mathchar 28946\relax}}\in\Omega$ such that ${\bm{\mathchar 28946\relax}}_{n}\rightarrow{\bm{\mathchar 28946\relax}}$, $\lim\inf_{n\rightarrow\infty}E({\bm{\mathchar 28946\relax}}_{n})\geq E({\bm{\mathchar 28946\relax}})$.

Appendix C Derivation of (22)

The derivation of (4.2) relies on the following Lemma from Carrillo et al., (2017).

where $\circ$ denotes function composition, i.e., $F$ is evaluated on the output of $K*\mu$. Then we have

Now it is ready to derive (4.2). In this case, $F=\log(\cdot)$. Let $F_{1}=\nabla\varphi_{\epsilon}*\left(F^{\prime}\circ(\varphi_{\epsilon}*\mu)\mu\right)$, $F_{2}=\left(F^{\prime}\circ(\varphi_{\epsilon}*\mu)\right)\nabla\varphi_{\epsilon}*\mu$. We use particles to approximate $\mu$, e.g., $\mu\approx\frac{1}{M}\sum_{i=1}^{M}\delta_{{\bm{\mathchar 28946\relax}}^{(i)}i}$. We have

Combing (C) and (C) gives the formula for $\operatorname{\mathbf{v}}$ in (4.2).