META-LEARNING FOR STOCHASTIC GRADIENT MCMC
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.
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.