Improved off-policy training of diffusion samplers

Marcin SenderaMinsu KimSarthak MittalPablo LemosLuca ScimecaJarrid Rector-BrooksAlexandre AdamYoshua BengioNikolay Malkin

article2024NeurIPS52 citations

Presents a unified benchmark and an effective replay-buffer exploration strategy for training diffusion models to sample from unnormalized energy densities via continuous generative flow networks.

Listen

Generating samples and estimating partition functions from unnormalized probability densities or energy functions is a foundational challenge in modern machine learning and computational science. While standard simulation-based Markov Chain Monte Carlo methods struggle with slow mode-mixing and prohibitive computational costs in high-dimensional multimodal spaces, amortized variational inference aims to train parametric models that generate samples rapidly. Diffusion models and continuous generative flow networks (GFlowNets) have emerged as powerful paradigms for this task. However, the existing research landscape suffers from benchmarking inconsistencies, differing architectures, unstated hyperparameters, and conflicting claims regarding algorithmic performance and sample efficiency.

The article systematically benchmarks and evaluates continuous-time diffusion-structured inference algorithms under a unified framework. It also introduces and tests a novel exploration mechanism designed to enhance sampling quality and mode coverage in off-policy training.

To establish reliable comparisons, the authors developed an open-source codebase unifying simulation-based variational approaches (such as the Path Integral Sampler) and reinforcement-learning-inspired continuous GFlowNets (such as Trajectory Balance). Using consistent neural architectures and shared training settings, the authors evaluated these models across diverse synthetic and applied benchmarks ranging from 2 to 1600 dimensions, including Gaussian mixture models, funnel distributions, multi-well potentials, Log-Gaussian Cox processes, and conditional latent posterior sampling for a variational autoencoder. The core algorithmic contribution augments continuous GFlowNets with a local search exploration strategy, executing parallel Metropolis-adjusted Langevin algorithm chains directly in target space and storing accepted high-probability samples in a prioritized replay buffer for subsequent training.

The findings provide crucial clarity on model behavior. First, baseline GFlowNet objectives without exploration exhibit severe mode collapse, but adding the proposed local search replay buffer substantially improves sample quality, mode coverage, and log-partition function accuracy across tasks. Second, incorporating a Langevin parametrization inductive bias—directly leveraging target energy gradients within the policy—consistently improves performance across both on-policy and off-policy models, despite introducing a two- to three-fold computational slowdown per iteration. Third, the analysis reveals that partial-trajectory optimization (Subtrajectory Balance) offers no consistent performance benefit over standard whole-trajectory balance while incurring noticeable computational overhead, directly contradicting claims in previous literature. Fourth, in conditional variational autoencoder experiments, standard Trajectory Balance struggles to learn conditional partition functions, whereas policy-only variance-reduction estimators (VarGrad) match or outperform simulation-based baselines.

These results demonstrate that the primary bottleneck in continuous amortized inference lies in exploration efficiency and gradient-guided credit assignment rather than partial-trajectory loss formulation. For engineering and scientific practitioners, adopting prioritized replay buffers with local search provides an effective, computationally lightweight safeguard against mode collapse. However, teams must weigh the high computational cost of Langevin gradient evaluations against the necessity for sharp mode resolution in high-dimensional tasks.

Practitioners deploying amortized diffusion samplers should utilize trajectory-level objectives (such as Trajectory Balance or VarGrad) combined with local search exploration buffers, while avoiding the added complexity of subtrajectory flow modeling. In discrete-time sampling with few time steps, practitioners should also parameterize and learn the forward diffusion variance rather than keeping it fixed, which prevents over-smoothing around target peaks. Further pilot investigations are recommended before scaling these methods to full expectation-maximization training loops in large-scale generative pipelines or complex physical inverse problems.

While the empirical conclusions are supported across multiple repeated trials and rigorous statistical significance tests, readers should note certain limitations. The benchmarks rely primarily on Gaussian discretization bridges, fixed step discretizations, and specific low-to-moderate dimensional potentials. Furthermore, high-dimensional benchmarks like Log-Gaussian Cox processes lack exact ground-truth partition functions, introducing minor baseline variance. Nonetheless, the comparative insights regarding exploration, replay buffers, and credit assignment remain highly robust.

CIELA Institute CIFAR Dreamfold Jagiellonian University Korea Advanced Institute of Science and Technology Mila Université de Montréal University of Edinburgh diffusion probabilistic models inference probabilistic inference diffusion samplers off-policy training continuous generative flow networks

arXiv: 2402.05098 GFNOrg/gfn-diffusion

Paper: GFlowNet Foundations, Yoshua Bengio et al. (2023). Its flow-conservation theory and sampling objectives establish the GFlowNet foundations needed to follow the source’s Trajectory Balance methods.
Paper: Score-Based Generative Modeling through Stochastic Differential Equations, Yang Song et al. (2021). Its continuous-time SDE and score-based sampling framework supplies the diffusion machinery the source adapts to unnormalized target densities.
Paper: Learning GFlowNets From Partial Episodes For Improved Convergence And Stability, Kanika Madan et al. (2023). Its subtrajectory-balance formulation provides the necessary context for the source’s direct comparison of partial- and whole-trajectory objectives.
Paper: Better Training of GFlowNets with Local Credit and Incomplete Trajectories, Ling Pan et al. (2023). Its analysis of local credit assignment and incomplete GFlowNet trajectories clarifies the credit-assignment alternatives assessed in the source.

