Improved off-policy training of diffusion samplers

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.

Paper

References (93)

11Transport meets Variational Inference: Controlled Monte Carlo Diffusions2023

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer FAJh7/10 · confidence 5/52024-07-09

Summary

The paper studies the problem of training diffusion models to sample from a target distribution. The contributions are summarized as follows: 1. A codebase is provided for the study of diffusion-based samplers, due to the issue of inconsistent experimental settings in previous research; 2. Exploration in the target space can be enhanced by GFlowNet-based off-policy training objectives and local search with the use of replay buffer. 3. Experimental results validate the effectiveness of the proposed approach.

Strengths

Sampling from a target distribution can be challenging in high-dimensional spaces, especially when the distribution of interest has many separated modes. This paper explores diffusion models to address this challenge. Unlike existing reverse KLD-based methods, such as PIS and DDS, this paper considers GFlowNet-based training objectives (e.g., trajectory balance, sub-trajectory balance), which enable off-policy training. This means that training trajectories are not necessarily from the current forward process, thus enhancing exploratory capability. Additionally, local search using a replay buffer can further enhance exploration in the target space. In general, the paper is well-written and well-organized. --- **After rebuttal:** I will increase my score to 7. Typos or incorrect writing should be corrected upon acceptance. ---

Weaknesses

Please see the below questions.

Questions

