Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Scalable Equilibrium Sampling with Sequential Boltzmann Generators

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

TL;DR

This work addresses scalable equilibrium sampling of molecular systems by extending Boltzmann Generators with Sequential Boltzmann Generators (SBG). SBG combines a Transformer-based, non-equivariant normalizing flow operating on all-atom Cartesian coordinates with inference-time continuous-time annealing via AIS/SMC and annealed Langevin dynamics to move samples toward the Boltzmann target $\\mu_{\\text{target}}(x)$. Key contributions include scalable training of softly SE(3)-equivariant proposals using TarFlow, an energy-adjusted proposal to improve ESS, and continuous-time sampling that leverages Jarzynski's equality for reweighting. Empirically, SBG achieves state-of-the-art performance on peptide systems up to hexapeptides and demonstrates scalable equilibrium sampling for decapeptide chignolin, outperforming CNF-based BGs and SNIS baselines and enabling substantial gains in sampling efficiency and coverage of metastable states.

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.

Scalable Equilibrium Sampling with Sequential Boltzmann Generators

TL;DR

This work addresses scalable equilibrium sampling of molecular systems by extending Boltzmann Generators with Sequential Boltzmann Generators (SBG). SBG combines a Transformer-based, non-equivariant normalizing flow operating on all-atom Cartesian coordinates with inference-time continuous-time annealing via AIS/SMC and annealed Langevin dynamics to move samples toward the Boltzmann target . Key contributions include scalable training of softly SE(3)-equivariant proposals using TarFlow, an energy-adjusted proposal to improve ESS, and continuous-time sampling that leverages Jarzynski's equality for reweighting. Empirically, SBG achieves state-of-the-art performance on peptide systems up to hexapeptides and demonstrates scalable equilibrium sampling for decapeptide chignolin, outperforming CNF-based BGs and SNIS baselines and enabling substantial gains in sampling efficiency and coverage of metastable states.

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.

Paper Structure

This paper contains 35 sections, 8 theorems, 75 equations, 20 figures, 10 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Background and Preliminaries
  3. Normalizing Flows
  4. Boltzmann Generators
  5. Sequential Boltzmann Generators
  6. Scaling Training of Boltzmann Generators
  7. Inference Time Scaling of Boltzmann Generators
  8. Experiments
  9. Results
  10. Scaling to Decapeptide
  11. Related Work
  12. Conclusion
  13. Alternate Paths
  14. Proposal Free Langevin Dynamics
  15. Proofs
  16. ...and 20 more sections

Key Result

Proposition 1

Given an $\mathrm{SE(3)}$-invariant $\mu_{\text{target}}(x)$, consider the decomposition of a data point $x \in {\mathbb{R}}^{n \times 3}$ into its constituent mean-free component, $\bar{x} \in {\mathbb{R}}^{(n-1) \times 3} \hookrightarrow {\mathbb{R}}^{n \times 3}$ and center of mass $c \in {\mathb

Figures (20)

  • Figure 1: SBG uses annealed Langevin dynamics to transport proposal flow samples towards towards the target distribution.
  • Figure 2: Samples generated by SBG on peptide systems ranging from 2 to 10 residues.
  • Figure 3: Energy histograms for baseline methods and SBG on alanine dipeptide dataset.
  • Figure 4: Energy distribution histograms for baseline ECNF++ (left) and SBG(right) on trialanine dataset.
  • Figure 5: Energy distribution histograms for baseline ECNF++ (left) and SBG(right) on alanine tetrapeptide dataset.
  • ...and 15 more figures

Theorems & Definitions (12)

  • Proposition 1
  • Proposition 2: albergo_nets_2024
  • Proposition 3
  • Proposition 3
  • proof
  • Lemma 1
  • proof
  • Proposition 3
  • proof
  • Proposition 4
  • ...and 2 more