Paper: Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching, Aaron J. Havens et al. (2025). It carries diffusion sampling for unnormalized energies forward with adjoint-based optimization and replay, extending the source’s exploration-focused sampler training.
Paper: NETS: A Non-equilibrium Transport Sampler, Michael Samuel Albergo et al. (2025). It continues off-policy learning for unnormalized-distribution sampling by adding a learned transport drift to annealed Langevin dynamics.

Cover for Improved off-policy training of diffusion samplers

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 (link) as a base for future work on diffusion models for amortized inference.

Topics

diffusion probabilistic models inference probabilistic inference variational approaches

Keywords

diffusion samplers off-policy training continuous generative flow networks stochastic control backward policy neural SDE hierarchical sampling variational inference Euler-Maruyama integration Langevin process

1 Introduction
2 Prior work
3 Setting: Diffusion-structured sampling
3.1 Euler-Maruyama hierarchical samplers
3.2 Euler-Maruyama samplers as GFlowNets
4 Exploration and credit assignment in continuous GFlowNets
4.1 Credit assignment methods
4.2 A new method for off-policy exploration with local search and replay buffer
5 Experiments
5.1 Tasks and baselines
5.2 Results
5.3 Extensions to general SDE learning problems
6 Conclusion
References
A Code and hyperparameters
B Target densities
B.1 Discrepancies in past work
C Additional results
C.1 Expanded unconditional sampling results
C.2 Variance-preserving noising process
C.3 Scalability study
D Experiment details
E Local search-guided GFlowNet
E.1 Local search algorithm
E.2 Ablation study for local search-guided GFlowNets

Knowls

Knowl 1 — Replay-buffer exploration with local MALA search
algorithm
The proposed local-search method combines exploratory forward sampling with replay of target-space samples refined by parallel Metropolis-adjusted Langevin (MALA). Its inputs are an unnormalized target density $R(x)=\exp[-E(x)]$ on $\mathbb{R}^d$ , a forward sampler $p_F$ , a backward trajectory policy $p_B$ , and a trajectory-balance training loss. For a complete trajectory $\tau=(x_0\to\cdots\to x_1)$ , the loss is $\mathcal{L}_{\mathrm{TB}}(\tau)=\left[\log\frac{Z_\theta p_F(\tau;\theta)}{R(x_1)p_B(\tau\mid x_1)}\right]^2$ , where $Z_\theta$ is a learned positive estimate of the normalizer and $p_F(\tau;\theta)$ and $p_B(\tau\mid x_1)$ are the forward and backward trajectory densities.

Maintain a FIFO buffer $D_{\mathrm{buffer}}$ of forward-generated terminal samples and a local-search buffer $D_{\mathrm{LS}}$ . Alternate training updates: on forward updates, sample exploratory trajectories from $p_F$ , store their endpoints in $D_{\mathrm{buffer}}$ , and update the sampler using their trajectory-balance losses. On replay updates, select terminal samples from $D_{\mathrm{LS}}$ , generate trajectories ending at those samples by sampling backward from $p_B$ , and update the sampler on those trajectories. Every 100 training iterations, seed $M$ parallel MALA chains with samples from $D_{\mathrm{buffer}}$ and use their post-burn-in states to refresh $D_{\mathrm{LS}}$ .

For a MALA chain at state $x$ , propose $y\sim q_\eta(\cdot\mid x)=\mathcal{N}(x+\eta\nabla\log R(x),2\eta I_d)$ and accept with probability $\min\{1,R(y)q_\eta(x\mid y)/[R(x)q_\eta(y\mid x)]\}$ . After each parallel transition, increase the shared step size $\eta$ by a factor of $1.1$ if the batch acceptance rate exceeds $0.574$ , or decrease it by a factor of $0.9$ if the rate is below $0.574$ . The unconditional-task defaults are $\eta_0=0.01$ , 200 MALA transitions, and 100 burn-in transitions; the conditional VAE task uses 500 transitions and 200 burn-in transitions. The reported training default is 25,000 iterations with batch size $M=300$ .

Replay samples are drawn using rank-based prioritization $p(x)\propto(0.01|D_{\mathrm{buffer}}|+\operatorname{rank}(x))^{-1}$ , where rank is based on $R(x)$ and higher-density samples receive higher priority. The FIFO capacity is 600,000 samples for unconditional tasks and 90,000 for the VAE task. The MALA inverse-temperature exponent is $1$ for unconditional tasks and for VAE trajectory-balance local search; the VAE VarGrad local search uses exponent $5$ .
Knowl 2 — Local search improves mode coverage, especially on the 25-mode Gaussian mixture
empirical result
Across five-run unconditional benchmarks, adding local search (LS) to exploratory trajectory-balance (TB) training improves all three reported metrics on the two-dimensional 25-mode Gaussian mixture (25GMM). The metrics are log-partition estimation error $\Delta\log Z$ , importance-weighted log-partition estimation error $\Delta\log Z^{\mathrm{RW}}$ , and empirical squared 2-Wasserstein distance $W_2^2$ ; lower values indicate better results. TB with exploration alone obtains $0.560\pm0.302$ , $0.422\pm0.320$ , and $3.61\pm1.41$ , respectively, while TB with exploration and LS obtains $0.171\pm0.013$ , $0.004\pm0.011$ , and $1.25\pm0.18$ .

