META-LEARNING FOR STOCHASTIC GRADIENT MCMC
Year: 2019

Abstract:
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 learning tasks with Bayesian fully connected neural networks, Bayesian convolutional neural networks and Bayesian recurrent neural networks, showing that the learned sampler out-performs generic, hand-designed SG-MCMC algorithms, and generalizes to different datasets and larger architectures.

heading: 1 INTRODUCTION
text: 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 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. ∗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).

heading: 2 BACKGROUND: A COMPLETE FRAMEWORK FOR SG-MCMC
text: 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.

heading: 3 META-LEARNING FOR SG-MCMC
text: 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.

heading: 3.1 EFFICIENT PARAMETERIZATION OF DIFFUSION AND CURL MATRICES
text: For neural networks, the dimensionality of θ can be at least tens of thousands. Thus, training and applying full D(z) and Q(z) matrices can cause a huge computational burden, let alone gradient computations required by Γ(z). To address this, we define the preconditioning matrices as follows:
Q(z) =
[ 0 −Qf (z)
Qf (z) 0
] , D(z) = [ 0 0 0 Df (z) ] ,
Qf (z) = diag[f φQ(z)], Df (z) = diag[αf φQ(z) f φQ(z) + f φD (z) + c], α, c > 0. (6)
Here f φD and f φQ are neural network parameterized functions that will be detailed in section 3.2, and c is a small positive constant. We chooseDf andQf to be diagonal for fast computation, although future work can explore low-rank matrix solutions. From Ma et al. (2015), our design has the unique stationary distribution π(θ) ∝ exp(−U(θ)) if f φD is non-negative for all z . We discuss the role of each precondition matrix for better intuition. The curl matrix Q(z) in (2) mainly controls the deterministic drift forces introduced by the energy gradient ∇θU(θ) (as seen in many HMC-like procedures and in eq. (5)). Usually, we only have access to the stochastic gradient ∇θ Ũ(θ) through data sub-sampling, and hereD(z) acts as the friction to counter for the associated noise that mainly affects the momentum p. This explains the design of the diffusion matrix D(z) that usesDf (z) to control the amount of friction and injected noise to the momentum. Furthermore,
Df (z) should also account for the pre-conditioning effect introduced by Qf (z), e.g, when the magnitude ofQf (z) is large, we need higher friction correspondingly. This explains the squared term f φQ(z) f φQ(z) inDf (z) design. The positive scaling constant α is heuristically selected following (Chen et al., 2014; Ma et al., 2015) (see appendix). Finally, the extra term Γ(z) = [Γθ(z),Γp(z)] is responsible for compensating the changes introduced by preconditioning matricesQ(z) andD(z).
The discretized dynamics of the state z = (θ,p) with step-size η and stochastic gradient∇θ Ũ(θ) are
pt+1 = (1− ηDf (zt))pt − ηQf (zt)∇θtŨ(θt) + ηΓp(zt) + , ∼ N (0, 2ηDf (zt)), θt+1 = θt + ηQf (ẑt)pt+1 + ηΓθ(ẑt), zt = [θt, pt], ẑt = [θt, pt+1].
(7)
We use a modified forward Euler discretization (Neal et al., 2011) here, and the computation graph of eq. (7) is visualized in the right part of Figure 1 (see appendix for SGHMC discretized updates). Again we see thatQf (z) is responsible for the acceleration of θ, and from the (1− ηDf (z)) term in the update equation of p, we see thatDf (z) controls the friction introduced to the momentum. Note that in the big-data setting, the noisy gradient is approximately Gaussian distributed with mean 0 and variance V (θ). Observing this, Ma et al. (2015) further suggested a correction scheme to counter for stochastic gradient noise, which samples the Gaussian noise ∼ N (0, 2ηDf (z)− η2B̃(θ)) with an empirical estimate of the variance B̃(θ) ≈ Qf (z)V (θ)QTf (z) instead. These corrections can be dropped when the discretization step-size η is small, therefore, we do not consider them in our experiments.

heading: 3.2 CHOICES OF INPUTS TO THE NEURAL NETWORKS
text: We now present detailed functional forms for f φQ and f φD . When designing these, our goal was to achieve a good balance between generalization power and computational efficiency. Recall that the curl matrixQ(z) mainly controls the drift of the dynamics, and the desired behavior is the fast traverse through low-density regions. One useful source of information to identify this is the energy function U(θ ).1 We also include the momentum pi to the inputs of f φQ , allowing theQ(z) matrix to observe the velocity information of the θi. We further add an offset β to Q(z) to prevent the vanishing of this matrix. Putting all of them together, we define the ith element of f φQ as
f φQ,i(z) = β + fφQ(U(θ), pi). (8)
The corresponding Γ(z) term requires both ∂θifφQ(U(θ), pi) and ∂pifφQ(U(θ), pi). The energy gradient ∂θiU(θ) also appears in (7), so it remains to compute ∂Uf φQ , which, along with ∂pifφQ(U(θ), pi), can be obtained by automatic differentiation (Abadi et al., 2015).
MatrixD(z) is responsible for the friction and the stochastic gradient noise, which are crucial for better exploration around high-density regions. Therefore, we also add the energy gradient ∂θiU(θ) to the inputs, meaning that the ith element of f φD is
f φD,i(z) = fφD (U(θ), pi, ∂θiU(θ)). (9)
By the construction of theD(z) matrix, the Γ vector only requires ∇pDf without computing any higher order information.
1The energy gradient ∇θU(θ) is also informative here, however, it requires expensive computation of the diagonal Hessian for Γ(z). For similar reasons we do not consider other higher order derivatives as inputs.
In practice, both U(θ) and ∂θiU(θ) are replaced by their stochastic estimates Ũ(θ) and ∂θiŨ(θ), respectively. To keep the scale of the inputs roughly the same across tasks, we rescale all the inputs using statistics computed by simulating the sampler with randomly initialized f φD and f φQ . When the computational budget is limited, we replace the exact gradient computation required by Γ(z) with finite difference approximations. We refer the reader to the appendix for details.

