[Paper Review] NeuTra-lizing Bad Geometry in Hamiltonian Monte Carlo Using Neural Transport

Matthew D. Hoffman, Pavel Sountsov+4 more|arXiv (Cornell University)|Mar 9, 2019

Markov Chains and Monte Carlo Methods28 references65 citations

TL;DR

NeuTra HMC learns a neural transport map to warp the target distribution, enabling faster Hamiltonian Monte Carlo by improving geometry and mixing; it often outperforms vanilla HMC in wall-clock time and effective-sample-size.

ABSTRACT

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.

Motivation & Objective

• Address unfavorable posterior geometry that slows HMC convergence and mixing.
• Develop a scalable transport-map approach using inverse autoregressive flows (IAFs) to warp the target space.
• Leverage variational inference to train the transport map that makes HMC more efficient.
• Evaluate NeuTra HMC on synthetic and real-world problems and compare to vanilla HMC.
• Explore applicability in amortized variational inference settings.

Proposed method

• Fit an inverse autoregressive flow (IAF) to minimize KL(q(θ) || p(θ)) via variational inference, defining a warp θ = fφ(z).
• Run Hamiltonian Monte Carlo in the warped z-space with target p(z) ∝ p(θ = f(z)) |∂f/∂z|, starting from samples drawn from q(z).
• Push z-space samples forward through f to obtain θ-space samples from p(θ).
• Use IAFs stacked in layers to achieve flexible, scalable transport maps; compute Jacobian determinants efficiently.
• Optionally employ amortized variational inference by conditioning the map on auxiliary inputs.
• Automate HMC hyperparameter tuning via Bayesian optimization to optimize mixing and convergence diagnostics.

Experimental results

Research questions

• RQ1 Does NeuTra HMC improve mixing and convergence speed over vanilla HMC on ill-conditioned targets?
• RQ2 How does a neural transport map (IAF-based) affect effective sample size and wall-clock efficiency of HMC?
• RQ3 Do neural transport maps help in both unconditional and conditional (amortized) posterior settings?
• RQ4 What are the trade-offs between transport-map training time and warmup/mixing benefits for NeuTra HMC?

Key findings

• NeuTra HMC significantly outperforms vanilla HMC in time to reach stationarity and in asymptotic ESS rates on evaluated problems.
• IAF-based transport maps better approximate target geometry than diagonal or tril matrices, improving exploration of challenging regions (e.g., funnel neck).
• In conditional models like DLGM and MNIST experiments, NeuTra HMC refined amortized posteriors, improving test NLL after integration with NeuTra steps.
• NeuTra introduces RMHMC-like adaptive behavior through the Jacobian of the transport map, achieving improved local geometry without full RMHMC complexity.
• Training time for the transport map can be offset by faster subsequent mixing, though trade-offs depend on problem geometry.