Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

FEAT: Free energy Estimators with Adaptive Transport

Jiajun He, Yuanqi Du, Francisco Vargas, Yuanqing Wang, Carla P. Gomes, José Miguel Hernández-Lobato, Eric Vanden-Eijnden

TL;DR

This work addresses the persistent challenge of estimating free energy differences by introducing FEAT, which learns adaptive non-equilibrium transports between endpoint distributions using stochastic interpolants. FEAT leverages escorted Jarzynski equality and controlled Crooks theorem to produce consistent, minimum-variance estimators, along with variational upper and lower bounds on ΔF, while unifying equilibrium and non-equilibrium methods under a single framework. The approach relies on learned transport fields and energy interpolants, supports both forward and backward estimators via FB RND, and includes a one-sided variant for absolute free energies. Empirical results across toy models, molecular systems, and quantum-field–theoretic settings show FEAT surpassing targeted FEP baselines, with favorable runtimes and robustness to discretization; the work also provides connections to TI, BAR, and MBAR, highlighting FEAT’s flexibility and potential impact in computational physics, chemistry, and biology.

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. Our PyTorch implementation is available at https://github.com/jiajunhe98/FEAT.

FEAT: Free energy Estimators with Adaptive Transport

TL;DR

This work addresses the persistent challenge of estimating free energy differences by introducing FEAT, which learns adaptive non-equilibrium transports between endpoint distributions using stochastic interpolants. FEAT leverages escorted Jarzynski equality and controlled Crooks theorem to produce consistent, minimum-variance estimators, along with variational upper and lower bounds on ΔF, while unifying equilibrium and non-equilibrium methods under a single framework. The approach relies on learned transport fields and energy interpolants, supports both forward and backward estimators via FB RND, and includes a one-sided variant for absolute free energies. Empirical results across toy models, molecular systems, and quantum-field–theoretic settings show FEAT surpassing targeted FEP baselines, with favorable runtimes and robustness to discretization; the work also provides connections to TI, BAR, and MBAR, highlighting FEAT’s flexibility and potential impact in computational physics, chemistry, and biology.

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. Our PyTorch implementation is available at https://github.com/jiajunhe98/FEAT.

Paper Structure

This paper contains 37 sections, 3 theorems, 104 equations, 4 figures, 8 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background and Motivation
  3. Free energy perturbation and Bennett acceptance ratio
  4. Jarzynski equality, Crooks fluctuation theorem, and their escorted variations
  5. Methods
  6. Learning a Transport with Stochastic Interpolants
  7. Estimating the Free Energy Difference with Escorted Jarzynski Equality
  8. Minimizing Variance with Non-equilibrium Bennett Acceptance Ratio
  9. Improving Accuracy and Efficiency of Free Energy Estimation with FB RND
  10. Limiting Cases of FEAT with Connections to Other Approaches
  11. One-sided FEAT
  12. Experiments
  13. Conclusions and Limitations
  14. Algorithm
  15. Extended review of free energy estimation methods
  16. ...and 22 more sections

Key Result

Corollary 3.1

Given $v^\psi_t$ and $U_t^\theta$, consider the forward and backward SDEs: where $\sigma_t\ge0$ and $\mu_a$ and $\mu_b$ denote the distributions associated with the energies $U_a$ and $U_b$, respectively. Define also the "corrected generalized work": Using the generalized work with correction, we have the same escorted Jarzynski equality as before: where $\hbox{\ialign{ \crcr \rightarrowfill@{

Figures (4)

  • Figure 1: Connection between our proposed algorithm and other free-energy estimation approaches.
  • Figure 2: (a) Two states of ALDP-T. $S_a$: $\phi \in (0, 2.15)$; $S_b$: $\phi \in [-\pi, 0]\cup[0, \pi)$; (b) (c) Estimators with \ref{['eq:jarzynski_finite', 'eq:bar_noneq']} and their dynamics along training in LJ-128 and ALDP-S.
  • Figure 3: Eight example lattice configurations. We can see that the lattices are either mostly positive or negative.
  • Figure 4: Umbrella samples and reweighted histogram. We denote the two umbrellas as $u_a$ and $u_b$.

Theorems & Definitions (8)

  • Corollary 3.1: Escorted Jarzynski Equality with imperfect boundary conditions
  • Proposition 3.2: Minimum variance non-equilibrium free energy estimator
  • Proposition 3.3: Discretized controlled Crooks theorem with imperfect boundary conditions
  • proof
  • proof
  • proof
  • proof
  • proof