Neural Non-Equilibrium Hamiltonian Monte Carlo for Corrected Boltzmann Sampling
Machine Learning
Summary
The authors present Neural Non-Equilibrium Hamiltonian Monte Carlo (NHMC), a method to efficiently sample from complex probability distributions by learning paths that move through probability space. NHMC first trains a model to propose paths between an easy starting distribution and the target distribution, then uses statistical corrections based on comparing forward and reverse paths to ensure accurate sampling. During evaluation, these corrections help estimate important quantities like normalizing constants and allow for acceptance testing in a Monte Carlo algorithm. The method works well when forward and reverse paths overlap sufficiently, but struggles otherwise, as shown in tests on physics-inspired models and molecular structures.
Authors
Moxian Qian
Abstract
Sampling from an unnormalized Boltzmann density requires proposals that move probability mass globally while retaining enough path-probability information for statistical correction. We introduce Neural Non-Equilibrium Hamiltonian Monte Carlo (NHMC), a train-then-correct learned Hamiltonian sampler. Starting from a tractable base distribution, NHMC learns stochastic Hamiltonian-style paths toward the target. Once training is complete, the learned proposal parameters are fixed; the proposal then generates complete paths and endpoint configurations, which are statistically corrected using the recorded non-equilibrium work. This dimensionless generalized work is determined by the probability ratio between the forward proposal path and a reverse reference path. During training, minimizing its mean reduces a path-space KL divergence and controls an upper bound on endpoint mismatch. During evaluation, the same quantity defines weights for self-normalized importance sampling on paths (path-SNIS), estimates normalizing constants or free-energy differences, and gives the acceptance ratio for path-space independent Metropolis-Hastings (path-IMH). The same forward-reverse laws also define a shared-bridge round-trip Metropolis kernel that acts directly on configurations and preserves the Boltzmann target. On double-well and finite-volume lattice $φ^4$ targets, the NHMC construction gives corrected estimates when path overlap is sufficient; when overlap is poor, weight degeneracy, low acceptance, and long autocorrelation expose proposal failure. We additionally report a molecular internal-coordinate feasibility study using an MD prior and learned-force path proposal.