Title: META-LEARNING FOR STOCHASTIC GRADIENT MCMC

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.

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 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.

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 (θ θ θ), θ θ θ ∈ R D , π(θ θ θ) ∝ exp(-U (θ θ θ)). In this paper, we focus on this sampling task with a Bayesian modeling set-up, i.e. given observed data D = {o o o n } N n=1 , we define a probabilistic model p(D, θ θ θ) = N n=1 p(o o o n |θ θ θ)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, o o o n = (x x x n , y y y n ), the prior p(θ θ θ) is a Gaussian N (θ θ θ; 0 0 0, λ -1 I I I), and the energy function is defined as

EQUATION

with (y y y, ŷ 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 z z = (θ θ θ, r r r) with auxiliary variables r r r, and sample from the augmented distribution π(z z z) ∝ exp (-H(z z z)), with the Hamiltonian H(z z z) = U (θ θ θ) + g(r r r) such that exp(-g(r r r))dr r r = C. Thus, marginalizing out the auxiliary variable r r 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:

EQUATION

with f f f (z z z) the deterministic drift, W W W (t) the Wiener process, and D D D(z z z) the diffusion matrix. As a simple example, Langevin dynamics considers z z z = θ θ θ, f f f (θ θ θ) = -∇ θ θ θ U (θ θ θ) and D D D(θ θ θ) = I. Then using forward Euler discretization with step-size η the update rule of the parameters is

θ θ θ t+1 = θ θ θ t -η∇ θ θ θt U (θ θ θ t ) + 2η t , t ∼ N (0 0 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:

EQUATION

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 f f (z z z) :

EQUATION

with Q Q Q(z z z) the curl matrix, D D D(z z z) the diffusion matrix and Γ Γ Γ(z z z) a correction term. Remarkably Ma et al. (2015) showed the completeness of their framework:

1. π(z z z) ∝ exp(-H(z z z)) is a stationary distribution of the SDE (2)+(5) for any pair of positive semi-definite matrix D D D(z z z) and skew-symmetric matrix Q Q Q(z z z); 2. for any Itô diffusion process that has the unique stationary distribution π(z z z), under mild conditions there exist D D D(z z z) and Q Q Q(z z 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 z z and the D D D(z z z), Q Q Q(z z 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 of D D D(z z z) and Q Q Q(z z 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 p 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 z z = (θ θ θ, p p p) (i.e. r r r = p p p), and the Hamiltonian is defined as H(z z z) = U (θ θ θ) + 1 2 p p p T p p p with identity mass matrix.

3.1. EFFICIENT PARAMETERIZATION OF DIFFUSION AND CURL MATRICES:

For neural networks, the dimensionality of θ θ θ can be at least tens of thousands. Thus, training and applying full D D D(z z z) and Q Q Q(z z z) matrices can cause a huge computational burden, let alone gradient computations required by Γ Γ Γ(z z z). To address this, we define the preconditioning matrices as follows:

EQUATION

Here f f f φ D and f f f φ Q are neural network parameterized functions that will be detailed in section 3.2, and c is a small positive constant. We choose D D D f and Q Q Q f 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 f f φ D is non-negative for all z z z.

We discuss the role of each precondition matrix for better intuition. The curl matrix Q Q Q(z z 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 ∇ Figure 1 : Comparing the computation graphs of SGHMC and the meta-learned sampler (with modified forward Euler discretization). Here F φ and G φ transformations are defined by Eq. ( 7).

D D D f (z z z) should also account for the pre-conditioning effect introduced by

Q Q Q f (z z z), e.g, when the magnitude of Q Q Q f (z z z

) is large, we need higher friction correspondingly. This explains the squared term

f f f φ Q (z z z) f f f φ Q (z z z) in D D D f (z z 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 z) = [Γ Γ Γ θ θ θ (z z z), Γ Γ Γ p p p (z z z)] is responsible for compensating the changes introduced by preconditioning matrices Q Q Q(z z z) and D D D(z z z).

The discretized dynamics of the state z z z = (θ θ θ, p p p) with step-size η and stochastic gradient ∇ θ θ θ Ũ (θ θ θ) are

EQUATION

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 that Q Q Q f (z z z) is responsible for the acceleration of θ θ θ, and from the (1 -ηD D D f (z z z)) term in the update equation of p p p, we see that D D D f (z z 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 0 0 and variance V V 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 0 0, 2ηD

D D f (z z z) -η 2 B(θ θ θ)) with an empirical estimate of the variance B(θ θ θ) ≈ Q Q Q f (z z z)V V V (θ θ θ)Q Q Q T f (z z z) instead.

These corrections can be dropped when the discretization step-size η is small, therefore, we do not consider them in our experiments.

3.2. CHOICES OF INPUTS TO THE NEURAL NETWORKS:

We now present detailed functional forms for f f f φ Q and f f f φ D . When designing these, our goal was to achieve a good balance between generalization power and computational efficiency. Recall that the curl matrix Q Q Q(z z 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 p i to the inputs of f f f φ Q , allowing the Q Q Q(z z z) matrix to observe the velocity information of the θ θ θ i . We further add an offset β to Q Q Q(z z z) to prevent the vanishing of this matrix. Putting all of them together, we define the i th element of

EQUATION

The corresponding Γ Γ Γ(z z z) term requires both

∂ θi f φ Q (U (θ θ θ), p i ) and ∂ pi f φ Q (U (θ θ θ), p i ). The en- ergy gradient ∂ θi U (θ θ θ) also appears in (7), so it remains to compute ∂ U f f f φ Q , which, along with ∂ pi f φ Q (U (θ θ θ), p i )

, can be obtained by automatic differentiation (Abadi et al., 2015) .

Matrix D D D(z z 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 ∂ θi U (θ θ θ) to the inputs, meaning that the

EQUATION

By the construction of the D D D(z z z) matrix, the Γ Γ Γ vector only requires ∇ p p p D D D f without computing any higher order information.

In practice, both U (θ θ θ) and ∂ θi U (θ θ θ) are replaced by their stochastic estimates Ũ (θ θ θ) and ∂ θi Ũ (θ θ θ), 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 f f φ D and f f f φ Q . When the computational budget is limited, we replace the exact gradient computation required by Γ Γ Γ(z z z) with finite difference approximations. We refer the reader to the appendix for details.

3.3. LOSS FUNCTION DESIGN FOR META-LEARNING:

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 q t (θ θ θ|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[q t ||π] 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): et al., 1999; Beal, 2003) . We further make the objective doubly stochastic: (1) the energy term is further approximated by its stochastic estimates Ũ (θ θ θ); (2) we use Monte Carlo variational inference (MCVI, Ranganath et al., 2014; Blundell et al., 2015) which estimates the lower-bound with samples θ θ θ k t ∼ q t (θ θ θ t |D), k = 1, ..., K. These samples {θ θ θ k t } 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:

L t VI (q t ) = -E qt [U (θ θ θ)] + H[q t ] (Jordan

EQUATION

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) = 1

T /τ T /τ s=1 q sτ (θ θ θ). 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

Θ Θ Θ k T,τ = {θ θ θ k sτ } T /τ s=1

obtained by thinning:

EQUATION

)

Gradient approximation We leverage the recently proposed Stein gradient estimator (Li & Turner, 2018) to estimate the intractable gradients ∇ φ log q t (θ θ θ) for cross-chain loss and ∇ φ log q(θ θ θ) for in-chain loss. Precisely, by the chain rule, we have

∇ φ log q t (θ θ θ) = ∇ φ θ θ θ∇ θ θ θ log q t (θ θ θ), so it remains to estimate the gradients G G G = (∇ θ θ θ 1 t log q t (θ θ θ 1 t ), . . . , ∇ θ θ θ K t log q t (θ θ θ K t )) T at the sampled locations {θ θ θ k t } K k=1 ∼ q t . The recipe first constructs a kernel matrix K K K with K K K ij = K(θ θ θ i t , θ θ θ j t )

, and then estimates the gradients by:

G G G ≈ -(K K K + λI I I) -1 ∇, K K K , where ∇, K K K ij = K k=1 ∂ θ θ θ k t (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(K 2 D + K 3 ) time (usually K D).

4. RELATED WORK:

Since the development of SGLD (Welling & Teh, 2011) 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 D D(z z z) and Q Q Q(z z 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. 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.

5.2. BAYESIAN FEEDFORWARD NEURAL NETWORKS:

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 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 f f φ D and f f f φ Q .

From the left panel, f f 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 D D(z z z) is rather interesting. Recall that D D D(z z z) is parametrized by both f f f φ D and f f f φ Q f f f φ Q (eq. ( 6)). Since Figure 4 (left) indicates that f f f φ Q is large in high energy regions, the amount of friction is adequate, so f f f φ D tends to reduce its output to maintain momentum (see the middle plot). By contrast, at low energy regions, f f f φ D increases to add friction in order to prevent divergence. The right panel visualizes the interactions between the momentum and the mean gradient -1 N ∇ θ θ θ 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 f f φ D returns large values when the signs of the momentum and the gradient differ.

5.3. BAYESIAN CONVOLUTIONAL NEURAL NETWORKS:

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.

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.

5.4. BAYESIAN RECURRENT NEURAL NETWORKS:

Lastly, we consider a more challenging setup: sequence modeling with Bayesian RNNs. Here a single datum is a sequence o o o n = {x x x 1 n , ..., x x x T n } and the log-likelihood is defined as log p(o o o n |θ θ θ) = T t=1 log p(x x x n t |x x x n 1 , . . . , x x x n t-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 f f φ Q to reduce the acceleration (denoted as NNSGHMC-s). Surprisingly, NNSGHMC-s 

6. CONCLUSIONS AND FUTURE WORK:

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:

EQUATION

)

where C C C is a positive definite matrix called friction coefficient. Thus, HMC's continuous-time dynamics is governed by the following SDE:

EQUATION

)

The discretized update rule (with simple Euler discretization) of HMC with step-size η is

EQUATION

)

If stochastic gradient ∇ Ũ (θ θ θ) is used, we need to replace the covariance matrix of with 2η(C C C -B B B) where B 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:

EQUATION

)

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 p

t = v v v t , l = η 2 ,k = (1 -ηC C C), θ θ θ t+1 = θ θ θ t + v v v t , v v v t+1 = kv v v t -l∇U (θ θ θ t ) + ˆ t ˆ t ˆ t , ˆ t ˆ t ˆ t ∼ N (0 0 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:

EQUATION

B FINITE DIFFERENCE APPROXIMATION FOR THE GAMMA VECTOR:

The main computational burden is the gradient computation required by Γ Γ Γ(z z z) vector. To simplify notations we will write

EQUATION

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-1 i 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:

EQUATION

Similarly, the Γ Γ Γ p p p term expands as

EQUATION

We further approximate

∂f φ Q ,i

∂U (θ θ θ) by the following

EQUATION

This only requires the storage of previous Q Q Q 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 , ∇ θ t i U (θ θ θ t ))

, thus, we have

EQUATION

Therefore the proposed finite difference method only requires one more forward passes to compute 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.

C DETAILS OF THE STEIN GRADIENT ESTIMATOR:

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 : R d×1 → R d ×1 be a differentiable multivariate test function which maps θ θ θ to a column vector h(θ θ θ) = [h 1 (θ θ θ), h 2 (θ θ θ), ..., h d (θ θ θ)] 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:

EQUATION

)

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 ) ∈ R d ×K , ∇ θ θ θ h = 1 K K k=1 ∇ θ θ θ k h(θ θ θ k ) ∈ R d ×d .

Then G is obtained by ridge regression (with || • || F the Frobenius norm of a matrix)

ĜStein:

EQUATION

which has an analytical solution

EQUATION

where

K := H T H, K ij = K(θ θ θ i , θ θ θ j ) := h(θ θ θ i ) T h(θ θ θ j ), ∇, K := KH T ∇ θ θ θ h, ∇, K ij = K k=1 ∇ θ θ θ k (j) K(θ θ θ i , θ θ θ k ).

Here θ θ θ k (j) denotes the j th 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(θ θ θ) T h(θ θ θ ) = 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; Biń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(K 3 ) 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 of D D D and Q Q Q 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 f f φ D and f f 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(y y y n |x x x n , θ θ θ) -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

EQUATION

where D train and D test are the dimensionality of θ θ θ in the training task and the test task, respectively. Importantly, for RNNs N represents the total sequence length, namely N = N data n=1 T n , where N data is the total number of sequences and T n is the sequence length for a datum x x x n . 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. as friction coefficient. The f f f φ Q and f f f φ D are defined by 2-layer MLPs with 10 hidden units. We set the offset values to 0 for both Q Q Q and D D D. Further, we scale up the output of D D D f (z z z) by 10 and its gradient input ∇U (θ θ θ) by 100. We scale up the energy input U (θ θ θ) to both f f f φ Q and f f 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 p p. After we finish all sub-epoch, we re-initialize the θ θ θ and p p p using replay techniques with probability 0.15. The sub-sample chain number for in-chain loss is set to 5.

F.3.1 THE REPLAY TECHNIQUE:

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.

F.3.2 ARCHITECTURE GENERALIZATION (NT):

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.

F.3.3 NT + SIGMOID GENERALIZATION:

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.

F.3.4 NT + DATASET GENERALIZATION:

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, 

F.4 BAYESIAN RNN:

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 sub-epochs with only cross chain loss. Each sub-epochs consists 70 iterations. We scale the D D D 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 f f φ D and f f 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.

G RNN DATASET DESCRIPTION:

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 f f φ Q and f f 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).

H ADDITIONAL PLOTS H.1 SHORT HORIZON PERFORMANCE COMPARISONS:

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 f f φ D for reference. 

:

The energy gradient ∇ θ θ θ U (θ θ θ) is also informative here, however, it requires expensive computation of the diagonal Hessian for Γ Γ Γ(z z z). For similar reasons we do not consider other higher order derivatives as inputs.

EXPERIMENTSWe 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.

Implementation follows the ESS function in the BEAST package http://beast.community.

Figures and Tables:

θ θ θ Ũ (θ θ θ) through data sub-sampling, and here D D D(z z z) acts as the friction to counter for the associated noise that mainly affects the momentum p p p. This explains the design of the diffusion matrix D D D(z z z) that uses D D D f (z z z) to control the amount of friction and injected noise to the momentum. Furthermore, the meta-learned sampler (NNSGHMC)

Figure 2: (Left) Sampler's bias measured by KL. (Middle) NNSGHMC trajectory plot on a 2D-Gaussian with manually injected gradient noise. (Right) SGHMC plot for the same settings.

Figure 3: Learning curves on test error (top) and negative test LL (bottom).

Figure 4: (Left) The contour plot of function f f f φ Q (Middle) The contour plot for f f f φ D for dimension 1 and 2 with fixed -∇ θ θ θ U (θ θ θ) (Right) The same plot for f f f φ D for dimension 2 and 3 with fixed energy.

Figure 5: Learning curves on test error (top) and negative test LL/100 (bottom).

Figure 7: (Left) The unrolled scheme of the meta sampler updates. Stop gradient operations are applied to the dashed arrows. (Right) A visualization of cross-chain in-chain training. The grey area represents samples across multiple chains, and we compute the cross chain loss for every 5 time steps. The purple area indicates the samples taken across time with sub-sampled chains 1 and 3. In this visualization the initial 15 samples are discarded for burn-in, and the thinning length is τ = 1 (effectively no thinning).

Figure 8: We only test the Network Generalization and Activation function generalization. The upper part indicates the test error plot and lower part are the negative test LL curve

The final performance for the samplers on MNIST, averaged over 10 independent runs.

Table content: <table><tr><td>Sampler</td><td>NT Err.</td><td>NT+AF Err</td><td>NT+Data Err</td><td colspan="2">NT NLL NT+AF NLL NT+Data NLL</td></tr><tr><td colspan="5">NNSGHMC 98.36±0.02 0.02 0.02% 97.72±0.02 0.02 0.02% 98.62±0.02 0.02 0.02% 640±6.25</td><td>875±3.19 3.19 3.19</td><td>230±3.23</td></tr><tr><td>SGHMC</td><td colspan="4">98.21±0.01% 97.72±0.01 0.01 0.01% 98.52±0.03% 705±3.44</td><td>929±2.95</td><td>246±5.43</td></tr><tr><td>SGLD</td><td colspan="4">98.27±0.02% 97.62±0.02% 98.54±0.01% 631±3.15</td><td>905±2.36</td><td>232±1.93</td></tr><tr><td>PSGLD</td><td colspan="4">98.31±0.02% 97.67±0.02% 98.60±0.02% 610±2.93 2.93 2.93</td><td>975±4.41</td><td>224±1.97 1.97 1.97</td></tr><tr><td colspan="6">5.1 SYNTHETIC EXAMPLE: SAMPLING GAUSSIAN VARIABLES WITH NOISY GRADIENTS</td></tr></table>

The best performance on CIFAR-10 over 200 epoch , averaged over 5 independent runs. All the samplers achieved the best performance after around 190 epochs.

Table content: <table><tr><td>Methods</td><td>NT Err.</td><td>NT+AF Err</td><td colspan="4">NT+Data Err NT NLL/100 NT+AF NLL/100 NT+Data NLL/100</td></tr><tr><td colspan="3">NNSGHMC 78.12±0.035 0.035 0.035% 74.41±0.11 0.11 0.11%</td><td>89.97±0.04 0.04 0.04%</td><td>68.88±0.15 0.15 0.15</td><td>79.55±0.057 0.057 0.057</td><td>15.66±0.28</td></tr><tr><td>SGHMC</td><td colspan="4">77.63±0.068% 73.68±0.17% 90.11±0.15 0.15 0.15% 70.39±0.27</td><td>82.08±0.48</td><td>15.26±0.07</td></tr><tr><td>SGLD</td><td colspan="3">76.38±0.085% 73.83±0.013% 89.34±0.06%</td><td>73.39±0.26</td><td>80.30±0.14</td><td>16.36±0.07</td></tr><tr><td>PSGLD</td><td>77.46±0.05%</td><td>73.16±0.1%</td><td>89.78±0.08%</td><td>69.89±0.2</td><td>83.70±0.21</td><td>15.72±0.04</td></tr><tr><td>Adam</td><td>70.94±0.10%</td><td>69.73±0.11%</td><td>85.44±0.14%</td><td>86.77±1.01</td><td>87.60±0.44</td><td>22.35±0.47</td></tr><tr><td>SGD-M</td><td>68.06±0.27%</td><td>68.76±0.17%</td><td>84.86±0.48%</td><td>99.12±1.06</td><td>90.35±0.29</td><td>23.64±0.73</td></tr></table>

Test NLL learning curve (with zoom-in for sampling methods) and the best performance. Santa and Adam results are fromChen et al. (2016) convergences in similar speeds as the original one, but with less amount of over-fitting and better final test NLL 8.40.

Table content: <table><tr><td>Method</td><td colspan="3">Piano Muse Nott JSB</td></tr><tr><td colspan="2">NNSGHMC 7.66</td><td colspan="2">7.27 3.37 8.49</td></tr><tr><td>SGHMC</td><td>7.65</td><td colspan="2">7.33 3.35 8.40</td></tr><tr><td>PSGLD</td><td>7.67</td><td colspan="2">7.48 3.28 8.42</td></tr><tr><td>Santa</td><td>7.6</td><td>7.2</td><td>3.39 8.46</td></tr><tr><td>Adam</td><td>8</td><td>7.56</td><td>3.7 8.51</td></tr><tr><td>Figure 6:</td><td/><td/><td/></tr></table>

The basic statistics for 4 RNN datasets, bold figure represents large difference compared to others. Size is the number of data point. Avg. Time is the averaged sequence and Energy scale is the rough scale of the train NLL when sampler converges. and 0.0018, respectively. For optimization methods, we use learning rates 0.002 and 0.003 for Adam and SGD-M respectively.

Table content: <table><tr><td/><td colspan="3">Piano Muse Nott</td><td>JSB</td></tr><tr><td>Size:train</td><td>87</td><td>524</td><td>694</td><td>229</td></tr><tr><td>Size:test</td><td>25</td><td>124</td><td>170</td><td>77</td></tr><tr><td>Avg. Time:train</td><td>872</td><td>467</td><td>254</td><td>60</td></tr><tr><td>Avg. Time:test</td><td>761</td><td>518</td><td>261</td><td>61</td></tr><tr><td colspan="2">Energy scale:train ≈ 7.2</td><td>≈ 7</td><td colspan="2">≈ 2.5 2.5 2.5 ≈ 7.8</td></tr><tr><td>0.2</td><td/><td/><td/><td/></tr></table>

Peer review (iclr2019):
 In the paper "Meta-Learning for Stochastic Gradient MCMC", the authors present a meta-learning approach to automatically design MCMC sampler based on Hamiltonian dynamics to mix faster on problems similar to the training problems. The approach is evaluated on simple multidimensional Gaussians, and Bayesian neural networks (including fully connected, convolutional, and recurrent networks). MCMC samplers in general, and Hamiltonian Monte Carlo sampler in particular, are very powerful tools to perform Bayesian inference in high-dimensional spaces. Combined with stochastic gradients, methods like Stochastic Gradient MCMC (SGMCMC), or Stochastic Gradient Langevin Dynamics (SGLD) have been successfully used to apply these methods in the large data regime, where only noisy estimates of the gradients are feasible. Even though, many different samplers exists, and they are provably correct (meaning they converge to the correct distribution), fast mixing and low auto-correlation within the chain can heavily depend on the problem at hand and the hyperparameters of the sampler used. The work presented here, uses the general framework for SG-MCMC samplers of Ma et al., parametrizes it with a neural network and learns its weights on representative training problems. The paper is well written, although occasional minor mistakes and typos can be found. It seems however, that the method is still quite laborious and some care needs to be taken to train the meta-sampler. The overall narrative is easy to follow, but could benefit from more detail in certain parts. In general, I argue for acceptance of the paper, but have the following questions/comments: - Below Eq. (7), an interpretation of the parametrizations Q_f and D_f is given. I greatly appreciate this, but the phrase *Q_f is responsible for the acceleration of .theta* is not really instructive. By definition, the change in .theta is mostly driven by the momentum p. Therefor, Q_f looks like an inverse mass (at least in the second line of (7)), but maybe that is not a very helpful analogy either. - at the beginning of section 3.2, the term *particles* is used. While I am fully aware of what that is supposed to mean, a reader less familiar with the topic could be confused, because there is no explanation of it. - It is unclear to me how the stochastic estimate .tilde{U}(.theta) in equation (10) is computed exactly. Is it estimated using the current mini-batch at time t, or is it estimated using a *holdout-test set*? - I was wondering how the correlation between the chains due to thinning for the In-chain loss affects the results. The text, does not address this at all. - The experiments are very thorough and I appreciate the comparison to the tuned baselines, but I am missing some details in the paper: (a) Did you tune the SGHMC method in Figure 2, as well? It is not mentioned in the text, and the sample path looks very volatile, which could indicate a poor combination of step length and noise. (b) How was the tuning of the base line methods performed? (c) Are the results in Figure 3 based on single runs, or do you show the mean over 10 independent runs (as in table 1). - The insets in Figure 6 are helpful, but I think you could shrink the *outer y axis* and have the inset in the top right corner instead. That way, the zoomed-out plot would show more details on its own.