PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yu Xie, Ludwig Winkler, Lixin Sun, Sarah Lewis, Adam E. Foster, José Jiménez Luna, Tim Hempel, Michael Gastegger, Yaoyi Chen, Iryna Zaporozhets, Cecilia Clementi, Christopher M. Bishop, Frank Noé

arXiv:2602.16634v2stat.MLcs.AIcs.LGphysics.bio-phphysics.chem-ph

TL;DR

Molecular diffusion samplers remove slow mixing but still struggle to estimate observables governed by rare equilibrium states. This paper steers pretrained diffusion models to generate biased ensembles and exactly reweights them, achieving fast and accurate rare-event thermodynamic estimation across toy and protein systems.

  • Problem

    Diffusion-model equilibrium samplers remove slow mixing but retain exponential sample complexity for observables dominated by low-probability states such as protein folding and binding.

  • Method

    The framework steers pretrained diffusion models to generate biased ensembles and recovers unbiased thermodynamics through direct reweighting or WHAM/MBAR.

  • Results

    Across toy systems, protein folding landscapes, and folding free energies, enhanced diffusion sampling converges faster, reduces sample requirements by one to several orders of magnitude, and operates within GPU minutes to hours per calculation.

  • Takeaways & Limitations

    The framework combines diffusion sampling with classical bias-and-reweighting methods to address both slow mixing and rare-state sampling within one approach.

  • Takeaways & Limitations

    The estimates depend on a sufficiently accurate pretrained equilibrium model, and model mismatch with physical thermodynamics propagates into reweighted results.

Abstract

from arXiv · show

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), MetaDiff (a batchwise analogue for metadynamics), and $Δ$G-Diff (free-energy differences via tilted ensembles). 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.

I. INTRODUCTION

Molecular simulation faces both slow mixing among long-lived states and exponentially difficult sampling of rare equilibrium states. Diffusion samplers address slow mixing, while enhanced diffusion sampling combines biased exploration with reweighting to target rare-event observables.

  • MD generates equilibrium ensembles and observables, but its effectiveness is limited by slow mixing and rare-state sampling.
  • At 300K, a folding free energy of −10 kcal/mol means only about 1 in 1.9 × 10^7 equilibrium samples is unfolded.
  • Enhanced sampling addresses rare states by sampling biased distributions and reweighting the resulting statistics to recover equilibrium quantities.
  • Diffusion equilibrium samplers remove slow-mixing bottlenecks by generating approximately independent equilibrium configurations, but rare-state estimation remains difficult.
  • High-dimensional biomolecular transitions remain challenging because useful bias coordinates may be unknown and slow relaxation processes can persist outside the biased variables.
  • The paper integrates enhanced sampling with diffusion models, enabling biased-ensemble methods and unbiasing within GPU minutes to hours when a suitable pretrained model is available.

II. BIASING AND UNBIASING DIFFUSION MODEL ENSEMBLES

The framework treats a pretrained diffusion model as an unbiased equilibrium sampler, steers it to generate biased ensembles, and reweights those samples to estimate equilibrium observables.

  • The method assumes a pretrained diffusion model whose output distribution p(x) serves as the unbiased equilibrium distribution.
  • The framework requires only that p(x) and its dimensionless energy representation exist, not that they can be explicitly evaluated.
  • The goal is to estimate unbiased expectations, including unfolded-state probabilities that determine folding free energies and require rare-event sampling.
  • Biased ensembles are generated by steering the diffusion model at inference time under energies u(x) + b_k(x).
  • Known bias potentials are then used to recover unbiased expectations; MBAR performs the reweighting when multiple biased ensembles are used.

A. Biased sampling from diffusion models

The biased-sampling procedure modifies reverse diffusion to target a tilted distribution while tracking weights and resampling particles to control statistical efficiency.

  • Given a pretrained model sampling p(x), the method targets a biased ensemble defined by a biasing potential b(x), without needing its normalization constant.
  • Steering can use score guidance, sequential Monte Carlo with Feynman-Kac potentials, or invariant correctors, with the paper employing Feynman-Kac Corrector sampling.
  • The implementation uses Feynman-Kac Corrector methodology to generate weighted samples from biased diffusion marginals.
  • Biased marginals are constructed by interpolating from a constant initial bias to the target bias through b_t.
  • The Kish effective sample size tracks weighted-batch efficiency, ranging from 1 when one sample carries all weight to n when weights are equal.
  • Particles are periodically resampled, especially when effective sample size falls below a threshold such as n/2, after which weights are refreshed.

B. Unbiasing to the equilibrium distribution

The method combines weighted samples from biased diffusion-model ensembles and reweights them to estimate unbiased equilibrium observables. Weighted MBAR provides asymptotically unbiased, minimum-variance estimates while accounting for steering weights and effective sample sizes.

  • The combined weights incorporate both within-ensemble steering importance weights and the factor that maps each biased ensemble back to the unbiased distribution.After terminal resampling, the steering weights are unity and only the unbiasing factor remains.
  • Weighted samples from multiple biased ensembles are aggregated and directly reweighted to estimate equilibrium expectations of arbitrary observables.The resulting estimator is asymptotically unbiased in the infinite-sample limit.
  • The reweighted samples support folding free energies and potential-of-mean-force profiles through membership functions, weighted histograms, or kernel density estimates.A folding membership function assigns folded states to 1 and unfolded states to 0, with optional intermediate values based on order parameters.
  • Unbiasing multiple biased ensembles with MBAR: MBAR combines biased ensembles into statistically optimal reweighting coefficients without requiring evaluation of log p(x).For K = 2 biased ensembles, MBAR reduces to the traditional Bennett acceptance ratio.
  • Per-sample effective masses are normalized to each ensemble’s effective sample size, allowing MBAR to down-weight windows with highly skewed particle weights.Terminal resampling recovers standard MBAR because all sample weights equal one.

C. Illustration of enhanced diffusion sampling

In an idealized double-well setting, enhanced diffusion sampling uses a bias that equalizes the minima and then reweights the biased samples. It converges faster than direct equilibrium sampling, with the advantage growing for larger free-energy differences.

  • A linear bias is chosen for each double-well potential to equalize the two energy minima in the biased distribution.The tested free-energy differences range from -2 to -14 kBT, with high-energy-state probabilities between 10^-1 and 10^-6.
  • Both direct and steered protocols converge to the exact free-energy difference in the large-sample limit.The steered protocol uses direct reweighting of the biased ensemble.
  • Enhanced diffusion sampling converges faster than equilibrium sampling across all inspected free-energy differences, and the performance gap increases with ΔG.
  • 10 to 100 samples suffice for convergence with enhanced diffusion sampling across all inspected ΔG values, whereas unbiased sampling requires a sample count that grows approximately exponentially with ΔG.Convergence requires both the empirical mean and standard deviation to be within 1 kcal/mol of the exact value across repetitions.
  • The idealized benchmark assumes a single perfectly known bias, while realistic applications require trying multiple biases and combining them with MBAR.The realistic setting therefore needs more samples and has a more complicated relationship between free-energy differences and convergence.

III. UMBRELLADIFF — UMBRELLA SAMPLING WITH DIFFUSION MODELS

