Improved off-policy training of diffusion samplers

Thank you for your comments and positive assessment of the paper.

Evaluation metrics

Thank you for your question regarding the different aspects of the methods' behaviours, as measured by various evaluation metrics (here: $\log \hat Z^{\rm RW}$ and $\log \hat Z$). We agree with your suggestions about the importance of one of them in measuring mode collapse behaviour and the other with the partition function approximation.

However, for measuring both mode collapse and mode coverege, we have also a few other options. If we have access to ground truth samples, sample-based metrics (such as those typically used for assessing generative models trained from data) are available. One of them is $\mathcal{W}_2^2$, which we already included in Appendix C.1.

Other options are, e.g., EUBO, ESS, MMD, and EMC, which are presented in the very interesting "Beyond ELBOs" paper you mentioned (which appeared after the submission deadline, but which we were already aware of). The question of which approach is the best in each metric is already hard. However, we noticed that $\log\hat Z^{\rm RW}$ results are better for the methods using Langevin Parametrization (e.g., together with Local Search), whereas $\log \hat Z$ performance is better when the method utilizes Local Search.

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.

DB objective

The DB objective fails to discover multiple modes even on the simplest task (25GMM) and thus performs significantly worse than SubTB, which is in turn worse than TB. We guess that this is due to poor credit assignment over long trajectories, as the learning signal to the policy is mediated by learning of the intermediate flow functions (densities).

The preprint you mention tests generalization abilities only in a single environment, which (1) is discrete, (2) has much shorter and variable-length trajectories, (3) is simple enough that all learning algorithms converge to the optimum when modes are not masked. We guess that all three of these limit the generalizability of those results to the diffusion case.

However, it would indeed be interesting, in future work:

