Scalable Inference-Time Annealing with Surrogate Likelihood Estimators

2026-05-29

AI Summary

Traditional molecular simulations are hindered by high computational costs, while existing generative models rely on intractable divergence terms during inference. This work proposes a scalable annealed sampling framework that integrates flow-based models with energy-based models. By constructing a surrogate likelihood estimator, the method progressively refines the flow model across a temperature ladder through annealing, entirely circumventing expensive divergence computations at inference time. Evaluated on alanine dipeptide and tripeptide systems, the approach achieves state-of-the-art sampling performance, significantly enhancing both the efficiency and scalability of Boltzmann distribution sampling for larger biomolecules.

Abstract

A long standing challenge in computational chemistry and biophysics is efficiently sampling the Boltzmann distribution of molecules. Advances in generative modeling have been proposed to address the limitations of conventional sampling techniques by eliminating the computational cost of simulation. A promising direction is iteratively finetuning diffusion models along a temperature ladder whereby training data is generated via importance sampling during inference-time annealing. Unfortunately, these methods require computing a divergence over the score field to estimate importance weights, rendering them intractable for larger systems. Here we present scalable inference-time annealing (SITA), which retrains flow-based models to generate samples at progressively lower temperatures using an energy-based model to facilitate fast surrogate likelihoods. We demonstrate state-of-the-art performance on both Alanine Dipeptide and Alanine Tripeptide while avoiding costly divergence terms. Our code is available at: https://github.com/countrsignal/sita.git

Problem

• Boltzmann distribution
• molecular sampling
• inference-time annealing
• importance weights
• score divergence

Innovation

• inference-time annealing
• surrogate likelihood
• flow-based models
• energy-based model
• Boltzmann distribution sampling