Scalable Equilibrium Sampling with Sequential Boltzmann Generators

2025-02-25

AI Summary

Efficient and scalable all-atom Cartesian-coordinate sampling for molecular thermodynamic equilibrium remains challenging. Method: We propose the Sequential Boltzmann Generator (SBG), the first non-equivariant invertible Transformer-based normalizing flow explicitly modeling atomic Cartesian coordinates. SBG integrates inference-time annealed Langevin dynamics with sequential Monte Carlo resampling, overcoming the high-variance and poor sample independence inherent in conventional importance sampling. Results: Our approach achieves, for the first time in Cartesian space, unbiased equilibrium sampling of peptides ranging from tri- to hexa-peptides. Evaluated across standard metrics—including sample fidelity, statistical independence, and conformational diversity—SBG outperforms existing Boltzmann generators and establishes new state-of-the-art performance for molecular systems. This work introduces a scalable, generalizable paradigm for thermodynamic modeling of complex peptide chains.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Deep Generative Models & Autoencoders
• Natural Language Processing: Generation

Application Category

• Search and Retrieval-Augmented AI: Agentic search
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs

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 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. In contrast to equivariant continuous flows of prior methods, we leverage exactly invertible non-equivariant architectures which are highly efficient both during sample generation and likelihood computation. As a result, this unlocks more sophisticated inference strategies beyond standard importance sampling. More precisely, as a second key improvement we perform inference-time scaling of flow samples using annealed Langevin dynamics which transports samples toward the target distribution leading to lower variance (annealed) importance weights which enable higher fidelity resampling with sequential Monte Carlo. 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.

Problem

• Scalable sampling of molecular states
• Equilibrium sampling in Cartesian coordinates
• Improving Boltzmann generators with non-equivariant Transformers

Innovation

• Sequential Boltzmann Generators
• Non-equivariant Transformer-based flow
• Anneal Langevin dynamics scaling