New flow-based sampler handles periodic crystals and gives accurate free energies for large ice models
Scientists present a new method that uses neural-network flows to sample equilibrium states of condensed matter while respecting the periodic nature of simulation boxes. The method combines continuous normalizing flows, which are invertible neural networks that map simple probability distributions to complicated ones, with Riemannian flow matching, a training objective adapted to the torus shape that represents periodic boundary conditions. This lets the model propose independent equilibrium configurations directly, instead of relying on long molecular dynamics or Monte Carlo simulations.
To make the idea practical for big systems, the authors design several physics-aware parts. They use a lattice-based prior distribution that places atoms near equilibrium lattice sites with a wrapped Gaussian (so positions respect periodicity). They fix the system center of mass to avoid trivial global drift and use a size-transferable, local neural architecture based on an equivariant transformer. The resulting model is called RFM-ET (Riemannian Flow Matching with an Equivariant Transformer). The team trained RFM-ET on a monatomic ice test system (the mW model) with up to 1,000 particles — more than four times larger than earlier flow-based benchmarks — while keeping a comparable computational budget.
A technical obstacle for continuous flows is that computing the exact change in log density during the flow is costly (scaling poorly with particle number). The authors use Hutchinson’s trace estimator, a stochastic trick that estimates the required trace in linear time. Because taking exponentials of noisy estimates introduces bias in the importance weights used for statistical reweighting, they apply a bias correction based on a second-order cumulant expansion. Practically, this means subtracting half the estimated variance of the stochastic trace from the log-weights before exponentiation. The paper shows this correction substantially reduces systematic error in free energy estimates when only a small number of stochastic probes are used.