Cyclical Stochastic Gradient MCMC for Bayesian Deep Learning

1 INTRODUCTION . Deep neural networks are often trained with stochastic optimization methods such as stochastic gradient decent ( SGD ) and its variants . Bayesian methods provide a principled alternative , which account for model uncertainty in weight space ( MacKay , 1992 ; Neal , 1996 ) , and achieve an automatic balance between model complexity and data fitting . Indeed , Bayesian methods have been shown to improve the generalization performance of DNNs ( Hernández-Lobato & Adams , 2015 ; Blundell et al. , 2015 ; Li et al. , 2016a ; Maddox et al. , 2019 ) , while providing a principled representation of uncertainty on predictions which is crucial for decision making . Approximate inference for Bayesian deep learning has typically focused on deterministic approaches , such as variational methods ( Hernández-Lobato & Adams , 2015 ; Blundell et al. , 2015 ) . By contrast , MCMC methods are now essentially unused for inference with modern deep neural networks , despite previously providing the gold standard of performance with smaller neural networks ( Neal , 1996 ) . Stochastic gradient Markov Chain Monte Carlo ( SG-MCMC ) methods ( Welling & Teh , 2011 ; Chen et al. , 2014 ; Ding et al. , 2014 ; Li et al. , 2016a ) provide a promising direction for a sampling based approach to inference in Bayesian deep learning . Indeed , it has been shown that stochastic methods , which use mini-batches of data , are crucial for finding weight parameters that provide good generalization in modern deep neural networks ( Keskar et al. , 2016 ) . However , SG-MCMC algorithms for inference with modern neural networks face several challenges : ( i ) In theory , SG-MCMC asymptotically converges to target distributions via a decreasing stepsize scheme , but suffers from a bounded estimation error in limited time ( Teh et al. , 2016 ; Chen et al. , 2015 ) . ( ii ) In practice , empirical successes have been reported by training DNNs in relatively short time ( Li et al. , 2016b ; Chen et al. , 2014 ; Gan et al. , 2016 ; Neelakantan et al. , 2016 ; Saatchi & Wilson , 2017 ) . For example , Saatchi & Wilson ( 2017 ) apply SG-MCMC to generative adversarial networks ( GANs ) to solve the mode collapse problem and capture diverse generation styles . However , the loss surface for DNNs is highly multimodal ( Auer et al. , 1996 ; Choromanska et al. , 2015 ) . In order for MCMC to be effective for posterior inference in modern neural networks , a crucial question remains : how do we make SG-MCMC efficiently explore a highly multimodal parameter space given a practical computational budget ? Several attempts have been made to improve the sampling efficiency of SG-MCMC . Stochastic Gradient Hamiltonian Monte Carlo ( SGHMC ) ( Chen et al. , 2014 ) introduces momentum to Langevin dynamics . Preconditioned stochastic gradient Langevin dynamics ( pSGLD ) ( Li et al. , 2016a ) adaptively adjusts the sampler ’ s step size according to the local geometry of parameter space . Though simple and promising , these methods are still inefficient at exploring multimodal distributions in practice . It is our contention that this limitation arises from difficulties escaping local modes when using the small stepsizes that SG-MCMC methods typically require . Note that the stepsize in SGMCMC controls the sampler ’ s behavior in two ways : the magnitude to deterministically drift towards high density regions wrt . the current stochastic gradient , and the level of injecting noise to randomly explore the parameter space . Therefore , a small stepsize reduces both abilities , resulting in a large numbers of iterations for the sampler to move across the modes . In this paper , we propose to replace the traditional decreasing stepsize schedule in SG-MCMC with a cyclical variant . To note the distinction from traditional SG-MCMC , we refer to this method as Cyclical SG-MCMC ( cSG-MCMC ) . The comparison is illustrated in Figure 1 . The blue curve is the traditional decay , while the red curve shows the proposed cyclical schedule . Cyclical SG-MCMC operates in two stages : ( i ) Exploration : when the stepsize is large ( dashed red curves ) , we consider this stage as an effective burn-in mechanism , encouraging the sampler to take large moves and leave the local mode using the stochastic gradient . ( ii ) Sampling : when the stepsize is small ( solid red curves ) , the sampler explores one local mode . We collect samples for local distribution estimation during this stage . Further , we propose two practical techniques to improve estimation efficiency : ( 1 ) a system temperature for exploration and exploitation ; ( 2 ) A weighted combination scheme for samples collected in different cycles to reflect their relative importance . This procedure can be viewed as SG-MCMC with warm restarts : the exploration stage provides the warm restarts for its following sampling stage . cSG-MCMC combines the advantages from ( 1 ) the traditional SG-MCMC to characterize the fine-scale local density of a distribution and ( 2 ) the cyclical schedule in optimization to efficiently explore multimodal posterior distributions of the parameter space . In limited time , cSG-MCMC is a practical tool to provide significantly better mixing than the traditional SG-MCMC for complex distributions . cSG-MCMC can also be considered as an efficient approximation to parallel MCMC ; cSG-MCMC can achieve similar performance to parallel MCMC with only a fraction of cost ( reciprocal to the number of chains ) that parallel MCMC requires . To support our proposal , we also prove the non-asymptotic convergence for the cyclical schedule . We note that this is the first convergence analysis of a cyclical stepsize algorithm ( including work in optimization ) . Moreover , we provide extensive experimental results to demonstrate the advantages of cSG-MCMC in sampling from multimodal distributions , including Bayesian neural networks and uncertainty estimation on several large and challenging datasets such as ImageNet . In short , cSG-MCMC provides a simple and automatic approach to inference in modern Bayesian deep learning , with promising results , and theoretical support . This work is a step towards enabling MCMC approaches in Bayesian deep learning . We release code at https : //github.com/ruqizhang/csgmcmc . 2 PRELIMINARIES : SG-MCMC WITH A DECREASING STEPSIZE SCHEDULE . SG-MCMC is a family of scalable sampling methods that enables inference with mini-batches of data . For a dataset D = { di } Ni=1 and a θ-parameterized model , we have the likelihood p ( D|θ ) and prior p ( θ ) . The posterior distribution is p ( θ|D ) ∝ exp ( −U ( θ ) ) , where U ( θ ) is the potential energy given by U ( θ ) = − log p ( D|θ ) − log p ( θ ) . When D is too large , it is expensive to evaluate U ( θ ) for all the data points at each iteration . Instead , SG-MCMC methods use a minibatch to approximate U ( θ ) : Ũ ( θ ) = −N ′ N ∑N ′ i=1 log p ( xi|θ ) − log p ( θ ) , where N ′ N is the size of minibatch . We recommend Ma et al . ( 2015 ) for a general review of SG-MCMC algorithms . We describe two SG-MCMC algorithms considered in this paper . SGLD & SGHMC Welling & Teh ( 2011 ) proposed Stochastic Gradient Langevin Dynamics ( SGLD ) , which uses stochastic gradients with Gaussian noise . Posterior samples are updated at the k-th step as : θk = θk−1−αk∇Ũ ( θk ) + √ 2αk k , where αk is the stepsize and k has a standard Gaussian distribution . To improve mixing over SGLD , Stochastic Gradient Hamiltonian Monte Carlo ( SGHMC ) ( Chen et al. , 2014 ) introduces an auxiliary momentum variable v. SGHMC is built upon HMC , with an additional friction term to counteract the noise introduced by a mini-batch . The update rule for posterior samples is : θk = θk−1 + vk−1 , and vk = vk−1−αk∇Ũ ( θk ) − ηvk−1 + √ 2 ( η − γ̂ ) αk k , where 1− η is the momentum term and γ̂ is the estimate of the noise . To guarantee asymptotic consistency with the true distribution , SG-MCMC requires that the step sizes satisfy the following assumption : Assumption 1 . The step sizes { αk } are decreasing , i.e. , 0 < αk+1 < αk , with 1 ) ∑∞ k=1 αk = ∞ ; and 2 ) ∑∞ k=1 α 2 k < ∞ . Without a decreasing step-size , the estimation error from numerical approximations is asymptotically biased . One typical decaying step-size schedule is αk = a ( b + k ) −γ , with γ ∈ ( 0.5 , 1 ] and ( a , b ) some positive constants ( Welling & Teh , 2011 ) . 3 CYCLICAL SG-MCMC . We now introduce our cyclical SG-MCMC ( cSG-MCMC ) algorithm . cSG-MCMC consists of two stages : exploration and sampling . In the following , we first introduce the cyclical step-size schedule , and then describe the exploration stage in Section 3.1 and the sampling stage in Section 3.2 . We propose an approach to combining samples for testing in Section F. Assumption 1 guarantees the consistency of our estimation with the true distribution in the asymptotic time . The approximation error in limited time is characterized as the risk of an estimator R = B2 +V , whereB is the bias and V is the variance . In the case of infinite computation time , the traditional SG-MCMC setting can reduce the bias and variance to zero . However , the time budget is often limited in practice , and there is always a trade-off between bias and variance . We therefore decrease the overall approximation error R by reducing the variance through obtaining more effective samples . The effective sample size can be increased if fewer correlated samples from different distribution modes are collected . For deep neural networks , the parameter space is highly multimodal . In practice , SG-MCMC with the traditional decreasing stepsize schedule becomes trapped in a local mode , though injecting noise may help the sampler to escape in the asymptotic regime ( Zhang et al. , 2017 ) . Inspired to improve the exploration of the multimodal posteriors for deep neural networks , with a simple and automatic approach , we propose the cyclical cosine stepsize schedule for SG-MCMC . The stepsize at iteration k is defined as : αk = α0 2 [ cos ( π mod ( k − 1 , dK/Me ) dK/Me ) + 1 ] , ( 1 ) where α0 is the initial stepsize , M is the number of cycles and K is the number of total iterations ( Loshchilov & Hutter , 2016 ; Huang et al. , 2017 ) . The stepsize αk varies periodically with k. In each period , αk starts at α0 , and gradually decreases to 0 . Within one period , SG-MCMC starts with a large stepsize , resulting in aggressive exploration in the parameter space ; as the stepsize is decreasing , SG-MCMC explores local regions . In the next period , the Markov chain restarts with a large stepsize , encouraging the sampler to escape from the current mode and explore a new area of the posterior . Related work in optimization . In optimization , the cyclical cosine annealing stepsize schedule has been demonstrated to be able to find diverse solutions in multimodal objectives , though not specifically different modes , using stochastic gradient methods ( Loshchilov & Hutter , 2016 ; Huang et al. , 2017 ; Garipov et al. , 2018 ; Fu et al. , 2019 ) . Alternatively , we adopt the technique to SGMCMC as an effective scheme for sampling from multimodal distributions .

Summary: This article presents cyclical stochastic gradient MCMC for Bayesian deep learning for inference in posterior distributions of network weights of Bayesian NNs. The posteriors of Bayesian NN weights are highly multi-modal and present difficulty for standard stochastic gradient MCMC methods. The proposed cyclical version periodically warm start the SG-MCMC process such that it can explore the multimodal space more efficiently.