# Efficient Training of Boltzmann Generators Using Off-Policy Log-Dispersion Regularization

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 35 of its 53 references that are in the map.
- First posted on arXiv 3 Feb 2026.

## Builds on (most cited first, 30 of 35)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Density estimation using Real NVP](https://mlsanity.com/paper/1605.08803/) (2016)
- [Neural Spline Flows](https://mlsanity.com/paper/1906.04032/) (2019)
- [Torsional Diffusion for Molecular Conformer Generation](https://mlsanity.com/paper/2206.01729/) (2022)
- [E(n) Equivariant Normalizing Flows](https://mlsanity.com/paper/2105.09016/) (2021)
- [Equivariant flow matching](https://mlsanity.com/paper/2306.15030/) (2023)
- [Flow Annealed Importance Sampling Bootstrap](https://mlsanity.com/paper/2208.01893/) (2022)
- [Normalizing Flows on Tori and Spheres](https://mlsanity.com/paper/2002.02428/) (2020)
- [AlphaFold Meets Flow Matching for Generating Protein Ensembles](https://mlsanity.com/paper/2402.04845/) (2024)
- [Improved sampling via learned diffusions](https://mlsanity.com/paper/2307.01198/) (2023)
- [Normalizing Flows are Capable Generative Models](https://mlsanity.com/paper/2412.06329/) (2024)
- [VarGrad: A Low-Variance Gradient Estimator for Variational Inference](https://mlsanity.com/paper/2010.10436/) (2020)
- [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)
- [OpenMM 8: Molecular Dynamics Simulation with Machine Learning Potentials](https://mlsanity.com/paper/2310.03121/) (2023)
- [Transferable Boltzmann Generators](https://mlsanity.com/paper/2406.14426/) (2024)
- [Resampling Base Distributions of Normalizing Flows](https://mlsanity.com/paper/2110.15828/) (2021)
- [Scalable Equilibrium Sampling with Sequential Boltzmann Generators](https://mlsanity.com/paper/2502.18462/) (2025)
- [SE(3) Equivariant Augmented Coupling Flows](https://mlsanity.com/paper/2308.10364/) (2023)
- [Progressive Inference-Time Annealing of Diffusion Models for Sampling from Boltzmann Densities](https://mlsanity.com/paper/2506.16471/) (2025)
- [Gradients should stay on Path: Better Estimators of the Reverse- and Forward KL Divergence for Normalizing Flows](https://mlsanity.com/paper/2207.08219/) (2022)
- [Temperature-Annealed Boltzmann Generators](https://mlsanity.com/paper/2501.19077/) (2025)
- [Progressive Tempering Sampler with Diffusion](https://mlsanity.com/paper/2506.05231/) (2025)
- [FALCON: Few-step Accurate Likelihoods for Continuous Flows](https://mlsanity.com/paper/2512.09914/) (2025)
- [Amortized Sampling with Transferable Normalizing Flows](https://mlsanity.com/paper/2508.18175/) (2025)
- [Scalable Normalizing Flows Enable Boltzmann Generators for Macromolecules](https://mlsanity.com/paper/2401.04246/) (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)
- [Efficient Regression-Based Training of Normalizing Flows for Boltzmann Generators](https://mlsanity.com/paper/2506.01158/) (2025)
- [Non-equilibrium Annealed Adjoint Sampler](https://mlsanity.com/paper/2506.18165/) (2025)

## Built on by (most cited first, 3 of 3)

- [Learning Structure, Energy, and Dynamics: A Survey of Artificial Intelligence for Protein Dynamics](https://mlsanity.com/paper/2604.25244/) (2026)
- [Mode Coverage in Normalizing Flow Boltzmann Generators via Log-Ratio Variation](https://mlsanity.com/paper/2609.09473/) (2026)
- [Looking Under the Streetlight: Evaluation in Generative Molecular Dynamics](https://mlsanity.com/paper/2610.12518/) (2026)

## Abstract (arXiv)

Sampling from unnormalized probability densities is a central challenge in computational science. Boltzmann generators are generative models that enable independent sampling from the Boltzmann distribution of physical systems at a given temperature. However, their practical success depends on data-efficient training, as both simulation data and target energy evaluations are costly. To this end, we propose off-policy log-dispersion regularization (LDR), a novel regularization framework that builds on a generalization of the log-variance objective. We apply LDR in the off-policy setting in combination with standard data-based training objectives, without requiring additional on-policy samples. LDR acts as a shape regularizer of the energy landscape by leveraging additional information in the form of target energy labels. The proposed regularization framework is broadly applicable, supporting unbiased or biased simulation datasets as well as purely variational training without access to target samples. Across all benchmarks, LDR improves both final performance and data efficiency, with sample efficiency gains of up to one order of magnitude.
