Stochastic gradient Markov chain Monte Carlo (SG-MCMC) has become increasingly popular for simulating posterior samples in large-scale Bayesian modeling. However, existing SG-MCMC schemes are not tailored to any specific probabilistic model, even a simple modification of the underlying dynamical system requires significant physical intuition. This paper presents the first meta-learning algorithm that allows automated design for the underlying continuous dynamics of an SG-MCMC sampler. The learned sampler generalizes Hamiltonian dynamics with state-dependent drift and diffusion, enabling fast traversal and efficient exploration of energy landscapes. Experiments validate the proposed approach on Bayesian fully connected neural network, Bayesian convolutional neural network and Bayesian recurrent neural network tasks, showing that the learned sampler outperforms generic, hand-designed SG-MCMC algorithms, and generalizes to different datasets and larger architectures. There is a resurgence of research interests in Bayesian deep learning BID8 Blundell et al., 2015; BID10 BID9 BID4 BID33 , 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 BID2 BID19 BID24 and continual learning BID28 . 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 BID46 . Recent advances in this area further introduced optimization techniques such as pre-conditioning BID1 BID30 , annealing (Ding et al., 2014 ) and adaptive learning rates BID16 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 BID17 and language modeling with recurrent neural networks BID5 . 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 BID36 Bengio et al., 1992; BID26 BID43 . 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 BID35 BID32 BID3 , and learning algorithms such as gradient descent (Andrychowicz et al., 2016; BID18 BID47 , Bayesian optimization BID5 and reinforcement learning (Duan et al., 2016; . Unfortunately, these advances cannot 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.• 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 BID16 . We have presented a meta-learning algorithm that can learn an SG-MCMC sampler on simpler tasks and generalizes to more complicated densities in high dimensions. Experiments on Bayesian MLPs, Bayesian CNNs and Bayesian RNNs confirmed the strong generalization of the trained sampler to the long-time horizon as well as across datasets and network architectures. Future work will focus on better designs for both the sampler and the meta-learning procedure. For the former, temperature variable augmentation as well as moving average estimation will be explored. For the latter, better loss functions will be proposed for faster training, e.g. by reducing the unrolling steps of the sampler during training. Finally, the automated design of generic MCMC algorithms that might not be derived from continuous Markov processes remains an open challenge.A COMPARING MOMENTUM SGD AND SGHMC Similar to the relationship between SGLD and SGD, SGHMC is closely related SGD with momentum (SGD-M). First in HMC, the state space is augmented with an additional momentum variable denoted as p p p ∈ R D . We assume an identity mass matrix associated with that momentum term. Then the corresponding drift f f f (θ θ θ, p p p) and diffusion matrix D D D are: DISPLAYFORM0 where C C C is a positive definite matrix called friction coefficient. Thus, HMC's continuous-time dynamics is governed by the following SDE: DISPLAYFORM1 The discretized update rule (with simple Euler discretization) of HMC with step-size η is DISPLAYFORM2 If stochastic gradient ∇Ũ (θ θ θ) is used, we need to replace the covariance matrix of with 2η(C C C −B B B) whereB B B is the variance estimation of the gradients.On the other hand, the update equations of SGD with momentum (SGD-M) are the following: DISPLAYFORM3 where k and l are called momentum discount factor and learning rate, respectively. Also we can rewrite the SGHMC update equations by setting ηp p DISPLAYFORM4 Thus, the discretized SGHMC updates can be viewed as the SGD-M update injected with carefully controlled Gaussian noise. Therefore, the hyperparameter of SGHMC can be heuristically chosen based on the experience of SGD-M and vice versa. BID27 showed that in practice, simple Euler discretization for HMC simulation might cause divergence, therefore advanced discretization schemes such as Leapfrog and modified Euler are recommended. We use modified Euler discretization in our implementation of SGHMC and the meta sampler, resulting in the following update: DISPLAYFORM5 DISPLAYFORM6 Due to the two-stage update of Euler integrator, at time t, we have f DISPLAYFORM7 ), which is not exactly the history from the previous time. Therefore we further approximate it using delayed estimate: DISPLAYFORM8 Similarly, the Γ Γ Γ p p p term expands as DISPLAYFORM9 We further approximate DISPLAYFORM10 ∂U (θ θ θ) by the following DISPLAYFORM11 This only requires the storage of previous Q Q Q matrix. However, DISPLAYFORM12 ∂pi requires one further forward pass to obtainf DISPLAYFORM13 Therefore the proposed finite difference method only requires one more forward passes to computê f f f t−1 φ D and instead, save 3 back-propagations. As back-propagation is typically more expensive than forward pass, our approach reduces running time drastically, especially when the sampler are applied to large neural network.Time complexity figures Every SG-MCMC method (including the meta sampler) requires ∇ θ θ θŨ (θ θ θ). The main burden is the forward pass and back-propagation through the D D D(z z z) and Q Q Q(z z z) matrices, where the latter one has been replaced by the proposed finite difference scheme. The time complexity is O(HD) for both forward pass and finite difference with H the number of hidden units in the neural network of the meta sampler. Parallel computation with GPUs improves real-time speed, indeed in our MNIST experiment the meta sampler spends roughly 1.5x time when compared with SGHMC.
This paper proposes a method to automate the design of stochastic gradient MCMC proposal using meta learning approach.  Prsents a meta-learning approach to automatically design MCMC sampler based on Hamiltonian dynamics to mix faster on problems similar to training problems Parameterizes diffusion and curl matrices by neural networks and meta-learn and optimize an sg-mcmc algorithm. 