Boltzmann Generators for Condensed Matter via Riemannian Flow Matching

TL;DR

This paper introduces Riemannian flow matching with Hutchinson's trace estimator to improve equilibrium sampling in condensed matter, enabling accurate free energy calculations for large systems.

Contribution

It develops a novel Riemannian normalizing flow framework tailored for periodic condensed matter systems, enhancing sampling efficiency and accuracy.

Findings

Successfully trained on large monatomic ice systems.

Achieved highly accurate free energy estimates.

Validated the approach's effectiveness for equilibrium sampling.

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.