Title: Generalizing Hamiltonian Monte Carlo with Neural Networks
Authors: Daniel Levy
Matt D. Hoffman
Jascha Sohl-Dickstein
Venue: ICLR 2018
Keywords: markov
chain
monte
carlo
sampling
posterior
deep
learning
hamiltonian
mcmc

Abstract: We present a general-purpose method to train Markov chain Monte Carlo kernels, parameterized by deep neural networks, that converge and mix quickly to their target distribution. Our method generalizes Hamiltonian Monte Carlo and is trained to maximize expected squared jumped distance, a proxy for mixing speed. We demonstrate large empirical gains on a collection of simple but challenging distributions, for instance achieving a 106x improvement in effective sample size in one case, and mixing when standard HMC makes no measurable progress in a second. Finally, we show quantitative and qualitative gains on a real-world task: latent-variable generative modeling. Python source code will be open-sourced with the camera-ready paper.

TL;DR: General method to train expressive MCMC kernels parameterized with deep neural networks. Given a target distribution p, our method provides a fast-mixing sampler, able to efficiently explore the state space.

Review discussion:
In this work, the authors propose a procedure for tuning the parameters of an HMC algorithm (I guess, if I have understood correctly).
I think this paper has a good and strong point: this work points out the difficulties in choosing properly the parameters in a HMC method (such as the step and the number of iterations in the leapfrog integrator, for instance). In the literature, specially in machine learning, there is ``fever’’ about HMC, in my opinion, partially unjustified.
If I have understood, your method is an adaptive HMC algorithm  where the parameters are updated online; or is the training  done in advance? Please, remark and clarify this point.
However, I have other additional comments:
- Eqs. (4) and (5) are quite complicated; I think a running toy example can help the interested reader.
- I suggest to compare the proposed method to other efficient methods that do not use the gradient information (in some cases as multimodal posteriors, the use of the gradient information can be counter-productive for sampling purposes), such as Multiple Try Metropolis (MTM) schemes
L. Martino, J. Read, On the flexibility of the design of Multiple Try Metropolis schemes, Computational Statistics, Volume 28, Issue 6, Pages: 2797-2823, 2013,
adaptive techniques,
H. Haario, E. Saksman, and J. Tamminen. An adaptive Metropolis algorithm. Bernoulli, 7(2):223–242, April 2001,
and component-wise strategies as Gibbs Sampling,
W. R. Gilks and P. Wild, Adaptive rejection sampling for Gibbs sampling, Appl. Statist., vol. 41, no. 2, pp. 337–348, 199.
At least, add a brief paragraph in the introduction citing and discussing this possible alternatives.
Thank you for your review and the pointer to references.
We wish to emphasize that our method is able, but not limited to, automatically tuning HMC parameters (which systems like Stan already have well-tested heuristics for). Our approach generalizes HMC, and is capable of learning proposal distributions that do not correspond to any tuned HMC proposal (but which can still be plugged into the Metropolis-Hastings algorithm to generate a valid MCMC algorithm).
Indeed, in our experiments, we find that our approach significantly outperforms well-tuned HMC kernels.
The training is done during the burn-in phase, and the trained sampler is then frozen. This is a common approach to adapting transition-kernel hyperparameters in the MCMC literature.
Regarding the references, we added those in the text. We also want to emphasize that all of these are complementary to and could be combined with our method.
For example, we could incorporate the intuition behind MTM by having several parametric operators and training each one when used.
Additionally, in the process of revisiting our experiments to compare against LAHMC, we empirically found that weighting the second term of our loss (the ‘burn-in’ term) could lead to even more improved auto-correlation and ESS on the diagnostic distributions. We therefore updated the paper and report the results obtained with slightly tuning that parameter (setting it to 0 or 1).
The paper proposed a generalized HMC by modifying the leapfrog integrator using neural networks to make the sampler to converge and mix quickly. Mixing is one of the most challenge problems for a MCMC sampler, particularly when there are many modes in a distribution. The derivations look correct to me. In the experiments, the proposed algorithm was compared to other methods, e.g., A-NICE-MC and HMC. It showed that the proposed method could mix between the modes in the posterior.
Although the method could mix well when applied to those particular experiments, it lacks theoretical justifications why the method could mix well.
Thank you very much for your review and comments. Guaranteeing mixing between modes is a fundamental (#P-Hard) problem.
As such, we do not hope to solve it in the general case.
Rather, we propose a method to greatly increase the flexibility and adaptability of a class of samplers which is already state of the art in many contexts. The relation between mixing time and expected square jump distance is thoroughly treated in [Pasarica & Gelman, 2010], and is the theoretical inspiration for our choice of training loss.
We further emphasize that, barring optimization issues, our method should always fare at least as well as HMC in terms of mixing.
Thank you once again, we have updated the text to more clearly discuss why our approach might be expected to lead to better mixing.
Additionally, in the process of revisiting our experiments to compare against LAHMC, we empirically found that weighting the second term of our loss (the ‘burn-in’ term) could lead to even more improved auto-correlation and ESS on the diagnostic distributions. We therefore updated the paper and report the results obtained with slightly tuning that parameter (setting it to 0 or 1).
The paper introduces a non-volume-preserving generalization of HMC whose transitions are determined by a set of neural network functions. These functions are trained to maximize expected squared jump distance.
This works because each variable (of the state space) is modified in turn, so that the resulting update is invertible, with a tractable transformation inspired by Dinh et al 2016.
Overall, I believe this paper is of good quality, clearly and carefully written, and potentially accelerates mixing in a state-of-the-art MCMC method, HMC, in many practical cases. A few downsides are commented on below.
The experimental section proves the usefulness of the method on a range of relevant test cases;
in addition, an application to a latent variable model is provided sec5.2.
Fig 1a presents results in terms of numbers of gradient evaluations, but I couldn't find much in the way of computational cost of L2HMC in the paper. I can't see where the number "124x" in sec 5.1 stems from. As a user, I would be interested in the typical computational cost of both "MCMC sampler training" and MCMC sampler usage (inference?), compared to competing methods. This is admittedly hard to quantify objectively, but just an order of magnitude would be helpful for orientation.
Would it be relevant, in sec5.1, to compare to other methods than just HMC, eg LAHMC?
I am missing an intuition for several things: eq7, the time encoding defined in Appendix C
Appendix Fig5, I cannot quite see how the caption claim is supported by the figure (just hardly for VAE, but not for HMC).
The number "124x ESS" in sec5.1 seems at odds with the number in the abstract, "50x".
# Minor errors
- sec1: "The sampler is trained to minimize a variation": should be maximize
"as well as on a the real-world"
- sec3.2 "and 1/2 v^T v the kinetic": "energy" missing
- sec4: the acronym L2HMC is not expanded anywhere in the paper
The sentence "We will denote the complete augmented...p(d)" might be moved to after "from a uniform distribution" in the same paragraph.
In paragraph starting "We now update x":
- specify for clarity: "the first update, which yields x' "/ "the second update, which yields x''  "
- "only affects $x_{\bar{m}^t}$": should be $x'_{\bar{m}^t}$  (prime missing)
- the syntax using subscript m^t is confusing to read; wouldn't it be clearer to write this as a function, eg "mask(x',m^t)"?
- inside zeta_2 and zeta_3, do you not mean $m^t" and $\bar{m}^t$ ?
- sec5: add reference for first mention of "A NICE MC"
- Appendix A:
- "Let's" -> "Let"
- eq12 should be x''=...
- Appendix C: space missing after "Section 5.1"
- Appendix D1: "In this section is presented" : sounds odd
- Appendix D3: presumably this should consist of the figure 5 ? Maybe specify.
We first and foremost want to thank you for your time and extremely valuable comments. We have uploaded a new version of the paper based on the feedback, and have addressed specific points below.
Clarification about 50x vs 124x:
We decided against advertising the 124x number as it is misleading considering that HMC completely failed on this task; the correct ratio was too large for us to experimentally measure.
As such, we reported the one for the Strongly-Correlated Gaussian. We clarified this in the text and detail that L2HMC can succeed when HMC fails.
Intuition on Eq 7.:
We define this reciprocal loss to encourage mixing across the entire state space. The second term corresponds exactly to Expected Square Jump Distance, which we want to maximize as a proxy for mixing. The first term discourages a particle from not-moving at all in a region of state space -- if d(x, x’) = 0, the first term would be infinite. We clarified that part in the text.
Time encoding:
Our operator L_\theta consists of the composition of M augmented leapfrog steps. For each of those leapfrog, the timestep t is provided as input to the networks Q, S and T.
Instead of providing it as a single scalar value, we provide it as a 2-d vector [cos(2 * pi * t / M), sin(2 * pi * t / M)].
Regarding samples in Fig5:
Sample quality and sharpness are inherently hard things to evaluate. Our observation was that many digits generated by L2HMC-DGLM look very sharp (Line 1 Column 2, Line 2 Column 8, Line 5 Column 2, Line 7 Columns 3 and 7…).
However, we will weaken the claim in the caption.
Comparison with LAHMC:
We compared our method to LAHMC on the evaluated energy functions. L2HMC significantly outperforms LAHMC on all tasks, for the same number of gradient evaluations. LAHMC is also unable to mix between modes in the MoG case. Results are reported in Appendix C.1.
We also note that L2HMC could be easily combined with LAHMC, by replacing the leapfrog integrator of LAHMC with the learned one of L2HMC.
In the process of revisiting our experiments to compare against LAHMC, we empirically found that weighting the second term of our loss (the ‘burn-in’ term) could lead to even more improved auto-correlation and ESS on the diagnostic distributions. We therefore updated the paper and report the results obtained with slightly tuning that parameter (setting it to 0 or 1).
Question about computation:
For the 2d-SCG case, on CPU, the training of the sampler took 160 seconds. The L2HMC overhead for sampling, with a batch-size of 200, was about 36%. This is negligible compared to an 106x improved ESS.  We also should note that for the latent generative model case, we train the sampler online with the same computations used to train everything else; in that case L2HMC and HMC perform the exact same number of gradient evaluation of the energy and thus requires no training budget.
Thank you once again for your valuable feedback, we hope this helps answer your questions!
We thank the reviewers for their valuable time and comments.
We updated the paper with the following modifications:
- Clarified some points and fixed typos pointed out by the reviewers.
- Added a ``Future Work section as well as additional relevant references.
- Added a comparison with Look-Ahead HMC (LAHMC; Sohl-Dickstein et al. 2014) in the Appendix.
Additionally, in the process of revisiting our experiments to compare against LAHMC, we empirically found that weighting the second term of our loss (the ‘burn-in’ term) could lead to even more improved auto-correlation and ESS on the diagnostic distributions. We therefore updated the paper and report the results obtained with slightly tuning that parameter (setting it to 0 or 1).

Meta-review:
role: PC
data: comment: This paper presents a learned inference architecture which generalizes HMC. It defines a parameterized family of MCMC transition operators which share the volume preserving structure of HMC updates, which allows the acceptance ratio to be computed efficiently. Experiments show that the learned operators are able to mix significantly faster on some simple toy examples, and evidence is presented that it can improve posterior inference for a deep latent variable model. This paper has not quite demonstrated usefulness of the method, but it is still a good proof of concept for adaptive extensions of HMC.



Decision: Accept (Poster)