• to study the regularizing effect of local-in-time objectives like DB in the diffusion sampler setting, as has been done in the standard diffusion setting in [Lai et al., "FP-Diffusion", ICML'23];
• to understand the meaning of these objectives in the continuous-time limit.

Missing results

We will include the missing results for DIS and DDS on Manywell. The main reason is that our (and, e.g., PIS's) Manywell setting is slightly different than the one used in the DIS paper. We are currently working on making DDS and DIS work on our Manywell version and plan to share those results before the end of the discussion period.

Missing reference

Thank you for pointing out this paper. We are aware of it and will happily include it in the revised version of our paper. (Note that that paper appeared only on June 11, after the NeurIPS submission date.)

Potential typos

Thank you for reading carefully! We fixed the small mistakes.

Thanks you again, and please let us know if we can provide any further information during the discussion period.

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.

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).

Thank you for your comments, in particular, for acknowledging the strength of our comprehensive benchmarking.

Regarding your questions about the benchmarks and their dimensionality, first, we kindly direct you to the response to all reviewers for discussion of the choice of target densities.

Second, the dimensions of the tasks studied are: 2, 10, 20 (with 784-dim conditioning vector), 32, and 1600. As just noted, this is in line with past work on diffusion samplers. But it is also worth noting that sampling from complex densities in high dimension without any prior knowledge (such as a pretrained model or known samples from the target) is generally an unsolved problem, which is perhaps why it has not been the subject of much study in this past work. Our methods, which use off-policy training, can be combined with prior knowledge, such as known target samples (such as by placing those samples in the replay buffer), as has already been demonstrated, e.g., in [42]. Simulation-based methods like PIS, DDS, etc. do not have this ability.

The use of diffusion samplers in settings with prior knowledge, while not the subject of this paper, is indeed a very interesting direction of research and likely important in more difficult applications.

Thank you again for your review. Please let us know if we can provide any other clarification that would help you in understanding and assessing the paper.

Thank you for your comments. Below we've tried to address what we believe to a number of misunderstandings.

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

First, we'd like to point out that this paper is not the first to connect Euler-Maruyama samplers of SDEs with GFlowNets (this was already done in [41] and a few other papers). The goal of this paper is to study methods that, for the first time, achieve results competitive with simulation-based methods on several continuous sampling benchmarks.

Next, we answer the rest of the points:

Applicability of diffusion samplers is limited

We respectfully disagree:

You are correct that there are many domains where we want to train a generative model from ground truth samples. That is an easier problem than sampling unnormalized densities, which is perhaps one reason for the greater attention it receives.

However, many important Bayesian inference problems require sampling distributions that are only available through unnormalized density functions, and sampling methods for such distributions are an active area of research. As noted by the other reviewers:

The studied problem of sampling from a distribution is an important issue with a long history in statistical inference (6eYp)

Enabling diffusion models to be efficiently applied to sampling from unnormalized probability distributions is a problem with high potential for impact (5SPJ).

A few such problems beyond Bayesian deep learning are simulations in statistical physics (e.g., [Nicoli et al., "Asymptotically unbiased estimation of physical observables with neural samplers", Physical Review E, 2020]), molecular dynamics problems (e.g., [Noé et al., "Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning", Science, 2019]), imaging inverse problems in many scientific fields, etc. In fact, it is these problems that have motivated the benchmarks that are adopted in the literature for entirely data-free inference methods such as ours (e.g., Manywell is a simple statistical physics model).

Neural SDE integration is slow compared to MCMC

This is inaccurate. The solution of a neural SDE requires as many neural net evaluations as the number of discretization steps chosen (for us, this is 100, in line with past work). This is far less expensive than the time to run MCMC chains to convergence on any nontrivial density.

Target distributions are simple

We kindly direct you to the response to all reviewers.

Why test on a Gaussian mixture?

Of course, you are correct that a Gaussian mixture density does not require a diffusion model to sample. However, 2D Gaussian mixture densities have been used as a benchmark in many past works ([41,85,87], among others) to test the ability to discover separated modes in an easily interpretable setting.

The proposed method seems to need an extra trained vae to perform sampling

This question appears to rest on a misunderstanding about what that experiment is showing. The goal there is to learn an encoder for a trained VAE, which is an established problem in the diffusion samplers literature (e.g., [87]). The goal is not to unconditionally sample new images, which could indeed be done by simply decoding noise samples with the pretrained decoder. Here, we need a pretrained decoder to provide the reward for the diffusion encoder that we train (the reward is the joint likelihood -- the product of the decoder likelihood and the prior density).

Comparing with GFlowNets for biological sequence design

There seems to be a misunderstanding here. Biological sequence design is a discrete sampling problem in which GFlowNets have indeed been used successfully. The goal of this paper is precisely to study the adaptation of GFlowNets and related methods to continuous sampling problems.

$\log Z$ estimation error

There may be some misunderstanding about the meaning of this metric. The metrics are variational lower bounds on $\log Z$, meaning that they are lower than the true log-partition function in expectation. The estimation error is a more meaningful metric than the raw value because the best possible value (0) is known and invariant to additive changes to the target log-density, which do not change the target distribution.

When the true $\log Z$ is not known, one can report the estimate, not the estimation error, as we indeed do in some of the experiments (LGCP in Table 1 and VAE in Table 2).

Thank you again for your feedback. Please let us know if we have misunderstood anything in your comments and we will be happy to provide further clarification. If we have sufficiently clarified these points and answered your questions, please consider updating the score.

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.

Thank you for your review. Below we will answer your questions and concerns.

Statistically significant improvement over the baselines

Firstly, we want to highlight that our method achieved comparable or better results to current SOTA models. In Tables 1 and 2, we highlighted the best results (red) and those that are not statistically-significantly worse according to the Welch test (blue). Whereas the baselines achieve good results usually in one or two tasks, our method and its variations are the best or among the best results in nearly all tasks simultanously.

Experiments on synthetic energy functions

We thank the reviewer for their suggestion on including the real-world sampling example. Note that the LGCP task is already a real-world example (it is derived from observations of pines in Finland). However, we recognize the importance of applications and will add more experiments (e.g., molecular conformations) in the revision.

We kindly direct you to the response to all reviewers for further discussion of the choice of densities.

Underwhelming results when compared to SOTA?

Simulation-based methods, such as PIS, are predominantly on-policy, making it challenging to explore low-energy regions in high-dimensional spaces and prone to mode collapse in complex densities (see Figure 1 in the Manywell task). In contrast, training a sampler using off-policy techniques allows for the utilization of powerful methods like local search, which enhances exploration and performance. This off-policy trick has been verified in various problems, including simple energy (25GMM), complex multimodal energies (Manywell), high-dimensional tasks (LGCP), and conditional posterior sampling tasks (VAE).

Thanks again for your feedback. Please let us know if you have any more questions and if any further clarifications would help you assess the paper.