Temperature-Annealed Boltzmann Generators

TL;DR

This paper introduces temperature-annealed Boltzmann generators that improve sampling efficiency for molecular systems by combining high-temperature training with reweighting to lower temperatures, outperforming existing methods.

Contribution

The authors develop a novel temperature-annealing approach for Boltzmann generators that reduces mode collapse and enhances sampling accuracy across complex molecular systems.

Findings

Achieves better sampling metrics than baseline methods.

Requires up to three times fewer energy evaluations.

Accurately resolves metastable states in large systems.

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.