FEAT: Free energy Estimators with Adaptive Transport (2504.11516v1)

Abstract: We present Free energy Estimators with Adaptive Transport (FEAT), a novel framework for free energy estimation -- a critical challenge across scientific domains. FEAT leverages learned transports implemented via stochastic interpolants and provides consistent, minimum-variance estimators based on escorted Jarzynski equality and controlled Crooks theorem, alongside variational upper and lower bounds on free energy differences. Unifying equilibrium and non-equilibrium methods under a single theoretical framework, FEAT establishes a principled foundation for neural free energy calculations. Experimental validation on toy examples, molecular simulations, and quantum field theory demonstrates improvements over existing learning-based methods.

Overview of "FEAT: Free energy Estimators with Adaptive Transport"

The paper introduces Free energy Estimators with Adaptive Transport (FEAT), a novel methodology for estimating free energy differences using machine learning. Free energy estimation is a fundamental problem across disciplines like chemistry, biology, and statistical mechanics. Existing methods, while powerful, often come with constraints such as the need for high overlap in distributions in equilibrium methods, or the challenge of navigating between non-equilibrium trajectories.

Key Contributions:

1. Unified Framework: FEAT provides a unified approach combining equilibrium-based methods and non-equilibrium trajectories, utilizing neural networks to model stochastic interpolants. These interpolants facilitate efficient estimation of free energy differences by exploiting both equilibrium and non-equilibrium conditions.
2. Use of Neural Networks: Neural networks are employed to learn mappings between states in a system. This approach significantly reduces variance in free energy estimation by exploiting recent advancements in neural samplers, particularly those involving flow-based transformations and score matching.
3. Robustness Against Errors: FEAT incorporates mechanisms to correct errors caused by imperfect boundary conditions. This allows it to perform well even when the learned energy functions do not exactly match the true energy landscapes at the endpoints.
4. Minimum-Variance Estimator: The paper extends the Bennett acceptance ratio to non-equilibrium scenarios, providing a minimum-variance estimator that yields high accuracy in free energy differences. This addresses a common limitation in existing methodologies like Thermodynamic Integration (TI), Free Energy Perturbation (FEP), and Targeted Functional Perturbation (TFEP).
5. Direct Connection to Crooks and Jarzynski Equalities: FEAT leverages the escorted Jarzynski equality, providing a theoretical grounding for the estimation process, with direct connections demonstrated to the distinguished fluctuation theorems in statistical mechanics.
6. Experimental Validation: Demonstrated across systems of increasing complexity, FEAT shows superior performance to both MBAR (multi-state Bennett acceptance ratio) and recent deep learning approaches like Targeted FEP with Flow Matching, Neural TI, and classical methods under challenging conditions (e.g., Lennard-Jones fluids, alanine di-peptide in solvent).

Architectural Approach:

• Stochastic Interpolants: FEAT constructs interpolants that smoothly transition between samples of two states. Using parameters such as $\alpha_t$ and $\beta_t$, these interpolants provide smoother and more accurate mappings.
• Neural Network Integration: Equipped with adaptive capabilities, FEAT neural networks learn mappings that reduce variance and respect equilibrium properties effectively.

Future Work:

Potential improvements in FEAT could focus on adapting this framework to broader problems, including multi-state free energy calculations or explicit solvent systems where the state space is significantly larger. Integration with transferable networks, analogous to works by \citet{klein2024transferable}, offers exciting prospects for improving its generalizability.

Conclusion and Implications:

FEAT marks a significant step in bridging equilibrium and non-equilibrium free energy calculation methods. It enhances computational efficiency, reduces estimation variance and increases robustness to errors, challenging and potentially surpassing classical simulation methods. As the field of machine learning continues to advance, so too does its potential to solve traditional problems in physical sciences, with FEAT being a promising intersection of these disciplines.