Recuit scalable en phase d'inférence avec des estimateurs de vraisemblance de substitution

Résumé

Un défi de longue date en chimie computationnelle et en biophysique est l'échantillonnage efficace de la distribution de Boltzmann des molécules. Des avancées en modélisation générative ont été proposées pour surmonter les limites des techniques d'échantillonnage conventionnelles en éliminant le coût computationnel de la simulation. Une direction prometteuse consiste à affiner itérativement des modèles de diffusion le long d'une échelle de température, où les données d'entraînement sont générées par échantillonnage d'importance lors du recuit en phase d'inférence. Malheureusement, ces méthodes nécessitent le calcul d'une divergence sur le champ de score pour estimer les poids d'importance, ce qui les rend irréalisables pour des systèmes plus grands. Nous présentons ici le recuit à l'inférence scalable (SITA), qui réentraîne des modèles basés sur les flots pour générer des échantillons à des températures progressivement plus basses, en utilisant un modèle basé sur l'énergie pour faciliter des vraisemblances de substitution rapides. Nous démontrons des performances de pointe à la fois sur l'alanine dipeptide et l'alanine tripeptide tout en évitant les termes de divergence coûteux. Notre code est disponible à l'adresse 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