Meta-Learning For Stochastic Gradient MCMC

1 INTRODUCTION . There is a resurgence of research interests in Bayesian deep learning ( Graves , 2011 ; Blundell et al. , 2015 ; Hernández-Lobato & Adams , 2015 ; Hernandez-Lobato et al. , 2016 ; Gal & Ghahramani , 2016 ; Ritter et al. , 2018 ) , which applies Bayesian inference to neural networks for better uncertainty estimation . It is crucial for e.g . better exploration in reinforcement learning ( Deisenroth & Rasmussen , 2011 ; Depeweg et al. , 2017 ) , resisting adversarial attacks ( Feinman et al. , 2017 ; Li & Gal , 2017 ; Louizos & Welling , 2017 ) and continual learning ( Nguyen et al. , 2018 ) . A popular approach to performing Bayesian inference on neural networks is stochastic gradient Markov chain Monte Carlo ( SG-MCMC ) , which adds properly scaled Gaussian noise to a stochastic gradient ascent procedure ( Welling & Teh , 2011 ) . Recent advances in this area further introduced optimization techniques such as pre-conditioning ( Ahn et al. , 2012 ; Patterson & Teh , 2013 ) , annealing ( Ding et al. , 2014 ) and adaptive learning rates ( Li et al. , 2016a ; Chen et al. , 2016 ) . All these efforts have made SG-MCMC highly scalable to many deep learning tasks , including shape and texture modeling in computer vision ( Li et al. , 2016b ) and language modeling with recurrent neural networks ( Gan et al. , 2017 ) . However , inventing novel dynamics for SG-MCMC requires significant mathematical work to ensure the sampler ’ s stationary distribution is the target distribution , which is less friendly to practitioners . Furthermore , many of these algorithms are designed as a generic sampling procedure , and the associated physical mechanism might not be best suited for sampling neural network weights . This paper aims to automate the SG-MCMC proposal design by introducing meta-learning techniques ( Schmidhuber , 1987 ; Bengio et al. , 1992 ; Naik & Mammone , 1992 ; Thrun & Pratt , 1998 ) . The general idea is to train a learner on one or multiple tasks in order to acquire common knowledge that generalizes to future tasks . Recent applications of meta-learning include learning to transfer knowledge to unseen few-shot learning tasks ( Santoro et al. , 2016 ; Ravi & Larochelle , 2017 ; Finn et al. , 2017 ) , and learning algorithms such as gradient descent ( Andrychowicz et al. , 2016 ; Li & Malik , 2017 ; Wichrowska et al. , 2017 ) , Bayesian optimization ( Chen et al. , 2017 ) and reinforcement learning ( Duan et al. , 2016 ; Wang et al. , 2016 ) . Unfortunately , these advances can not be directly transferred to the world of MCMC samplers , as a naive neural network parameterization of the transition kernel does not guarantee the posterior distribution to be the stationary distribution of the sampler . ∗Equal contribution †Work done at the University of Cambridge We present to the best of our knowledge the first attempt towards meta-learning an SG-MCMC algorithm . Concretely , our contribution includes : • An SG-MCMC sampler that extends Hamiltonian dynamics with learnable diffusion and curl matrices . Once trained , the sampler can generalize to different datasets and architectures . • Extensive evaluation of the proposed sampler on Bayesian fully connected neural networks , Bayesian convolutional neural networks and Bayesian recurrent neural networks , with comparisons to popular SG-MCMC schemes based on e.g . Hamiltonian Monte Carlo ( Chen et al. , 2014 ) and pre-conditioned Langevin dynamics ( Li et al. , 2016a ) . 2 BACKGROUND : A COMPLETE FRAMEWORK FOR SG-MCMC . Consider sampling from a target density π ( θ ) that is defined by an energy function : U ( θ ) , θ ∈ RD , π ( θ ) ∝ exp ( −U ( θ ) ) . In this paper , we focus on this sampling task with a Bayesian modeling set-up , i.e . given observed dataD = { on } Nn=1 , we define a probabilistic model p ( D , θ ) = ∏N n=1 p ( on|θ ) p ( θ ) , and we want samples from the target density defined as posterior distribution π ( θ ) = p ( θ|D ) . We use Bayesian neural networks as an illustrating example , in this case , on = ( xn , yn ) , the prior p ( θ ) is a Gaussian N ( θ ; 0 , λ−1I ) , and the energy function is defined as U ( θ ) = − N∑ n=1 log p ( yn|xn , θ ) − log p ( θ ) = N∑ n=1 ` ( yn , NNθ ( xn ) ) + λ||θ||22 , ( 1 ) with ` ( y , ŷ ) usually defined as the ` 2 loss for regression or the cross-entropy loss for classification . A typical MCMC sampler constructs a Markov chain with a transition kernel , and corrects the proposed samples with Metropolis-Hastings ( MH ) rejection steps . Some of these methods , e.g . Hamiltonian Monte Carlo ( HMC ) ( Duane et al. , 1987 ; Neal et al. , 2011 ) , further augment the state space as z = ( θ , r ) with auxiliary variables r , and sample from the augmented distribution π ( z ) ∝ exp ( −H ( z ) ) , with the Hamiltonian H ( z ) = U ( θ ) + g ( r ) such that ∫ exp ( −g ( r ) ) dr = C. Thus , marginalizing out the auxiliary variable r will not affect the stationary distribution π ( θ ) ∝ exp ( −U ( θ ) ) . For deep learning tasks , the observed dataset D often contains thousands , if not millions , of instances , making MH rejection steps computationally prohibitive . Fortunately this is mitigated by SG-MCMC , whose transition kernel is implicitly defined by a stochastic differential equation ( SDE ) that leaves the target density invariant ( Welling & Teh , 2011 ; Ahn et al. , 2012 ; Patterson & Teh , 2013 ; Chen et al. , 2014 ; Ding et al. , 2014 ) . Such a Markov process is called Itô diffusion governed by the continuous-time SDEs : dz = f ( z ) dt+ √ 2D ( z ) dW ( t ) , ( 2 ) with f ( z ) the deterministic drift , W ( t ) the Wiener process , and D ( z ) the diffusion matrix . As a simple example , Langevin dynamics considers z = θ , f ( θ ) = −∇θU ( θ ) andD ( θ ) = I . Then using forward Euler discretization with step-size η the update rule of the parameters is θt+1 = θt − η∇θtU ( θt ) + √ 2η t , t ∼ N ( 0 , I ) . ( 3 ) Stochastic gradient Langevin dynamics ( SGLD , Welling & Teh , 2011 ) proposed an approximation to ( 3 ) , by replacing the exact gradient∇θU ( θ ) with an estimate using a mini-batch of datapoints : ∇θ Ũ ( θ ) = − N M M∑ m=1 ∇θ log p ( om|θ ) −∇θ log p ( θ ) , o1 , ... , oM ∼ D. ( 4 ) Therefore SGLD can be viewed as a stochastic gradient descent ( SGD ) that adds in a properly scaled Gaussian noise term . Similarly , SGHMC ( Chen et al. , 2014 ) is closely related to momentum SGD ( see appendix ) . Furthermore , MH rejection steps are usually dropped in SG-MCMC when a carefully selected discretization step-size is in use . Therefore SG-MCMC has the same computational complexity as many stochastic optimization algorithms , making it highly scalable for sampling posterior distributions of neural network weights conditioned on big datasets . Ma et al . ( 2015 ) derived a framework of SG-MCMC samplers using advanced statistical mechanics ( Yin & Ao , 2006 ; Shi et al. , 2012 ) , which explicitly parameterizes the drift f ( z ) : f ( z ) = − [ D ( z ) +Q ( z ) ] ∇zH ( z ) + Γ ( z ) , Γi ( z ) = d∑ j=1 ∂ ∂zj ( Dij ( z ) +Qij ( z ) ) , ( 5 ) withQ ( z ) the curl matrix , D ( z ) the diffusion matrix and Γ ( z ) a correction term . Remarkably Ma et al . ( 2015 ) showed the completeness of their framework : 1. π ( z ) ∝ exp ( −H ( z ) ) is a stationary distribution of the SDE ( 2 ) + ( 5 ) for any pair of positive semi-definite matrixD ( z ) and skew-symmetric matrixQ ( z ) ; 2. for any Itô diffusion process that has the unique stationary distribution π ( z ) , under mild conditions there existD ( z ) andQ ( z ) matrices such that the process is governed by ( 2 ) + ( 5 ) . As a consequence , the construction of an SG-MCMC algorithm reduces to defining the state-space z and the D ( z ) , Q ( z ) matrices . Indeed Ma et al . ( 2015 ) also cast existing SG-MCMC samplers within the framework , and proposed an improved version of SG-Riemannian-HMC . In general , an appropriate design of these two matrices leads to significant improvements in mixing as well as reduction of sample bias ( Li et al. , 2016a ; Ma et al. , 2015 ) . However , this design has been historically based on strong physical intuitions from statistical mechanics ( Duane et al. , 1987 ; Neal et al. , 2011 ; Ding et al. , 2014 ) . Therefore , it can still be difficult for practitioners to understand and engineer the sampling method that is best suited to their machine learning tasks . In the next section , we will describe our recipe on meta-learning an SG-MCMC sampler of the form ( 2 ) + ( 5 ) . Before the presentation , we emphasize that the completeness result of the framework is beneficial for our meta-learning task . On the one hand , as meta-learning searches the best algorithm for a given set of tasks , it is crucial that the search space is large enough to contain many useful candidates . On the other hand , some form of “ correctness ” guarantee is often required to achieve better generalization to test tasks that might not be very similar to the training tasks . Ma et al . ( 2015 ) ’ s completeness result indicates that our proposed method searches SG-MCMC samplers in the biggest subset of all Itô diffusion processes such that each instance is a valid posterior sampler . Therefore , the proposed meta-learning algorithm has the best from both worlds , indeed our experiments show that the learned sampler is superior to a number of other baseline SG-MCMC methods . 3 META-LEARNING FOR SG-MCMC . This section presents a meta-learning approach to learn an SG-MCMC proposal from data . Our aim is to design an appropriate parameterization ofD ( z ) andQ ( z ) , so that the sampler can be trained on simple tasks with a meta-learning procedure , and generalize to more complicated densities . For simplicity , we only augment the state-space by one extra variable p called momentum ( Duane et al. , 1987 ; Neal et al. , 2011 ) , although generalization to e.g . thermostat variable ( Ding et al. , 2014 ) is straightforward . Thus , the augmented state-space is z = ( θ , p ) ( i.e . r = p ) , and the Hamiltonian is defined as H ( z ) = U ( θ ) + 12p Tp with identity mass matrix .

Summary: This paper proposes a novel method to perform meta-learning for stochastic gradient MCMC. They utilize a general family of SDEs that guarantees preservation of the target density with somewhat loose constraint on the drift and diffusion functions (from Ma et al. (2015)). Then, they propose learning these functions on a set of training tasks and evaluating on unseen, different tasks, in a meta-learning fashion.