heading: 3.3 LOSS FUNCTION DESIGN FOR META-LEARNING
text: Another challenge is to design a meta-learning procedure for the sampler to encourage faster convergence and low bias on test tasks. To achieve these goals we propose two loss functions that we named as the cross-chain loss and the in-chain loss. From now on we consider the discretized dynamics and define qt(θ|D) as the marginal distribution of the random variable θ at time t.
Cross-chain loss We introduce cross-chain loss that encourages the sampler to converge faster. Since the sampler is guaranteed to have the unique stationary distribution π(θ) ∝ exp(−U(θ)), fast convergence means that KL[qt||π] is close to zero when t is small. Therefore this KL-divergence becomes a sensible objective to minimize, which is equivalent to maximizing the variational lowerbound (or ELBO): LtVI(qt) = −Eqt [U(θ)] + H[qt] (Jordan et al., 1999; Beal, 2003). We further make the objective doubly stochastic: (1) the energy term is further approximated by its stochastic estimates Ũ(θ); (2) we use Monte Carlo variational inference (MCVI, Ranganath et al., 2014; Blundell et al., 2015) which estimates the lower-bound with samples θkt ∼ qt(θt|D), k = 1, ...,K. These samples {θkt } K,T k=1,t=1 are obtained by simulating K parallel Markov chains with the sampler, and the cross-chain loss is defined by accumulating the lower-bounds through time:
Lcross-chain = 1
T T∑ t=1 LtVI({θkt }Kk=1), LtVI({θkt }Kk=1) = − 1 K K∑ k=1 [ Ũ(θkt ) + log qt(θ k t |D) ] . (10)
By minimizing this objective, we can improve the convergence of the sampler, especially at the early times of the Markov chain. The objective also takes the sampler bias into account because the two distributions will match when the KL-divergence is minimized.
In-chain loss For very big neural networks, simulating multiple Markov chains is prohibitively expensive. The issue is mitigated by thinning that collects samples for every τ step (after burn-in), which effectively draws samples from the averaged distribution q̄(θ|D) = 1bT/τc ∑bT/τc s=1 qsτ (θ). The in-chain loss is, therefore, defined as the ELBO evaluated at the averaged distribution q̄, which is then approximated by Monte Carlo with samples ΘkT,τ = {θksτ} bT/τc s=1 obtained by thinning:
Lin-chain = 1
K K∑ k=1 LkVI ( ΘkT,τ ) , LkVI ( ΘkT,τ ) = − 1 bT/τc bT/τc∑ s=1 [ Ũ(θksτ ) + log q̄(θ k sτ |D) ] . (11)
Gradient approximation We leverage the recently proposed Stein gradient estimator (Li & Turner, 2018) to estimate the intractable gradients ∇φ log qt(θ) for cross-chain loss and ∇φ log q̄(θ) for in-chain loss. Precisely, by the chain rule, we have ∇φ log qt(θ) = ∇φθ∇θ log qt(θ), so it remains to estimate the gradients G = (∇θ1t log qt(θ 1 t ), . . . ,∇θKt log qt(θ K t ))
T at the sampled locations {θkt }Kk=1 ∼ qt. The recipe first constructs a kernel matrix K with K ij = K(θit, θ j t ), and then
estimates the gradients by G ≈ −(K + λI )−1〈∇,K〉, where 〈∇,K〉ij = ∑K k=1 ∂θkt (j)K(θ k t , θ i t). In our experiments, we use the RBF kernel, and the corresponding gradient estimator has a simple analytic form that can be computed efficiently in O(K2D +K3) time (usually K D).

heading: 4 RELATED WORK
text: Since the development of SGLD (Welling & Teh, 2011), SG-MCMC has been increasingly popular for posterior sampling on big data. In detail, Chen et al. (2014) scaled up HMC with stochastic gradients, Ding et al. (2014) further augmented the state space with an auxiliary temperature variable, and Springenberg et al. (2016) improved robustness through scale adaptation. The SG-MCMC
extensions to Riemannian Langevin dynamics and HMC (Girolami & Calderhead, 2011) have also been proposed (Patterson & Teh, 2013; Ma et al., 2015). Our proposed sampler architecture further generalizes SG-Riemannian-HMC as it decouples the design of D(z) and Q(z) matrices, and the detailed functional form of these two matrices are also learned from data.
Our approach is closely related to the recent line of work on learning optimization algorithms. Specifically, Andrychowicz et al. (2016) trained a recurrent neural network (RNN) based optimizer that transfers to similar tasks with supervised learning. Later Chen et al. (2017) generalized this approach to Bayesian optimization (Brochu et al., 2010; Snoek et al., 2012) which is gradient-free. We do not use RNNs in our approach as it cannot be represented within the framework of Ma et al. (2015). We leave the combination of learnable RNN proposals to future work. Also Li & Turner (2018) presented an initial attempt to meta-learn an approximate inference algorithm, which simply combined the stochastic gradient and the Gaussian noise with a neural network. Thus the stationary distribution of that sampler (if it exists) is only an approximation to the exact posterior. On the other hand, the proposed sampler (with η → 0) is guaranteed to be correct by the complete framework (Ma et al., 2015). Very recently Wu et al. (2018) discussed that short-horizon meta-objectives for learning optimizers can cause a serious issue for long-time generalization. We found this bias is less severe in our approach, again due to the fact that the learned sampler is provably correct.
Recent research also considered improving HMC with a trainable transition kernel. Salimans et al. (2015) improved upon vanilla HMC by introducing a trainable re-sampling distribution for the momentum. Song et al. (2017) parameterized the HMC transition kernel with a trainable invertible transformation called non-linear independent components estimation (NICE, Dinh et al., 2014), and train it with Wasserstein adversarial training (Arjovsky et al., 2017). Levy et al. (2018) generalized HMC by augmenting the state space with a binary direction variable, and they parameterized the transition kernel with a non-volume preserving invertible transformation that is inspired by the real-valued non-volume preserving (RealNVP) flows (Dinh et al., 2017). The sampler is trained with the expected squared jump distance (Pasarica & Gelman, 2010). We note that adversarial training is less reliable for high dimensional data, thus it is not considered in this paper. Also, the jump distance does not explicitly take the sampling bias and convergence speed into account. More importantly, the purpose of these approaches is to directly improve the HMC-like sampler on the target distribution, and with NICE/RealNVP parametrization it is difficult to generalize the sampler to densities of different dimensions. In contrast, our goal is to learn an SG-MCMC sampler that can later be transferred to sample from different Bayesian neural network posterior distributions, which will typically have different dimensionality and include tens of thousands of random variables.

