# Temperature-Annealed Boltzmann Generators

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

## Key facts

- Cited by 15 papers in the map (data through 14 Oct 2026).
- Builds on 29 of its 50 references that are in the map.
- First posted on arXiv 31 Jan 2025.

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

- [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)
- [Graph neural networks for materials science and chemistry](https://mlsanity.com/paper/2208.09481/) (2022)
- [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)
- [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)
- [Normalizing Flows on Tori and Spheres](https://mlsanity.com/paper/2002.02428/) (2020)
- [Iterated Denoising Energy Matching for Sampling from Boltzmann Densities](https://mlsanity.com/paper/2402.06121/) (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)
- [Continual Repeated Annealed Flow Transport Monte Carlo](https://mlsanity.com/paper/2201.13117/) (2022)
- [Diffusion Generative Flow Samplers: Improving learning signals through partial trajectory optimization](https://mlsanity.com/paper/2310.02679/) (2023)
- [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)
- [Transferable Boltzmann Generators](https://mlsanity.com/paper/2406.14426/) (2024)
- [Stochastic Optimal Control for Collective Variable Free Sampling of Molecular Transition Paths](https://mlsanity.com/paper/2207.02149/) (2022)
- [SE(3) Equivariant Augmented Coupling Flows](https://mlsanity.com/paper/2308.10364/) (2023)
- [Transport meets Variational Inference: Controlled Monte Carlo Diffusions](https://mlsanity.com/paper/2307.01050/) (2023)
- [On the Universality of Volume-Preserving and Coupling-Based Normalizing Flows](https://mlsanity.com/paper/2402.06578/) (2024)
- [Efficient mapping of phase diagrams with conditional Boltzmann Generators](https://mlsanity.com/paper/2406.12378/) (2024)
- [A theoretical perspective on mode collapse in variational inference](https://mlsanity.com/paper/2410.13300/) (2024)
- [Iterated Energy-based Flow Matching for Sampling from Boltzmann Densities](https://mlsanity.com/paper/2408.16249/) (2024)
- [Transition Path Sampling with Improved Off-Policy Training of Diffusion Path Samplers](https://mlsanity.com/paper/2405.19961/) (2024)
- [Designing losses for data-free training of normalizing flows on Boltzmann distributions](https://mlsanity.com/paper/2301.05475/) (2023)
- [Adaptive Annealed Importance Sampling with Constant Rate Progress](https://mlsanity.com/paper/2306.15283/) (2023)
- [Conditional Normalizing Flows for Active Learning of Coarse-Grained Molecular Representations](https://mlsanity.com/paper/2402.01195/) (2024)
- [TRADE: Transfer of Distributions between External Conditions with Normalizing Flows](https://mlsanity.com/paper/2410.19492/) (2024)

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

- [Underdamped Diffusion Bridges with Applications to Sampling](https://mlsanity.com/paper/2503.01006/) (2025)
- [Progressive Inference-Time Annealing of Diffusion Models for Sampling from Boltzmann Densities](https://mlsanity.com/paper/2506.16471/) (2025)
- [Trust Region Constrained Measure Transport in Path Space for Stochastic Optimal Control and Inference](https://mlsanity.com/paper/2508.12511/) (2025)
- [Bridge Matching Sampler: Scalable Sampling via Generalized Fixed-Point Diffusion Matching](https://mlsanity.com/paper/2603.00530/) (2026)
- [Learning Boltzmann Generators via Constrained Mass Transport](https://mlsanity.com/paper/2510.18460/) (2025)
- [Energy-Weighted Flow Matching: Unlocking Continuous Normalizing Flows for Efficient and Scalable Boltzmann Sampling](https://mlsanity.com/paper/2509.03726/) (2025)
- [Annealing in variational inference mitigates mode collapse: A theoretical study on Gaussian mixtures](https://mlsanity.com/paper/2602.12923/) (2026)
- [Efficient Training of Boltzmann Generators Using Off-Policy Log-Dispersion Regularization](https://mlsanity.com/paper/2602.03729/) (2026)
- [Discrete diffusion samplers and bridges: Off-policy algorithms and applications in latent spaces](https://mlsanity.com/paper/2602.05961/) (2026)
- [AI-based Methods for Simulating, Sampling, and Predicting Protein Ensembles](https://mlsanity.com/paper/2509.17224/) (2025)
- [Learning Structure, Energy, and Dynamics: A Survey of Artificial Intelligence for Protein Dynamics](https://mlsanity.com/paper/2604.25244/) (2026)
- [Autoregressive Boltzmann Generators](https://mlsanity.com/paper/2606.27361/) (2026)
- [Mode Coverage in Normalizing Flow Boltzmann Generators via Log-Ratio Variation](https://mlsanity.com/paper/2609.09473/) (2026)
- [Iterative Exact Discrete Guidance for Energy-Based Sampling](https://mlsanity.com/paper/2609.37043/) (2026)
- [Looking Under the Streetlight: Evaluation in Generative Molecular Dynamics](https://mlsanity.com/paper/2610.12518/) (2026)

## Abstract (arXiv)

Efficient sampling of unnormalized probability densities such as the Boltzmann distribution of molecular systems is a longstanding challenge. Next to conventional approaches like molecular dynamics or Markov chain Monte Carlo, variational approaches, such as training normalizing flows with the reverse Kullback-Leibler divergence, have been introduced. However, such methods are prone to mode collapse and often do not learn to sample the full configurational space. Here, we present temperature-annealed Boltzmann generators (TA-BG) to address this challenge. First, we demonstrate that training a normalizing flow with the reverse Kullback-Leibler divergence at high temperatures is possible without mode collapse. Furthermore, we introduce a reweighting-based training objective to anneal the distribution to lower target temperatures. We apply this methodology to three molecular systems of increasing complexity and, compared to the baseline, achieve better results in almost all metrics while requiring up to three times fewer target energy evaluations. For the largest system, our approach is the only method that accurately resolves the metastable states of the system.