UmbrellaDiff applies umbrella sampling to steered diffusion models by generating biased ensembles across a reaction coordinate and removing the bias with weighted MBAR. Its independent sampling avoids slow orthogonal relaxation and kinetic trapping that can hinder traditional MD-based umbrella sampling.

  • Umbrella sampling setup: Umbrella sampling estimates a marginal density and free-energy profile along a user-chosen reaction coordinate ξ.The coordinate may be one-dimensional, such as protein end-to-end distance or fraction of native contacts, with multidimensional extensions described as straightforward.
  • Umbrella sampling setup: UmbrellaDiff uses K biased ensembles centered at prescribed ξ values, commonly with harmonic restraints of stiffness κk.Each ensemble is generated by diffusion-model steering and later unbiased with weighted MBAR.
  • Bias design and diagnostics: Neighboring umbrella centers are typically spaced within one or two σ to obtain approximately 10–30% overlap along ξ.The bias parameters can be adapted using pilot samples and empirical variances to equalize window widths.
  • Bias design and diagnostics: UmbrellaDiff should monitor effective sample sizes, neighboring-window overlap, connected overlap matrices, MBAR uncertainty estimates, and stability under increasing sample size.These diagnostics support adaptive changes to window centers, force constants, sampling allocation, or bridging windows.
  • UmbrellaDiff versus MD umbrella sampling: Unlike conventional MD umbrella sampling, UmbrellaDiff independently generates steered samples from noisy states, avoiding the need to equilibrate slow orthogonal degrees of freedom between windows.
  • UmbrellaDiff versus MD umbrella sampling: UmbrellaDiff avoids kinetic trapping from hidden orthogonal barriers because diffusion-model iid sampling prevents slow intra-window relaxation.The method can recover unbiased free-energy profiles when off-path metastable states contribute substantially, whereas traditional umbrella sampling can show hysteresis and slow convergence.

IV. METADIFF — METADYNAMICS WITH DIFFUSION MODELS

MetaDiff adapts metadynamics to steered diffusion models by updating a bias from batched, independently sampled biased ensembles. Because each update targets a well-defined thermodynamic state, weighted MBAR can provide online unbiased free-energy estimates.

  • Method: MetaDiff samples iid weighted configurations from the equilibrium distribution of the current bias before each update.Unlike biased MD trajectories, each steering invocation targets the current biased ensemble directly.
  • Method: Batch size 1 recovers standard metadynamics applied to steered diffusion sampling.The batchwise update is therefore a direct generalization of single-hill metadynamics.
  • Method: MetaDiff deposits a Gaussian kernel per sample, distributing a fixed per-iteration bias mass across the batch.The bias remains a sum of Gaussians, while the total deposited mass is independent of batch size.
  • Theory: At stationarity, the well-tempered bias satisfies u_Ξ(ξ) = −(γ/(γ −1))b_∞(ξ)+C.For γ = ∞, the update reduces to standard metadynamics; typical γ values are approximately 5–20.
  • Illustration: 1.77·10^4-sample batches explored the second minimum after the second batch and immediately produced an accurate free-energy estimate for a −13.8 kBT double well.The unbiased diffusion sampler using 1.11·10^5 samples failed to sample the second minimum.

V. ∆G-DIFF: FREE ENERGY DIFFERENCES BETWEEN TWO STATES

ΔG-Diff estimates free-energy differences between two states by steering diffusion samplers through tilted ensembles and combining the resulting weighted samples with MBAR. In protein folding tests, steering substantially reduced sample-complexity growth across stable proteins, though applicability depended on usable overlap and reliable references.

  • Method: ΔG-Diff estimates ΔG_AB between configurational states A and B using a progress coordinate ξ(x) ∈ [0,1].The coordinate indicates how far a configuration has advanced from state A to state B.
  • Tilted ensembles: Linear tilts favor state A for a > 0 and state B for a < 0, creating ensembles that improve sampling of both states.The method uses centered linear tilts in reduced units.
  • Algorithm: ΔG-Diff adds tilt values until both endpoints dominate and neighboring ensembles meet an overlap threshold, then refines samples to a target ESS.Weighted MBAR subsequently unbiases all tilts to estimate the target a = 0 ensemble.
  • Method: A single tilted ensemble can suffice when a good guess of ΔG is available, and unlike classical MD-based enhanced sampling, ΔG-Diff does not require overlapping distributions in coordinate space.The steered ensembles are combined with MBAR to estimate the free-energy value.
  • Protein folding: 18 proteins remained after excluding 8 systems, and approximately 1,000 particles produced close agreement between enhanced-sampling and converged unbiased ΔG estimates.Excluded systems included five with RMSD ranges beyond 2 nm and three with unreliable unconverged references.
  • Protein folding: Steered sampling showed much weaker sample-count scaling than unbiased sampling as folding stability increased, demonstrating reduced sample complexity for stable proteins.Unbiased sampling required exponentially more samples with ΔG, whereas steered sampling scaled much more weakly over the tested range.
  • Conclusion: The study reports that UmbrellaDiff and MetaDiff reduce sample requirements by one to several orders of magnitude and make protein folding free energies scale sub-exponentially in ΔG.These conclusions concern the reported toy-system and protein experiments.

