Boltzmann Generators for Condensed Matter via Riemannian Flow Matching

2026-02-10

AI Summary

This work addresses the challenges of inefficient equilibrium sampling and inaccurate free energy estimation in condensed-phase systems by proposing a continuous normalizing flow method that incorporates periodic structural constraints. The approach constructs a Boltzmann generator via Riemannian flow matching—a technique applied here for the first time to condensed-phase systems—and integrates Hutchinson’s stochastic trace estimator with a cumulant-expansion-based bias correction scheme to enable thermodynamically consistent and efficient reweighting. Demonstrated on a monoatomic water model, the method successfully trains the largest generative model to date for such systems and achieves high-accuracy free energies without requiring multi-stage estimation protocols.

Abstract

Sampling equilibrium distributions is fundamental to statistical mechanics. While flow matching has emerged as scalable state-of-the-art paradigm for generative modeling, its potential for equilibrium sampling in condensed-phase systems remains largely unexplored. We address this by incorporating the periodicity inherent to these systems into continuous normalizing flows using Riemannian flow matching. The high computational cost of exact density estimation intrinsic to continuous normalizing flows is mitigated by using Hutchinson's trace estimator, utilizing a crucial bias-correction step based on cumulant expansion to render the stochastic estimates suitable for rigorous thermodynamic reweighting. Our approach is validated on monatomic ice, demonstrating the ability to train on systems of unprecedented size and obtain highly accurate free energy estimates without the need for traditional multistage estimators.

Problem

• equilibrium sampling
• condensed matter
• periodicity
• free energy estimation
• generative modeling

Innovation

• Riemannian flow matching
• continuous normalizing flows
• equilibrium sampling
• free energy estimation
• Hutchinson's trace estimator