heading: 5 EXPERIMENTS
text: We evaluate the meta-learned SG-MCMC sampler, which is referred to as NNSGHMC or the meta sampler in the following. Detailed test set-ups are reported in the appendix. The code is available at https://github.com/WenboGong/MetaSGMCMC.

heading: 5.1 SYNTHETIC EXAMPLE: SAMPLING GAUSSIAN VARIABLES WITH NOISY GRADIENTS
text: We first consider sampling Gaussian variables to demonstrate fast convergence and low bias of the meta sampler. To mimic stochastic gradient settings, we manually inject Gaussian noise with unit variance to the gradient as suggested by (Chen et al., 2014). The training density is a 10D Gaussian with randomly generated diagonal covariance matrix, and the test density is a 20D Gaussian. For evaluation, we simulate K = 50 parallel chains for T = 12, 000 steps. Then we follow Ma et al. (2015) to evaluate the sampler’s bias measured by the KL divergence from the empirical estimate to the ground truth. Results are visualized on the left panel of Figure 2, showing that the meta sampler both converges much faster and achieves lower bias compared to SGHMC. The effective sample size2 for SGHMC and NNSGHMC are 22 and 59, again indicating better efficiency of the meta sampler. For illustration purposes, we also plot in the other two panels the trajectory of samples by simulating NNSGHMC (middle) and SGHMC (right) on a 2D Gaussian for a fixed amount of time ηT . This confirms that the meta sampler explores more efficiently and is less affected by the injected noise.

heading: 5.2 BAYESIAN FEEDFORWARD NEURAL NETWORKS
text: Next, we consider Bayesian neural network classification on MNIST data with three generalization tests: network architecture generalization (NT), activation function generalization (AF) and dataset generalization (Data). In all tests, the sampler is trained with a 1-hidden layer multi-layer perceptron (MLP) (20 units, ReLU activation) as the underlying model for the target distribution π(θ). We also report long-time horizon generalization results, meaning that the simulation time steps in test time are much longer than that of training (cf. Andrychowicz et al., 2016). Algorithms in comparison include SGLD (Welling & Teh, 2011), SGHMC (Chen et al., 2014) and preconditioned SGLD (PSGLD, Li et al., 2016a). Note that PSGLD uses RMSprop-like preconditioning techniques (Tieleman & Hinton, 2012) that require moving average estimates of the gradient’s second moments. Therefore the underlying dynamics of PSGLD cannot be represented within our framework (6). Thus we mainly focus on comparisons with SGLD and SGHMC, and leave the PSGLD results as reference. The discretization step-sizes for the samplers are tuned on the validation dataset for each task.
Architecture generalization (NT) In this test we use the trained sampler to draw samples from the posterior distribution of a 2-hidden layer MLP with 40 units and ReLU activations. Figure 3 shows the learning curves of test error and negative test log-likelihood (NLL) for 100 epochs, where the final performance is reported in Table 1. Overall NNSGHMC achieves the fastest convergence even when compared with PSGLD. It has the lowest test error compared to SGLD and SGHMC. NNSGHMC’s final test LL is on par with SGLD and slightly worse than PSGLD, but it is still better than SGHMC.
2Implementation follows the ESS function in the BEAST package http://beast.community.
Architecture + Activation function generalization (NT+AF) Next we replace NT’s test network’s activation function with sigmoid and re-run the same test as before. Again results in Figure 3 and Table 1 show that NNSGHMC converges faster than others for both test error and NLL. It also achieves the best NLL results among all samplers, and the same test error as SGHMC.
Architecture + Dataset generalization (NT+Data) In this test we split MNIST into training task (classifying digits 0-4) and test task (digits 5-9). The meta sampler is trained with the smaller MLP, and it is evaluated on the task with the larger MLP with NT’s architecture. Thus, the meta sampler is trained without any knowledge of the test task’s training and test data. From Figure 3, we see that NNSGHMC, although a bit slower at the start, catches up quickly and proceeds to lower error. The difference between these samplers NLL results is marginal, and NNSGHMC is on par with PSGLD.
Learned strategies For better intuition, we visualize in Figure 4 the contours of f φD and f φQ . From the left panel, f φQ has learned a nearly linear strategy w.r.t. the energy and small variations w.r.t. the momentum. This enables the sampler for fast traversal through low density (high energy) regions and better exploration at high density (low energy) area.
The strategy learned for the diffusion matrix D(z) is rather interesting. Recall that D(z) is parametrized by both f φD and f φQ f φQ (eq. (6)). Since Figure 4 (left) indicates that f φQ is large in high energy regions, the amount of friction is adequate, so f φD tends to reduce its output to maintain momentum (see the middle plot). By contrast, at low energy regions, f φD increases to add friction in order to prevent divergence. The right panel visualizes the interactions between the momentum and the mean gradient − 1N∇θU(θ) at a fixed energy level. This indicates that the meta sampler has learned a strategy to prevent overshoot by producing large friction, indeed f φD returns large values when the signs of the momentum and the gradient differ.