The gains are not uniform across targets or metrics. On the 32-dimensional Manywell distribution, TB with exploration changes from $4.01\pm0.05$ , $2.68\pm0.06$ , and $6.15\pm0.02$ to $4.57\pm2.13$ , $0.19\pm0.29$ , and $5.66\pm0.05$ after adding LS: the importance-weighted error and Wasserstein distance improve, but the ordinary log-partition error does not. On the 10-dimensional Funnel, TB with exploration plus LS gives $0.653\pm0.025$ , $0.285\pm0.099$ , and $21.9\pm4.0$ , compared with $0.749\pm0.015$ , $0.226\pm0.138$ , and $21.3\pm4.0$ without LS, a mixed change across metrics. The paper’s page-7 Manywell sample projections visually illustrate the associated mode-coverage benefit: replay-assisted local search retains modes that are missing from some baseline samples.
Knowl 3 — Common benchmark tasks and evaluation protocol
experimental setup
The study benchmarks amortized diffusion samplers on four unconditional energy distributions and one conditional posterior. The 25GMM is a two-dimensional mixture with 25 modes centered on the Cartesian grid $\{-10,-5,0,5,10\}^2$ , with component variance $0.3$ . The 10-dimensional Funnel has $x_0\sim\mathcal{N}(0,9)$ and $x_{1:9}\mid x_0\sim\mathcal{N}(0,\exp(x_0)I)$ . Manywell is a 32-dimensional product of 16 identical two-dimensional double-well distributions. The Log-Gaussian Cox process (LGCP) task has dimension 1600. The conditional task samples a 20-dimensional latent posterior of a pretrained VAE given an MNIST image.

The primary comparison trains Gaussian forward policies with 100 Euler–Maruyama steps, $\Delta t=0.01$ , and a fixed forward variance matched to a discretized Brownian-bridge backward process. The diffusion-rate parameter $\sigma^2$ is 5 for 25GMM and LGCP and 1 for Funnel and Manywell. The compared methods include SMC and GGNS, simulation-driven methods DIS, DDS, and PIS, and diffusion GFlowNets trained with TB, VarGrad, or FL-SubTB, with exploration, Langevin parametrization, and local search tested as enhancements. A common neural architecture is used for the simulation-driven and GFlowNet methods.

The log-partition metrics use 2,000 forward trajectories: one estimates the expectation of the log importance weight, and the importance-weighted version uses the log of the summed trajectory importance weights. Both are lower-bound estimators of $\log Z$ ; the importance-weighted estimate approaches $\log Z$ with increasing sample count. Unconditional results also report sample-based $W_2^2$ . The true normalizer is known analytically for the unconditional tasks except LGCP. Main benchmark table entries report means and standard deviations over five runs.
Knowl 4 — Langevin parametrization injects target-gradient information into the sampler
model/method
The Langevin parametrization (LP) changes the drift of a diffusion sampler to $u(x_t,t;\theta)=\mathrm{NN}_1(x_t,t;\theta)+\mathrm{NN}_2(t;\theta)\nabla E(x_t)$ . Here $x_t\in\mathbb{R}^d$ is the current state, $t\in[0,1]$ is time, $E$ is the target energy, $\mathrm{NN}_1$ outputs a vector-valued learned correction, and $\mathrm{NN}_2$ outputs a scalar scale. The energy-gradient term gives the policy direct information about the target at intermediate sampling steps, rather than only providing reward information through a training loss.

The experiments find LP beneficial for off-policy GFlowNet training on most tasks or metrics, including when added to TB or FL-SubTB, but its per-iteration computation is reported to be roughly 2–3 times higher because the energy gradient must be evaluated at sampling steps. The authors conjecture, without proving, that when $\mathrm{NN}_1$ is independent of $x_t$ , LP and the forward-looking state-flow parametrization used by FL-SubTB have the same short-subtrajectory discrepancy asymptotics as $\Delta t\to0$ .
Knowl 5 — Exploration helps mode discovery; FL-SubTB does not consistently outperform TB
empirical result
Adding noise to the forward policy during off-policy training improves discovery of separated modes, particularly on 25GMM. The reported exploration factor is $0.2$ (additional variance) and is decayed linearly over the first half of training. The paper’s page-8 exploration plot shows that exploration should be reduced over training: persistent high exploration can impair the sampler’s later allocation of modeling capacity to high-likelihood regions, while slowly decaying exploration gives the best reported strategy.

The study also compares complete-trajectory TB with subtrajectory balance (SubTB), which learns state flows and enforces balance on partial trajectories, including a forward-looking variant that incorporates intermediate target energy. In these diffusion experiments, FL-SubTB provides no consistent improvement over TB, and its added state-flow training entails greater computational cost. The VarGrad objective, which minimizes the minibatch variance of the trajectory-balance log ratio, gives results broadly similar to TB. These comparisons qualify earlier claims that partial-trajectory credit assignment reliably improves diffusion samplers.
Knowl 6 — Conditional VAE sampling is competitive with PIS when using VarGrad and LP
empirical result
For conditional sampling of the 20-dimensional latent posterior of a pretrained MNIST VAE, the authors report the ordinary and importance-weighted log-partition lower bounds, $\widehat{\log Z}$ and $\widehat{\log Z}^{\mathrm{RW}}$ ; less negative values are better. PIS obtains $-102.54\pm0.437$ and $-47.753\pm2.821$ , while baseline TB obtains $-162.73\pm35.55$ and $-61.407\pm17.83$ . VarGrad alone obtains $-102.54\pm0.934$ and $-46.502\pm1.018$ .