- In terms of Table 1, which proposed methods perform best according to $\log Z^{LB}$ and $\log Z^{RW}$? Both evaluation metrics reflect different aspects of performance. For example, $\log Z^{LB}$ should indicate mode collapse behaviour, with lower values suggesting more serious mode collapse (e.g., PIS+LP: 13.19 vs. TB+Expl: 4.01, as illustrated in Figure 1)? In contrast, lower $\log Z^{RW}$ values suggest a closer approximation to the ground-truth $\log Z$. - In terms of C2 variance-preserving noising process, - As far as I undertand from the paper, GFlowNet uses PIS architectures, where the initial state is a point mass at 0. How did you design it under VP settings, where the marginal distribution $p_{t}^{ref}$ is an invariant Gaussian distribution, i.e., the initial state is Gaussian distributed, not a point mass at 0? - In VP settings, we gradually add more and more noises to data (data --> Gaussian: $\beta$ increases). However, the sampling process starts with simple Gaussians and ends with samples, so should $\beta$ decrease, when Gaussian --> data? - I am curious if the authors have ever tried the detailed balance objective, due to its superior generalization capability (https://arxiv.org/pdf/2407.03105, different settings though)? - The results for DIS and DDS over Manywell are still missing compared to the previous version. However, the DIS paper does include the results. Any reasons? - Missing reference - Beyond ELBOs: A Large-Scale Evaluation of Variational Methods for Sampling. Both papers share the same goal. It would be better to discuss it in the paper. Potential typos: - Eq.7: To define the reverse KL, a distribution is missing between $\int$ and $\log$? Why does $\mu_{0}(x_{0})$ appear before $d x_{\Delta t}$? - Eq.8: $d x_{0}$ is missing? - Eq.10 & 11: $p_{target}(x_{1})$ --> $R(x_{1})$? - Line 181: $Z_{\theta}$ is missing in the equation? - Eq.12: $f(x_{m})$ --> $f(x_{m \Delta t})$, as well as $f(x_{n})$? - Line 209: $I_{d}$ is missing in $p_{t}(x_{t})$? - Line 290 & 294: $\log$ is missing before $\int$, and $P_{F}(\tau | x_{1})$ --> $P_{F}(\tau | x_{0})$?

Rating

7

Confidence

5

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes. Limitations were included.

Reviewer FAJh2024-08-12

Thank you for your detailed responses, and I apologize for the delayed reply. ### Evaluation metrics - I agree that if ground-truth samples are available, then we can use sample-based metrics for evaluation. - I am still a bit confused. For example, if we look at TB + Expl. + LS (ours), we see that $\Delta \log Z = 0.171$ and $\Delta \log Z^{\mathrm{RW}} = 0.004$, which are the best in the column. We might conclude that it performs the best? It gives the lowest approximation error for $\log Z$, while suffering from more serious mode collapse? ### VP setting > Good question and observations! As you noticed, we used the PIS model architectures, but changed the noising (backward) process to add increasing levels of noise and made the initial state Gaussian-distributed and not a point mass at 0 (which can be thought of as a fixed first transition from an abstract initial state to a sample from the Gaussian). The exact procedure is presented in Appendix C.2. - Do you mean that the policy at the first step, from $s_{0} = ((0, 0), 0)$ to (x_{1}, 1) is constrained to be a unit Gaussian, i.e., $P_{F}((x_{1}, 1) | ((0, 0), 0)) := P(x_{1})$, while subsequent steps are conditional Gaussians with a known variance, i.e., $P_{F}(\cdot | (x_{t}, t))$? > In VP settings, we gradually add more and more noises to data (data --> Gaussian: $\beta$ increases). However, the sampling process starts with simple Gaussians and ends with samples, so should $\beta$ decrease, when Gaussian --> data? - In terms of the equation below Eq.(16), we see that $\beta$ increases from $0.01$ to $4$, implying that we are adding more and more noise for the forward policy. This approach works for diffusion generative modeling. However, the sampling process starts with simple Gaussians and ends with samples, so shouldn't $\beta$ decrease (i.e., $\beta$ decreases from $4$ to $0.01$)? This is my previous concern. [Updated] Thank you for providing new results.

Authorsrebuttal2024-08-12

Thank you for following up. **On the evaluation metrics:** we are not sure to understand your question. In fact, "TB + Expl. + LS (ours)" shows **less** serious mode collapse and correspondingly lower $\log Z$ estimation error. **On the VP setting:** > Do you mean that the policy at the first step ... is constrained to be a unit Gaussian? Yes, precisely as you wrote (assuming integer indexing of time steps). > However, the sampling process starts with simple Gaussians and ends with samples, so shouldn't $\beta$ decrease? We have mistakenly written the noise schedule in equation (16) **in reverse** (i.e., the time indexing convention for diffusion models trained from data, not the one for diffusion samplers). In fact, $\beta$ decreases linearly from 4 to 0.01 (as we have just confirmed in the code that was run for the experiments). Thank you again for reading so carefully.

Reviewer FAJh2024-08-12

**On the evaluation metrics:** for example, TB+Expl: 4.01 vs. TB + Expl. + LS (ours): 4.57, together with Figure 1, where you mentioned that our proposed replay buffer with local search is capable of preventing mode collapse. I might conclude that TB+Expl suffers from more serious mode collapse, with giving a lower value? Other points are clear to me now. Thank you.

Authorsrebuttal2024-08-12

That task (25GMM) is easy enough that none of the methods in the last group (using either LS or LP) are exhibiting mode collapse -- they sample points around each of the 25 modes. Note that: - a sampler that finds 9 of 25 modes (the inner 3-by-3 square) will have a $\log Z^{\rm RW}$ estimation error of about $-\log\frac{9}{25}\approx0.44$ (as TB and VarGrad without any extensions do); - a sampler that finds only one mode will have an error of about $-\log\frac{1}{25}\approx1.40$, and two modes $-\log\frac{2}{25}\approx1.10$ (and the methods with error >1 do indeed usually find one or two modes).

Reviewer FAJh2024-08-12

Thank you for taking the time to clarify. I am willing to increase my score to 7. Typos or incorrect writing should be corrected upon acceptance.

Authorsrebuttal2024-08-12

Much appreciated, and thank you yet again for your detailed reading and finding all the small mistakes that we missed.

Reviewer fepi3/10 · confidence 3/52024-07-12

Summary

This paper proposes an off-policy diffusion-based sampler training method to match a target distribution and a corresponding exploration strategy and credit assignment to improve it.

Strengths

1. The proposed idea of this paper is interesting, which connects the Euler-Maruyama sampler and GFlowNets.

Weaknesses

1. Although the authors mention that traditional MCMC have high cost in sampling, the proposed method based on neural sde seems to still have this problem. To the reviewer’s knowledge, the solving procedure of neural sde is time-consuming as well. 2. The experimental target distribution also seems relatively simple. In the reviewer’s opinion, for GMM, we can first sample a mode according to the weights of different modes and then obtain a sample in this mode. Hence, it seems unnecessary to use complex model like diffusion. 3. In many real-world applications like image generation, the pdf (may be unnormalized) of the target distribution is unavailable and we can only achieve data samples from the target distribution. Hence, the application scenarios of the proposed model are limited. 4. Besides, as mentioned in the conditional sampling case, the proposed method seems to need an extra trained vae to perform sampling. However, the vae can directly do the image generation. In that case, what is the real contribution of the proposed method?

Questions

1. The biological sequence design seems more appropriate to be the validation benchmark for the proposed model, which is also considered in GFlowNets. So the reviewer wonders how the proposed method is compared with GFlowNets in such tasks. 2. Could the authors explain why they use the log-partition function estimation error as the metrics rather than the log-partition function itself? Similar to MLE (Maximum Likelihood Estimation), the model can be considered better with higher $\log Z$.

Rating

3

Confidence

3

Soundness

2

Presentation

2

Contribution

1

Limitations

See Weakness 3.

Reviewer fepi2024-08-09

Thanks for the detailed responses. However, the reviewer still considers it is necessary to do some experiments on high-dimensional real-world applications to show the applicability of the proposed method or make enough theoretical contribution like DDS anyway. Otherwise, it is hard to judge whether the proposed method is useful and efficient. In that case, the reviewer decide to reject this paper and wish the authors could find appropriate complex real-world applications for their method later. Moreover, I guess a potential application for your method is the robot locomotion task, which is considered in CFlowNets [1]. [1] Li Y, Luo S, Wang H, et al. Cflownets: Continuous control with generative flow networks[J]. arXiv preprint arXiv:2303.02430, 2023.

Authorsrebuttal2024-08-09

Thank you for following up. As pointed out in [41, Section 3.1], the referenced paper [1] makes some serious errors in math -- incorrectly doing a change of variables in an integral -- that make the algorithm unsound. Formulating this task as a continuous sampling problem requires access to the Jacobian of the environment’s transition function, which isn’t available. As noted in the initial response, we have demonstrated the effectiveness of our sampling method on a standard collection of tasks in the diffusion samplers literature. If you see a task of higher difficulty that is standard in the field, please let us know.

Reviewer 6eYp4/10 · confidence 2/52024-07-13

Summary

This paper focuses on the problem of sampling with distributions defined by a black-box and unnormalized energy function. This work provides a comprehensive review of existing works, including both variational methods and policy-based methods, and offers a codebase and benchmark to replicate and evaluate the existing works. Additionally, this work proposes a method to improve existing policy-based methods via local search and a replay buffer.

Strengths

1. The studied problem of sampling from a distribution is an important issue with a long history in statistical inference. The paper provides a good review of recent works on this topic by leveraging diffusion models. The codebase that unifies existing methods is certainly useful to the community for continuing research on this topic. 2. The experiments are comprehensive in baselines, including not just diffusion-based methods but also classical MCMC algorithms. The results clearly show the advantages of diffusion-based methods and the techniques proposed in this work.

Weaknesses

1. It appears to me that this work only tests the algorithm on relatively simple and manually-constructed scenarios. Are there any real and important applications within the field? I am not very familiar with this field, but I think that only conducting experiments on synthetic datasets makes this topic less practical. I believe the main advantage of the diffusion-based method over classical methods is in modeling complex distributions, making experiments on synthetic examples less meaningful. 2. Additionally, the tested scenarios are all low-dimensional cases. I wonder how this algorithm performs on high-dimensional cases, such as when the energy function is learned through neural networks. For example, is it possible to apply this algorithm to image generation where the energy function is represented by an image classifier? Testing the algorithm on high-dimensional tasks like these would provide a better understanding of its scalability and practicality in more complex and realistic settings.

Questions

Is there any benchmark in this field involving a real application and complex high-dimensional distribution?

Rating

4

Confidence

2

Soundness

2

Presentation

3

Contribution

3

Limitations

See my comments on the weakness above.

Reviewer 5SPJ6/10 · confidence 3/52024-07-17

Summary

The paper presents a variety of improvements to off-policy strategies for training diffusion models to sample from unnormalized densities. Equation 13. These include maintaining a replay buffer (obtained with Langevin sampling) to enable efficient off-policy exploration and incorporating an inductive bias into the neural network which estimates the SDE drift term. They also present a software library containing unified implementation of these techniques and e.g. diffusion model training.

Strengths

- Enabling diffusion models to be efficiently applied to sampling from unnormalized probability distributions is a problem with high potential for impact - Thorough experimental analysis and comparison of different alterations to the training procedure. On most problems considered, the authors' contributions are necessary to achieve good results with trajectory balance. - The contribution of a software library could be valuable to the community.

Weaknesses

- The results are not overwhelming - although the proposed contributions are helpful compared to a basic version of TB, there is only one modeling task in Tables 1-2 (25GMM) where they provide a statistically significant improvement over the baselines. - The experiments are on synthetic energy functions and MNIST VAE. Including more real-world data or models would be informative.

Questions

See weaknesses. My main concern is the underwhelming results when compared to DIS, DDS, PIS, PIS+LP. Is there any reason why the proposed method(s) should be preferred to these? Or e.g. expected to scale better to more complex problems?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

Adequately addressed

Authorsrebuttal2024-08-12

Additional DDS+DIS results

As promised, here are the results for DDS and DIS on Manywell. First, we clarify the difference between the task as considered in DIS and in this work, noting that both have been studied in other papers as well and neither choice seems obviously better. | | Ours | DIS (hardest version) | |---|---|---| | Dimension | 32 | 50 | | No. of double wells | 16 | 5 | | No. of modes | $2^{16}$ | $2^5$ | We consider Manywell as a product of 16 double well distributions having unnormalized density given by $R(x,y) = \exp(-x^4 + 6x^2 + 0.5x - 0.5y^2)$, following [Midgley et al., "Flow annealed importance sampling bootstrap", ICLR'23], whereas in the DIS paper, the authors consider the more general case of the distribution being a product of some number of double well distributions. However, despite their considering this problem in higher dimensionality, they are stacking a smaller number of double wells (no more than 5), which makes the number of modes in distribution much smaller. We ran DDS and DIS with our setting (after searching for the best hyperparameters, DDS was still somewhat unstable and often encounters an explosion of the loss and mode collapse, in which case we take results from the preceding checkpoint): | Algorithm | $\Delta\log Z$ | $\Delta\log Z^{\rm RW}$ | ${\cal W}_2^2$ | |---|---|---|---| | DIS (4 runs) | 10.52 ± 1.02 | 3.05 ± 0.46 | 5.98 ± 0.46 | | DDS (3 runs) | 7.36 ± 2.43 | 0.23 ± 0.05 | 5.71 ± 0.16 | | VarGrad + LP + LS (ours) | 4.11 ± 0.45 | 0.02 ± 0.21 | 5.30 ± 0.02 | We see that our best method outperforms both DIS and DDS in both metrics, but both are comparable with the various baselines in Table 1 of the paper. We hope you find this comparison helpful. Thank you again for your time.

Authorsrebuttal2024-08-13

Dear Reviewer 5SPJ, This message is to ask if you have any final questions and if our response has affected your assessment of the paper. Thank you, The authors.

Authorsrebuttal2024-08-13

Dear Reviewer 6eYp, This message is to ask if you have any final questions and if our response has affected your assessment of the paper. Thank you, The authors.

Reviewer 5SPJ2024-08-13

Thank you for the rebuttal. After considering it, I believe that this paper does present a valuable contribution through its thorough comparisons and the generally-good performance of the proposed method. I still believe that it would be strengthened by experiments on more complex/high-dimensional problems but given that, as the author's point out in the rebuttal, the evaluations tasks are standard in the continuous sampling community, I believe they are sufficient for acceptance. I have raised my score to 6.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC