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

TL;DR

This paper introduces PITA, a novel diffusion-based sampling framework that combines annealing and smoothing techniques to efficiently generate samples from complex Boltzmann distributions, enabling equilibrium sampling of molecular systems.

Contribution

PITA is the first method to successfully perform equilibrium sampling of molecular systems using diffusion models with inference-time annealing and PDE-based techniques.

Findings

Enables equilibrium sampling of N-body systems and peptides.

Reduces energy function evaluations significantly.

First to achieve molecular system sampling with diffusion models.

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