heading: 5.3 BAYESIAN CONVOLUTIONAL NEURAL NETWORKS
text: Following the setup of BNN MNIST experiments, we also test our algorithm on convolutional neural networks (CNNs) for CIFAR-10 (Krizhevsky, 2009) classification, again with three generalization tasks (NT, AF and Data). The meta sampler is trained using a smaller CNN with two convolutional layers (3× 3× 3× 8 and 3× 3× 8× 8) and one fully connected (fc) layer (50 hidden units). ReLU activations, and max-pooling operators of size 2 are applied after each convolutional layer. The meta sampler is trained using 100 “meta-epochs”, where each “meta-epoch” has 5 data epochs. At the beginning of each “meta-epoch”, a “replay” technique inspired by experience replay (Lin, 1993) is utilized (see appendix). The discretization step-sizes are tuned on a validation dataset for each task.
Architecture generalization (NT) The test CNN has two convolutional layers (3× 3× 3× 16 and 3× 3× 16× 16) and one fc layer (100 hidden units), resulting in roughly 4× dimenality of θ. Figure 5 shows that the meta sampler achieves the fastest learning at the first 10 epochs, and continues to have better performance in both test accuracy and NLL. Interestingly, PSGLD slows down quickly after 3 epochs, and it converges to a worse answer. The best performance over 200 epochs is shown in Table 2, where the meta sampler is a clear winner in both accuracy and NLL. This demonstrates that our sampler indeed converges faster and has found a better posterior mode.
Architecture + Activation function generalization (NT+AF) We use the same CNN architecture as in NT but replace the ReLU activations with sigmoid. Figure 5 and Table 2 show that the meta sampler again has better convergence speed and the best final performance.
Architecture + Dataset generalization (NT+Data) We split CIFAR-10 according to labels 0-4 as the training task and 5-9 as the test task. We also used the same CNN architecture as in NT. From Figure 5 and Table 2, the meta sampler consistently achieves the fastest convergence speed. It also achieves similar accuracy as SGHMC, but it has slightly worse test NLL compared to SGHMC.

heading: 5.4 BAYESIAN RECURRENT NEURAL NETWORKS
text: Lastly, we consider a more challenging setup: sequence modeling with Bayesian RNNs. Here a single datum is a sequence on = {x1n, ...,xTn} and the log-likelihood is defined as log p(on|θ) =∑T t=1 log p(x n t |xn1 , . . . ,xnt−1, θ), with each of the conditional densities produced by a gated recurrent unit (GRU) network (Cho et al., 2014). We consider four polyphonic music datasets for this task: Piano-midi (Piano) as training data, and Nottingham (Nott), MuseData (Muse) and JSB chorales (JSB) for evaluation. The meta sampler is trained on a small GRU with 100 hidden states. At test time, we follow Chen et al. (2016) and set the step-size to η = 0.001. We found SGLD significantly under-performs, so instead, we report the performances of two optimizers, Adam (Kingma & Ba, 2014) and Santa (Chen et al., 2016), taken from Chen et al. (2016). Again, these two optimizers use moving average schemes which are out of the scope of our framework, so we mainly compare the meta sampler with SGHMC and leave the others as references.
The meta sampler is tested on the four datasets using 200 unit GRU. So for Piano this corresponds to architecture generalization only. From Figure 6 we see that the meta sampler achieves faster convergence compared to SGHMC, at the same time it achieves similar speed as Santa at early stages. All the samplers achieve best results close to Santa on Piano. The meta sampler successfully generalizes to the other three datasets, demonstrating faster convergence than SGHMC consistently, and better final performance on Muse. Interestingly, the meta sampler’s final results on Nott and JSB are slightly worse than other samplers. Presumably, these two datasets are very different from Muse and Piano, therefore, the energy landscape is less similar to the training density (see appendix). Specifically, JSB is a dataset with much shorter sequences. And in this case, SGHMC also exhibits over-fitting but to a smaller degree. Therefore, we further test the meta sampler on JSB without the offset β in f φQ to reduce the acceleration (denoted as NNSGHMC-s). Surprisingly, NNSGHMC-s
convergences in similar speeds as the original one, but with less amount of over-fitting and better final test NLL 8.40.

heading: 6 CONCLUSIONS AND FUTURE WORK
text: 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.

heading: ACKNOWLEDGMENTS
text: We thank Shixiang Gu, Mark Rowland and Cheng Zhang for comments on the manuscript. We also appreciate Changyou Chen for providing the experiment results of Bayesian RNN (Chen et al., 2016). Wenbo Gong is supported by the CSC-Cambridge Trust Scholarship.

