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

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.

• The paper introduces a novel framework that embeds bias potentials into the reverse diffusion process for efficient rare event sampling.
• The methodology adapts umbrella sampling, metadynamics, and free energy difference estimation within diffusion models to achieve order-of-magnitude sample reductions.
• Numerical benchmarks on biomolecular systems demonstrate reliable reconstruction of free energy landscapes even in high-stability proteins where unbiased methods fail.

Enhanced Diffusion Sampling with Diffusion Models for Rare Event and Free Energy Calculations

Introduction and Problem Statement

The challenge of rare-event sampling is a central bottleneck in molecular simulation, particularly in the context of biomolecular dynamics, where thermodynamically relevant observables—such as folding free energies or binding affinities—are dominated by low-probability states that are rarely visited in unbiased trajectories. Traditional MD-based enhanced sampling methods introduce bias potentials to circumvent the slow-mixing problem but require careful biasing along presumed reaction coordinates and may remain limited by dynamical trapping in orthogonal slow modes. Recently, diffusion models (DMs), especially in the form of models like BioEmu, have emerged as equilibrium samplers capable of generating approximately independent (i.i.d.) samples, thus eliminating correlation-induced slow mixing. However, even with efficient decorrelation, the estimation of rare state observables remains exponentially costly for large free energy gaps.

The paper "Enhanced Diffusion Sampling: Efficient Rare Event Sampling and Free Energy Calculation with Diffusion Models" (2602.16634) addresses this outstanding sample-complexity problem by extending the classical bias-and-reweighting paradigm of enhanced sampling to pretrained diffusion models. The approach enables efficient biased exploration of configuration space and exact recovery of equilibrium observables through reweighting. The framework is instantiated through three core methodologies: UmbrellaDiff (umbrella sampling), MetaDiff (metadynamics), and $\Delta$G-Diff (free energy difference estimation), each adapted to the diffusion model context.

Methodological Framework

Enhanced diffusion sampling generalizes classical enhanced sampling methods—typically built around biased MD trajectories and reweighting—to diffusion model-based generative samplers. The procedure consists of two stages: (1) generating samples from biased ensembles via controlled steering during diffusion model inference, and (2) unbiasing these samples, either by direct reweighting (single-window) or statistically optimal estimators like MBAR (multi-window), to yield unbiased thermodynamic estimates.

The key innovation is embedding bias potentials into the reverse (denoising) process of the pretrained diffusion model. Bias is introduced via score modification or, more formally, through the Feynman-Kac Corrector (FKC)—a stochastic differential equation approach that augments the reverse process with a bias-dependent drift term and computes an associated importance weight trajectory for each sampled path. This produces efficient sampling from the biased distribution even when the exact density of the generative model is intractable.

The importance weights are combined post hoc across one or more biased ensembles via MBAR, providing minimum-variance, asymptotically unbiased estimates of equilibrium observables for arbitrary choices of the bias potential.

Figure 1: Enhanced sampling via diffusion models accelerates estimation of rare-event probabilities by introducing bias potentials and exact reweighting; comparison of unbiased and biased sampling shows the dramatic reduction in sample requirements for increasing free energy gap $\Delta G$.

Algorithms: UmbrellaDiff, MetaDiff, and $\Delta$G-Diff

UmbrellaDiff

Analogous to classical umbrella sampling, UmbrellaDiff introduces a set of harmonic or otherwise localized biases along a user-defined coordinate (e.g., fraction of native contacts or rmsd). Each bias (window) steers a batch of generation trajectories, and the biased samples are reweighted via MBAR to reconstruct the potential of mean force (PMF) along the chosen coordinate. Importantly, the iid nature of sampling in this context removes the necessity of kinetic equilibration between windows, circumventing the slow orthogonal mixing that hampers traditional umbrella sampling.

MetaDiff

MetaDiff adapts the metadynamics paradigm to diffusion models; bias "hills" are deposited in collective variable (CV) space based on ensembles generated via steering, typically using Gaussian kernels centered on sample CV values and with adaptive tempering schemes. Each batch defines a thermodynamic state, with biases cumulatively updated across iterations. Samples from each batch are aggregated via MBAR, providing an on-the-fly unbiased PMF estimate. Unlike MD-based metadynamics, which suffers from nonequilibrium bias deposition, MetaDiff produces equilibrium samples at each iteration, permitting real-time uncertainty quantification and early stopping.