With exploration and LP, VarGrad reaches $-99.472\pm0.259$ and $-46.574\pm0.736$ ; adding local search gives $-99.783\pm0.312$ and $-46.245\pm0.543$ . Thus the strongest reported VarGrad configurations are competitive with or better than PIS on these estimates, although the ordinary log-bound changes slightly when LS is added. LP and LS improve TB relative to its baseline but do not close the gap as effectively as VarGrad. The authors hypothesize that TB’s difficulty partly reflects learning a condition-dependent partition-function estimator, which VarGrad does not require.
Knowl 7 — Learning the forward variance helps samplers with coarse time discretization
empirical result
The main benchmarks fix the forward policy variance, but the implementation also permits learning it. In a 10-step 25GMM experiment, the fixed-variance model gives a log-partition estimate of $-1.67$ , compared with $-0.62$ for the learned-variance model. The page-9 sampling visualization shows why: fixed variance adds substantial Gaussian noise near the final step and blurs the mixture components, whereas the learned policy reduces its variance around the target peaks near $t=1$ . Across varying numbers of steps, the learned-variance samplers achieve similar quality with fewer discretization steps.

The implementation additionally supports optimizing the backward noising process and using varying backward noise schedules. In a Manywell comparison of three representative methods, a variance-preserving backward process gives results broadly similar to the Brownian backward process. The authors’ results therefore support learning forward variance for coarse discretizations, while not establishing a general advantage for one backward schedule.
Knowl 8 — Replay-buffer capacity and sampling strategy affect local-search quality
empirical result
An ablation on 32-dimensional Manywell compares local-search replay buffers of 30,000, 60,000, and 600,000 samples, with five independent runs per capacity. The reported $(\Delta\log Z,\Delta\log Z^{\mathrm{RW}},W_2^2)$ values are $(4.41\pm0.10,2.73\pm0.15,6.17\pm0.02)$ , $(4.06\pm0.05,2.38\pm0.38,6.14\pm0.04)$ , and $(4.57\pm2.13,0.19\pm0.29,5.66\pm0.05)$ , respectively. The largest buffer improves the importance-weighted error and Wasserstein distance, consistent with better mode coverage, but does not improve every metric monotonically. The page-24 sample visualizations likewise show more captured modes with increasing buffer capacity.

The authors also report that rank-prioritized replay converges faster than uniform replay, and that adapting the MALA step size toward an acceptance rate of $0.574$ performs better than holding the step size fixed at $0.01$ in their ablations. These results motivate retaining diverse past high-density samples rather than relying only on current sampler trajectories.
Knowl 9 — High-dimensional scaling exposes an accuracy–cost tradeoff
empirical result
A Manywell scaling study evaluates dimensions 8, 32, 128, and 512. The authors report that LP’s computational overhead increases with dimension but is important for performance, FL-SubTB with LP has still higher overhead, and local search has relatively low computational overhead when run intermittently. Accuracy does not automatically improve from adding LS at high dimension: at dimensions 128 and 512, TB+LS has $(\Delta\log Z,\Delta\log Z^{\mathrm{RW}})$ errors of $(458.7,139.3)$ and $(1626,1077)$ , while TB+LP+LS gives $(66.6,14.9)$ and $(326,209)$ . TB+LP alone gives $(46.4,14.0)$ and $(259,169)$ at those dimensions, respectively, so adding LS to LP is not uniformly accuracy-improving there.

Runs at dimensions 128 and 512 were sometimes limited to 12 hours; runs at dimensions 8 and 32 completed in under 3 hours on an RTX8000 GPU. Separately, 200 MALA steps for local search took 6.04 seconds on average over five Manywell trials using an Intel Xeon Gold 6338 CPU and an NVIDIA RTX 4090 GPU. The evidence supports low overhead for occasional local-search updates, not uniform scaling of sampling quality with dimension.
Knowl 10 — Benchmark definitions and prior results have reproducibility caveats
limitation
The authors identify inconsistencies that limit direct comparison with earlier diffusion-sampling results. The Funnel target used in this study has $x_0\sim\mathcal{N}(0,9)$ ; they note that earlier work used variance 1 instead, and that the discrepancy propagated to later studies. For the 1600-dimensional LGCP, there is no reliable common ground-truth $\log Z$ : prior sources used values 512.6 and 501.8, and the authors caution that long-run SMC estimates can remain noisy in such high dimensions. They also report that the published FL-SubTB implementation they examined did not reproduce the earlier reported results, and that LGCP training did not converge within the allotted time.

These caveats mean that some cross-paper performance claims are not directly verifiable from the published benchmark values. More broadly, the paper reports empirical comparisons rather than new theorems; its proposed continuous-time relationship between LP and forward-looking SubTB remains a conjecture.

