Towards More Theoretically-Grounded Particle Optimization Sampling for Deep Learning

Many deep-learning based methods such as Bayesian deep learning ( DL ) and deep reinforcement learning ( RL ) have heavily relied on the ability of a model being able to efficiently explore via Bayesian sampling . Particle-optimization sampling ( POS ) is a recently developed technique to generate high-quality samples from a target distribution by iteratively updating a set of interactive particles , with a representative algorithm the Stein variational gradient descent ( SVGD ) . Though obtaining significant empirical success , the non-asymptotic convergence behavior of SVGD remains unknown . In this paper , we generalize POS to a stochasticity setting by injecting random noise in particle updates , called stochastic particle-optimization sampling ( SPOS ) . Notably , for the first time , we develop non-asymptotic convergence theory for the SPOS framework , characterizing convergence of a sample approximation w.r.t . the number of particles and iterations under both convex- and noncovexenergy-function settings . Interestingly , we provide theoretical understanding of a pitfall of SVGD that can be avoided in the proposed SPOS framework , i.e. , particles tend to collapse to a local mode in SVGD under some particular conditions . Our theory is based on the analysis of nonlinear stochastic differential equations , which serves as an extension and a complementary development to the asymptotic convergence theory for SVGD such as ( Liu , 2017 ) . With such theoretical guarantees , SPOS can be safely and effectively applied on both Bayesian DL and deep RL tasks . Extensive results demonstrate the effectiveness of our proposed framework . 1 INTRODUCTION . Recent development of deep-learning techniques have required the ability of an algorithm to efficiently explore some particular space ( e.g. , parameter space ) via Bayesian sampling , due to high model complexity of modern deep models . Meanwhile , recent years have seen the development of scalable sampling methods such as stochastic gradient MCMC ( SG-MCMC ) ( Welling & Teh , 2011 ; Chen et al. , 2014 ; Ding et al. , 2014 ; Chen et al. , 2015 ) and Stein variational gradient descent ( SVGD ) ( Liu & Wang , 2016 ) to facilitate big-data analysis . SG-MCMC is a family of scalable Bayesian sampling algorithms relying on Itó diffusions , linear stochastic differential equations ( SDE ) with appropriately designed coefficients such that the corresponding stationary distributions match a target distribution . One potential issue is that samples might be highly correlated partially due to the nature of Markov chains , leading to low sample efficiency , an undesirable property of SG-MCMC . SVGD , on the other hand , belongs to the family of particle-based sampling methods , which optimize a set of interacting particles to minimize some distance metric between the target distribution and the distribution induced by the particles . By optimizing the distance measure , one maintains an optimal set of particles at each time . Recent development of SVGD has shown that the underlying mathematical principle is based on a family of nonlinear SDEs , in the sense that coefficients of the SDE depend on the current density of the particles . Though achieving numerous practical successes ( Liu & Wang , 2016 ; Feng et al. , 2017 ; Liu et al. , 2017 ; Haarnoja et al. , 2017 ; Liu & Zhu , 2018 ) , little theory has been developed to understand the convergence property of the algorithm . A recent theoretical development has interpreted SVGD as a special type of gradient flow in the space of probability measures , and developed theory to disclose its asymptotic convergence behavior ( Liu , 2017 ) . Recently , Chen et al . ( 2018 ) unified SG-MCMC and SVGD by proposing a particle-optimization sampling framework that interprets both as Wasserstein gradient flows ( WGFs ) . Generally speaking , a WGF is a partial differential equation ( PDE ) defined on the space of probability measures , describing the evolution of a density function over time . In ( Chen et al. , 2018 ) , the authors define a WGF by combining the corresponding Fokker-Planck equations for both SG-MCMC and SVGD , and solve it with deterministic particle approximations . However , due to the diffusion nature , deterministicparticle approximation leads to a hard-to-control error , challenging for theoretical analysis . Based on the unified framework in ( Chen et al. , 2018 ) , we propose to solve WGFs with stochastic particle approximation , leading to stochastic particle-optimization sampling ( SPOS ) . The idea is instead of solving the WGF with an uncontrollable deterministic approximation for a diffusion term , we solve it stochastically by injecting random noise to the particle-update equations . Remarkably , for the first time , we develop nonasymptotic convergence theory for the family of SPOS algorithms , considering both convex- and nonconvex-energy functions . Different from existing theory for SGMCMC based algorithms ( Teh et al. , 2016 ; Vollmer et al. , 2016 ; Chen et al. , 2015 ; Raginsky et al. , 2017 ; Zhang et al. , 2017 ; Xu et al. , 2018 ) , our development relies on the theory of nonlinear SDEs , which is more involved and less explored in literature . Particularly , we adopt tools from granular media equations ( Malrieu , 2003 ; Cattiaux et al. , 2008 ) to develop nonasymptotic error bounds in terms of the 1-Wasserstein distance . Please refer to Section M in the Supplementary Material ( SM ) for detailed distinctions of our work to existing work . Within our theoretical framework , we provide a formal theoretical understanding of a pitfall of SVGD , i.e. , particles tend to collapse to one point under some particular conditions ; whereas this can be avoided in the proposed SPOS framework due to the injected random noise . Our theory is general enough to be applied for various deep-learning tasks , including Bayesian deep learning and Bayesian exploration in reinforcement learning . Our extensive experimental results well suggest advantages of our framework compared to existing methods . 2 PRELIMINARIES . This section introduces necessary preliminaries , along with notations used in this paper . For the sake of clarity , through out the paper , we use bold letters to denote variables in continuous-time diffusions and model definitions , e.g. , θt in ( 1 ) defined below ( indexed by “ time ” t ) . By contrast , normal unbold letters are used to denote parameters in algorithms ( discrete solutions of continuous-time diffusions ) , e.g. , θ ( i ) k in ( 3 ) below ( indexed by “ iteration ” k ) . For conciseness , all the proofs as well as some extra experimental results are presented in the SM . Discussion on the complexity of our algorithm is also included in Section L of the SM . 2.1 STOCHASTIC GRADIENT MCMC . In Bayesian sampling , we aim to generate random samples from a posterior distribution p ( θ|X ) ∝ p ( X|θ ) p ( θ ) , where θ ∈ Rd represents the model parameter with a prior distribution p ( θ ) , and X , { xq } Nq=1 represents the observed data with likelihood p ( X|θ ) = ∏ q p ( xq |θ ) . Define the potential en- ergy as : U ( θ ) , − log p ( X|θ ) − log p ( θ ) = − ∑N q=1 ( log p ( xq |θ ) + 1N log p ( θ ) ) , ∑N q=1 Uq ( θ ) . SG-MCMC algorithms belong to diffusion-based sampling methods , where a continuous-time diffusion process is designed such that its stationary distribution matches the target posterior distribution . The diffusion process is driven by a specific stochastic differential equation ( SDE ) . For example , in stochastic gradient Langevin dynamic ( SGLD ) , the SDE endows the following form : dθt = −β−1F ( θt ) dt+ √ 2/βdWt , ( 1 ) where F ( θ ) , ∇θU ( θ ) = ∑N q=1∇θUq ( θ ) , ∑N q=1 Fq ( θ ) ; t is the time index ; β > 0 is the temperature parameter ; andWt ∈ Rd is a d-dimensional Brownian motion . More instances of SDEs corresponding to other SG-MCMC algorithms can be defined by specifying different forms of F and potentially other diffusion coefficients . We focus on SGLD and ( 1 ) in this paper , and refer interested readers to ( Ma et al. , 2015 ) for more detailed description of general SG-MCMC algorithms . Denote the probability density function of θt in ( 1 ) as νt , and a ·b , a > b for two vectors a and b . It is known that νt is characterized by the following Fokker-Planck ( FP ) equation ( Risken , 1989 ) : ∂tνt = ∇θ · ( β−1νtF ( θ ) + β−1∇θνt ) . ( 2 ) where the stationary distribution ν∞ equals to our target distribution p ( θ|X ) according to Chiang & Hwang ( 1987 ) . SGLD generates samples from p ( θ|X ) by numerically solving the SDE ( 1 ) . For scalability , it replaces F ( θk ) in each iteration with an unbiased evaluation by randomly sampling a subset of X , i.e . F ( θk ) is approximated by : Gk , NBk ∑ q∈Ik Fq ( θk ) , where Ik is a random subset of [ 1 , 2 , · · · , N ] with size Bk in each iteration . Based on the above settings , SGLD uses the Euler method with stepsize hk to numerically solve ( 1 ) and obtains the update equation : θk+1 = θk − β−1Gkhk + √ 2β−1hkξk , ξk ∼ N ( 0 , I ) . 2.2 STEIN VARIATIONAL GRADIENT DESCENT . Different from SG-MCMC , SVGD is a deterministic particle-optimization algorithm that is able to generate samples from a target distribution . In the algorithm , a set of particles interact with each other , driving them to high density locations in the parameter space while keeping them far away from each other with repulsive force . The update equations of the particles follow the fastest descent direction of the KL-divergence between current empirical distribution of the particles and the target distribution , on an RKHS induced by a kernel function κ ( · , · ) ( Liu & Wang , 2016 ) . Formally , Liu & Wang ( 2016 ) derived the following updating rules for the particles { θ ( i ) k } Mi=1 at the k-th iteration with stepsize hk and G ( i ) k , N Bk ∑ q∈Ik Fq ( θ ( i ) k ) : θ ( i ) k+1 = θ ( i ) k + hk M M∑ j=1 [ κ ( θ ( j ) k , θ ( i ) k ) G ( i ) k +∇θ ( j ) k κ ( θ ( j ) k , θ ( i ) k ) ] , ∀i ( 3 ) where the first term in the bracket encourages particles to locate on high density modes , and the second term serves as repulsive force that pushes away different particles . Different from SG-MCMC , only particles at the current iteration , { θ ( i ) k } Mi=1 , are used to approximate the target distribution . 2.3 PARTICLE-OPTIMIZATION BASED SAMPLING METHODS . SG-MCMC and SVGD , though look closely related , behave very differently in terms of algorithms , e.g. , stochastic and noninteractive versus deterministic and interactive particle updates . Recently , Chen et al . ( 2018 ) proposed a deterministic particle-optimization framework that unifies both SGMCMC and SVGD . Specifically , the authors viewed both SG-MCMC and SVGD as Wasserstein gradient flows ( WGFs ) on the space of probabilistic measures , and derived several deterministic particle-optimization techniques for particle evolutions , like what SVGD does . For SG-MCMC , the FP equation ( 2 ) for SGLD is a special type of WGFs . Together with an interpretation of SVGD as a special case of the Vlasov equation in nonlinear PDE literature , Chen et al . ( 2018 ) proposed a general form of PDE to characterize the evolution of the density for the model parameter θ , denoted as νt at time t with ν∞ matching our target ( posterior ) distribution , i.e. , ∂tνt = ∇θ · ( νtβ −1F ( θ ) + νt ( K ∗ νt ( θ ) ) + β−1∇θνt ) , ( 4 ) where K is a function controlling the interaction of particles in the PDE system . For example , in SVGD , Chen et al . ( 2018 ) showed that K and K ∗ νt ( θ ) endow the following forms : K ( θ , θ′ ) , F ( θ′ ) κ ( θ′ , θ ) −∇θ′κ ( θ′ , θ ) and K ∗ νt ( θ ) , ∫ K ( θ , θ′ ) νt ( θ′ ) dθ′ ( 5 ) where κ ( · , · ) is a kernel function such as the RBF kernel . In the following , we introduce a new unary function K ( θ ) = exp ( −‖θ‖ 2 η2 ) , thus κ ( θ , θ ′ ) can be rewritten as κ ( θ , θ′ ) = K ( θ − θ′ ) . Hence , ( 4 ) with K defined in ( 5 ) can be equivalently rewritten as : ∂tνt = ∇θ · ( νtβ −1F ( θ ) + νt ( EY∼νtK ( θ − Y ) F ( Y ) −∇K ∗ νt ( θ ) ) + β−1∇θνt ) , ( 6 ) where Y is a random sample from νt but independent of θ . Importantly , Proposition 1 ( Chen et al . ( 2018 ) ) The stationary distribution of ( 6 ) equals to our target distribution , which means ν∞ ( θ ) = p ( θ|X ) . Chen et al . ( 2018 ) proposed to solve ( 4 ) numerically with deterministic particle-optimization algorithms such as the blob method . Specifically , the continuous density νt is approximated by a set of M particles { θ ( i ) t } Mi=1 that evolve over time t , i.e . νt ≈ 1M ∑M i=1 δθ ( i ) t ( θ ) , where δ θ ( i ) t ( θ ) = 1 if θ = θ ( i ) t and 0 otherwise . Note ∇θνt in ( 4 ) is no longer a valid definition when adopting particle approximation for νt . Consequently , ∇θνt needs nontrivial approximations , e.g. , by discrete gradient flows or blob methods proposed in ( Chen et al. , 2018 ) . We omit the details here for simplicity .

Summary: This paper considers the problem of Bayesian inference using particle optimization sampler. Similarly to SGLD, authors propose Stochastic Particle Optimization Sampler (SPOS), augmenting Stein Variational Gradient Descent (SVGD) with diminishing Gaussian noise, replacing the hard-to-compute term of the Chen et al. (2018) formulation. Various theoretical results are given.