Scalable Equilibrium Sampling with Sequential Boltzmann Generators

• URL: http://arxiv.org/abs/2502.18462v1
• Date: Tue, 25 Feb 2025 18:59:13 GMT
• Title: Scalable Equilibrium Sampling with Sequential Boltzmann Generators
• Authors: Charlie B. Tan, Avishek Joey Bose, Chen Lin, Leon Klein, Michael M. Bronstein, Alexander Tong,
• Abstract summary: We extend the Boltzmann generator framework and introduce Sequential Boltzmann generators with two key improvements.<n>The first is a highly efficient non-equivariant Transformer-based normalizing flow operating directly on all-atom Cartesian coordinates.<n>We demonstrate the first equilibrium sampling in Cartesian coordinates of tri, tetra, and hexapeptides that were so far intractable for prior Boltzmann generators.
• Score: 60.00515282300297
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Flow Perturbation to Accelerate Unbiased Sampling of Boltzmann distribution [2.103187931015573]
  Flow-based generative models have been employed for sampling the Boltzmann distribution, but their application is hindered by the computational cost of obtaining the Jacobian of the flow. We introduce the flow perturbation method, which incorporates optimized perturbations into the flow. By reweighting trajectories generated by the perturbed flow, our method achieves unbiased sampling of the Boltzmann distribution with orders of magnitude speedup.
  arXiv  Detail & Related papers  (2024-07-15T12:29:17Z)
• Transferable Boltzmann Generators [6.36887173111677]
  We propose a first framework for Boltzmann Generators that are transferable across chemical space. We show that our proposed architecture enhances the efficiency of Boltzmann Generators trained on single molecular systems.
  arXiv  Detail & Related papers  (2024-06-20T15:50:12Z)
• Iterated Denoising Energy Matching for Sampling from Boltzmann Densities [109.23137009609519]
  Iterated Denoising Energy Matching (iDEM) iDEM alternates between (I) sampling regions of high model density from a diffusion-based sampler and (II) using these samples in our matching objective. We show that the proposed approach achieves state-of-the-art performance on all metrics and trains $2-5times$ faster.
  arXiv  Detail & Related papers  (2024-02-09T01:11:23Z)
• Equivariant Flow Matching with Hybrid Probability Transport [69.11915545210393]
  Diffusion Models (DMs) have demonstrated effectiveness in generating feature-rich geometries. DMs typically suffer from unstable probability dynamics with inefficient sampling speed. We introduce geometric flow matching, which enjoys the advantages of both equivariant modeling and stabilized probability dynamics.
  arXiv  Detail & Related papers  (2023-12-12T11:13:13Z)
• SE(3) Equivariant Augmented Coupling Flows [16.65770540017618]
  Coupling normalizing flows allow for fast sampling and density evaluation. Standard coupling architecture precludes endowing flows that operate on the Cartesian coordinates of atoms.
  arXiv  Detail & Related papers  (2023-08-20T20:49:15Z)
• Adaptive Annealed Importance Sampling with Constant Rate Progress [68.8204255655161]
  Annealed Importance Sampling (AIS) synthesizes weighted samples from an intractable distribution. We propose the Constant Rate AIS algorithm and its efficient implementation for $alpha$-divergences.
  arXiv  Detail & Related papers  (2023-06-27T08:15:28Z)
• Equivariant flow matching [0.9208007322096533]
  We introduce equivariant flow matching, a new training objective for equivariant continuous normalizing flows (CNFs) Equivariant flow matching exploits the physical symmetries of the target energy for efficient, simulation-free training of equivariant CNFs. Our results show that the equivariant flow matching objective yields flows with shorter integration paths, improved sampling efficiency, and higher scalability compared to existing methods.
  arXiv  Detail & Related papers  (2023-06-26T19:40:10Z)
• NAG-GS: Semi-Implicit, Accelerated and Robust Stochastic Optimizer [45.47667026025716]
  We propose a novel, robust and accelerated iteration that relies on two key elements. The convergence and stability of the obtained method, referred to as NAG-GS, are first studied extensively. We show that NAG-arity is competitive with state-the-art methods such as momentum SGD with weight decay and AdamW for the training of machine learning models.
  arXiv  Detail & Related papers  (2022-09-29T16:54:53Z)
• Sampling in Combinatorial Spaces with SurVAE Flow Augmented MCMC [83.48593305367523]
  Hybrid Monte Carlo is a powerful Markov Chain Monte Carlo method for sampling from complex continuous distributions. We introduce a new approach based on augmenting Monte Carlo methods with SurVAE Flows to sample from discrete distributions. We demonstrate the efficacy of our algorithm on a range of examples from statistics, computational physics and machine learning, and observe improvements compared to alternative algorithms.
  arXiv  Detail & Related papers  (2021-02-04T02:21:08Z)
• Gaussianization Flows [113.79542218282282]
  We propose a new type of normalizing flow model that enables both efficient iteration of likelihoods and efficient inversion for sample generation. Because of this guaranteed expressivity, they can capture multimodal target distributions without compromising the efficiency of sample generation.
  arXiv  Detail & Related papers  (2020-03-04T08:15:06Z)
• Stochastic Normalizing Flows [2.323220706791067]
  We show that normalizing flows can be used to learn the transformation of a simple prior distribution. We derive an efficient training procedure by which both the sampler's and the flow's parameters can be optimized end-to-end. We illustrate the representational power, sampling efficiency and correctness of SNFs on several benchmarks including applications to molecular sampling systems in equilibrium.
  arXiv  Detail & Related papers  (2020-02-16T23:29:32Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.