Improved off-policy training of diffusion samplers
Marcin Sendera, Minsu Kim, Sarthak Mittal, Pablo Lemos, Luca Scimeca, Jarrid Rector-Brooks, Alexandre Adam, Yoshua Bengio, Nikolay Malkin
Abstract
We study the problem of training diffusion models to sample from a distribution with a given unnormalized density or energy function. We benchmark several diffusion-structured inference methods, including simulation-based variational approaches and off-policy methods (continuous generative flow networks). Our results shed light on the relative advantages of existing algorithms while bringing into question some claims from past work. We also propose a novel exploration strategy for off-policy methods, based on local search in the target space with the use of a replay buffer, and show that it improves the quality of samples on a variety of target distributions. Our code for the sampling methods and benchmarks studied is made public at https://github.com/GFNOrg/gfn-diffusion as a base for future work on diffusion models for amortized inference.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Cited by top-tier papers28
- Trajectory Balance with Asynchrony: Decoupling Exploration and Learning for Fast, Scalable LLM Post-TrainingBrian R. Bartoldson, Siddarth Venkatraman, James Diffenderfer, Moksh Jain et al.NeurIPS 2025 · 34 citations
- DriftLite: Lightweight Drift Control for Inference-Time Scaling of Diffusion ModelsYinuo Ren, Wenhao Gao, Lexing Ying, Grant M. Rotskoff et al.ICLR 2026 · 24 citations
- Trust Region Constrained Measure Transport in Path Space for Stochastic Optimal Control and InferenceDenis Blessing, Julius Berner, Lorenz Richter, Carles Domingo-Enrich et al.NeurIPS 2025 · 24 citations
- Markovian Flow Matching: Accelerating MCMC with Continuous Normalizing FlowsAlberto Cabezas, Louis Sharrock, Christopher NemethNeurIPS 2024 · 14 citations
- On scalable and efficient training of diffusion samplersMinkyu Kim, Kiyoung Seong, Dongyeop Woo, Sungsoo Ahn et al.NeurIPS 2025 · 11 citations
Builds on29
- Denoising Diffusion Probabilistic ModelsJonathan Ho, Ajay Jain, Pieter AbbeelNeurIPS 2020 · 35,902 citations
- High-Resolution Image Synthesis with Latent Diffusion ModelsRobin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser et al.CVPR 2022 · 13,123 citations
- Score-Based Generative Modeling through Stochastic Differential EquationsYang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar et al.ICLR 2021 · 1,270 citations
- Maximum Likelihood Training of Score-Based Diffusion ModelsYang Song, Conor Durkan, Iain Murray, Stefano ErmonNeurIPS 2021 · 958 citations
- Tackling the Generative Learning Trilemma with Denoising Diffusion GANsZhisheng Xiao, Karsten Kreis, Arash VahdatICLR 2022 · 726 citations
Related papers
- Adaptive teachers for amortized samplersMinsu Kim, Sanghyeok Choi, Taeyoung Yun, Emmanuel Bengio et al.ICLR 2025
- Discrete Diffusion Samplers and Bridges: Off-Policy Algorithms and Applications in Latent SpacesArran Carter, Sanghyeok Choi, Kirill Tamogashev, Víctor Elvira et al.ICML 2026 · 1 citation
- Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative ModelsSiddarth Venkatraman, Mohsin Hasan, Minsu Kim, Luca Scimeca et al.ICML 2025
- Flow Sampling : Learning to Sample from Unnormalized Densities via Denoising Conditional ProcessesAaron Havens, Brian Karrer, Neta ShaulICML 2026 · 2 citations
- Local Search GFlowNetsMinsu Kim, Taeyoung Yun, Emmanuel Bengio, Dinghuai Zhang et al.ICLR 2024 · 59 citations
