FEAT: Free energy Estimators with Adaptive Transport

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

arXiv.org Machine Learning 

FEA T leverages learned transports implemented via stochastic interpolants and provides consistent, minimum-variance estimators based on escorted Jarzyn-ski 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, FEA T 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. 1 Introduction Estimating free energy is fundamental across machine learning (appearing as normalization factors and the model evidence), statistical mechanics (partition functions), chemistry, and biology (Chipot and Pohorille, 2007; Leli ` evre et al., 2010; Tuckerman, 2023). The free energy is expressed as: F = k BT log Z, Z = null Ωexp( βU (x))dx (1) where Ω R d, U: Ω R is the energy function, assumed to be such that Z <, and β = 1 /k BT combines the Boltzmann constant k B and temperature T . Rather than calculating F directly, one typically estimates the free energy difference between systems (or states) S a and S b with energies U a and U b, which is essential for biological conformational changes, ligand-macromolecule binding, and chemical reaction mechanisms (Wang et al., 2015): F = F b F a = k BT log Z b Z a (2) This computational challenge has driven numerous approaches. Zwanzig (1954) reformulated the problem as importance sampling, where one system serves as the proposal, enabling free energy difference estimation via Monte Carlo sampling. This free energy perturbation (FEP) method, however, suffers from high variance when the energies U a and U b of systems S a and S b differ significantly, particularly in high-dimensional spaces. The authors contributed equally to this work. The order is randomly assigned and will be randomly reshuffled in each version of the paper to reflect this equal contribution.