Figure 2: PMF reconstruction with MetaDiff in a double-well potential—metadynamics iterations rapidly recover the free energy landscape across large energy gaps inaccessible to unbiased diffusion sampling.

$\Delta$G-Diff

$\Delta$G-Diff formulates free energy differences between two macrostates as a stratification problem over a progress coordinate, employing a series of tilted (linear bias) ensembles to span the entire relevant region (e.g., folded/unfolded). A diagnostic criteria ensures sufficient overlap and dominance of both end-states across the windowed ensembles. The aggregate is reweighted via MBAR to optimally estimate $\Delta G$. Notably, the sample complexity for estimating $\Delta G$ via $\Delta$G-Diff grows sub-exponentially with the gap, in stark contrast to the exponential scaling found with unbiased estimators.

Numerical Results and Sample Efficiency

Quantitative benchmarks in model systems and realistic protein folding scenarios demonstrate order-of-magnitude reductions in required sample counts for rare-event estimation. In particular, for folding free energies exceeding $5$ kcal/mol, unbiased diffusion model sampling would require up to $10^7$ samples—mapping to years of GPU time for moderate proteins—whereas enhanced diffusion sampling closes the gap in minutes to hours.

Figure 3: Catastrophic failure rates versus sample size: for large $\Delta G$, unbiased sampling never yields unfolded states at tractable batch sizes; steered sampling (blue) eliminates catastrophic failures system-wide.

Figure 4: Mean absolute error (MAE) in $\Delta G$ estimation as a function of sample size; the advantage of steered sampling (blue) over unbiased (grey) is pronounced, especially for more stable proteins where rare states are essentially invisible in equilibrium sampling.

Notably, for a panel of 18 ProThermDB proteins, comparison of MAE and catastrophic failure as a function of sample count confirms that steered sampling both converges much faster and is robust for high-stability systems, in which unbiased diffusion model sampling produces zero unfolded samples (no effective estimate) until prohibitively large batch sizes.

Theoretical and Practical Implications

The proposed framework demonstrates that combining independent sampling via generative models with adaptive bias and reweighting constructions from statistical physics enables practical rare-event thermodynamics calculations that are, in principle, extensible to arbitrarily high energy gaps—provided the generative model has sufficient support and fidelity across the relevant configurational space.

A key formal claim: the rare-state bottleneck can be closed (in terms of practical compute) for observables dominated by exponentially suppressed regions, under the assumption of a sufficiently expressive and accurately trained diffusion model.

The methodology relaxes critical constraints of traditional enhanced sampling:

• No requirement for inter-window configurational overlap in phase space, only overlap in (possibly low-dimensional) bias coordinate space.
• No necessity for connected MD trajectories or explicit bridging, as each diffusion model sample is effectively independent.
• Robust handling of orthogonal slow modes and hidden barriers, which are decoupled from the sampling process.

On the other hand, the approach inherits all limitations of importance sampling: inaccuracies in the generative model lead to systematic deviation in equilibrium estimates, and highly aggressive biases can still result in weight degeneracy if there is insufficient overlap.

Outlook and Speculative Directions

The authors suggest several avenues for algorithmic and theoretical extension, including:

• Automated bias optimization and adaptive windowing using MBAR diagnostics.
• Incorporation of direct energy evaluation in energy-based or hybrid DMs for even more flexible biasing and improved estimator properties.
• Extensions to path reweighting and rare-event kinetics by marrying trajectory-sampling generators with bias and reweighting procedures.

The generality of the framework also points toward broad application in chemical physics, materials science, and any domain with a scalable generative equilibrium sampler and rare-event observables of interest.

Conclusion

Enhanced diffusion sampling, as articulated in this work, realizes the long-standing goal of making rare-event thermodynamics tractable for complex, high-dimensional systems by merging modern generative modeling with the mature framework of enhanced sampling. The practical outcome is a drastic reduction in sample complexity for rare-state estimation at no cost to equilibrium correctness—provided sufficient model accuracy and careful biasing. This paradigm is likely to become a central component of future simulation-based workflows for biophysics, chemistry, and beyond.