Temperature-Annealed Boltzmann Generators

2025-01-31

AI Summary

Efficient sampling from the Boltzmann distribution of molecular systems remains challenging—conventional approaches (e.g., molecular dynamics, random sampling, or end-to-end training) suffer from incomplete conformational coverage and mode collapse. Method: This paper proposes the Temperature-Annealed Variational Normalizing Flow (TAVNF) framework, introducing a novel training paradigm: high-temperature initialization → progressive annealing → reweighted reverse KL divergence optimization, enabling smooth distributional transition and lossless recovery of metastable states. Contribution/Results: TAVNF achieves state-of-the-art performance across three molecular systems: energy evaluation error is reduced by up to 3× compared to baselines; on the largest system, it is the first method to fully resolve all metastable conformations—thereby overcoming a fundamental bottleneck in complex molecular sampling.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Calibration & Uncertainty Quantification
• Multiagent Systems: Multiagent Learning

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Security and Privacy: Large-scale security measurements

Abstract

Efficient sampling of unnormalized probability densities such as the Boltzmann distribution of molecular systems is a longstanding challenge. Next to conventional approaches like molecular dynamics or Markov chain Monte Carlo, variational approaches, such as training normalizing flows with the reverse Kullback-Leibler divergence, have been introduced. However, such methods are prone to mode collapse and often do not learn to sample the full configurational space. Here, we present temperature-annealed Boltzmann generators (TA-BG) to address this challenge. First, we demonstrate that training a normalizing flow with the reverse Kullback-Leibler divergence at high temperatures is possible without mode collapse. Furthermore, we introduce a reweighting-based training objective to anneal the distribution to lower target temperatures. We apply this methodology to three molecular systems of increasing complexity and, compared to the baseline, achieve better results in almost all metrics while requiring up to three times fewer target energy evaluations. For the largest system, our approach is the only method that accurately resolves the metastable states of the system.

Problem

• Complex Probability Distribution Sampling
• Molecular Distributions
• Overcoming Limitations of Traditional Methods

Innovation

• Temperature Annealing
• Boltzmann Generator
• Complex Probability Distribution Sampling