VI. DISCUSSION

The framework transfers bias-and-reweighting enhanced sampling to diffusion samplers, addressing rare-state estimation after diffusion models remove slow mixing. Its scope includes equilibrium free energies and broader rare-event observables, while limitations include model mismatch, overlap failure, integrator dependence, and the current focus on equilibrium properties.

  • Core contribution: Diffusion samplers address slow mixing, while steering pretrained models and exact reweighting address rare-state sampling for converged free-energy estimation.The framework supports direct reweighting, WHAM, and MBAR.
  • Implications for enhanced sampling: Independent samples remove dynamical trapping within windows, leaving overlap as a statistical requirement diagnosable with MBAR/WHAM tools.Bias variables therefore need not correspond to slow dynamical modes.
  • Limitations: Reweighted estimates inherit pretrained-model mismatch, and insufficient overlap or aggressive bias can cause weight degeneracy and degraded effective sample sizes.Five of 26 candidate proteins were excluded because steering entered low-density regions and produced extreme importance weights.
  • Limitations: Reported wall-clock costs depend quantitatively on the integrator and step count because FKC steering requires SDE integrators.SDE schemes typically require more denoising steps than ODE schemes for equivalent fidelity.
  • Limitations and future directions: The present work focuses on equilibrium properties; dynamical observables require additional structure such as path reweighting or consistent generative dynamical models.Related transport models may have stationary distributions that do not coincide with the desired equilibrium distribution without corrections.
  • Scope: The formulation is presented as applicable beyond biomolecular folding to materials, soft matter, and condensed-phase chemistry when iid equilibrium samplers exist.The supported target is rare-event statistics in systems where equilibrium sampling is otherwise available.

Appendix A: BioEmu Folding Free Energies Details

The protein benchmark evaluates steered diffusion sampling on filtered ProThermDB systems, showing faster and more reliable folding-free-energy estimation than unbiased sampling, especially for stable proteins.

  • Sampling protocol: Steering used a clamped linear collective-variable potential with a per-system estimated optimal slope and independent batches pooled for analysis.Clipping bounds prevent the potential from becoming unbounded.
  • Dataset and filtering: 26 ProThermDB proteins were selected, but 8 were excluded because of difficult steering or unreliable unfolded-state reference estimates, leaving 18 systems.The retained systems span predicted unbiased |∆Gfold| values from 1.1 to 5.7 kcal/mol.
  • Accuracy: Steered sampling achieves sub-kcal/mol accuracy at n ≈100-1,000 depending on protein stability.Mean absolute error is measured relative to each method’s own converged pooled estimate.
  • Scaling with stability: Steered sampling outperforms unsteered sampling at ∆G ≈3.2 kcal/mol, with the advantage increasing for larger ∆G.Unbiased sampling requires exponentially more samples as ∆G increases.
  • Rare-state coverage: Steered sampling eliminates catastrophic failures at small sample sizes across the benchmark, whereas unbiased sampling needs orders of magnitude more particles for very stable proteins.Catastrophic failure is defined as drawing no unfolded sample.
Loading 2602.16634v2…