Paper: Diffusion Generative Flow Samplers: Improving learning signals through partial trajectory optimization (ICLR-2024)

Reviewer comment (Experiments):
**Limited Experiments:** I consider the experiments presented in the paper to be its primary weakness. While the theory is well-presented and comprehensible, the experiments reported in the paper appear to fall short in terms of comparison with more sophisticated baseline methods, particularly in the context of flow-based samplers.

Author response:
**Regarding experiment limitation.** We understand the reviewer's point about the limitations in our experiments. To address this, we have carried out further experiments, which we hope will answer the reviewer's concerns. The details are provided below.

Reviewer comment (Theory):
On the bottom of page 5, the authors state, "On the other hand, notice that our objectives in Equation 9 or Equation 14 do not require the training samples to follow any particular distribution (only have full support)." While I believe that in the asymptotic regime of infinitely many samples this holds, it may not be true for practical cases with a finite number of samples. In such cases, amortized variational inference schemes might encounter issues, as discussed in recent works [4]. I'm interested in whether this assumption can be relaxed or if the authors can provide arguments for the general validity of the 'full support' assumption in their case.

Author response:
> On the bottom of page 5, the authors state, "On the other hand, notice that our objectives in Equation 9 or Equation 14 do not require the training samples to follow any particular distribution (only have full support)." While I believe that in the asymptotic regime of infinitely many samples this holds, it may not be true for practical cases with a finite number of samples. In such cases, amortized variational inference schemes might encounter issues, as discussed in recent works [4]. I'm interested in whether this assumption can be relaxed or if the authors can provide arguments for the general validity of the 'full support' assumption in their case.

Your statement is grounded in prior GFlowNet theory [10, 11]. These theories propose that when a GFlowNet achieves zero training loss across the board, it is assured to sample objects with probability proportional to the given reward (density) function. In essence, the learned probabilistic model is capable of generating samples that perfectly align with the target density. For additional details, we direct you to Section 3.2 in [10] and Section 3.4 in [11]. In our paper, we are actually making the assumption that, given a neural network with infinite capacity and training data encompassing the entire space, the network can be trained to achieve zero training loss. As per GFlowNet theory, under the conditions that the GFlowNet model possesses sufficient capacity and the training trajectories are fully supportive, the GFlowNet can be accurately trained to sample. However, in practical scenarios, both two assumptions are not fully met: neural network architectures have limited expressiveness, and training data coverage is incomplete. When these assumptions are relaxed, it is no longer certain that our model will still learn to sample correctly. In real-world applications, our aim is to update the model efficiently with limited training data. A key insight, therefore, is to ensure that the training trajectories cover significant modes of the target distribution, reflecting the most important aspects despite data limitations.

Reviewer comment (Experiments):
In the discussion about Variance Gradient Updates (section 3.3A), the authors cite works like Roederer et al. (2017), which reported that the variance of the gradient does not necessarily vanish even when the optimal distribution is achieved. However, in Refs. [1,2], authors use information-theoretical arguments to demonstrate that the variance can indeed approach zero when the target distribution is perfectly learned, using the so-called path gradient method for normalizing flows. It would be interesting to see a comparison of DGFS with such a method (specifically designed for continuous normalizing flows in [2]). Additionally, given the discussion in Sec. 3.3A, it would be beneficial to include these references and expand the discussion.

Author response:
> In the discussion about Variance Gradient Updates (section 3.3A), the authors cite works like Roederer et al. (2017), which reported that the variance of the gradient does not necessarily vanish even when the optimal distribution is achieved. However, in Refs. [1,2], authors use information-theoretical arguments to demonstrate that the variance can indeed approach zero when the target distribution is perfectly learned, using the so-called path gradient method for normalizing flows. 
> It would be interesting to see a comparison of DGFS with such a method (specifically designed for continuous normalizing flows in [2]). 
> Additionally, given the discussion in Sec. 3.3A, it would be beneficial to include these references and expand the discussion.

Thank you for bringing up the important work of [1,2], which we have indeed missed in the initial writing. Similar to the analysis in Roeder et al. (2017), [1,2] propose to decompose the total derivative into the summation of a path-gradient term and a score term, and only conisder the path-gradient term when calculating loss objectives. The path-gradient term indeed has zero gradient variance at the optimal solution. We think this is very related to our work, thus we have added reference and discussion in the updated draft.

We agree with the reviewer's suggestion for doing experimental comparison, so we implement the path-gradient version of neural ODE algorithm in [2] with PyTorch. We train the neural ODE to learn to sample from the unnormalized densities in the benchmark in our submission. On the manywell task, this method obtains a log normalizing factor bias of 2.577, which is outperformed by DGFS. We also find it challenging to scale to hard tasks. We are performing more experiments on this algorithm and will add this baseline into the final version of our work. For completeness we provide our anonymous implementation in [this](https://file.io/vEQo2YFry6vv) link.

