Temperature-Annealed Boltzmann Generators

Henrik Schopmans, Pascal Friederich

TL;DR: We can efficiently sample molecular systems using normalizing flows with reweighting-based temperature-annealing.

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.

Lay Summary:

Understanding how molecules behave and interact is key to breakthroughs in areas like drug discovery. Computer simulations of physical systems are like virtual microscopes that offer insights into a system's behavior without requiring expensive lab experiments. However, to cover all relevant interactions and processes, the movement of the atoms in a system typically has to be simulated for a very long time, requiring large amounts of computational resources.

Instead of simulating the system over time, variational sampling methods train generative machine learning models to directly match the probability distribution of the physical system. This is a promising approach to make computer simulations more efficient. However, such variational methods often show a problem called mode collapse, where only a small fraction of the system's behavior is learned by the machine learning model.

Our work presents a simple and effective fix: we start training the model under easier conditions, simulating a higher-temperature environment where the mode collapse problem disappears. Then, we gradually "cool" the system down, guiding the model toward realistic behaviors at the target temperature. This approach allows efficient exploration of the atomistic behavior of molecular systems, without mode collapse.

Link To Code: https://github.com/aimat-lab/TA-BG

Primary Area: Applications->Chemistry, Physics, and Earth Sciences

Keywords: boltzmann generators, normalizing flows, molecular systems, variational sampling

Submission Number: 6832