Scalable Equilibrium Sampling with Sequential Boltzmann Generators

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

Keywords: Normalizing Flows, Boltzmann Generators, Annealed Importance Sampling

TL;DR: We scale up Boltzmann Generators using large all atom transformers and apply non-equilibrium transport to the proposal samples.

Abstract: Scalable sampling of molecular states in thermodynamic equilibrium is a long-standing challenge in statistical physics. Boltzmann generators tackle this problem by pairing powerful normalizing flows (NF) with importance sampling to obtain statistically independent samples under the target distribution. In this paper, we extend the Boltzmann generator framework and introduce Sequential Boltzmann Generators (SBG) with two key improvements. The first is a highly efficient non-equivariant Transformer-based normalizing flow operating directly on all-atom Cartesian coordinates. The second is inference time scaling of flow samples with non-equilibrium transport towards the target distribution and reweighting through a novel application of Sequential Monte Carlo (SMC). SBG is more computationally efficient as no explicit equivariance constraint is encoded in the NF, but is instead softly introduced via data augmentations. Further, SBG improves sample quality by applying SMC along the transport path. SBG achieves state-of-the-art performance w.r.t. all metrics on molecular systems, demonstrating the first equilibrium sampling in Cartesian coordinates of tri, tetra, and hexapeptides that were so far intractable for prior Boltzmann generators.

Submission Number: 97