Neural Thermodynamic Integration: Free Energies from Energy-based Diffusion Models

TL;DR

Neural TI tackles the TI bottleneck by learning a time-dependent Hamiltonian $U_t^\theta$ that interpolates between interacting and non-interacting systems, and by employing a denoising diffusion model to sample all intermediate ensembles. It formulates TI along the diffusion time with a score-based energy function, enabling free-energy differences to be estimated from a single reference calculation via $\beta\Delta F = -\int_0^1 \mathrm{d}t\, \langle \partial_t U_t^\theta\rangle_t$. Demonstrated on Lennard-Jones fluids, Neural TI yields accurate excess chemical potentials and continuous $\mu$–density relations across gas–liquid transitions, capturing large free-energy changes up to about $200\,k_B T$. The approach offers a transferable, scalable framework for TI in molecular systems, with future extensions to more complex intramolecular and electrostatic interactions through suitable energy-function design. Overall, energy-based diffusion with TI provides a practical alternative to traditional TI and flow-based methods for estimating free energies in condensed-phase systems.

Abstract

Thermodynamic integration (TI) offers a rigorous method for estimating free-energy differences by integrating over a sequence of interpolating conformational ensembles. However, TI calculations are computationally expensive and typically limited to coupling a small number of degrees of freedom due to the need to sample numerous intermediate ensembles with sufficient conformational-space overlap. In this work, we propose to perform TI along an alchemical pathway represented by a trainable neural network, which we term Neural TI. Critically, we parametrize a time-dependent Hamiltonian interpolating between the interacting and non-interacting systems, and optimize its gradient using a score matching objective. The ability of the resulting energy-based diffusion model to sample all intermediate ensembles allows us to perform TI from a single reference calculation. We apply our method to Lennard-Jones fluids, where we report accurate calculations of the excess chemical potential, demonstrating that Neural TI reproduces the underlying changes in free energy without the need for simulations at interpolating Hamiltonians.

Figures (9)

• Figure 1: Schematic summary of the proposed approach. We interpolate between the target, $\mathcal{H}_0$, and latent, $\mathcal{H}_1$, Hamiltonians with a time-dependent potential $U_t^\theta$. During sampling, the normalizing constant of the target can be estimated via thermodynamic integration. Particles whose separation from their closest neighbor is less than $0.85\sigma$ are colored red, the rest are colored blue. This color-coding illustrates that in the ideal gas (right) there are many colliding particles and as the LJ potential is turned on (left) the particles do not overlap anymore.
• Figure 2: The soft-core LJ potential $U_t^\text{LJ}$ in Equation \ref{['eq:soft_LJ']} for various values of $t \in [0,1]$. Note that for larger values of $t$, particles can get closer to each other without experiencing strong repulsive forces. This is necessary since in the diffusion process it is inevitable that particles get close to each other as $t$ increases.
• Figure 3: Radial distribution functions as predicted by Monte Carlo simulations (gray) and a diffusion model (red) trained on densities $\rho \in \{0.19,0.37,0.56,0.74,0.93\}$. Note that the model reconstructs $g(r)$ across the the gas-liquid phase transition.
• Figure 4: Distribution of the number of particles in the grand canonical ensemble at different chemical potentials from GCMC simulations (gray) and estimated by thermodynamic integration with a trained diffusion model (red). The vertical blue lines denote the canonical ensembles that the diffusion model was trained on.
• Figure 5: Expected density as a function of the chemical potential (left) and estimates of $\mu_\text{ex}$ as a function of the expected density (right). The left plot suggests a that a gas-liquid phase transition takes place at $\beta\mu \approx -28$. The blue lines denote the canonical ensembles that the diffusion model was trained on.