Schaalbare Inferentietijd-Annealing met Surrogaat Aannemelijkheidsschatters

Samenvatting

Een al lang bestaande uitdaging in computationele chemie en biofysica is het efficiënt bemonsteren van de Boltzmannverdeling van moleculen. Vooruitgang in generatieve modellering is voorgesteld om de beperkingen van conventionele bemonsteringstechnieken aan te pakken door de rekenkosten van simulatie te elimineren. Een veelbelovende richting is het iteratief finetunen van diffusiemodellen langs een temperatuurladder, waarbij trainingsdata worden gegenereerd via importantiebemonstering tijdens het annealen op het moment van inferentie. Helaas vereisen deze methoden het berekenen van een divergentie over het scoreveld om importantiegewichten te schatten, wat ze onhandelbaar maakt voor grotere systemen. Hier presenteren we schaalbaar annealen tijdens inferentie (SITA), dat op stroming gebaseerde modellen hertraint om monsters te genereren bij progressief lagere temperaturen met behulp van een op energie gebaseerd model om snelle surrogate waarschijnlijkheden mogelijk te maken. We demonstreren state-of-the-art prestaties op zowel Alanine Dipeptide als Alanine Tripeptide, terwijl we kostbare divergentietermen vermijden. Onze code is beschikbaar op https://github.com/countrsignal/sita.git

English

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