Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Boltzmann Generators for Condensed Matter via Riemannian Flow Matching

Emil Hoffmann, Maximilian Schebek, Leon Klein, Frank Noé, Jutta Rogal

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.

Boltzmann Generators for Condensed Matter via Riemannian Flow Matching

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.
Paper Structure (13 sections, 25 equations, 3 figures, 3 tables)

This paper contains 13 sections, 25 equations, 3 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Methods
  3. Implementation
  4. Experimental Results
  5. Discussion
  6. Appendix
  7. Geodesic on the flat torus
  8. Systems
  9. 3N-dim Gaussian prior: Einstein-Crystal
  10. Monoatomic-Water
  11. Second-Order Cumulant Expansion for Hutchinson’s Estimator
  12. ODE Integration
  13. Implementation details

Figures (3)

  • Figure 1: Left: Radial distribution function of the cubic ice system with 1000 particles as obtained from MD and RFM-ET. Center: Energy histograms of the cubic ice system with 1000 (top) and 216 (bottom) particles as obtained from MD and RFM-ET. Right: The 216 particles (blue) and 1000 particles (light blue) cubic ice systems and their training times with our method and local coupling flows (LCF) (schebek_scalable_2025).
  • Figure 2: Left: Effective sample size computed for the models described in the text, including results for models evaluated on the same system size as used during training, as well as for models transferred to larger system sizes than those seen during training. Error bars are shown where visible. Right: The deviation to the reference free energy of a 512 particles mW system in the cubic ice phase, computed with a varying number of Hutchinson probes, with and without bias correction. The shaded area corresponds to the accuracy needed for precise phase separation between hexagonal and cubic ice. Errors for our method in all plots were estimated over $10$k samples obtained from three models trained with different seeds.
  • Figure 3: ESS against the number of Hutchinson probes used at each integration step. Dotted lines correspond to values obtained with exact divergence calculations.