Enhanced Diffusion Sampling: Efficient Rare Event Sampling and Free Energy Calculation with Diffusion Models

TL;DR

The paper tackles the dual challenge of slow mixing and rare-state estimation in molecular dynamics by marrying diffusion-model samplers with classical enhanced-sampling biasing and exact reweighting. It introduces three algorithms—UmbrellaDiff, ΔG-Diff, and MetaDiff—that steer pretrained diffusion models to biased ensembles and recover unbiased thermodynamics via MBAR/WHAM, enabling accurate folding free-energy calculations and rare-event statistics on GPU-scale runtimes. Key contributions include a general steering framework for diffusion models, novel implementations of umbrella sampling and metadynamics in this context, and demonstrated efficiency gains for protein folding systems using BioEmu. This work provides a practical pathway to routinely compute rare-state observables and free energies with diffusion-model samplers, broadening the applicability of data-driven equilibrium ensembles in biomolecular modeling.

Abstract

The rare-event sampling problem has long been the central limiting factor in molecular dynamics (MD), especially in biomolecular simulation. Recently, diffusion models such as BioEmu have emerged as powerful equilibrium samplers that generate independent samples from complex molecular distributions, eliminating the cost of sampling rare transition events. However, a sampling problem remains when computing observables that rely on states which are rare in equilibrium, for example folding free energies. Here, we introduce enhanced diffusion sampling, enabling efficient exploration of rare-event regions while preserving unbiased thermodynamic estimators. The key idea is to perform quantitatively accurate steering protocols to generate biased ensembles and subsequently recover equilibrium statistics via exact reweighting. We instantiate our framework in three algorithms: UmbrellaDiff (umbrella sampling with diffusion models), $Δ$G-Diff (free-energy differences via tilted ensembles), and MetaDiff (a batchwise analogue for metadynamics). Across toy systems, protein folding landscapes and folding free energies, our methods achieve fast, accurate, and scalable estimation of equilibrium properties within GPU-minutes to hours per system -- closing the rare-event sampling gap that remained after the advent of diffusion-model equilibrium samplers.

Figures (6)

• Figure 1: Enhanced sampling with diffusion models efficiently computes probabilities of rare eventsa) Unbiased diffusion model probability density $p(x)$ and corresponding energy $-\ln p(x)$. In this example the free energy difference between the two minima is set to $\Delta G = -7\,k_BT$. b) Biased density $p(x) \mathrm{e}^{-b(x)}$ using a linear tilt $b(x) = \beta x$ with $\beta$ chosen to achieve $\Delta G_{\mathrm{biased}} = 0\,k_BT$. c) Estimates of $\Delta G$ with an unbiased diffusion model and enhanced sampling using a tilting potential that sets $\Delta G_{\mathrm{biased}} = 0\,k_BT$. Mean (solid line) and 95% confidence interval (shaded area) are shown. d) Number of samples required to get an estimate within 1 kcal/mol of the exact $\Delta G$ value and with a standard deviation of 1 kcal/mol, as a function of $\Delta G$.
• Figure 2: UmbrellaDiff bypasses kinetic trap problem in traditional umbrella sampling. a) Simple two-state 2D potential. Umbrella sampling is performed using 8 umbrellas that bias $\xi = x$ using regular umbrella sampling with Langevin dynamics (orange) and UmbrellaDiff (blue). Contours show 95% of the sample density of individual umbrellas. b) Free energy as a function of $x$ using traditional umbrella sampling (orange) and UmbrellaDiff (blue). Both methods approximate the true potential of mean force (grey). c) As a) but with a third high-probability state (top right). Traditional umbrella sampling (orange) does not sample the third state as it is off-path and long simulation trajectories in each umbrella would be required to sample this state. UmbrellaDiff (blue) samples the equilibrium distribution conditioned on the umbrella potential and therefore samples the third state without issues. d) As b). Traditional umbrella sampling estimates a wrong free energy due to failing to sample the third state (orange), UmbrellaDiff estimates the correct free energy (blue).
• Figure 3: MetaDiff illustration on 1-dimensional double-well potential. Estimated PMF using iterations of metadynamics at a free energy difference of $\Delta G=-13.8 k_BT$.
• Figure 4: $\Delta$G-Diff illustration and protein folding results with BioEmu: a) Processes with large free energy differences (top left, black), e.g., protein folding, sample almost exclusively one state in equilibrium (top left, blue). $\Delta$G-Diff creates a series of tilted biased potentials until both states are dominantly sampled in at least one ensemble each (top left, bottom right), and these ensembles can be combined with MBAR to yield an estimate of $\Delta$G. b) Comparison of folding free energies $\Delta$G sampled by unsteered BioEmu and $\Delta$G-Diff with single optimal tilt. c) Sampling efficiency as a function of $\Delta$G: unsteered sampling (grey) requires exponentially many samples in $\Delta$G, while the number of samples required by $\Delta$G-Diff shows much weaker scaling with $\Delta$G. d) Sample folded and unfolded state of a stable protein domain of fibronectin (PDB code 1ttg).
• Figure 5: Per-system catastrophic failure rate as a function of sample size for 18 ProThermDB proteins, sorted by reference $\Delta G$. Catastrophic failure is defined as the probability that a random draw of $n$ particles contains no unfolded sample (FNC $< 0.5$). Blue: steered sampling; grey: unbiased sampling. Steered sampling eliminates catastrophic failures at small sample sizes across all systems, while unbiased sampling requires orders of magnitude more particles for very stable proteins.