heading: A COMPARING MOMENTUM SGD AND SGHMC
text: 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 ∈ RD. We assume an identity mass matrix associated with that momentum term. Then the corresponding drift f (θ,p) and diffusion matrixD are:
f (θ,p) = [ 0 I −I −C ] [ ∇U(θ) p ] , D = [ 0 0 0 C ] , (12)
where C is a positive definite matrix called friction coefficient. Thus, HMC’s continuous-time dynamics is governed by the following SDE:
dθ = pdt, dp = −∇U(θ)dt−Cpdt+ √ 2CdW (t). (13)
The discretized update rule (with simple Euler discretization) of HMC with step-size η is
θt+1 = θt + ηpt,
pt+1 = (1− ηC)pt − η∇U(θt) + t, t ∼ N (0, 2ηC).
(14)
If stochastic gradient ∇Ũ(θ) is used, we need to replace the covariance matrix of with 2η(C − B̂) where B̂ is the variance estimation of the gradients.
On the other hand, the update equations of SGD with momentum (SGD-M) are the following:
θt+1 = θt + vt, vt+1 = kvt − l∇U(θt). (15)
where k and l are called momentum discount factor and learning rate, respectively. Also we can rewrite the SGHMC update equations by setting ηpt = vt, l = η2,k = (1− ηC),
θt+1 = θt + vt,
vt+1 = kvt − l∇U(θt) + ̂t , ̂t ∼ N (0, 2l(1− k)).
(16)
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.
Neal et al. (2011) 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:
pt+1 = (1− ηC)pt − η∇U(θt) + t, θt+1 = θt + ηpt+1,
t ∼ N (0, 2ηC). (17)

