# NETS: A Non-Equilibrium Transport Sampler

arXiv: https://arxiv.org/abs/2410.02711 · HTML: https://mlsanity.com/paper/2410.02711/

## Key facts

- Cited by 69 papers in the map (data through 14 Oct 2026).
- Builds on 25 of its 44 references that are in the map.
- First posted on arXiv 3 Oct 2024.

## Builds on (most cited first, 25 of 25)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Stochastic Interpolants: A Unifying Framework for Flows and Diffusions](https://mlsanity.com/paper/2303.08797/) (2023)
- [Diffusion Schr"odinger Bridge with Applications to Score-Based Generative Modeling](https://mlsanity.com/paper/2106.01357/) (2021)
- [Path Integral Sampler: a stochastic control approach for sampling](https://mlsanity.com/paper/2111.15141/) (2021)
- [Denoising Diffusion Samplers](https://mlsanity.com/paper/2302.13834/) (2023)
- [Flow-based generative models for Markov chain Monte Carlo in lattice field theory](https://mlsanity.com/paper/1904.12072/) (2019)
- [An optimal control perspective on diffusion-based generative modeling](https://mlsanity.com/paper/2211.01364/) (2022)
- [Flow Annealed Importance Sampling Bootstrap](https://mlsanity.com/paper/2208.01893/) (2022)
- [Iterated Denoising Energy Matching for Sampling from Boltzmann Densities](https://mlsanity.com/paper/2402.06121/) (2024)
- [Action Matching: Learning Stochastic Dynamics from Samples](https://mlsanity.com/paper/2210.06662/) (2022)
- [Annealed Flow Transport Monte Carlo](https://mlsanity.com/paper/2102.07501/) (2021)
- [Continual Repeated Annealed Flow Transport Monte Carlo](https://mlsanity.com/paper/2201.13117/) (2022)
- [Beyond ELBOs: A Large-Scale Evaluation of Variational Methods for Sampling](https://mlsanity.com/paper/2406.07423/) (2024)
- [Improved off-policy training of diffusion samplers](https://mlsanity.com/paper/2402.05098/) (2024)
- [Learning Interpolations between Boltzmann Densities](https://mlsanity.com/paper/2301.07388/) (2023)
- [Stochastic normalizing flows as non-equilibrium transformations](https://mlsanity.com/paper/2201.08862/) (2022)
- [Liouville Flow Importance Sampler](https://mlsanity.com/paper/2405.06672/) (2024)
- [A Unified Perspective on Natural Gradient Variational Inference with Gaussian Mixture Models](https://mlsanity.com/paper/2209.11533/) (2022)
- [Posterior Sampling with Denoising Oracles via Tilted Transport](https://mlsanity.com/paper/2407.00745/) (2024)
- [Harmonic Path Integral Diffusion](https://mlsanity.com/paper/2409.15166/) (2024)

## Built on by (most cited first, 30 of 69)

- [Feynman-Kac Correctors in Diffusion: Annealing, Guidance, and Product of Experts](https://mlsanity.com/paper/2503.02819/) (2025)
- [Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching](https://mlsanity.com/paper/2504.11713/) (2025)
- [Sequential Controlled Langevin Diffusions](https://mlsanity.com/paper/2412.07081/) (2024)
- [Scalable Equilibrium Sampling with Sequential Boltzmann Generators](https://mlsanity.com/paper/2502.18462/) (2025)
- [Debiasing Guidance for Discrete Diffusion with Sequential Monte Carlo](https://mlsanity.com/paper/2502.06079/) (2025)
- [LEAPS: A discrete neural sampler via locally equivariant networks](https://mlsanity.com/paper/2502.10843/) (2025)
- [No Trick, No Treat: Pursuits and Challenges Towards Simulation-free Training of Neural Samplers](https://mlsanity.com/paper/2502.06685/) (2025)
- [Progressive Inference-Time Annealing of Diffusion Models for Sampling from Boltzmann Densities](https://mlsanity.com/paper/2506.16471/) (2025)
- [Consistent Sampling and Simulation: Molecular Dynamics with Energy-Based Diffusion Models](https://mlsanity.com/paper/2506.17139/) (2025)
- [Trust Region Constrained Measure Transport in Path Space for Stochastic Optimal Control and Inference](https://mlsanity.com/paper/2508.12511/) (2025)
- [DriftLite: Lightweight Drift Control for Inference-Time Scaling of Diffusion Models](https://mlsanity.com/paper/2509.21655/) (2025)
- [Tilt Matching for Scalable Sampling and Fine-Tuning](https://mlsanity.com/paper/2512.21829/) (2025)
- [Training Neural Samplers with Reverse Diffusive KL Divergence](https://mlsanity.com/paper/2410.12456/) (2024)
- [Adjoint Schr"odinger Bridge Sampler](https://mlsanity.com/paper/2506.22565/) (2025)
- [MDNS: Masked Diffusion Neural Sampler via Stochastic Optimal Control](https://mlsanity.com/paper/2508.10684/) (2025)
- [Solving Inverse Problems via Diffusion-Based Priors: An Approximation-Free Ensemble Sampling Approach](https://mlsanity.com/paper/2506.03979/) (2025)
- [Accelerated Parallel Tempering via Neural Transports](https://mlsanity.com/paper/2502.10328/) (2025)
- [RNE: plug-and-play diffusion inference-time control and energy-based training](https://mlsanity.com/paper/2506.05668/) (2025)
- [FEAT: Free energy Estimators with Adaptive Transport](https://mlsanity.com/paper/2504.11516/) (2025)
- [Proximal Diffusion Neural Sampler](https://mlsanity.com/paper/2510.03824/) (2025)
- [Value Gradient Guidance for Flow Matching Alignment](https://mlsanity.com/paper/2512.05116/) (2025)
- [Progressive Tempering Sampler with Diffusion](https://mlsanity.com/paper/2506.05231/) (2025)
- [Amortized Sampling with Transferable Normalizing Flows](https://mlsanity.com/paper/2508.18175/) (2025)
- [Complexity Analysis of Normalizing Constant Estimation: from Jarzynski Equality to Annealed Importance Sampling and beyond](https://mlsanity.com/paper/2502.04575/) (2025)
- [Scaling of Stochastic Normalizing Flows in SU(3) lattice gauge theory](https://mlsanity.com/paper/2412.00200/) (2024)
- [On scalable and efficient training of diffusion samplers](https://mlsanity.com/paper/2505.19552/) (2025)
- [Learning Boltzmann Generators via Constrained Mass Transport](https://mlsanity.com/paper/2510.18460/) (2025)
- [Non-equilibrium Annealed Adjoint Sampler](https://mlsanity.com/paper/2506.18165/) (2025)
- [Energy-Weighted Flow Matching: Unlocking Continuous Normalizing Flows for Efficient and Scalable Boltzmann Sampling](https://mlsanity.com/paper/2509.03726/) (2025)
- [Iterative Importance Fine-tuning of Diffusion Models](https://mlsanity.com/paper/2502.04468/) (2025)

## Abstract (arXiv)

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.