Free Energy Surface Sampling via Reduced Flow Matching

Abstract: Sampling the free energy surface, namely, the distribution of collective variables (CVs), is a crucial problem in statistical physics, as it underpins a better understanding of chemical reactions and conformational transitions. Traditional methods for free energy surface sampling involve simulation in high-dimensional configuration space and projecting the resulting configurations onto the CV space. To reduce the computational costs of such sampling, we propose FES-FM, a reduced flow matching (FM) method for free energy sampling (FES). We train a dynamical transport map in the CV space, thereby enabling direct sampling of the free energy surface. For many-particle systems, we construct a prior distribution based on the Hessian at a local minimum of the potential, which ensures both rotation-translation invariance and physically meaningful configurations. We evaluate the proposed method across a variety of potential functions and collective variables. Comparative experiments demonstrate that our approach drastically reduces computational costs while delivering superior accuracy per unit sampling time.

• The paper introduces FES-FM, a variational framework for direct sampling from free energy surfaces by learning reduced-space neural transport maps.
• It leverages collective variables to bypass expensive high-dimensional simulations and employs physics-informed loss functions for accurate sampling.
• Numerical experiments on systems like the Müller-Brown potential demonstrate reduced computational cost and improved fidelity via Wasserstein metrics.

Free Energy Surface Sampling via Reduced Flow Matching

Introduction and Motivation

Sampling the free energy surface (FES) along collective variables (CVs) is central to statistical physics, molecular simulation, and chemical physics. Traditional approaches require extensive simulations in the full, high-dimensional configuration space—often governed by the Boltzmann distribution—and subsequent projection onto the CV manifold. This methodology is computationally limiting for rare-event systems and complex molecular assemblies, where costly exploration of high-dimensional landscapes is unavoidable. The paper "Free Energy Surface Sampling via Reduced Flow Matching" (2605.00337) introduces FES-FM, a variational framework that trains a dynamical transport map directly in CV space, thereby enabling efficient direct sampling from the FES and eliminating the need for high-dimensional simulations during generation.

Methodological Framework

Flow Matching and Model Reduction

FES-FM is conceptually rooted in recent advances in flow matching for generative modeling and Boltzmann sampling, in particular, methods that construct ODE-based neural transport maps interpolating between a tractable prior and a physical target measure. Whereas classical flow-matching methods operate in the high-dimensional configuration space, FES-FM leverages the CV map $\xi: \mathbb{R}^n \rightarrow \mathbb{R}^d$ ($d \ll n$) to project the dynamics and distributions to the low-dimensional CV manifold. The objective is to learn a time-dependent velocity field $u(y, t)$ specified directly in CV space, driving the evolution of samples from a chosen prior toward the FES measure.

This reduced-space approach expunges the prohibitive cost of high-dimensional dynamics during FES sample generation. At training time, the method still utilizes importance-weighted estimators and non-equilibrium reweighting—aligned with the Jarzynski equality and annealed importance sampling—but the essential innovation is learning generative dynamics for the CV distribution itself.

Training Losses and Implementation

The training protocol proceeds in two stages:

1. Warm-up: A high-dimensional transport map $b_{\theta_0}$ is approximately learned to mitigate variance in Jarzynski estimators, a necessity for accurate expectations under non-equilibrium measures.
2. Main Training: With $b_{\theta_0}$ fixed, three neural networks are fitted: a CV-space velocity field $u_{\theta_1}$, a local mean force estimator $v_{\theta_1}$, and a normalizing constant estimator $c_{\theta_1}$. The loss comprises two components: (i) an $L^2$ regression for the local mean force and (ii) a physics-informed term enforcing the reduced-space Liouville/transport equation. All expectations are estimated via non-equilibrium importance-sampled trajectories, and model reduction is structurally supported by the coarea formula and explicit conditioning on CV manifolds.

Sampling Workflow Comparison

The difference in sampling workflows between FES-FM and high-dimensional methods such as NETS-P is foundational: NETS-P generates configurations in high dimensions and then projects to CVs, whereas FES-FM projects the prior onto CV space first and then evolves dynamics in the reduced space.

