# On scalable and efficient training of diffusion samplers arXiv: https://arxiv.org/abs/2505.19552 ยท HTML: https://mlsanity.com/paper/2505.19552/ ## Key facts - Cited by 10 papers in the map (data through 14 Oct 2026). - Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models. - Builds on 34 of its 49 references that are in the map. - First posted on arXiv 26 May 2025. ## Placed in - Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/) ## Builds on (most cited first, 30 of 34) - [Exploration by Random Network Distillation](https://mlsanity.com/paper/1810.12894/) (2018) - [E(n) Equivariant Graph Neural Networks](https://mlsanity.com/paper/2102.09844/) (2021) - [Equivariant Diffusion for Molecule Generation in 3D](https://mlsanity.com/paper/2203.17003/) (2022) - [Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation](https://mlsanity.com/paper/2106.04399/) (2021) - [GFlowNet Foundations](https://mlsanity.com/paper/2111.09266/) (2021) - [Trajectory balance: Improved credit assignment in GFlowNets](https://mlsanity.com/paper/2201.13259/) (2022) - [The Primacy Bias in Deep Reinforcement Learning](https://mlsanity.com/paper/2205.07802/) (2022) - [Path Integral Sampler: a stochastic control approach for sampling](https://mlsanity.com/paper/2111.15141/) (2021) - [Bayesian Structure Learning with Generative Flow Networks](https://mlsanity.com/paper/2202.13903/) (2022) - [Denoising Diffusion Samplers](https://mlsanity.com/paper/2302.13834/) (2023) - [An optimal control perspective on diffusion-based generative modeling](https://mlsanity.com/paper/2211.01364/) (2022) - [Learning GFlowNets from partial episodes for improved convergence and stability](https://mlsanity.com/paper/2209.12782/) (2022) - [Flow Annealed Importance Sampling Bootstrap](https://mlsanity.com/paper/2208.01893/) (2022) - [Sample-Efficient Optimization in the Latent Space of Deep Generative Models via Weighted Retraining](https://mlsanity.com/paper/2006.09191/) (2020) - [Iterated Denoising Energy Matching for Sampling from Boltzmann Densities](https://mlsanity.com/paper/2402.06121/) (2024) - [GFlowNets and variational inference](https://mlsanity.com/paper/2210.00580/) (2022) - [A theory of continuous generative flow networks](https://mlsanity.com/paper/2301.12594/) (2023) - [Better Training of GFlowNets with Local Credit and Incomplete Trajectories](https://mlsanity.com/paper/2302.01687/) (2023) - [Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching](https://mlsanity.com/paper/2504.11713/) (2025) - [NETS: A Non-Equilibrium Transport Sampler](https://mlsanity.com/paper/2410.02711/) (2024) - [Diffusion Generative Flow Samplers: Improving learning signals through partial trajectory optimization](https://mlsanity.com/paper/2310.02679/) (2023) - [Beyond ELBOs: A Large-Scale Evaluation of Variational Methods for Sampling](https://mlsanity.com/paper/2406.07423/) (2024) - [Generative Augmented Flow Networks](https://mlsanity.com/paper/2210.03308/) (2022) - [Local Search GFlowNets](https://mlsanity.com/paper/2310.02710/) (2023) - [Improved off-policy training of diffusion samplers](https://mlsanity.com/paper/2402.05098/) (2024) - [Sequential Controlled Langevin Diffusions](https://mlsanity.com/paper/2412.07081/) (2024) - [Transferable Boltzmann Generators](https://mlsanity.com/paper/2406.14426/) (2024) - [Scalable Equilibrium Sampling with Sequential Boltzmann Generators](https://mlsanity.com/paper/2502.18462/) (2025) - [Underdamped Diffusion Bridges with Applications to Sampling](https://mlsanity.com/paper/2503.01006/) (2025) - [No Trick, No Treat: Pursuits and Challenges Towards Simulation-free Training of Neural Samplers](https://mlsanity.com/paper/2502.06685/) (2025) ## Built on by (most cited first, 10 of 10) - [Trust Region Constrained Measure Transport in Path Space for Stochastic Optimal Control and Inference](https://mlsanity.com/paper/2508.12511/) (2025) - [Bridge Matching Sampler: Scalable Sampling via Generalized Fixed-Point Diffusion Matching](https://mlsanity.com/paper/2603.00530/) (2026) - [Learning Boltzmann Generators via Constrained Mass Transport](https://mlsanity.com/paper/2510.18460/) (2025) - [Energy-Weighted Flow Matching: Unlocking Continuous Normalizing Flows for Efficient and Scalable Boltzmann Sampling](https://mlsanity.com/paper/2509.03726/) (2025) - [Reinforced sequential Monte Carlo for amortised sampling](https://mlsanity.com/paper/2510.11711/) (2025) - [Enhanced Sampling for Efficient Learning of Coarse-Grained Machine Learning Potentials](https://mlsanity.com/paper/2510.11148/) (2025) - [Efficient Training of Boltzmann Generators Using Off-Policy Log-Dispersion Regularization](https://mlsanity.com/paper/2602.03729/) (2026) - [Onsager-Machlup Posterior Transport for Deep Gaussian Processes](https://mlsanity.com/paper/2605.23434/) (2026) - [Stop the Sampler! Classifier-Based Adaptive Stopping for Sampling Kernels](https://mlsanity.com/paper/2606.16073/) (2026) - [Mode Collapse Is Cheap to Detect: A Ground-Truth-Free Pre-Flight Check for Neural Samplers](https://mlsanity.com/paper/2609.26272/) (2026) ## Abstract (arXiv) We address the challenge of training diffusion models to sample from unnormalized energy distributions in the absence of data, the so-called diffusion samplers. Although these approaches have shown promise, they struggle to scale in more demanding scenarios where energy evaluations are expensive and the sampling space is high-dimensional. To address this limitation, we propose a scalable and sample-efficient framework that properly harmonizes the powerful classical sampling method and the diffusion sampler. Specifically, we utilize Monte Carlo Markov chain (MCMC) samplers with a novelty-based auxiliary energy as a Searcher to collect off-policy samples, using an auxiliary energy function to compensate for exploring modes the diffusion sampler rarely visits. These off-policy samples are then combined with on-policy data to train the diffusion sampler, thereby expanding its coverage of the energy landscape. Furthermore, we identify primacy bias, i.e., the preference of samplers for early experience during training, as the main cause of mode collapse during training, and introduce a periodic re-initialization trick to resolve this issue. Our method significantly improves sample efficiency on standard benchmarks for diffusion samplers and also excels at higher-dimensional problems and real-world molecular conformer generation.