NETS: A Non-equilibrium Transport Sampler

We introduce the Non-Equilibrium Transport Sampler (NETS), an algorithm for sampling from unnormalized probability distributions. NETS builds on non-equilibrium sampling strategies that transport a simple base distribution into the target distribution in finite time, as pioneered in Neal's annealed importance sampling (AIS). In the continuous-time setting, this transport is accomplished by evolving walkers using Langevin dynamics with a time-dependent potential, while simultaneously evolving importance weights to debias their solutions following Jarzynski's equality. The key innovation of NETS is to add to the dynamics a learned drift term that offsets the need for these corrective weights by minimizing their variance through an objective that can be estimated without backpropagation and provably bounds the Kullback-Leibler divergence between the estimated and target distributions. NETS provides unbiased samples and features a tunable diffusion coefficient that can be adjusted after training to maximize the effective sample size. In experiments on standard benchmarks, high-dimensional Gaussian mixtures, and statistical lattice field theory models, NETS shows compelling performances.

This paper introduces a way to learn how to sample from a distribution with density $\rho_t$ in a statistically unbiased way, given access only to the potential $U_t$. It does so by learning a neural network that augments standard powerful Monte Carlo algorithms in a principled way that has intrinsic connections to the equations governing non-equilibrium thermodynamics. The method is contextualized and numerically compared to other recent approaches in the literature.