Coverage note — No substantial contributed material was omitted; detailed neural-network layer specifications and command-line reproduction settings are implementation details rather than separate findings.

References

1.Adam, A., Coogan, A., Malkin, N., Legin, R., Perreault-Levasseur, L., Hezaveh, Y., and Bengio, Y. Posterior samples of source galaxies in strong gravitational lenses with score-based priors. arXiv preprint arXiv:2211.03812, 2022.
2.Agrawal, A. and Domke, J. Amortized variational inference for simple hierarchical models. Neural Information Processing Systems (NeurIPS), 2021.
3.Albergo, M. S., Kanwar, G., and Shanahan, P. E. Flow-based generative models for Markov chain Monte Carlo in lattice field theory. Physical Review D, 100(3):034515, 2019.
4.Atanackovic, L., Tong, A., Wang, B., Lee, L. J., Bengio, Y., and Hartford, J. DynGFN: Towards bayesian inference of gene regulatory networks with GFlowNets. Neural Information Processing Systems (NeurIPS), 2023.
5.Bandeira, A. S., Maillard, A., Nickl, R., and Wang, S. On free energy barriers in Gaussian priors and failure of cold start MCMC for high-dimensional unimodal distributions. Philosophical transactions. Series A, Mathematical, physical, and engineering sciences, 381, 2022.
6.Bengio, E., Jain, M., Korablyov, M., Precup, D., and Bengio, Y. Flow network based generative models for non-iterative diverse candidate generation. Neural Information Processing Systems (NeurIPS), 2021.
7.Bengio, Y., Lahlou, S., Deleu, T., Hu, E. J., Tiwari, M., and Bengio, E. GFlowNet foundations. Journal of Machine Learning Research, 24(210):1–55, 2023.
8.Berner, J., Richter, L., and Ullrich, K. An optimal control perspective on diffusion-based generative modeling. arXiv preprint arXiv:2211.01364, 2022.
9.Berner, J., Richter, L., Sendera, M., Rector-Brooks, J., and Malkin, N. From discrete-time policies to continuous-time diffusion samplers: Asymptotic equivalences and faster training. arXiv preprint arXiv:2501.06148, 2025.
10.Buchner, J. Nested sampling methods. arXiv preprint arXiv:2101.09675, 2021.
11.Burda, Y., Grosse, R. B., and Salakhutdinov, R. Importance weighted autoencoders. International Conference on Learning Representations (ICLR), 2016.
12.Chen, T., Liu, G.-H., and Theodorou, E. A. Likelihood training of Schrödinger bridge using forward-backward SDEs theory. International Conference on Learning Representations (ICLR), 2022.
13.Chopin, N. A sequential particle filter method for static models. Biometrika, 89(3):539–552, 2002.
14.Cornish, R., Caterini, A., Deligiannidis, G., and Doucet, A. Relaxing bijectivity constraints with continuously indexed normalising flows. International Conference on Machine Learning (ICML), 2020.
15.De Bortoli, V. Convergence of denoising diffusion models under the manifold hypothesis. Transactions on Machine Learning Research (TMLR), 2022.
16.Del Moral, P., Doucet, A., and Jasra, A. Sequential Monte Carlo samplers. Journal of the Royal Statistical Society Series B: Statistical Methodology, 68(3):411–436, 2006.
17.Deleu, T., Góis, A., Emezue, C., Rankawat, M., Lacoste-Julien, S., Bauer, S., and Bengio, Y. Bayesian structure learning with generative flow networks. Uncertainty in Artificial Intelligence (UAI), 2022.
18.Deleu, T., Nouri, P., Malkin, N., Precup, D., and Bengio, Y. Discrete probabilistic inference as control in multi-path environments. Uncertainty in Artificial Intelligence (UAI), 2024.
19.Dinh, L., Sohl-Dickstein, J., and Bengio, S. Density estimation using Real NVP. International Conference on Learning Representations (ICLR), 2017.
20.Duane, S., Kennedy, A., Pendleton, B. J., and Roweth, D. Hybrid Monte Carlo. Physics Letters B, 195(2):216–222, 1987.
21.Föllmer, H. An entropy approach to the time reversal of diffusion processes. pp. 156–163, 1985.
22.Gao, C., Isaacson, J., and Krause, C. i-flow: High-dimensional integration and sampling with normalizing flows. Machine Learning: Science and Technology, 1(4):045023, 2020.
23.Grathwohl, W., Chen, R. T., Bettencourt, J., Sutskever, I., and Duvenaud, D. FFJORD: Free-form continuous dynamics for scalable reversible generative models. International Conference on Learning Representations (ICLR), 2019.
24.Grenander, U. and Miller, M. I. Representations of knowledge in complex systems. Journal of the Royal Statistical Society: Series B (Methodological), 56(4):549–581, 1994.
25.Halton, J. H. Sequential Monte Carlo. In Mathematical Proceedings of the Cambridge Philosophical Society, volume 58, pp. 57–78. Cambridge University Press, 1962.
26.Harrison, J., Willes, J., and Snoek, J. Variational Bayesian last layers. International Conference on Learning Representations (ICLR), 2024.
27.Hernández-Lobato, J. M. and Adams, R. Probabilistic backpropagation for scalable learning of Bayesian neural networks. International Conference on Machine Learning (ICML), 2015.
28.Ho, J., Jain, A., and Abbeel, P. Denoising diffusion probabilistic models. Neural Information Processing Systems (NeurIPS), 2020.
29.Hoffman, M., Sountsov, P., Dillon, J. V., Langmore, I., Tran, D., and Vasudevan, S. NeuTra-lizing bad geometry in Hamiltonian Monte Carlo using neural transport. arXiv preprint arXiv:1903.03704, 2019.
30.Hoffman, M. D., Blei, D. M., Wang, C., and Paisley, J. W. Stochastic variational inference. Journal of Machine Learning Research (JMLR), 14:1303–1347, 2013.
31.Hoffman, M. D., Gelman, A., et al. The No-U-Turn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo. Journal of Machine Learning Research (JMLR), 15(1):1593–1623, 2014.
32.Holdijk, L., Du, Y., Hooft, F., Jaini, P., Ensing, B., and Welling, M. Stochastic optimal control for collective variable free sampling of molecular transition paths. Neural Information Processing Systems (NeurIPS), 2023.
33.Hu, E. J., Malkin, N., Jain, M., Everett, K., Graikos, A., and Bengio, Y. GFlowNet-EM for learning compositional latent variable models. International Conference on Machine Learning (ICML), 2023.
34.Hu, E. J., Jain, M., Elmoznino, E., Kaddar, Y., Lajoie, G., Bengio, Y., and Malkin, N. Amortizing intractable inference in large language models. International Conference on Learning Representations (ICLR), 2024.
35.Izmailov, P., Vikram, S., Hoffman, M. D., and Wilson, A. G. What are Bayesian neural network posteriors really like? International Conference on Machine Learning (ICML), 2021.
36.Jain, M., Bengio, E., Hernandez-Garcia, A., Rector-Brooks, J., Dossou, B. F., Ekbote, C. A., Fu, J., Zhang, T., Kilgour, M., Zhang, D., et al. Biological sequence design with gflownets. International Conference on Machine Learning (ICML), 2022.
37.Jang, H., Kim, M., and Ahn, S. Learning energy decompositions for partial inference of GFlowNets. International Conference on Learning Representations (ICLR), 2024.
38.Jing, B., Corso, G., Chang, J., Barzilay, R., and Jaakkola, T. Torsional diffusion for molecular conformer generation. Neural Information Processing Systems (NeurIPS), 2022.
39.Kim, M., Ko, J., Zhang, D., Pan, L., Yun, T., Kim, W., Park, J., and Bengio, Y. Learning to scale logits for temperature-conditional GFlowNets. arXiv preprint arXiv:2310.02823, 2023.
40.Kim, M., Yun, T., Bengio, E., Zhang, D., Bengio, Y., Ahn, S., and Park, J. Local search GFlowNets. International Conference on Learning Representations (ICLR), 2024.
41.Kingma, D. P. and Welling, M. Auto-encoding variational Bayes. International Conference on Learning Representations (ICLR), 2014.
42.Lahlou, S., Deleu, T., Lemos, P., Zhang, D., Volokhova, A., Hernández-Garcıa, A., Ezzine, L. N., Bengio, Y., and Malkin, N. A theory of continuous generative flow networks. International Conference on Machine Learning (ICML), 2023.
43.Lemos, P., Malkin, N., Handley, W., Bengio, Y., Hezaveh, Y., and Perreault-Levasseur, L. Improving gradient-guided nested sampling for posterior inference. arXiv preprint arXiv:2312.03911, 2023.
44.Madan, K., Rector-Brooks, J., Korablyov, M., Bengio, E., Jain, M., Nica, A., Bosc, T., Bengio, Y., and Malkin, N. Learning GFlowNets from partial episodes for improved convergence and stability. International Conference on Machine Learning (ICML), 2022.
45.Malkin, N., Jain, M., Bengio, E., Sun, C., and Bengio, Y. Trajectory balance: Improved credit assignment in gflownets. Neural Information Processing Systems (NeurIPS), 2022.
46.Malkin, N., Lahlou, S., Deleu, T., Ji, X., Hu, E., Everett, K., Zhang, D., and Bengio, Y. GFlowNets and variational inference. International Conference on Learning Representations (ICLR), 2023.
47.Máté, B. and Fleuret, F. Learning interpolations between Boltzmann densities. Transactions on Machine Learning Research (TMLR), 2023.
48.Mittal, S., Bracher, N. L., Lajoie, G., Jaini, P., and Brubaker, M. A. Exploring exchangeable dataset amortization for bayesian posterior inference. In ICML 2023 Workshop on Structured Probabilistic Inference {&} Generative Modeling, 2023.
49.Møller, J., Syversveen, A., and Waagepetersen, R. Log Gaussian Cox processes. Scandinavian Journal of Statistics, 25(3):451–482, 1998. ISSN 0303-6898.
50.Nichol, A. and Dhariwal, P. Improved denoising diffusion probabili1stic models. International Conference on Machine Learning (ICML), 2021.
51.Nicoli, K. A., Nakajima, S., Strodthoff, N., Samek, W., Müller, K.-R., and Kessel, P. Asymptotically unbiased estimation of physical observables with neural samplers. Physical Review E, 101(2):023304, 2020.
52.Noé, F., Olsson, S., Köhler, J., and Wu, H. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(6457):eaaw1147, 2019.
53.Nüsken, N. and Richter, L. Solving high-dimensional Hamilton–Jacobi–Bellman PDEs using neural networks: perspectives from the theory of controlled diffusions and measures on path space. Partial Differential Equations and Applications, 2(4):48, 2021.
54.Øksendal, B. Stochastic Differential Equations: An Introduction with Applications. Springer, 2003.
55.Pan, L., Malkin, N., Zhang, D., and Bengio, Y. Better training of GFlowNets with local credit and incomplete trajectories. International Conference on Machine Learning (ICML), 2023.
56.Pillai, N. S., Stuart, A. M., and Thiéry, A. H. Optimal scaling and diffusion limits for the langevin algorithm in high dimensions. The Annals of Applied Probability, 22(6), December 2012.
57.Radev, S. T., Mertens, U. K., Voss, A., Ardizzone, L., and Köthe, U. Bayesflow: Learning complex stochastic models with invertible neural networks. IEEE transactions on neural networks and learning systems, 33(4):1452–1466, 2020.
58.Ranganath, R., Gerrish, S., and Blei, D. Black box variational inference. Artificial Intelligence and Statistics (AISTATS), 2014.
59.Rector-Brooks, J., Madan, K., Jain, M., Korablyov, M., Liu, C.-H., Chandar, S., Malkin, N., and Bengio, Y. Thompson sampling for improved exploration in GFlowNets. arXiv preprint arXiv:2306.17693, 2023.
60.Rezende, D. and Mohamed, S. Variational inference with normalizing flows. International Conference on Machine Learning (ICML), 2015.
61.Rezende, D. J., Mohamed, S., and Wierstra, D. Stochastic backpropagation and approximate inference in deep generative models. International Conference on Machine Learning (ICML), 2014.
62.Richter, L., Boustati, A., Nüsken, N., Ruiz, F. J. R., and Ömer Deniz Akyildiz. VarGrad: A low-variance gradient estimator for variational inference. Neural Information Processing Systems (NeurIPS), 2020.
63.Richter, L., Berner, J., and Liu, G.-H. Improved sampling via learned diffusions. International Conference on Learning Representations (ICLR), 2023.
64.Roberts, G. O. and Rosenthal, J. S. Optimal scaling of discrete approximations to langevin diffusions. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 60(1):255–268, 1998.
65.Roberts, G. O. and Tweedie, R. L. Exponential convergence of Langevin distributions and their discrete approximations. Bernoulli, pp. 341–363, 1996.
66.Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models. Conference on Computer Vision and Pattern Recognition (CVPR), 2021.
67.Särkkä, S. and Solin, A. Applied stochastic differential equations. Cambridge University Press, 2019.
68.Shen, M. W., Bengio, E., Hajiramezanali, E., Loukas, A., Cho, K., and Biancalani, T. Towards understanding and improving GFlowNet training. International Conference on Machine Learning (ICML), 2023.
69.Skilling, J. Nested sampling for general Bayesian computation. Bayesian Analysis, 1(4):833 – 859, 2006. doi: 10.1214/06-BA127. URL https://doi.org/10.1214/06-BA127.
70.Sohl-Dickstein, J., Weiss, E. A., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequilibrium thermodynamics. International Conference on Machine Learning (ICML), 2015.
71.Song, Y., Durkan, C., Murray, I., and Ermon, S. Maximum likelihood training of score-based diffusion models. Neural Information Processing Systems (NeurIPS), 2021.
72.Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. International Conference on Learning Representations (ICLR), 2021.
73.Tiapkin, D., Morozov, N., Naumov, A., and Vetrov, D. Generative flow networks as entropy-regularized RL. arXiv preprint arXiv:2310.12934, 2023.
74.Tripp, A., Daxberger, E., and Hernández-Lobato, J. M. Sample-efficient optimization in the latent space of deep generative models via weighted retraining. Neural Information Processing Systems (NeurIPS), 2020.
75.Tzen, B. and Raginsky, M. Neural stochastic differential equations: Deep latent Gaussian models in the diffusion limit. arXiv preprint arXiv:1905.09883, 2019.
76.Tzen, B. and Raginsky, M. Theoretical guarantees for sampling and inference in generative models with latent diffusions. Conference on Learning Theory (CoLT), 2019.
77.van Krieken, E., Thanapalasingam, T., Tomczak, J., van Harmelen, F., and ten Teije, A. A-NeSI: A scalable approximate method for probabilistic neurosymbolic inference. Neural Information Processing Systems (NeurIPS), 2023.
78.Vargas, F., Grathwohl, W., and Doucet, A. Denoising diffusion samplers. International Conference on Learning Representations (ICLR), 2023.
79.Vargas, F., Padhy, S., Blessing, D., and Nüsken, N. Transport meets variational inference: Controlled Monte Carlo diffusions. International Conference on Learning Representations (ICLR), 2024.
80.Vincent, P. A connection between score matching and denoising autoencoders. Neural computation, 23(7):1661–1674, 2011.
81.Wu, H., Köhler, J., and Noé, F. Stochastic normalizing flows. Neural Information Processing Systems (NeurIPS), 2020.
82.Xiao, Z., Kreis, K., and Vahdat, A. Tackling the generative learning trilemma with denoising diffusion GANs. International Conference on Leraning Representations (ICLR), 2022.
83.Zhang, D., Malkin, N., Liu, Z., Volokhova, A., Courville, A., and Bengio, Y. Generative flow networks for discrete probabilistic modeling. International Conference on Machine Learning (ICML), 2022.
84.Zhang, D., Chen, R. T. Q., Malkin, N., and Bengio, Y. Unifying generative models with GFlowNets and beyond. arXiv preprint arXiv:2209.02606, 2023.
85.Zhang, D., Rainone, C., Peschl, M., and Bondesan, R. Robust scheduling with GFlowNets. International Conference on Learning Representations (ICLR), 2023.
86.Zhang, D., Chen, R. T. Q., Liu, C.-H., Courville, A., and Bengio, Y. Diffusion generative flow samplers: Improving learning signals through partial trajectory optimization. International Conference on Learning Representations (ICLR), 2024.
87.Zhang, Q. and Chen, Y. Diffusion normalizing flow. Neural Information Processing Systems (NeurIPS), 2021.
88.Zhang, Q. and Chen, Y. Path integral sampler: a stochastic control approach for sampling. International Conference on Learning Representations (ICLR), 2022.
89.Zhu, Y., Wu, J., Hu, C., Yan, J., Hsieh, C.-Y., Hou, T., and Wu, J. Sample-efficient multi-objective molecular optimization with GFlowNets. Neural Information Processing Systems (NeurIPS), 2023.
90.Zimmermann, H., Lindsten, F., van de Meent, J.-W., and Naesseth, C. A. A variational perspective on generative flow networks. Transactions on Machine Learning Research (TMLR), 2023.