Reviewer comment (Novelty):
The concept of injecting information between a sequence of transformations strongly reminds me of simulated annealing and annealed importance sampling (as seen in Stochastic Normalizing Flow [4,5]). Could the authors expand on this, highlighting the significant similarities and differences, if any? There appears to be a notable overlap worth exploring.

Author response:
> The concept of injecting information between a sequence of transformations strongly reminds me of simulated annealing and annealed importance sampling (as seen in Stochastic Normalizing Flow [4,5]). Could the authors expand on this, highlighting the significant similarities and differences, if any? There appears to be a notable overlap worth exploring.

We agree that modeling complex distribution through a sequence of transformations is an old yet powerful idea. Representative works include annealed importance sampling [12] and sequential Monte Carlo sampler [13], followed by many learning augmented variants [4,5]. These methods explicitly define a series of densities by linearly interpolating between a reference distribution (usually a standard Gaussian) and the target density in the log probability scale. These methods then design different kinds of transformations between these intermediate level; these transformations are usually complex and are based on MCMC kernels, sometimes involving normalizing flows and rejection-acceptance steps. On the other hand, in our work and other diffusion-based sampling methods, we do not explicitly define the densities in the intermediate level. Instead, we define clear conditional probability for the transition between adajacent levels. In our case, they are all simple Gaussian distributions. We are thus not able to know the exact density functions in the intermediate levels, and thus we use an amortized way to approximate them (Eq. 11-12). In summary, the former kind of modeling explicitly define the intermediate densities and implicitly define the transformations, while the latter modeling does the opposite.

Reviewer comment (Experiments):
I'm interested in seeing how the lower bound estimator for the log partition function used in this work compares to the asymptotically unbiased estimator for the partition function proposed in [7,8].

Author response:
> I'm interested in seeing how the lower bound estimator for the log partition function used in this work compares to the asymptotically unbiased estimator for the partition function proposed in [7,8].

Thank you for pointing out references [7, 8], which we have now included in our revised draft. The asymptotically unbiased estimator in these works relates to Eq. 17 in the DGFS paper, in that both utilize importance sampling to estimate the normalizing factor. In [7, 8], the proposal is a learned normalizing flow. In our work, the proposal is the forward process 
$\mathcal{Q}$ described in Eq. 1, and our aim is to calculate the integration of the target process $\mathcal{P}$, which has the same value to the integration of the unnormalized density function. We have incorporated a discussion about these papers in Section 3.2 of our updated draft.

Reviewer comment (Experiments):
As pointed out in the **Weaknesses** section, I'm curious about how the results for the many-well task compare to [3], which has reported notable results in data-free training for multimodal target densities for quantum chemical tasks as well as the many-well problem.

Author response:
> As pointed out in the Weaknesses section, I'm curious about how the results for the many-well task compare to [3], which has reported notable results in data-free training for multimodal target densities for quantum chemical tasks as well as the many-well problem.
 
Thank you for bringing up this work. We do have include this work in our original writing, and we point to Appendix C.5 for our results.

Reviewer comment (Experiments):
While I acknowledge the significant contribution and potential impact of the paper, I believe additional robust experimental evidence and a broader discussion, as outlined above, would strengthen its acceptance. Adding new baselines, particularly [2,3], to the comparison in Figure 2 and Table 1 would be a valuable addition to the paper.

Author response:
**Regarding experiment limitation.** We understand the reviewer's point about the limitations in our experiments. To address this, we have carried out further experiments, which we hope will answer the reviewer's concerns. The details are provided below.

Reviewer comment (Writing):
Page 7: to the best of my knowledge, what the authors call "Hamilton Monte Carlo" is more often found in the literature as "Hamiltonian Monte Carlo".

Author response:
**Regarding minor points.** Thank you for pointing out these points. We have fixed them in the revision.

Reviewer comment (Novelty):
Page 7: In the related work section, the authors might consider citing some of the references listed below, being relevant within the context of NF-based samplers. Notably, some of these references have shown how to derive an asymptotically unbiased estimator for the normalizing constant [7,8,9], also known as the partition function in the realm of Boltzmann distributions. This is a quantity of interest for the present work as well.

Author response:
**Regarding related works.** We appreciate the reviewer's suggestion to include additional references. We have read these papers and agree that they are very related to our paper, therefore we have added a discussion about prior works [1-3, 5-9] in the related work section in the updated draft. We believe these additions have enriched our writing by providing more descriptions especially on normalizing flow based methods.