Source-linked AI summary
Generalizing Hamiltonian Monte Carlo with Neural Networks
Daniel Levy, Matthew D. Hoffman, Jascha Sohl-Dickstein
TL;DR
Sampling high-dimensional targets known only up to a normalizing constant is difficult. This paper trains neural-network MCMC kernels that generalize HMC using an expected-squared-jumped-distance objective. It reports gains on challenging distributions and a latent-variable generative-modeling task.
Problem
Sampling high-dimensional distributions known only up to a normalizing constant is generally difficult, despite its importance for learning and inference.
Source: §1 · p. 1Method
The method trains expressive neural-network MCMC kernels that generalize HMC, retaining a tractable Metropolis-Hastings correction and optimizing expected squared jumped distance.
Source: §1 · p. 2, §4.2 · p. 5Results
The authors report significant gains over HMC where it performs poorly, plus improved log-likelihood and greater posterior-sample complexity on a latent-variable generative model.
Source: §1 · p. 2Takeaways & Limitations
For analytically specified targets, the method is presented as a fast-mixing sampler intended for black-box use in sampling-intensive domains.
Source: §7 · p. 10Proposal updates must remain invertible with tractable Jacobians, and each sub-update must avoid nonlinear dependence on the variables it updates.
Source: §4.1 · p. 4
Abstract
from arXiv · showhide
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. We release an open source TensorFlow implementation of the algorithm.
1 INTRODUCTION
Sampling high-dimensional distributions known only up to a normalizing constant is difficult, and standard MCMC performance depends strongly on the proposal. The paper introduces a trainable neural-network generalization of HMC and reports empirical gains on difficult distributions and a latent-variable generative-modeling task.
1 INTRODUCTION
- High-dimensional distributions known only up to a normalizing constant arise in several fields, and sampling from them is critical for learning and inference.The introduction names protein folding, physics simulations, and machine learning as examples.Source: §1 · p. 1
- MCMC convergence and mixing depend strongly on choosing an effective proposal, while detailed balance can encourage random-walk behavior and slow exploration.It can also be difficult in practice to identify burn-in and determine when samples are uncorrelated.Source: §1 · p. 1
- HMC can struggle with multimodal targets, ill-conditioned energy landscapes, rapidly changing gradients, and movement across energy levels.Its volume-preserving dynamics make transitions across energy levels difficult, and crossing low-density regions between modes is unlikely.Source: §1 · p. 1
- The method generalizes HMC with parametric functions, retains a tractable Metropolis-Hastings correction, and is trained to improve mixing.Its parameterization can reduce to standard HMC and emulate within-trajectory tempering and diagonal mass matrices.Source: §1 · p. 2
- The paper evaluates its method on distributions where HMC struggles and on training latent-variable generative models.The introduction frames the approach as returning an exact sampler with convergence and mixing guarantees through a tractable Metropolis-Hastings step.Source: §1 · p. 2
- The authors report significant empirical gains on distributions where HMC performs poorly.They also report improved model log-likelihood and greater complexity in posterior samples for a latent-variable generative model.Source: §1 · p. 2
2 RELATED WORK
Prior work has adapted proposals, applied machine-learning models to MCMC, and used learned transformations as proposals. The paper positions its approach as combining expressive transformations with gradient information and an explicit mixing objective.
2 RELATED WORK
- Pasarica and Gelman (2010) motivate tuning proposal hyperparameters by maximizing expected squared jumped distance as a principled, practical approach.Other HMC adaptation has tuned step size or leapfrog-step count, while Salimans et al. (2015) proposed a variational-inference-based scheme.Source: §2 · p. 2
- Machine-learning methods have been used to learn MCMC proposals or energy gradients, while Boltzmann machines have supported energy approximations and hand-designed proposals.Examples include kernel methods, Restricted and semi-Restricted Boltzmann machines, and physics applications.Source: §2 · p. 2
- Song et al. (2017) use adversarially trained volume-preserving transformations as MCMC proposals, but the approach lacks gradient information and relies on potentially unstable adversarial training.The passage also notes possible mode collapse and difficulty mixing across energy levels.Source: §2 · p. 2
- RNVP uses invertible layers with tractable Jacobians and motivates the paper’s extension of the HMC leapfrog integrator.Each layer modifies only a subset of variables using a function of the remaining variables.Source: §2 · p. 2
3 BACKGROUND
The background introduces MCMC and Metropolis-Hastings before describing HMC as an augmented-state proposal based on approximately integrated Hamiltonian dynamics. It explains the properties that support valid acceptance steps and long-distance moves.
3.1 MCMC METHODS AND METROPOLIS-HASTINGS
- MCMC constructs a Markov chain whose stationary distribution is the target; detailed balance is a sufficient condition for stationarity.The chain must also be irreducible and aperiodic, and detailed balance equates probability flow between states in both directions.Source: §3.1 · p. 3
- Metropolis-Hastings turns a proposal into a detailed-balance transition through an accept/reject rule, but cautious proposals or low acceptance can slow mixing.For continuous spaces, HMC addresses this by proposing long moves that stay roughly on iso-probability contours.Source: §3.1 · p. 3
3.2 HAMILTONIAN MONTE CARLO
- HMC augments the state with independent momentum and approximately integrates Hamiltonian dynamics, moving along approximate iso-probability contours.The target is represented through potential energy, while momentum contributes kinetic energy.Source: §3.2 · p. 3
- Leapfrog is the integrator typically used to simulate the Hamiltonian dynamics in HMC.A single time step consists of the described leapfrog updates.Source: §3.2 · p. 3
- Leapfrog integration and a momentum flip form an involutive transformation, while the leapfrog sub-updates are volume-preserving shears.These properties simplify the Metropolis-Hastings-Green acceptance probability.Source: §3.2 · p. 3
4 L2HMC: TRAINING MCMC SAMPLERS
L2HMC learns a parametric leapfrog operator from an energy function, generalizing HMC while retaining a tractable Metropolis-Hastings accept/reject step. Its design and training target fast mixing, fast burn-in, and movement across energy levels and modes.
4.1 AUGMENTING HMC
- HMC can traverse long distances but may mix slowly across energy levels, struggle with ill-conditioned covariances, and fail to move between modes without additional tricks.These limitations motivate desiderata including fast mixing and burn-in, energy-level exploration, and mode mixing.Source: §4.1 · p. 4
4.1.2 UPDATE STEPS
- The learned updates preserve invertibility and tractable Jacobian determinants while allowing non-volume-preserving transformations; momentum scaling may aid movement through low-density zones, and gradient scaling may improve conditioning.The functions Q, S, and T use shared-weight multilayer perceptrons in the experiments.Source: §4.1 · p. 4, §4.1.2 · p. 5
4.1.1 STATE SPACE
- The augmented state combines position x, standard-normal momentum v, and a uniformly drawn binary direction d; fixed random masks divide coordinates into complementary subsets for each step.Each mask has half its entries set to 0 and half to 1.Source: §4.1.1 · p. 4
4.1.2 UPDATE STEPS
- Each modified leapfrog step updates momentum using position and the energy gradient, updates complementary masked position subsets, then updates momentum again.Learned translations and rescalings modify the updates; setting the relevant functions to zero recovers the standard leapfrog updates.Source: §4.1.2 · pp. 4–5
4.1.3 MCMC TRANSITIONS
- MCMC transitions apply the learned leapfrog operator with its Metropolis-Hastings acceptance probability, then resample momentum and direction; setting Q, S, and T to zero reduces the parameterization to standard HMC.Each transition is stated to satisfy detailed balance with respect to the target distribution.Source: §4.1.3 · p. 5
4.2 LOSS AND TRAINING PROCEDURE
- Minimizing lag-one autocorrelation is equivalent to maximizing expected squared jumped distance, but that objective alone can leave some regions with little movement.A reciprocal loss term is added to optimize typical and worst-case behavior, strongly penalizing states where the sampler cannot move effectively.Source: §4.2 · p. 5
- The loss trains for both mixing on target-distribution draws and fast burn-in from initialization draws, with λb controlling the burn-in regularization strength.The scale parameter λ captures the problem's characteristic length scale.Source: §4.2 · pp. 5–6
- Each training iteration requires one network pass and can be batched; the procedure applies to learnable operators with tractable Jacobian determinants.The algorithm takes an energy function, its gradient, and an initial distribution over the augmented state space as inputs.Source: §4.2 · p. 6
5 EXPERIMENTS
The experiments test L2HMC on challenging energy landscapes and latent-variable generative modeling. They report faster mixing, higher data likelihood, and richer posterior behavior than the compared approaches.
5.1 VARIED COLLECTION OF ENERGY FUNCTIONS
- L2HMC learns a diagonal inertia tensor for an ill-conditioned Gaussian and adapts to a rotated, non-axis-aligned Gaussian despite using element-wise transformations.The rotated Gaussian has variances [10^2, 10^-2] and is rotated by π/4.Source: §5.1 · p. 6
- L2HMC mixes faster than well-tuned HMC and A-NICE-MC on toy distributions, achieves more than 106× improved ESS on SCG, and mixes across MoG modes where HMC gets stuck.The MoG comparison uses a low-density region between modes.Source: Figure 1 · p. 7, §5.1 · p. 7
- On the two-mode Gaussian diagnostic, A-NICE-MC struggles to mix because its volume-preserving transformation maps only a fraction of the large mode to the small one, while L2HMC traverses between modes.The clusters are closer than in the cited mog2 setup, making this a diagnostic of volume-preserving transformations.Source: §5.1 · p. 7
5.2 LATENT-VARIABLE GENERATIVE MODEL
- For latent-variable modeling, L2HMC replaces HMC in the posterior-sampling setup and is jointly trained with the generative model; experiments use dynamically binarized MNIST.The jointly trained sampler performs the same number of energy-function gradient computations as HMC.Source: §5.2 · p. 8
5.2.1 SAMPLE QUALITY AND DATA LIKELIHOOD
- The decoder trained with L2HMC generates sharper samples than the compared HMC and vanilla-VAE methods.The paper cautions that higher log-likelihood does not necessarily correspond to better samples.Source: §5.2.1 · p. 8
- At a similar number of gradient computations, L2HMC-DLGM achieves higher training and held-out data log-likelihood than HMC-DLGM.The paper estimates log-likelihood using annealed importance sampling.Source: §5.2.1 · p. 8, Figure 2 · p. 8
5.2.2 INCREASED EXPRESSIVITY OF THE POSTERIOR
- L2HMC switches among posterior modes where the approximate posterior remains on one reconstructed digit, and its posterior samples show a non-Gaussian latent-space shape.The visualization uses 512 parallel chains, each run for 20,000 Metropolis-Hastings steps, and projects samples along the highest-variance direction.Source: §5.2.2 · p. 9, Figure 3 · p. 9
6 FUTURE WORK
Future work could extend the training loss to target additional autocorrelations and sample statistics, and explore combining the method with other MCMC approaches.
6 FUTURE WORK
- The loss targets lag-one autocorrelation; proposed extensions target higher lags and faster autocorrelation decay in energy, position, or learned sample statistics.Learned statistics could combine aspects of this method with the adversarial approach of Song et al. (2017).Source: §6 · p. 9
- The authors propose exploring combinations with minibatch HMC, shadow Hamiltonians, gradient preconditioning, alternative accept-reject rules, non-canonical dynamics, AIS variants, discrete-state extensions, and alternative integrators.These are proposed directions, not evaluated combinations.Source: §6 · p. 9
- They also suggest combining the approach with methods that do not use gradient information, including multiple-try, adaptive, and component-wise strategies.The multiple-try idea would use several parametric operators, training each when used.Source: §6 · p. 9
7 CONCLUSION
The conclusion presents a general method for training expressive neural-network MCMC kernels, while the appendices describe reverse updates and Jacobian computation for the learned operator.
7 CONCLUSION
- Given a target known up to a constant, the method aims to provide a fast-mixing sampler that can explore state space efficiently.The authors hope it can be used where sampling is a bottleneck, including protein folding and physics simulations.Source: §7 · p. 10
A REVERSE LEAPFROG OPERATOR
- A reverse leapfrog step inverts the corresponding forward equations, with steps applied from M down to 1 to invert Lθ.The appendix introduces the updates for a single time step with direction d = −1.Source: A REVERSE LEAPFROG OPERATOR · p. 13
B DETERMINANT OF THE JACOBIAN
- For the forward operator, the appendix states that its Jacobian can be computed; i labels an intermediate ζ value at step t, and d denotes direction.The displayed Jacobian expression is given in the passage, but its terms are not sufficiently clear to summarize further.Source: B DETERMINANT OF THE JACOBIAN · p. 13
C.1 IMPLEMENTATION DETAILS
The implementation uses neural networks for the sampler, specifies ESS estimation and training details, and reports that L2HMC outperforms LAHMC across the evaluated tasks.
C.1 IMPLEMENTATION DETAILS
- The Q, S, and T networks use two ReLU hidden layers with 10 units, or 100 for 50-d ICG; MoG training anneals temperature but evaluation does not use it.The networks are trained with Adam at learning rate 10^-3 for 5,000 iterations with batch size 200.Source: §C.1 · p. 13
C.2 AUTO-CORRELATION AND ESS
- ESS estimation truncates the autocorrelation sum when autocorrelation falls below 0.05.The section introduces definitions of autocorrelation and ESS, but their displayed formulas are not included in the supplied passages.Source: §C.2 · p. 14
C.3 COMPARISON WITH LAHMC
- L2HMC largely outperforms LAHMC on all tasks, while LAHMC cannot mix between modes on the MoG task.Table 1 reports ESS for a fixed number of gradient evaluations.Source: §C.3 · p. 14, Table 1 · p. 14
D.1 TRAINING ALGORITHM
- For the training algorithm, Kθ(·|x) denotes the next-state distribution after sampling momentum and direction and applying the Metropolis-Hastings step.Source: §D.1 · p. 14
D.2 IMPLEMENTATION DETAILS OF L2HMC-DGLM
The latent-variable implementation describes sampler and model initialization, network architecture, and the AIS evaluation setup; the associated figure caption reports sharper mean activations from the L2HMC-DGLM decoder.
D.2 IMPLEMENTATION DETAILS OF L2HMC-DGLM
- The sampler, decoder, and approximate posterior are initialized, then updated using minibatches and separate decoder, ELBO, and sampler objectives.The latent energy is defined from the negative log likelihood and negative log prior.Source: §D.2 · p. 15
- The L2HMC-DGLM decoder is described as producing sharper mean activations.Source: Figure 5 · p. 15
- The x and v networks use two hidden layers of 200 ReLU units, and the sampler’s loss scale λ is the approximate posterior’s standard deviation.The auxiliary image is first passed through a two-layer softplus network with 512 hidden units.Source: §D.2 · pp. 14–15
- AIS evaluation runs 20 chains per data point for 10,000 annealing steps with 10 leapfrogs per step, using a step size chosen for 0.65 acceptance.Source: §D.2 · p. 15
D.3 MNIST SAMPLES
The section presents generated samples from VAE, HMC-DGLM, and L2HMC-DGLM.
D.3 MNIST SAMPLES
- Samples from VAE, HMC-DGLM, and L2HMC-DGLM are shown for comparison.Source: §D.3 · p. 15