Source-linked AI summary
NeuTra-lizing Bad Geometry in Hamiltonian Monte Carlo Using Neural Transport
Matthew Hoffman, Pavel Sountsov, Joshua V. Dillon, Ian Langmore, Dustin Tran, Srinivas Vasudevan
TL;DR
Unfavorable posterior geometry can make HMC require many expensive target and gradient evaluations to mix. NeuTra HMC learns IAF transport maps using variational inference and runs HMC in warped space. The authors report consistent improvements over HMC across synthetic and real problems, often by an order of magnitude.
Problem
Unfavorable target geometry can require many target-log-probability and gradient evaluations for HMC to mix between faraway states.
Source: §2.1 · p. 2Method
NeuTra HMC uses inverse autoregressive flows parameterized by neural networks and fit using variational inference to build transport maps for HMC.
Source: §1 · p. 1Results
NeuTra HMC consistently outperforms HMC on synthetic and real problems, often by an order of magnitude.
Source: §1 · p. 1Takeaways & Limitations
Using IAFs instead of affine flows often dramatically improves mixing speed, especially on posteriors often found in hierarchical Bayesian models.
Source: §5 · p. 8If transport maps fail to capture target geometry adequately, NeuTra could slow mixing in the tails.
Source: §5 · p. 8
Abstract
from arXiv · showhide
Hamiltonian Monte Carlo is a powerful algorithm for sampling from difficult-to-normalize posterior distributions. However, when the geometry of the posterior is unfavorable, it may take many expensive evaluations of the target distribution and its gradient to converge and mix. We propose neural transport (NeuTra) HMC, a technique for learning to correct this sort of unfavorable geometry using inverse autoregressive flows (IAF), a powerful neural variational inference technique. The IAF is trained to minimize the KL divergence from an isotropic Gaussian to the warped posterior, and then HMC sampling is performed in the warped space. We evaluate NeuTra HMC on a variety of synthetic and real problems, and find that it significantly outperforms vanilla HMC both in time to reach the stationary distribution and asymptotic effective-sample-size rates.
1. Introduction
Unfavorable target geometry can make HMC require many costly evaluations to mix. NeuTra HMC uses IAF-based transport maps and HMC in warped space; the paper reports consistent gains over HMC and improvements over earlier transport-map approaches.
1. Introduction
- Unfavorable target geometry can require many log-probability and gradient evaluations before HMC mixes between faraway states.The problem is especially costly when successive chain samples are highly correlated.Source: §1 · p. 1
- NeuTra uses a transport map of inverse autoregressive flows parameterized by neural networks and fit using variational inference.The map warps the sampling space before HMC is run.Source: §1 · p. 1
- The approach improves on earlier transport-map MCMC with more powerful, scalable IAF maps and gradient-based HMC.Earlier polynomial-regression maps could be too expensive for high-dimensional problems.Source: §1 · p. 1
- The authors also adapt the strategy to train variational autoencoders.Source: §1 · p. 1
- NeuTra HMC consistently outperforms HMC across varied problems, often by an order of magnitude.The evaluation covers synthetic and real problems.Source: §1 · p. 1
2. Neural Transport MCMC
NeuTra MCMC combines Hamiltonian Monte Carlo with learned invertible flows to make sampling more effective when the target has unfavorable geometry. The method learns a transport map, samples in the transformed space, and can also use maps trained in an amortized setting.
2.1. Hamiltonian Monte Carlo
- HMC augments parameters with Gaussian momentum, simulates dynamics using reversible, volume-preserving leapfrog updates, and accepts or rejects the resulting proposal using Metropolis probabilities.Accurate discretization keeps energy changes small and proposals more likely to be accepted.Source: §2.1 · p. 2
- Unfavorable geometry can force HMC to use small steps in high-curvature directions and many steps in low-curvature directions, while heavy or light tails can slow mixing further.HMC is most efficient on roughly isotropic distributions with roughly Gaussian tails.Source: §2.1 · p. 2
2.2. Normalizing Flows and Variational Inference
- Normalizing flows use invertible transformations to turn a simple base distribution into a more expressive approximation, with variational training minimizing KL divergence to the target.The ELBO can be estimated using samples, the Jacobian determinant, and the unnormalized target density.Source: §2.2 · p. 2
- IAFs parameterize shifted and scaled coordinates with neural networks of earlier inputs, enabling parallel computation and a lower-triangular Jacobian.Stacking flows increases map expressiveness while preserving an overall Jacobian formed from the component Jacobians.Source: §2.2 · pp. 2–3
2.3. Neural-Transport MCMC
- NeuTra combines HMC with IAF maps to scale transport-map MCMC to higher-dimensional problems, where polynomial maps may require restrictive independence assumptions.The authors identify HMC’s gradient information and IAFs’ scalability and flexibility as improvements over earlier transport-map MCMC.Source: §2.3 · p. 3
- NeuTra fits an IAF transport map by variational KL minimization, runs HMC on the Jacobian-adjusted target in z-space, then maps samples forward to θ-space.This uses the variational approximation to accelerate MCMC while allowing MCMC to correct its mismatch with the target.Source: §2.3 · p. 3
2.3.1. NeuTra in the amortized setting
- NeuTra also supports amortized inference: a conditional map can take auxiliary inputs x, allowing the approach to work with maps trained across multiple inputs.The cost of fitting a conditional variational distribution can be amortized over many values of x.Source: §2.3.1 · p. 3
3. Related Work
The related work situates NeuTra alongside geometric, adaptive, and neural approaches to improving MCMC. It contrasts NeuTra’s learned transport maps with RMHMC’s metric-based adaptation and other learned-kernel or variational-inference hybrids.
3.1. Riemannian Manifold HMC
- RMHMC can make rapid progress with an appropriate position-dependent metric, but its implicit integrator and metric calculations can be expensive and model-specific, especially in high dimensions.Rapidly changing metrics may still require small step sizes.Source: §3.1 · p. 4
- NeuTra’s continuous-time dynamics are equivalent to RMHMC with a metric defined by the transport-map Jacobian, potentially retaining local adaptation with lower implementation and computational complexity.The potential benefit depends on learning a good explicit map; Figure 1 illustrates locally adaptive effective step sizes in the tails.Source: §3.1 · p. 4, Figure 1 · p. 4
3.2. Learned MCMC Kernels
- Other learned MCMC methods tune kernels or integrators, while NeuTra uses variational approximations to speed MCMC; the cited alternatives may sacrifice symplectic stability or lack strong convergence guarantees.L2HMC uses a neural-network-parameterized generalization of leapfrog, while Song et al. use adversarial training.Source: §3.2 · pp. 4–5
3.3. Deep Generative Models
- Related deep-generative-model work includes MCMC-based variational bounds for VAE training, HMC step-size tuning from mean-field variational scales, and a transport-map objective whose mixing effect is unclear.Titsias’s objective maximizes the last chain sample’s log density while ignoring an intractable entropy term.Source: §3.3 · p. 5
4. Experiments
Experiments cover synthetic Gaussian and funnel targets, German-credit logistic regression, and an MNIST variational autoencoder. NeuTra HMC often improves mixing and reaches low bias quickly, while results vary with the target and the cost of training the transport map.
4.1.1. Transport maps and training procedure
- Experiments use two synthetic targets, German-credit sparse logistic regression, and an MNIST variational autoencoder, with IAF and diagonal and triangular affine transport maps.The synthetic targets include a 100-dimensional Gaussian whose eigenvalues span six orders of magnitude and a 100-dimensional funnel; transport maps are trained for 5000 steps.Source: §4 · p. 5, §4.1 · p. 5, §4.1.1 · p. 5
4.1.2. HMC sampler hyperparameters
- Sampler tuning uses Bayesian optimization to balance low potential scale reduction and high ESS per gradient evaluation, using the worst componentwise values.All tested chains reached R̂ < 1.1; because chains start from an underdispersed initial distribution, the reported R̂ values are lower bounds.Source: §4.1.2 · p. 6
4.1.3. Results
- On the funnel, NeuTra reaches the bias estimator’s noise floor quickly, while non-neural maps make little progress and their chains mix and warm up slowly.For the Gaussian, the diagonal preconditioner converges slowly, but its low computational cost lets it overtake NeuTra before IAF training finishes; the triangular preconditioner performs best.Source: §4.1.3 · p. 7
- For sparse logistic regression, the target is sufficiently well behaved that non-neural transport maps also reach the bias estimator’s noise floor.When bias requirements are not stringent, the authors note that a simpler preconditioner—or, for some problems, variational samples alone—may be preferable.Source: §4.1.3 · pp. 7–8
- In all cases except the Gaussian, NeuTra significantly outperforms non-neural transport maps in asymptotic ESS rates, often by over an order of magnitude.Figure 4 reports componentwise chain statistics after discarding the first 500 of 1000 samples; the authors attribute the advantage to simpler geometry and fewer costly leapfrog steps.Source: Figure 4 · p. 7, §4.1.3 · p. 8
4.2. Conditional Target Distributions
- For dynamically binarized MNIST, NeuTra HMC refines the approximate posterior using the existing IAF map without adding model parameters.The procedure initializes a short NeuTra HMC chain from the approximate posterior for each minibatch; additional model evaluations at each leapfrog step significantly reduce training speed.Source: §4.2 · p. 8, Table 1 · p. 8
5. Discussion
NeuTra HMC warps target geometry with IAFs trained using variational inference; IAFs often improve mixing over affine flows, especially for hierarchical Bayesian posteriors. If the maps miss target geometry, mixing in the tails could slow.
5. Discussion
- NeuTra HMC accelerates HMC by nonlinearly warping target geometry with inverse autoregressive flows trained using variational inference.Source: §5 · p. 8
- IAFs instead of affine flows often dramatically improve mixing speed, especially on posteriors often found in hierarchical Bayesian models.Source: §5 · p. 8
- If maps fail to capture target geometry adequately, NeuTra could slow mixing in the tails; architectures and regularization strategies are proposed for exploration.Source: §5 · p. 8