Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
This paper introduces scalable inference-time annealing (SITA), a novel approach that retrains flow-based models to sample the Boltzmann distribution more efficiently by leveraging an energy-based model for fast surrogate likelihood estimation. By iteratively fine-tuning diffusion models along a temperature ladder and generating training data through importance sampling, SITA circumvents the computationally expensive divergence calculations typically required in larger systems. The method achieves state-of-the-art performance on the sampling of Alanine Dipeptide and Alanine Tripeptide, demonstrating significant improvements in efficiency without sacrificing accuracy.
SITA achieves state-of-the-art sampling efficiency for molecular distributions while eliminating costly divergence calculations, revolutionizing how we approach computational chemistry.
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