Source-linked AI summary
Iterated Denoising Energy Matching for Sampling from Boltzmann Densities
Tara Akhound-Sadegh, Jarrid Rector-Brooks, Avishek Joey Bose, Sarthak Mittal, Pablo Lemos, Cheng-Hao Liu, Marcin Sendera, Siamak Ravanbakhsh, Gauthier Gidel, Yoshua Bengio, Nikolay Malkin, Alexander Tong
TL;DR
Sampling Boltzmann targets is challenging when target samples are scarce and existing methods face scaling or mode-coverage constraints. iDEM iteratively trains a diffusion sampler using an energy-based matching objective and model samples. It performs strongly across benchmarks at lower training and memory cost, and is the first energy-based method reported to scale to LJ-55.
Problem
Boltzmann sampling is difficult when target samples are scarce, while classical methods scale poorly in high-dimensional physical systems and reverse-KL objectives can miss modes.
Source: §1 · p. 2Method
iDEM alternates diffusion-model sampling with a stochastic matching objective on the energy, using model samples to iteratively train the sampler.
Source: §1 · p. 2Results
Takeaways & Limitations
The approach extends energy-based training to the challenging LJ-55 particle system while offering high mode coverage and strong performance across multiple benchmarks.
Source: §6 · p. 9The method assumes the energy and its cheaply computable gradient are known, while the partition function is unavailable.
Source: §3 · p. 3
Abstract
from arXiv · showhide
Efficiently generating statistically independent samples from an unnormalized probability distribution, such as equilibrium samples of many-body systems, is a foundational problem in science. In this paper, we propose Iterated Denoising Energy Matching (iDEM), an iterative algorithm that uses a novel stochastic score matching objective leveraging solely the energy function and its gradient -- and no data samples -- to train a diffusion-based sampler. Specifically, iDEM alternates between (I) sampling regions of high model density from a diffusion-based sampler and (II) using these samples in our stochastic matching objective to further improve the sampler. iDEM is scalable to high dimensions as the inner matching objective, is simulation-free, and requires no MCMC samples. Moreover, by leveraging the fast mode mixing behavior of diffusion, iDEM smooths out the energy landscape enabling efficient exploration and learning of an amortized sampler. We evaluate iDEM on a suite of tasks ranging from standard synthetic energy functions to invariant $n$-body particle systems. We show that the proposed approach achieves state-of-the-art performance on all metrics and trains $2-5\times$ faster, which allows it to be the first method to train using energy on the challenging $55$-particle Lennard-Jones system.
1. Introduction
Sampling Boltzmann distributions is difficult when target samples are scarce and classical or learned alternatives face computational or mode-coverage constraints. iDEM trains a diffusion sampler from the energy and its gradient, achieving strong benchmark performance while scaling to LJ-55.
1. Introduction
- The target is a many-body Boltzmann distribution whose energy is invariant to rotations, reflections, and particle permutations, and whose modes the sampler should cover..Source: §1 · p. 1
- Monte Carlo and molecular-dynamics methods are computationally expensive and scale poorly in high-dimensional physical systems, while Boltzmann-generator efficacy depends on initial samples, flow expressivity, and importance-weight fidelity.Reverse-KL variational objectives can also seek modes rather than explore the full energy landscape.Source: §1 · p. 2
- The motivating gap is scalable sampling from the energy and its gradient while achieving high mode coverage.Source: §1 · p. 2
- iDEM alternates model sampling with simulation-free stochastic regression on the energy, using reverse-SDE samples to amortize sampling and explore the energy landscape.Its objective uses energy and its gradient, and the method can use model samples with an exploratory off-policy scheme rather than requiring target samples.Source: §1 · p. 2
- iDEM is competitive with or exceeds prior state-of-the-art approaches and neural sampler baselines, with lower training and memory costs and energy-based training on LJ-55.Experiments include synthetic Gaussian mixtures and invariant double-well and Lennard-Jones particle systems.Source: §1 · p. 2
2. Background and preliminaries
The target is specified through an energy even when its normalized density is unavailable. The section reviews importance sampling and diffusion models, including why standard denoising score matching is not directly tractable for an unsamplable target.
2. Background and preliminaries
- Table 1 compares approaches by whether they are MCMC-free and off-policy, their forward simulation steps during training, and whether they backpropagate through time.The caption gives comparison criteria, not the approaches' outcomes.Source: Table 1 · p. 3
- The target density is proportional to exp(−E(x)); its partition function is generally intractable, so the energy gives access to an unnormalized density rather than an evaluable exact density.Source: §2 · p. 3
2.1. Classical sampling methods
- Importance sampling estimates target expectations using samples from q(x) weighted by exp(−E(x))/q(x), but finding a good proposal is challenging for high-dimensional or separated multimodal targets.The optimal proposal depends roughly on the observable times the target density.Source: §2.1 · p. 3
2.2. Denoising diffusion
- Diffusion models learn the score of a noised distribution for reverse-time generation, but denoising score matching requires samples from p0, which is intractable for the Boltzmann setting considered next.The forward process progressively adds noise to samples from p0.Source: §2.2 · p. 3
3. ITERATED DENOISING ENERGY MATCHING
iDEM addresses the unavailable target score and the lack of informative training locations with a two-level iterative scheme. It combines energy-based score matching with a diffusion sampler that supplies points for subsequent training.
3. ITERATED DENOISING ENERGY MATCHING
- iDEM alternates between energy-based score matching and reverse-SDE sampling to address unavailable target scores and uninformative training locations.The method assumes the energy and its gradient are available, while the normalizing constant and exact target sampling are not.Source: §3 · pp. 3–4
3.1. Denoising diffusion with a Boltzmann target (C1)
- DEM estimates the convolved target score by Monte Carlo and trains a score network using the energy and its gradient, enabling simulation-free, off-policy training.Unlike classical diffusion regression, this approach does not require target samples and can use training points from any distribution.Source: §3.1 · p. 4, Figure 2 · p. 5
- The Monte Carlo score estimator is consistent, with bias decaying as O(1/K), while low-density regions require larger K for accurate estimates.The paper reports that experiments on a Gaussian mixture show decreasing bias and mean squared error as K increases.Source: §3.1 · p. 5
3.2. Amortized sampling with a diffusion sampler (C2)
- The outer loop uses the reverse SDE to generate informative points from a mass-covering prior and stores batches in a replay buffer for inner-loop training.The sampler is fixed during this sampling phase, so gradients are not backpropagated through the SDE solver; repeated inner and outer loops improve the sampler and buffer samples.Source: §3.2 · p. 6
3.3. Incorporating symmetries in iDEM
- For n-body systems, iDEM incorporates SE(3) × S_n symmetry by using zero-center-of-mass Gaussian noise and an equivariant diffusion sampler.The zero-center-of-mass construction handles translation invariance, while the standard normal is already rotation- and permutation-invariant.Source: §3.3 · p. 6
4. Experimental results
The experiments assess iDEM on a 40-mode GMM and three particle potentials, comparing it with sampler baselines. iDEM performs strongly across sample-quality measures and trains faster than FAB; ablations examine estimator scaling, diffusion time, and use of model samples.
4. Experimental results
- The evaluation covers a 40-GMM and DW-4, LJ-13, and LJ-55 particle potentials, with PIS, DDS, FAB, and pDEM as baselines.The benchmarks span dimensions and include equivariant particle systems; PIS and DDS are diffusion-based, while FAB combines AIS samples with flow training.Source: §4 · pp. 6–7
4.1. Main results
- iDEM matches or outperforms all baselines on NLL for GMM, DW-4, and LJ-13, and achieves the best NLL on LJ-55, where PIS and DDS train unstably.It also outperforms considered baselines on W2 and TV; Table 2 reports NLL, ESS, and W2.Source: §4.1 · p. 7, Table 2 · p. 8
- On LJ-13 and LJ-55, iDEM most accurately models the true energy; its separation from FAB is especially large on LJ-55.Figure 4 compares ground-truth energy histograms with energies of generated samples; DDS is omitted from both plots and PIS from LJ-55 because they diverge.Source: §4.1 · p. 7, Figure 4 · p. 7
- iDEM trains about 4× faster than FAB on LJ-13 and LJ-55 and 1.8× faster on GMM and DW-4; it is also faster than PIS and DDS.pDEM takes less time but has training-stability costs: 2 of 3 runs diverged on LJ-55.Source: §4.1 · pp. 7–8, Table 3 · p. 8
4.2. Ablation experiments
- On GMM, increasing K reduces DEM estimator bias and MSE; the bias approaches zero with an empirical asymptotic decay of O(1/K).This supports DEM consistency and empirically validates Proposition 1.Source: §4.2 · p. 8, Figure 5 · p. 8
- On GMM, log-MSE decreases as K increases but rises as diffusion time t approaches 1, where the estimator has higher bias for the same K.The passage attributes the higher bias to diffusion moving farther from the energy modes.Source: §4.2 · p. 8, Figure 5 · p. 8
- In the iDEM ablation against pure off-policy pDEM, samples from sθ improve DW-4 energy histograms, yield only minor improvements on LJ-13, while pDEM fails to learn in 2 of 3 LJ-55 runs.Source: §4.2 · p. 8
5. Related work
Related work spans Monte Carlo and variational approaches, equivariant flows, diffusion-based neural samplers, and generative flow networks. The paper distinguishes iDEM by its simulation-free training objective.
MCMC and variational approximations
- Monte Carlo methods are regarded as sampling gold standards but can be expensive and slow to converge, while reverse-KL variational methods can be mode-seeking.Hybrid flow–MCMC approaches, including FAB and CRAFT, combine these approaches and have shown empirical benefits.Source: §5 · pp. 8–9
Equivariant flows and Boltzmann generators
- Equivariant normalizing flows have been used for Boltzmann sampling, with SE(3)-equivariant generative models spanning robotics, molecular modeling, and protein generation.The cited work also includes flow-based proposal distributions for molecular dynamics.Source: §5 · p. 9
Neural samplers
- Diffusion-based neural samplers exploit fast mode mixing to amortize MCMC, but their objectives require simulation, unlike iDEM.The cited approaches include Berner et al. (2022), Vargas et al. (2023), and Zhang & Chen (2022).Source: §5 · p. 9
GFlowNets
- Continuous generative flow networks support off-policy learning from sampled states or trajectories without differentiating through the simulation that produced them.iDEM is described as simulation-free training for generative flow networks of a particular structure.Source: §5 · p. 9
6. Conclusion
The conclusion presents iDEM as an iterative, amortized diffusion approach to Boltzmann sampling and reports strong benchmark performance and faster training. It also identifies bias and sample variance in the DEM objective as areas for future work.
6. Conclusion
- iDEM uses a novel stochastic matching loss in an inner loop to train a diffusion sampler for amortized sampling from Boltzmann distributions.The sampler proposes informative samples that are used to further accelerate training.Source: §6 · p. 9
- iDEM is reported to offer high mode coverage, state-of-the-art performance on multiple benchmarks, and significantly faster training than previous approaches.The conclusion identifies iDEM as the first approach scalable to energy-based training on LJ-55.Source: §6 · p. 9
- The DEM objective is biased and may be affected by sample variance; reducing variance and speeding up outer-loop SDE simulation are future directions.Source: §6 · p. 9
Impact statement
The impact statement situates amortized Boltzmann sampling in statistical modeling and applications, highlighting molecular design as a motivation and immediate application for iDEM.
Impact statement
- Amortized sampling from Boltzmann densities is relevant to statistical modeling, including high-dimensional Bayesian posteriors, and to applications.Source: Impact statement · p. 9
- The statement highlights molecular design, with possible applications in drug and material discovery, as a motivation and immediate application for iDEM.Source: Impact statement · p. 9
- The authors foresee no immediate negative impacts but encourage due caution to prevent potential misuse.Source: Impact statement · p. 9
Contribution statement
The contribution statement describes how the idea, experiments, theory, writing, and experiment design were distributed among the authors.
Contribution statement
- J.R., N.M., and A.T. initially conceived a stochastic off-policy continuous regression objective, while A.J.B., T.A., G.G., and A.T. independently approached the problem through active inference and a bi-level iterative scheme.Source: Contribution statement · p. 9
- Experiments were primarily led by T.A., J.R., and A.T. for iDEM, and by J.R., S.M., P.L., C.H.L., and M.S. for ablations and baselines.Source: Contribution statement · p. 9
- G.G., A.J.B., and N.M. led theory development; A.J.B. and N.M. drove writing, and all authors contributed to experiment design.Source: Contribution statement · p. 9
A. Proofs of propositions
The appendix establishes score-estimation and equivariance results, extends iDEM to other noising processes, relates its objective to flow matching, and reviews sampling methods and their trade-offs.
A. Proofs of propositions
- Under sub-Gaussian conditions, Proposition 1 bounds the Monte Carlo score estimator’s bias with high probability when K is sufficiently large.The proof estimates the convolved energy score using Gaussian perturbations and a concentration argument.Source: A. Proofs of propositions · p. 14
- For a G-invariant density and Gaussian sampling, Proposition 2 states that the Monte Carlo score estimator is G-equivariant.The proof uses the invariance of the density and equivariance of its gradient.Source: A. Proofs of propositions · p. 15
B. iDEM for non-VE noising processes
- A linear-drift noising SDE can be rescaled to a variance-exploding SDE with an equivalent score; matrix coefficients are also identified as a possible extension.Using the original or rescaled state as the network input is an implementation choice; the paper uses the original state for numerical stability.Source: B. iDEM for non-VE noising processes · p. 15
C. iDEM and flow matching
- iDEM and flow matching express the convolved score differently, while the flow-matching expectation for a Boltzmann target is intractable and prior work approximates it with importance weighting.The authors speculate about directly estimating the regression target at x_t and reusing generated points for simulation-free training.Source: C. iDEM and flow matching · pp. 15–16
D. Sampling with MCMC
- Metropolis-Hastings can converge under detailed balance but may converge slowly in high dimensions, where annealing is often needed for mode coverage.HMC uses energy gradients and is more efficient than MH, but still struggles with mode coverage and is sensitive to tuning; MALA is another gradient-based alternative.Source: §D.1 · p. 16, §D.2 · p. 17, §D.3 · p. 17
D.4. Annealed importance sampling
- Annealed importance sampling combines a sequence of intermediate distributions with invariant Markov transitions to reduce importance-sampling variance while achieving high mode coverage.The sequence is weighted to construct the AIS estimator.Source: §D.4 · p. 17
D.5. Boltzmann generators
- Boltzmann generators need target samples for forward-KL training, whereas SMC progresses through intermediate distributions and can yield high-quality samples.SMC can also estimate the log-partition function from average log-importance weights.Source: §D.5 · p. 17, §D.6 · p. 18
D.7. Nested Sampling
- Nested sampling replaces low-likelihood live points with higher-likelihood prior samples, then reweights the removed points to estimate the partition function and draw samples.The prior volume contracts approximately at each iteration and termination occurs once it is sufficiently small.Source: §D.7 · p. 18
E. Related simulation-based and simulation-free diffusion-like samplers
The section contrasts samplers by their training and sampling costs, then gives the experimental tuning and implementation setup used for evaluation.
E. Related simulation-based and simulation-free diffusion-like samplers
- Simulation-based SDE samplers require numerical integration during training and store trajectories, while FAB uses AIS samples and needs likelihood computation.For the SDE methods, the stated training cost scales linearly with the number of integration steps and memory is O(Ld); FAB’s memory is O(L + d).Source: E. Related simulation-based and… · p. 19
- DEM’s inner loop is simulation-free, while sampling from its prior requires simulation only at inference.FAB’s sampling can be cheaper with a standard normalizing flow, at the cost of less architectural flexibility.Source: E. Related simulation-based and… · p. 19
F.1. Experimental Setup
- Hyperparameters were grid-searched for each task and method, with three random seeds; FAB on GMM, DW-4, and LJ-13 used previously reported settings.Networks used Adam and NVIDIA A100 GPUs with 40GB VRAM.Source: §F.1 · pp. 19–20
- The setup tunes iDEM’s noise schedule and learning rate by task, with LJ-55 using K = 100 regression-target samples and a 40GB A100-constrained FAB batch size of 8.The task-specific settings also vary the architectures and integration or noise parameters used by the baselines.Source: §F.1 · p. 20
F.2. Metrics reported in Table 2 and Table 5
This section defines the reported metrics and evaluation procedures, including qualifications on NLL estimation, timing, and the particle-system test data.
F.2. Metrics reported in Table 2 and Table 5
- NLL is evaluated with a shared OT-CFM model, but LJ-55 values are excluded for PIS, DDS, and pDEM when retraining the CFM fails to improve over random initialization.For LJ-13 and LJ-55, some sampler outputs degrade the trained CFM’s NLL; reported aggregates exclude values deemed divergent.Source: §F.2 · pp. 20–21
- ESS is estimated from model-proposal log importance weights and normalized by test-set size.The weights use exp(−E(x))/p_model(x).Source: §F.2 · p. 21
- Total variation uses 200 bins per dimension for the low-dimensional GMM, but interatomic-distance distributions for larger equivariant tasks to avoid exponential bin growth.The latter comparison indicates how well sampled energies match the test data.Source: §F.2 · p. 21
- The log-partition-function estimate is a lower bound, so the largest estimate is preferred for equivariant tasks; 2-Wasserstein distance uses Euclidean distance between empirical samples.The discrete transport optimization uses the Hungarian algorithm.Source: §F.2 · p. 21
F.3. Timing experiment setup
- Training time is measured as iterations to convergence multiplied by iterations per second, excluding evaluation computations.Convergence is determined from training loss.Source: §F.3 · p. 22
F.4.1. GAUSSIAN MIXTURE MODEL
- The GMM evaluation uses 40 equally weighted Gaussians in two dimensions, with means uniformly distributed over the −40 to 40 box.Its test set contains 1,000 samples.Source: §F.4.1 · p. 22
F.4.2. DW-4
- DW-4 is a four-particle, two-dimensional double-well system evaluated against Klein et al. (2023b) MCMC samples, which the authors acknowledge are not perfect ground truth.They omit an earlier test set because it was generated from a single MCMC chain and considered biased.Source: §F.4.2 · p. 22
F.4.3. LENNARD-JONES POTENTIAL
- LJ-13 has 39 dimensions and LJ-55 has 165; the score explodes as interatomic distance approaches zero, making higher-dimensional tasks challenging.Evaluation uses Klein et al. (2023b) MCMC samples rather than an earlier dataset the authors consider biased.Source: §F.4.3 · p. 23, Figure 7 · p. 23
G.1. Additional metrics
The additional results compare energy histograms and interatomic distances, alongside supplementary TV, log-partition-function, and NLL metrics.
G.1. Additional metrics
- Table 5 reports TV and log Z, while Table 6 compares native NLL estimates with NLL from a retrained flow-matching model.The paper reports means and standard deviations over three seeds for Table 5; Table 6 describes the native estimators used by FAB and PIS.Source: §G.1 · p. 23, Table 5 · p. 23, Table 6 · p. 24
- iDEM and FAB match the DW-4 ground-truth energy histogram similarly, while PIS and DDM perform noticeably worse.The figure compares each method’s generated energies with the ground-truth histogram.Source: §G.1 · p. 23, Figure 8 · p. 24
G.2. Interatomic distances
- Across DW-4, LJ-13, and LJ-55, iDEM’s interatomic distances agree most with ground truth, and baseline differences grow with dataset complexity.The comparison is reported for model-generated samples and the ground-truth system.Source: §G.2 · p. 24
G.3. Further ablations
Among the score estimators examined, only the LogSumExp estimate behaved well and approached the true score. Further tests show that estimator bias grows with dimension and noise, while clipping supports good samples even in a 10,000-dimensional task.
G.3. Further ablations
- The LogSumExp estimate alone behaved well and approached the true score; the Jensen estimate retained large MSE, while the ratio estimate produced NaN or infinite values below 500 MC samples.The ratio estimate also requires exponentiated energies directly, which can cause large variance and numerical instability.Source: §G.3 · pp. 24–25, Figure 10 · p. 26
- Increasing K lowered log-bias and improved DW-4 performance, while bias also increased farther from data for fixed K.The authors report that more MC samples are needed for proper score estimation.Source: §G.3 · p. 25, Figure 11 · p. 26
- In the 40-GMM, bias and MSE increased with dimension and noise level, but mean-score cosine similarity stayed near 1, indicating mainly magnitude error.The authors hypothesize that early magnitude overestimation can be manageable if later estimates are accurate and steps are small enough.Source: §G.3 · p. 26, Figure 12 · p. 27
- In 10,000 dimensions, the score estimator produced good samples across clipping values from 10^4 to 10^8.The estimator was used directly rather than through a network; the figure compares true and generated distributions over the first two principal components.Source: §G.3 · pp. 26–27, Figure 13 · p. 27