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

2025-06-19

AI Summary

Efficient sampling from unnormalized Boltzmann densities—particularly in molecular simulation—remains a longstanding challenge, as existing diffusion-based samplers fail under scale limitations. Method: This paper introduces Progressive Inference-Time Annealing (PITA), a novel framework that jointly integrates temperature annealing and diffusion smoothing via dual-path interpolation. PITA leverages Feynman–Kac partial differential equations for modeling and sequential Monte Carlo (SMC) for inference-time dynamic adaptive annealing, enabling end-to-end equilibrium sampling of N-body systems and peptides in Cartesian coordinates. Results: On benchmark molecular systems—including alanine dipeptide and tripeptides—PITA substantially reduces energy function evaluations, overcomes the scalability bottleneck of current diffusion samplers, and achieves, for the first time, efficient, unbiased equilibrium sampling across small molecules to peptides.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Reasoning under Uncertainty: Probabilistic Programming
• Machine Learning: Calibration & Uncertainty Quantification

Application Category

• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs

Abstract

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

Problem

• Efficient sampling from unnormalized Boltzmann densities
• Scaling diffusion-based samplers for molecular systems
• Reducing energy evaluations in equilibrium sampling

Innovation

• Progressive Inference-Time Annealing (PITA) framework
• Combines annealing and diffusion smoothing techniques
• Uses Feynman-Kac PDE with Sequential Monte Carlo