heading: B FINITE DIFFERENCE APPROXIMATION FOR THE GAMMA VECTOR
text: The main computational burden is the gradient computation required by Γ(z) vector. To simplify notations we will writeD = D(z) andQ = Q(z). From the parametrization of theQ andD matrices in eq. (6), for θ,p ∈ RD we have Γ(z) = [Γθ ,Γp ]. For the first term Γθ , we have
Γθ,i = −∇p ·Qi,: = − ∂fφQ,i
∂pi . (18)
Due to the two-stage update of Euler integrator, at time t, we have f t−1φQ,i = fφQ,i(U(θt−1), p t−1 i ), f̂ t−1φQ,i = fφQ,i(U(θt−1), p t i) and f t−1 φD,i = fφD,i(U(θt−1), p t−1 i ,∇θt−1i U(θt−1)). Thus a proper finite
difference method requires fφQ,i(U(θt), p t−1 i ), which is not exactly the history from the previous time. Therefore we further approximate it using delayed estimate:
∂f tφQ,i ∂pti ≈ f̂ t−1φQ,i − f t−1 φQ,i pti − p t−1 i ⇒ Γtθ ≈ − Q̂ t−1 −Qt−1 pt − pt−1 . (19)
Similarly, the Γp term expands as
Γp,i = ∇ · [D +Q]i,:
= ∂fφQ,i ∂θi + ∂fφD,i ∂pi + 2αfφQ,i ∂fφQ,i ∂pi
= ∂fφQ,i
∂U(θ)
∂U(θ) ∂θi + ∂fφD,i ∂pi + 2αfφQ,i ∂fφQ,i ∂pi .
(20)
We further approximate ∂fφQ,i
∂U(θ) by the following
∂fφQ,i ∂U(θ) ≈
f tφQ,i − f̂ t−1 φQ,i
U(θt)− U(θt−1) (21)
This only requires the storage of previousQ matrix. However, ∂fφD,i∂pi requires one further forward pass to obtain f̂ t−1φD,i = fφD,i(U(θt), p t−1 i ,∇θtiU(θt)), thus, we have
∂fφD,i ∂pi
≈ f tφD,i − f̂ t−1 φD,i
pti − p t−1 i
⇒ Γp ≈ Qt − Q̂
t−1
U(θt)− U(θt−1) ∇θU(θt) +
f tφD − f̂ t−1 φD
pt − pt−1 + 2αf tφQ
Q̂ t−1 −Qt−1
pt − pt−1 .
(22)
Therefore the proposed finite difference method only requires one more forward passes to compute 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∇θ Ũ(θ). The main burden is the forward pass and back-propagation through the D(z) and Q(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.

heading: C DETAILS OF THE STEIN GRADIENT ESTIMATOR
text: For a distribution q(θ) that is implicitly defined by a generative procedure, the density q(θ) is often intractable. Li & Turner (2018) derived the Stein gradient estimator that estimates G = (∇θ1 log q(θ1), · · · ∇θK log q(θK))T on samples θ1, ..., θK ∼ q(θ). There are two different ways to derive this gradient estimator, here we briefly introduce one of them, and refer the readers to Li & Turner (2018) for details.
We start by introducing Stein’s identity (Stein, 1972; 1981; Gorham & Mackey, 2015; Liu et al., 2016). Let h : Rd×1 → Rd′×1 be a differentiable multivariate test function which maps θ to a column vector h(θ) = [h1(θ), h2(θ), ..., hd′(θ)]
T. One can use integration by parts to show the following Stein’s identity when a boundary condition lim||θ||→∞ q(θ)h(θ) = 0 is assumed for the test function:
Eq[h(θ)∇θ log q(θ)T +∇θh(θ)] = 0, ∇θh(θ) = (∇θh1(θ), · · · ,∇θhd′(θ))T ∈ Rd ′×d. (23)
This boundary condition holds for almost any test function if q has sufficiently fast-decaying tails (e.g. Gaussian tails). Li & Turner (2018) proposed the Stein gradient estimator for ∇θ log q(θ) by
inverting a Monte Carlo (MC) version of Stein’s identity (23):
− 1 K
HG ≈ ∇θh, H = ( h(θ1), · · · ,h(θK) ) ∈ Rd ′×K , ∇θh = 1
K K∑ k=1 ∇θkh(θk) ∈ Rd ′×d.
Then G is obtained by ridge regression (with || · ||F the Frobenius norm of a matrix)
ĜSteinV := arg min Ĝ∈RK×d
||∇θh + 1
K HĜ||2F +
η
K2 ||Ĝ||2F , η ≥ 0, (24)
which has an analytical solution
ĜSteinV = −(K + ηI)−1〈∇,K〉, (25)
where K := HTH, Kij = K(θi, θj) := h(θi)Th(θj),
〈∇,K〉 := KHT∇θh, 〈∇,K〉ij = K∑ k=1 ∇θk(j)K(θi, θk).
Here θk(j) denotes the jth element of vector θk. One can show that the RBF kernel satisfies Stein’s identity (Liu et al., 2016). In this case h(θ) = K(θ, ·), d′ = +∞ and by the reproducing kernel property, h(θ)Th(θ ′) = 〈K(θ, ·),K(θ ′, ·)〉H = K(θ,θ ′). Li & Turner (2018) also show that the Stein gradient estimator can be obtained by minimizing a Monte Carlo estimate of the kernelized Stein discrepancy (Chwialkowski et al., 2016; Liu et al., 2016).
The kernel choice It is well-known for kernel methods that a better choice of the kernel can greatly improve the performance. However, optimal kernels are often problem specific, and they are generally difficult to obtain. Recently, a popular approach for kernel design is to compose a simple kernel (e.g. RBF kernel) on features extracted from a deep neural network. Representative work include deep kernel learning for Gaussian processes (Wilson et al., 2016), and adversarial approaches to learn kernel parameters (Li et al., 2017; Bińkowski et al., 2018). Unfortunately, both approaches do not scale very well to our application as θ has at least tens of thousands of dimensions. Furthermore, they both considered kernel learning for observed data, while in our case θ is a latent variable to be inferred. Therefore it remains a research question on how to learn kernels on latent variables efficiently, and addressing this question is out of the scope of the paper. Instead, we follow Liu & Wang (2016); Li & Turner (2018) to use RBF kernel for the gradient estimator. Other kernels can be trivially adapted to our method. We expect even better performance if an optimal kernel is in use, but we leave the investigation to future work.
Time complexity figures During meta sampler training, the Stein gradient estimator requires the kernel matrix inversion which is O(K3) for cross-chain training. In practice, we only run a few parallel Markov chains K = 20 ∼ 50, thus, this will not incur huge computation cost. For in-chain loss the computation can also be reduced with proper thinning schemes.
D IMPLEMENTATION DETAILS OF THE TRAINING LOSS
We visualize on the left panel of Figure 7 the unrolled computation scheme. We apply truncated back-propagate through time (BPTT) to train the sampler. Specifically, we manually stop the gradient flow through the input ofD andQ matrices to avoid computing higher order gradients.
We also illustrate cross-chain in-chain training on the right panel of Figure 7. Cross-chain training encourages both fast convergence and low bias, provided that the samples are taken from parallel chains. On the other hand, in-chain training encourages sample diversity inside a chain. In practice, we might consider thinning the chains when performing in-chain training. Empirically this improves the Stein gradient estimator’s accuracy as the samples are spread out. Computationally, this also prevents inverting big matrices for the Stein gradient estimator, and reduces the number of backpropagation operations. Another trick we applied is parallel chain sub-sampling: if all the chains are used, then there is less encouragement of singe chain mixing, since the parallel chain samples can be diverse enough already to give reasonable gradient estimate.
E INPUT PRE-PROCESSING
One potential challenge is that for different tasks and problem dimensions, the energy function, momentum and energy gradient can have very different scales and magnitudes. This affects the meta sampler’s generalization, for example, if training and test densities have completely different energy scales, then the meta sampler is likely to produce wrong strategies. This is especially the case when the meta sampler is generalized to much bigger networks or to very different datasets.
To mediate this issue, we propose to pre-process the inputs to both f φD and f φQ networks to make it at similar scale as those in training task. Recall that the energy function is U(θ) = − ∑N n=1 log p(yn|xn, θ)− log p(θ) where the prior log p(θ) is often an isotropic Gaussian distribution. Thus the energy function scale linearly w.r.t both the dimensionality of θ and the total number of observations N . Often the energy function is further approximated using mini-batches of M datapoints. Putting them together, we propose pre-processing the energy as
U(θ) = 1
M M∑ m=1 log p(ym|xm, θ) + Dtrain NDtest log p(θ) (26)
where Dtrain and Dtest are the dimensionality of θ in the training task and the test task, respectively. Importantly, for RNNs N represents the total sequence length, namely N = ∑Ndata n=1 Tn, where Ndata is the total number of sequences and Tn is the sequence length for a datum xn. We also define M accordingly. The momentum and energy gradient magnitudes are estimated by simulating a randomly initialized meta sampler for short iterations. With these statistics we normalize both the momentum and the energy gradient to have roughly zero mean and unit variance.

heading: F EXPERIMENT SETUP

heading: F.1 TOY EXAMPLE
text: We train our meta sampler on a 10D uncorrelated Gaussian with mean (3, ..., 3) and randomly generated covariance matrix. We do not set any offset and additional frictions, i.e. α = 0 and β = 0. The noise estimation matrix B̃ are set to be 0 for both meta sampler and SGHMC. To mimic stochastic gradient, we manually inject Gaussian noise with zero mean and unit variance into ∇θ Ũ(θ) = ∇θU(θ)+ , ∼ N (0, I ). The functions f φD and f φQ are represented by 1-hidden-layer MLPs with 40 hidden units. For training task, the meta sampler step size is 0.01. The initial positions are drawn from Uniform([0, 6]D). We train our sampler for 100 epochs and each epochs consists 4 x 100 steps. For every 100 steps, we updates theQ andD matrices using Adam optimizer with learning rate 0.0005. Then we continue the updated sampler with last position and momentum until 4
sub-epochs are finished. We re-initialize the momentum and position. We use both cross-chain and in-chain losses. The Stein Gradient estimator uses RBF kernel with bandwidth chosen to be 0.5 times the median-heuristic estimated value. We unroll the Markov Chain for 20 steps before we manually stop the gradient. For cross-chain training, we take sampler across chain for each 2 time steps. For in-Chain, we discard initial 50 points for burn-in and sub-sample the chain with batch size 5. We thin the samples for every 3 steps. For both training and evaluation, we run 50 parallel Markov Chains.
The test task is to draw samples from a 20D correlated Gaussian with with mean (3, ..., 3) and randomly generated covariance matrix. The step size is 0.025 for both meta sampler and SGHMC. To stabilize the meta sampler we also clamp the output values of f φQ within [−5, 5]. The friction matrix for SGHMC is selected as I .

heading: F.2 BAYESIAN MLP MNIST
text: In MNIST experiment, we apply input pre-processing on energy function as in (26) and scale energy gradient by 70. Also, we scale up f φD by 50 to account for sum of stochastic noise. The offset α is selected as 0.01η as suggested by Chen et al. (2014), where η = √ lr N with lr the per-batch
learning rate. We also turn off the off-set and noise estimation, i.e. β = 0 and B̃ = 0. We run 20 parallel chains for both training and evaluation. We only adopt the cross chain training with thinning samplers of 5 times step. We also use the finite difference technique during evaluation to speed-up computations.

heading: F.2.1 ARCHITECTURE GENERALIZATION (NT)
text: We train the meta sampler on a smaller BNN with architecture 784-20-10 and ReLU activation function, then test it on a larger one with architecture 784-40-40-10. In both cases the batch size is 500 following Chen et al. (2014). Both f φD and f φQ are parameterized by 1-hidden-layer MLPs with 10 units. The per-batch learning rate is 0.007. We train the sampler for 100 epochs and each one consists of 7 sub-epochs. For each sub-epoch, we run the sampler for 100 steps. We re-initialize θ and momentum after each epoch. To stabilize the meta sampler in evaluation, we first run the meta sampler with small per-batch learning rate 0.0085 for 3 data epochs and clamp theQ values. After, we increase the per-batch learning rate to 0.018 with clipped f φQ . The learning rate for SGHMC is 0.01 for all times. For SGLD and PSGLD, they are 0.2 and 1.4× 10−3 respectively. These step-sizes are tuned on MNIST validation data.

heading: F.2.2 NT + ACTIVATION FUNCTION GENERALIZATION
text: We modify the test network’s activation function to sigmoid. We use almost the same settings as in network generalization tests, except that the per-batch learning rates are tuned again on validation data. For the meta sampler and SGHMC, they are 0.18 and 0.15. For SGLD and PSGLD, they are 1 and 1.3× 10−2.

heading: F.2.3 NT + DATASET GENERALIZATION
text: We train the meta sampler on ReLU network with architecture 784-20-5 to classify images 0-4, and test the sampler on ReLU network 784-40-40-5 to classify images 5-9. The settings are mostly the same as in network architecture generalization for both training and evaluation. One exception is again the per-batch learning rate for PSGLD, which is tuned as 1.3× 10−3. Note that even though we use the same per-batch learning rate as before, the discretization step-size is now different due to smaller training dataset, thus, α will be automatically adjusted accordingly.

heading: F.3 BAYESIAN CONVOLUTIONAL NEURAL NETWORK ON CIFAR-10
text: CIFAR-10 dataset contains 50,000 training images with 10 labels and 10,000 test images. We train our meta sampler using smaller CNN classifier with two convolutional layer (3 × 3 × 3 × 8 and 3× 3× 8× 8, no padding) and one fc layer of 50 hidden units. Therefore the dimensionality of θ is 15, 768. The training sampler discretization step-size η is √ 0.0007 50000. and scaling term is α = 0.005 η . To make it analogous to optimization methods, we call 0.0007 as per-batch learning rate and 0.005
as friction coefficient. The f φQ and f φD are defined by 2-layer MLPs with 10 hidden units. We set the offset values to 0 for bothQ andD. Further, we scale up the output ofDf (z) by 10 and its gradient input ∇U(θ) by 100. We scale up the energy input U(θ) to both f φQ and f φD by 5. We train our meta sampler using 100 “meta epoch” with 5 data epoch and 500 batch size. Within each ”meta epoch”, we repeat the following computation for 10 times: we run 50 parallel chains using the meta sampler for 50 iterations (0.5 dataset epoch), compute the loss function, and update the meta sampler’s parameters using Adam. We manually stop the gradient after 20 iterations. Then we start the next sub-epoch using the last θ and p. After we finish all sub-epoch, we re-initialize the θ and p using replay techniques with probability 0.15. The sub-sample chain number for in-chain loss is set to 5.

heading: F.3.1 THE REPLAY TECHNIQUE
text: Experience replay (Lin, 1993) is a technique broadly used in reinforcement learning literature. Inspired by this, in Bayesian CNN experiments we train the meta sampler in a similar way, and we found this replay technique particularly useful for more complicated dataset like CIFAR-10.
At the beginning of each “meta epoch”, each chain is initialized either with a specific state randomly chosen from a replay pool, or with a random state sampled from a Gaussian distribution. We use a pre-defined replay probability to control the replay strategy. The replay pool is updated after each sub-epoch, and it has a queue-like data structure of constant size, so that the old states are replaced by the new ones. Therefore, this replay technique is useful for both short-time and long-time horizon generalization. On one hand, the meta sampler can continue with previous states, allowing it to accommodate long-time horizon behavior. On the other hand, due to non-zero probability of random restart, the meta sampler can learn a better strategy for fast convergence. Therefore with this replay technique, the sampler can observe both burn-in and roughly-converged behavior, and this balance is controlled by the replay probability.

heading: F.3.2 ARCHITECTURE GENERALIZATION (NT)
text: For architecture generalization, the test CNN has two convolutional layer (3 × 3 × 3 × 16 and 3 × 3 × 16 × 16, no padding) and one fully connected layer with 100 hidden units. Thus, the dimensionality of θ is 61,478, roughly 4 times of the training dimension. We run 20 parallel chains in test time. We split the 50,000 training images into 45,000 training and 5,000 validation images, and tune the discretization step-size of each sampling and optimization methods on the validation set for 80 epochs. For test, we run the tuned samplers/optimizers for 200 data epoch (roughly 40 times longer than training) to ensure convergence. For the meta sampler, the per-batch rate is 0.003. For SGHMC, the per-batch is also 0.003 with friction coefficient 0.01. For SGLD, the per-batch learning rate is 0.15. PSGLD uses 1.3× 10−3 as learning rate and 0.99 as moving average term. For optimization methods, we use learning rate 0.002 for Adam and 0.003 for SGD-M. The momentum term is 0.9. To prevent overfitting, we use weight penalty with coefficient 0.001.

heading: F.3.3 NT + SIGMOID GENERALIZATION
text: The test CNN has same architecture as in NT, except that it replaces all ReLU activation functions with sigmoid activations. We fix all other parameters for sampling method and only re-tune the step sizes using same setup as in NT. The per-batch rate for meta sampler, SGHMC, SGLD and PSGLD are 0.1, 0.03, 0.5 and 0.005 respectively. For optimization methods, the step size for Adam and SGD-M are 0.002 and 0.03 respectively.

heading: F.3.4 NT + DATASET GENERALIZATION
text: We split the CIFAR-10 training and test dataset according to the labels. We use training data with labels 0-4 for meta sampler training, training data with labels 5-9 for test CNN training, and test data with labels 5-9 for test CNN evaluation. Thus, the meta sampler has no access to the test task’s training and test data during sampler training. We train our sampler using the same scaling terms as in NT but reduce the discretization step-size to 0.0005. The rest setup is the same as in NT.
We use the same test CNN architecture and ReLU activation as in NT, and tune the learning rate using validation data. The step size for the meta sampler, SGHMC, SGLD and PSGLD are 0.0015, 0.005,
0.2 and 0.0018, respectively. For optimization methods, we use learning rates 0.002 and 0.003 for Adam and SGD-M respectively.

heading: F.4 BAYESIAN RNN
text: The Piano data is selected as the training task, which is further split into training, validation and test subsets. We use batch-size 1, meaning that the energy and the gradient are estimated on a single sequence. The meta sampler uses similar neural network architectures as in MNIST tests. The training and evaluation per-batch learning rate for all the samplers is set to be 0.001 following Chen et al. (2016). We train the meta sampler for 40 epochs with 7 sub-epochs with only cross chain loss. Each sub-epochs consists 70 iterations. We scale theD output by 20 and set α = 0.002η , where η is defined in the same way as before. We use zero offset during training, i.e. β = 0. We apply input pre-processing for both f φD and f φQ . To prevent divergence of the meta sampler at early training stage. We also set the constant of c = 100 to the fφD . For dataset generalization, we tune the off-set value based on Piano validation set and transfer the tuned setting β = −1.5 to the other three datasets. For Piano architecture generalization, we do not tune any hyper-parameters including β and use exactly same settings as training. Exact gradient is used in RNN experiments instead of computing finite differences.

heading: G RNN DATASET DESCRIPTION
text: We list some data statistics in Table 3 which roughly indicates the similarity between datasets. Piano dataset is the smallest in terms of data number, however, the averaged sequence length is the largest. Muse dataset is similar to Piano in sequence length and energy scale but much larger in terms of data number. On the other hand, Nott dataset has very different energy scale compared to the other three. This potentially makes the generalization much harder due to inconsistent energy scale fed into f φQ and f φD . For JSB, we notice a very short sequence length on average, therefore the GRU model is more likely to over-fit. Indeed, some algorithms exhibits significant over-fitting behavior on JSB dataset compared to other data (Santa is particularly severe).

heading: H ADDITIONAL PLOTS

heading: H.1 SHORT HORIZON PERFORMANCE COMPARISONS
text: We also run the samplers using the same settings as in MNIST experiments for a short period of time (500 iterations). We also compare to other optimization methods including momentum SGD (SGD-M) and Adam. We use the same per-batch learning rate for SGD-M and SGHMC as in MNIST experiment. For Adam, we use 0.002 for ReLU and 0.01 for Sigmoid network.
The results are shown in Figure 8. Meta sampler and Adam achieves the fastest convergence speed. This again confirms the faster convergence of the meta sampler especially at initial stages. We also provide additional contour plots (Figure 9) for MNIST experiments to demonstrate the strategy learned by f φD for reference.

Review summaries:
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.