NETS: A Non-Equilibrium Transport Sampler (2410.02711v3)

Abstract: We propose an algorithm, termed the Non-Equilibrium Transport Sampler (NETS), to sample from unnormalized probability distributions. NETS can be viewed as a variant of annealed importance sampling (AIS) based on Jarzynski's equality, in which the stochastic differential equation used to perform the non-equilibrium sampling is augmented with an additional learned drift term that lowers the impact of the unbiasing weights used in AIS. We show that this drift is the minimizer of a variety of objective functions, which can all be estimated in an unbiased fashion without backpropagating through solutions of the stochastic differential equations governing the sampling. We also prove that some these objectives control the Kullback-Leibler divergence of the estimated distribution from its target. NETS is shown to be unbiased and, in addition, has a tunable diffusion coefficient which can be adjusted post-training to maximize the effective sample size. We demonstrate the efficacy of the method on standard benchmarks, high-dimensional Gaussian mixture distributions, and a model from statistical lattice field theory, for which it surpasses the performances of related work and existing baselines.

• The paper presents a novel sampler that integrates annealed Langevin dynamics with a learnable drift to effectively reduce variance in importance weights.
• It establishes an unbiased sampling method with a tunable diffusion coefficient, underpinned by Fokker-Planck equations and Jarzynski's equality.
• Empirical results on high-dimensional Gaussian mixtures and lattice field models demonstrate improved effective sample size and overall sampling efficiency.

Overview of "NETS: A Non-Equilibrium Transport Sampler"

The paper introduces the Non-Equilibrium Transport Sampler (NETS), a novel algorithm aimed at sampling from unnormalized probability distributions. It extends annealed importance sampling (AIS) via Jarzynski's equality by incorporating a learned drift term in the stochastic differential equation (SDE) used for non-equilibrium sampling. This enhancement effectively reduces the variance in importance weights typical of AIS, which can thereby improve sampling efficiency.

Algorithmic Contributions

The authors present several key contributions:

• Integrated Transport and Annealed Dynamics: The NETS algorithm employs annealed Langevin dynamics augmented with learnable transport. The added drift is shown to be the minimizer of objective functions, thus allowing a reduction in the effect of unbiasing weights.
• Unbiased Nature and Tunable Diffusion: The method is theoretically proven to be unbiased and features a tunable diffusion coefficient. This tunability can be adjusted post-training to optimize the effective sample size.
• Off-Policy Learning and KL Control: A significant theoretical contribution is the demonstration that one of the objective functions, grounded in Physics-Informed Neural Networks (PINNs), can be estimated without bias. This objective also governs the Kullback-Leibler (KL) divergence between the estimated and target distribution, ensuring efficient learning without initial data dependency.

Theoretical Foundations

The theoretical core of NETS is built around the manipulation of Fokker-Planck equations and Jarzynski's equality, establishing the unbiased nature of the sampler. By identifying and learning an optimal drift function, the authors manage to control the lag between evolving samples and the target distribution, a common issue in conventional AIS where large variance in weights is typically encountered.

Empirical Validation

The efficacy of NETS is demonstrated across several domains:

• Standard Benchmarks: Evaluations on high-dimensional Gaussian mixtures and lattice field theory models reveal that NETS outperforms existing baselines, showing improved sample quality and effective sample size.
• Performance Metrics: The authors highlight the ability of NETS to adaptively enhance sampling performance by tuning the diffusion coefficient, distinguishing it from traditional ergodic methods.

Implications and Speculations

NETS represents a significant step forward in the sampling domain, particularly for distributions that are not log-concave. The integration of learnable transport with annealed dynamics offers a flexible and robust solution to a longstanding challenge in computational statistical mechanics and Bayesian inference frameworks. Theoretical underpinnings regarding drift estimation pave the way for future exploration in adaptive sampling techniques, potentially influencing the development of more efficient generative models.

Future Directions

Several promising research directions emerge from this work:

• Complex System Simulations: Further exploration into complex systems might leverage NETS for unbiased sampling in higher dimensions, particularly in physical simulations and inverse problems.
• Algorithmic Hybridization: Integrating NETS with other sampling methods such as Sequential Monte Carlo (SMC) to harness their respective strengths could provide enhanced performance in even more challenging sampling tasks.

In conclusion, the NETS algorithm, by addressing the critical issues of bias and variance in the sampling process, offers substantial contributions to both theoretical and practical aspects of statistical sampling. Its introduction marks a valuable addition to the algorithmic toolkit available for tackling high-dimensional, non-log-concave target distributions.