# Progressive Inference-Time Annealing of Diffusion Models for Sampling from Boltzmann Densities

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

## Key facts

- Cited by 26 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 31 of its 49 references that are in the map.
- First posted on arXiv 19 Jun 2025.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

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

- [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)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [E(n) Equivariant Graph Neural Networks](https://mlsanity.com/paper/2102.09844/) (2021)
- [Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions](https://mlsanity.com/paper/2209.11215/) (2022)
- [Equivariant Flows: Exact Likelihood Generative Learning for Symmetric Densities](https://mlsanity.com/paper/2006.02425/) (2020)
- [Reduce, Reuse, Recycle: Compositional Generation with Energy-Based Diffusion Models and MCMC](https://mlsanity.com/paper/2302.11552/) (2023)
- [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)
- [Equivariant flow matching](https://mlsanity.com/paper/2306.15030/) (2023)
- [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)
- [Action Matching: Learning Stochastic Dynamics from Samples](https://mlsanity.com/paper/2210.06662/) (2022)
- [Feynman-Kac Correctors in Diffusion: Annealing, Guidance, and Product of Experts](https://mlsanity.com/paper/2503.02819/) (2025)
- [Normalizing Flows are Capable Generative Models](https://mlsanity.com/paper/2412.06329/) (2024)
- [Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching](https://mlsanity.com/paper/2504.11713/) (2025)
- [NETS: A Non-Equilibrium Transport Sampler](https://mlsanity.com/paper/2410.02711/) (2024)
- [Beyond ELBOs: A Large-Scale Evaluation of Variational Methods for Sampling](https://mlsanity.com/paper/2406.07423/) (2024)
- [Target Score Matching](https://mlsanity.com/paper/2402.08667/) (2024)
- [Sequential Controlled Langevin Diffusions](https://mlsanity.com/paper/2412.07081/) (2024)
- [Transferable Boltzmann Generators](https://mlsanity.com/paper/2406.14426/) (2024)
- [Scalable Equilibrium Sampling with Sequential Boltzmann Generators](https://mlsanity.com/paper/2502.18462/) (2025)
- [No Trick, No Treat: Pursuits and Challenges Towards Simulation-free Training of Neural Samplers](https://mlsanity.com/paper/2502.06685/) (2025)
- [Composition and Control with Distilled Energy Diffusion Models and Sequential Monte Carlo](https://mlsanity.com/paper/2502.12786/) (2025)
- [Temperature-Annealed Boltzmann Generators](https://mlsanity.com/paper/2501.19077/) (2025)
- [Progressive Tempering Sampler with Diffusion](https://mlsanity.com/paper/2506.05231/) (2025)

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

- [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)
- [FALCON: Few-step Accurate Likelihoods for Continuous Flows](https://mlsanity.com/paper/2512.09914/) (2025)
- [CREPE: Controlling Diffusion with Replica Exchange](https://mlsanity.com/paper/2509.23265/) (2025)
- [Discrete Feynman-Kac Correctors](https://mlsanity.com/paper/2601.10403/) (2026)
- [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)
- [Control Variate Score Matching for Diffusion Models](https://mlsanity.com/paper/2512.20003/) (2025)
- [Flow Sampling: Learning to Sample from Unnormalized Densities via Denoising Conditional Processes](https://mlsanity.com/paper/2605.03984/) (2026)
- [Diffusion Path Samplers via Sequential Monte Carlo](https://mlsanity.com/paper/2601.21951/) (2026)
- [Discrete Adjoint Schr"odinger Bridge Sampler](https://mlsanity.com/paper/2602.08243/) (2026)
- [Protein Language Model Embeddings Improve Generalization of Implicit Transfer Operators](https://mlsanity.com/paper/2602.11216/) (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)
- [A unified perspective on fine-tuning and sampling with diffusion and flow models](https://mlsanity.com/paper/2605.00229/) (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)
- [Conditional Diffusion Sampling](https://mlsanity.com/paper/2605.04013/) (2026)
- [Autoregressive Boltzmann Generators](https://mlsanity.com/paper/2606.27361/) (2026)
- [Energy Guided Geometric Flow Matching](https://mlsanity.com/paper/2509.25230/) (2025)
- [Active Learning for Gaussian Process Regression Under Self-Induced Boltzmann Weights](https://mlsanity.com/paper/2605.10654/) (2026)
- [Scalable Inference-Time Annealing with Surrogate Likelihood Estimators](https://mlsanity.com/paper/2605.31498/) (2026)
- [Neural Non-Equilibrium Hamiltonian Monte Carlo for Corrected Boltzmann Sampling](https://mlsanity.com/paper/2607.15682/) (2026)
- [Mode Coverage in Normalizing Flow Boltzmann Generators via Log-Ratio Variation](https://mlsanity.com/paper/2609.09473/) (2026)

## Abstract (arXiv)

Sampling efficiently from a target unnormalized probability density remains a core challenge, with relevance across countless high-impact scientific applications. A promising approach towards this challenge is the design of amortized samplers that borrow key ideas, such as probability path design, from state-of-the-art generative diffusion models. However, all existing diffusion-based samplers remain unable to draw samples from distributions at the scale of even simple molecular systems. In this paper, we propose Progressive Inference-Time Annealing (PITA), a novel framework to learn diffusion-based samplers that combines two complementary interpolation techniques: I.) Annealing of the Boltzmann distribution and II.) Diffusion smoothing. PITA trains a sequence of diffusion models from high to low temperatures by sequentially training each model at progressively higher temperatures, leveraging engineered easy access to samples of the temperature-annealed target density. In the subsequent step, PITA enables simulating the trained diffusion model to procure training samples at a lower temperature for the next diffusion model through inference-time annealing using a novel Feynman-Kac PDE combined with Sequential Monte Carlo. Empirically, PITA enables, for the first time, equilibrium sampling of N-body particle systems, Alanine Dipeptide, and tripeptides in Cartesian coordinates with dramatically lower energy function evaluations. Code available at: https://github.com/taraak/pita