Research map / Probabilistic and scientific modeling / Variational and generative sampling
On scalable and efficient training of diffusion samplers
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes SGDS, which pairs an MCMC Searcher (AIS/MALA/MD with RND novelty rewards) that collects off-policy samples with a diffusion Learner trained by trajectory balance on mixed on- and off-policy data, plus periodic re-initialization to counter primacy bias.
Problem
Training diffusion samplers to draw from unnormalized energy distributions without data is sample-inefficient and suffers mode collapse in high-dimensional problems with expensive energy evaluations.
Approach
A gradient-informed MCMC Searcher with RND-based novelty bonuses generates off-policy samples for a GFlowNet-style trajectory-balance diffusion Learner; the Learner is periodically re-initialized (keeping log Z and the replay buffer) to counter primacy bias.
Key results
- On Manywell-128, SGDS reaches an EUBO-ELBO gap of 70.35 with 20M energy calls, versus 94.28 for TB+LP at 300M calls.
- On LJ-55, SGDS achieves an EUBO-ELBO gap of 33.01 with 45K energy calls, versus 145-323 for iDEM variants at much higher budgets (iDEM at 120M calls for W2 17.17 vs SGDS 4.25).
- On peptides, SGDS attains the lowest EUBO-ELBO gap on dipeptide, tripeptide and tetrapeptide (17.11, 5.60, 26.60 x10^3), outperforming MLE and PIS/TB baselines.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use SGDS instead of the strongest listed baseline | Diffusion samplers · Manywell-64 | EUBO - ELBO gap: 44.75 ± 1.39 -> 16.26 ± 0.22 | Helped |
| Use SGDS instead of the strongest listed baseline | Diffusion samplers · Manywell-128 | EUBO - ELBO gap: 103.45 ± 0.88 -> 70.35 ± 4.31 | Helped |
| Use SGDS instead of PIS | Diffusion samplers · LJ-13 | EUBO - ELBO gap: 2.04 ± 0.32 -> 1.53 ± 0.25 | Helped |
| Use SGDS instead of PIS | Diffusion samplers · LJ-55 | EUBO - ELBO gap: 45.53 ± 5.58 -> 33.01 ± 0.93 | Helped |
| Use a second SGDS training round instead of stopping after round 1 | SGDS · Manywell-32 | EUBO - ELBO gap: 4.11 ± 0.05 -> 3.85 ± 0.03 | Helped |
| Continue to a second SGDS round | SGDS · 40GMM | EUBO - ELBO gap: 0.73 ± 0.02 -> 0.73 ± 0.05 | No clear effect |
| Allocate the full 20M energy-call budget to AIS and train an MLE Learner (ablation) | Manywell-128 amortized learning · Manywell-128 | Amortized learning performance | Hurt |
| Continue training from first-round weights instead of re-initializing the Learner (ablation) | SGDS · Manywell-128 | EUBO - ELBO gap | Hurt |
| Re-initialize the learned log Z parameter at the start of round 2 (ablation) | SGDS · Manywell-128 | Training stability and performance | Hurt |
| Remove the RND novelty reward from later-round Searcher sampling (ablation) | SGDS · Manywell-128 | EUBO - ELBO gap | Hurt |
Surprising findings
- Off-policy Searcher samples from AIS with 20M energy calls alone did not yield a good amortized Learner under MLE training, underscoring the need for the on/off-policy TB mixture.
- Several LP-based baselines (PIS+LP, TB+LP) show large absolute metric values or divergent training on LJ potentials, so LP methods were omitted there.
Limitations
- Searcher requires gradient evaluations and parallel MCMC chains, adding energy-call cost that is counted in the budget.
- Multi-round training and hyperparameters (RND weight, buffer size, re-initialization schedule) are task-specific.
- Peptide experiments are limited to vacuum small peptides with TorchANI potential.
Experimental setup
Datasets: 40GMM Manywell-32/64/128 LJ-13 LJ-55 Alanine dipeptide/tripeptide/tetrapeptide reference ensembles
Benchmarks: 40GMM Manywell-32 Manywell-64 Manywell-128 LJ-13 LJ-55 Alanine dipeptide Alanine tripeptide Alanine tetrapeptide
Metrics: ELBO EUBO EUBO-ELBO gap energy calls Wasserstein-2 energy histogram W2
Baselines: PIS TB TB+Expl+LS FL-SubTB+LP GAFN AT+LP iDEM Adjoint Sampling MLE
Builds on: trajectory balance GFlowNet iDEM path integral sampler random network distillation
Closest papers by meaning
- Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching (ICML 2025 · similarity 0.91)
- Consistent Sampling and Simulation: Molecular Dynamics with Energy-Based Diffusion Models (NeurIPS 2025 · similarity 0.91)
- Scalable Discrete Diffusion Samplers: Combinatorial Optimization and Statistical Physics (ICLR 2025 · similarity 0.90)
- Transition Path Sampling with Improved Off-Policy Training of Diffusion Path Samplers (ICLR 2025 · similarity 0.90)
- Improved off-policy training of diffusion samplers (NeurIPS 2024 · similarity 0.90)
- Progressive Inference-Time Annealing of Diffusion Models for Sampling from Boltzmann Densities (NeurIPS 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.