Fokker-Planck Score Learning: Efficient Free-Energy Estimation under Periodic Boundary Conditions

TL;DR

The paper tackles the challenge of efficiently estimating free-energy profiles under periodic boundary conditions in molecular simulations. It introduces Fokker-Planck Score Learning, a physics-informed score-based diffusion framework that maps PBC simulations to a Brownian particle in a periodic potential and learns the nonequilibrium steady-state score to reconstruct the PMF. By leveraging the analytic non-equilibrium steady-state solution and enforcing periodicity, the method achieves faster convergence and lower variance than traditional umbrella-sampling approaches, as demonstrated on a 1D toy model and a lipid bilayer permeation system. Key ideas include energy-based diffusion models to ensure a conservative score and Fourier-feature periodic parametrization to simplify training, with potential extensions to higher dimensions and atomistic simulations. Overall, the approach provides a principled, data-efficient route to PMF reconstruction under periodic boundaries, with practical impact for membrane permeation, ligand binding, and materials design.

Abstract

Accurate free-energy estimation is essential in molecular simulation, yet the periodic boundary conditions (PBC) commonly used in computer simulations have rarely been explicitly exploited. Equilibrium methods such as umbrella sampling, metadynamics, and adaptive biasing force require extensive sampling, while non-equilibrium pulling with Jarzynski's equality suffers from poor convergence due to exponential averaging. Here, we introduce a physics-informed, score-based diffusion framework: by mapping PBC simulations onto a Brownian particle in a periodic potential, we derive the Fokker-Planck steady-state score that directly encodes free-energy gradients. A neural network is trained on non-equilibrium trajectories to learn this score, providing a principled scheme to efficiently reconstruct the potential of mean force (PMF). On benchmark periodic potentials and small-molecule membrane permeation, our method is up to one order of magnitude more efficient than umbrella sampling.

Figures (4)

• Figure 1: We consider a computer simulation with periodic boundary conditions, where the red particle is pulled by a constant external driving force $f$ through a conservative potential or potential of mean force, $U(x)$. The steady-state solution of the Fokker--Planck equation of a Brownian particle in a periodic potential, $p^\mathrm{ss}$, informs the score of our diffusion model. The diffusion model interpolates between the original non-equilibrium system at constant flux $J$ and a trivial system at rest, shown left to right, respectively. Denoising of the diffusion model allows us to efficiently reconstruct $U(x)$ by exploiting the structure of $p^\mathrm{ss}$.
• Figure 2: Learning the energy landscape of a driven 1D Langevin toy model. (a) Reference free energy profile. (b) Non-equilibrium steady-state distributions $p^\mathrm{ss}$ of the toy model under various constant driving forces $f$. (c) Inferred free energy landscapes $\hat{U}_\theta$ obtained from the diffusion model, averaged over twenty independent trials for each biasing force; shaded regions denote $\pm$ standard deviation. (d) Mean absolute error (MAE) between the inferred and reference free energy profiles as a function of the driving force $f$. Shaded regions denote $\pm$ standard deviation over 20 independent trials. (e--h) Interpolation between latent space ($\tau=1$) and data space ($\tau=0$) in case biasing force $f=13$, illustrating: (e, g) Training objective $\nabla U^\theta_\text{eff}(x_\tau, \tau)$ and learned equilibrium score $\nabla U^\theta_\tau(x_\tau)$, respectively. (f, h) Learned negative log-density $U^\theta_\text{eff}(x_\tau, \tau)$ and $U^\theta_\tau(x_\tau)$, respectively.
• Figure 3: Lipid bilayer system and its characterization. (a) Illustration of the simulated system, depicting a coarse-grained Martini POPC lipid bilayer with a permeating solute (N1 bead type). (b) Reference free-energy profile for N1 bead permeation, determined with MBAR via high-resolution umbrella sampling with a window spacing of $\delta z = 0.02nm$. (c) Position-dependent diffusion coefficient, $D(z)$, of the N1 bead along the bilayer normal, derived from molecular dynamics simulations with $f=4kJ\per mol\per nm$.
• Figure 4: Performance evaluation of Fokker--Planck Score Learning for reconstructing the free-energy profile of N1 bead permeation through a POPC lipid bilayer. (a) Mean absolute error (MAE) versus computational wall time for Fokker--Planck Score Learning (green) and umbrella sampling with MBAR (orange). (b--d) Comparison of reconstructed free-energy profiles. (b) Fokker--Planck Score Learning using approximately $128.9ns$ of MD simulation data. (c) MBAR with comparable wall time to (b). (d) MBAR with a comparable MAE to (b). (e--g) Similar comparison as (b--d), but with Fokker--Planck Score Learning utilizing approximately $325ns$ of MD simulation data. (e) Fokker--Planck Score Learning. (f) MBAR with comparable wall time to (e). (g) MBAR with a comparable MAE to (e). Shaded regions indicate $\pm$ standard deviations.