Citation

MLA

Sendera, M., et al. “Improved Off-policy Training of Diffusion Samplers”. Advances in Neural Information Processing Systems, vol. 37, 2024, pp. 81016–45, https://proceedings.neurips.cc/paper_files/paper/2024/file/93c936b9e492def9c00782cab79dbc6d-Paper-Conference.pdf.

APA

Sendera, M., Kim, M., Mittal, S., Lemos, P., Scimeca, L., Rector-Brooks, J., Adam, A., Bengio, Y., & Whitammer, E. S. (2024). Improved off-policy training of diffusion samplers. Advances in Neural Information Processing Systems, 37, 81016–81045. https://proceedings.neurips.cc/paper_files/paper/2024/file/93c936b9e492def9c00782cab79dbc6d-Paper-Conference.pdf

Chicago

Sendera, M., M. Kim, S. Mittal, et al. 2024. “Improved Off-policy Training of Diffusion Samplers”. Advances in Neural Information Processing Systems 37: 81016–45. https://proceedings.neurips.cc/paper_files/paper/2024/file/93c936b9e492def9c00782cab79dbc6d-Paper-Conference.pdf.

Harvard

Sendera, M. et al. (2024) “Improved off-policy training of diffusion samplers”, Advances in Neural Information Processing Systems. Curran Associates, Inc., pp. 81016–81045. Available at: https://proceedings.neurips.cc/paper_files/paper/2024/file/93c936b9e492def9c00782cab79dbc6d-Paper-Conference.pdf.

