Scalable Equilibrium Sampling with Sequential Boltzmann Generators

Charlie B. Tan, Joey Bose, Chen Lin, Leon Klein, Michael M. Bronstein, Alexander Tong

TL;DR: We scale up Boltzmann Generators using large all-atom transformers and apply sequential Monte Carlo to improve sampling.

Abstract: Scalable sampling of molecular states in thermodynamic equilibrium is a long-standing challenge in statistical physics. Boltzmann generators tackle this problem by pairing normalizing flows with importance sampling to obtain uncorrelated samples under the target distribution. In this paper, we extend the Boltzmann generator framework with two key contributions, denoting our framework Sequential Boltzmann Generators (SBG). The first is a highly efficient Transformer-based normalizing flow operating directly on all-atom Cartesian coordinates. In contrast to the equivariant continuous flows of prior methods, we leverage exactly invertible non-equivariant architectures which are highly efficient during both sample generation and likelihood evaluation. This efficiency unlocks more sophisticated inference strategies beyond standard importance sampling. In particular, we perform inference-time scaling of flow samples using a continuous-time variant of sequential Monte Carlo, in which flow samples are transported towards the target distribution with annealed Langevin dynamics. SBG achieves state-of-the-art performance w.r.t. all metrics on peptide systems, demonstrating the first equilibrium sampling in Cartesian coordinates of tri-, tetra- and hexa-peptides that were thus far intractable for prior Boltzmann generators.

Lay Summary: Simulating the behavior of molecules is essential for science but takes a huge amount of time. We built a machine learning model that can generate realistic molecular structures much faster than traditional methods. It learns from existing data and improves its guesses using a process inspired by physics. Our method scales to larger molecules than ever before, which could help accelerate drug and material discovery.

Application-Driven Machine Learning: This submission is on Application-Driven Machine Learning.

Link To Code: https://github.com/charliebtan/transferable-samplers

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

Keywords: Normalizing Flows, Boltzmann Generators, Annealed Importance Sampling

Submission Number: 5008