Figure 2: FES-FM evolves CV-space dynamics, avoiding high-dimensional simulation during generation, while NETS-P projects high-dimensional samples after simulation.

Prior Construction: Hessian-Informed Harmonic Distribution

A key obstacle in molecular systems is the construction of priors that respect physical symmetries (translation, rotation) and yield meaningful configurations under the CV map. The authors propose a Hessian-informed harmonic prior centered at a local energy minimum and equivariant under $\mathrm{E}(3)$. This prior incorporates the local curvature (Hessian eigenstructure), eliminates unphysical global degrees of freedom, and constrains sampling to meaningful, stable molecular configurations. This approach overcomes the limitations of mean-free Gaussian priors, which produce nonphysical or numerically unstable samples, particularly in many-body systems.

Numerical Results and Empirical Analysis

Benchmark: M\"uller-Brown Potential

The M\"uller-Brown surface is used to validate the capacity to learn nonlinear CVs and FESs. The projected CV distribution from FES-FM accurately reproduces the true FES along the transition path and across metastable states. This demonstrates the method's ability to capture rare-event statistics and complicated nonlinear reductions.

Figure 3: M\"uller-Brown potential, learned CV, and the empirical distribution of FES-FM samples matching the reference distribution.

High-Dimensional Systems

In high-dimensional double-well systems up to $d \ll n$0, FES-FM achieves at least an order-of-magnitude reduction in sample generation time with improved or comparable Wasserstein distance to the reference FES, leading to a substantial improvement in accuracy per unit time compared to high-dimensional flow-matching baselines.

Many-Particle Systems

The method is further evaluated on synthetic molecular systems, including a three-particle system in $d \ll n$1 (angle CV) and a four-particle system in $d \ll n$2 (dihedral angle CV), with physical priors constructed by the Hessian-informed harmonic strategy. FES-FM achieves both lower computational cost and improved sample fidelity, as measured by Wasserstein metrics of the generated CV distribution. Histograms reveal close agreement between FES-FM output and ground truth, whereas projection-based baselines exhibit wider variance.

Figure 1: Comparison of sampled CV distributions for NETS-P and FES-FM in many-particle systems, illustrating the statistical accuracy and stability of FES-FM in low-dimensional CV space.

Theoretical and Practical Implications

FES-FM demonstrates that carefully constructed model reduction—where neural generative models and transport maps are trained directly in physically meaningful CV coordinates—can bypass the immense costs of high-dimensional trajectory simulation. This result challenges the orthodoxy in free energy estimation and rare event simulation, where biasing and enhanced sampling are usually implemented in configuration space. It connects advances in generative modeling (flow matching, normalizing flows, score-based methods) with physical simulation, emphasizing that model reduction is not only computationally efficient but can offer improved accuracy for tasks where the distribution over CVs is the core scientific observable.

The Hessian-informed prior ensures $d \ll n$3 invariance and physically correct support, a critical issue for molecular applications. The framework, being modular, is extensible to realistic all-atom models and force fields, where the separation between CVs and configuration space is even more pronounced.

Limitations and Future Directions

One remaining limitation is the necessity of high-dimensional NETS-based simulation during training, specifically for reweighting via the Jarzynski equality. The ultimate efficiency and scope of the approach would be further enhanced if training could be performed entirely in the reduced space, perhaps by constructing rigorous estimators for conditional expectations on CV manifolds or leveraging new developments in score-based/intrinsic manifold generative modeling.

Another future direction is coupling FES-focused sampling directly with automatic or learned identification of optimal CVs. This could enable fully end-to-end generative frameworks for complex molecular and material systems.

Conclusion

The FES-FM framework establishes a practical and principled approach for direct sampling from free energy surfaces in reduced CV space without full-dimensional simulation during generation. The introduction of Hessian-informed harmonic priors for many-particle systems ensures alignment with physical constraints. The reported numerical results substantiate the claim of superior accuracy per sampling time and lower computational burden compared to conventional high-dimensional approaches. This methodology thus offers a significant advancement in scalable, physically meaningful generative modeling for statistical physics and molecular simulation (2605.00337).