Vancouver

1. Sendera M, Kim M, Mittal S, Lemos P, Scimeca L, Rector-Brooks J, Adam A, Bengio Y, Whitammer ES (2024) Improved off-policy training of diffusion samplers. In: Advances in Neural Information Processing Systems. Curran Associates, Inc., pp 81016–81045

BibTeX

@inproceedings{sendera2024improved,
  title = {Improved off-policy training of diffusion samplers},
  author = {Sendera, Marcin and Kim, Minsu and Mittal, Sarthak and Lemos, Pablo and Scimeca, Luca and Rector-Brooks, Jarrid and Adam, Alexandre and Bengio, Yoshua and Whitammer, Esmeralda S.},
  year = {2024},
  booktitle = {Advances in Neural Information Processing Systems},
  publisher = {Curran Associates, Inc.},
  volume = {37},
  pages = {81016-81045},
  url = {https://proceedings.neurips.cc/paper_files/paper/2024/file/93c936b9e492def9c00782cab79dbc6d-Paper-Conference.pdf}
}

Metadata:DOI registry

Source Code

This paper has an official code repository available. Click below to access the source code.

View Repository

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: Published with permission

Improved off-policy training of diffusion samplers

Abstract

Topics

Keywords

Table of Contents

Knowls

Knowl 1 — Replay-buffer exploration with local MALA search

Knowl 2 — Local search improves mode coverage, especially on the 25-mode Gaussian mixture

Knowl 3 — Common benchmark tasks and evaluation protocol

Knowl 4 — Langevin parametrization injects target-gradient information into the sampler

Knowl 5 — Exploration helps mode discovery; FL-SubTB does not consistently outperform TB

Knowl 6 — Conditional VAE sampling is competitive with PIS when using VarGrad and LP

Knowl 7 — Learning the forward variance helps samplers with coarse time discretization

Knowl 8 — Replay-buffer capacity and sampling strategy affect local-search quality

Knowl 9 — High-dimensional scaling exposes an accuracy–cost tradeoff

Knowl 10 — Benchmark definitions and prior results have reproducibility caveats

References

Citation

BibTeX

Source Code

Access the Paper