Paper: Improved sampling via learned diffusions (ICLR-2024)

Reviewer comment (Novelty):
Although the author claims the framework is new. The underlying math theory and contributions have existed long before.

Author response:
*The underlying math theory has existed before.* We agree that the theoretical *tools* used in our novel framework are rather classical (e.g., Girsanov theorem, Ito formula, stochastic optimal control, divergences, optimality PDEs and corresponding BSDEs of SBs, etc.). However, to the best of our knowledge, we are the first to leverage these tools to derive the general likelihood for the (continuous-time) path space setting (see Proposition 2.3). This result is one of the key ingredients to deriving our methods. Only through this result, it becomes possible to consider other divergences, such as the log-variance divergence, which has not been applied to the sampling problem before and—according to our numerical results—yields substantial improvements. 

    We also think that it is more of an advantage than a disadvantage to explain and develop modern machine learning methods (leading to state-of-the-art results) with known mathematical tools that are well-studied and mature.

Reviewer comment (Writing):
To further improve the readability, authors should clearly define the log-variance loss and outline the algorithm pseudocode in the appendix.

Author response:
*Outline the algorithm in the appendix.* We added the pseudocode to the appendix in Section A.6. Thank you for making this suggestion, which greatly improves the presentation of our work.

* *Provide more details and analysis on the log-variance loss /  Clearly define the log-variance loss.* Thank you for this hint and suggestion. In the newly added algorithm environment, we explicitly define the empirical version of the log-variance loss.
In particular, we just compute the empirical variance $\frac{1}{m-1}\sum_{i=1}^m \left(x_i - \frac{1}{m}\sum_{j=1}^m x_j\right)^2$. In PyTorch, this can be done with the command `torch.var`. Note that Bessel's correction is not important since it only rescales the loss. We theoretically introduce the log-variance divergence (and the corresponding log-variance loss) in Definition 2.4 and Equation (8) in the paper.  

    We also added further empirical evaluation of the variance-reducing effect of the log-variance loss in Figure 5 in the appendix. This supports our analysis in Proposition 2.5 (robustness at the solution) and Remark A.1 in the appendix (control variate interpretation). Moreover, we refer to [1, 2] for further analysis.

Reviewer comment (Experiments):
PIS work shows the network that has information of target distribution gradients can perform much better. Is it also true with new log-variance loss?

Author response:
*Networks with target information perform better*. Thank you for this remark. We can confirm the observation made in the PIS paper that it usually helps to incorporate $p_\mathrm{target}$ (or its unnormalized version $\rho$) in the neural network. 

     However, different from the KL-based loss, the log-variance loss yields a derivative-free method when omitting the target in the parametrization. More precisely, it can be used for black-box problems where we have only zeroth-order information from $p_\mathrm{target}$. We added a corresponding discussion and experiment in Section A.7 in the appendix. Based on a novel extension of our algorithm leveraging partial trajectory optimization, we can show that we can achieve similar results even in the derivative-free setting.

Reviewer comment (Theory):
I am a bit confused discussion of stop gradient, does it lead to a biased gradient estimation? Is it the same trick used in Chen 2021a?

Author response:
*Does "stop gradient" lead to a biased gradient estimation?* Thank you for the question. We now clarify this point in our manuscript in Appendix A.6. Detaching the control $w$ in the computational graph of the log-variance loss does *not* lead to a biased gradient estimation. This comes from the fact that the log-variance divergence actually corresponds to a family of divergences parametrized by the function $w$. Specifically, one can optimize the controls $u$ and $v$ for (almost) arbitrary choices of $w$, all of them yielding unbiased and valid losses. As outlined in the text (and motivated by the control variate interpretation in Remark A.1 in the appendix), we usually choose $w = u$ in practice and detach $u$ from the computational graph (since we do not want to take the derivative w.r.t. to $w$). For the KL-based loss, one *must* set $w=u$, and detaching would lead to a meaningless loss. Since the work in Chen (2021a) is based on the KL divergence, they need to rely on a costly, alternating training scheme when detaching the control.

Reviewer comment (Reproducibility):
Can authors provide more descriptions and analysis on log-variance and its computations? How do authors estimate the mean, which is required to calculate the variance?

Author response:
*Provide more details and analysis on the log-variance loss /  Clearly define the log-variance loss.* Thank you for this hint and suggestion. In the newly added algorithm environment, we explicitly define the empirical version of the log-variance loss.
In particular, we just compute the empirical variance $\frac{1}{m-1}\sum_{i=1}^m \left(x_i - \frac{1}{m}\sum_{j=1}^m x_j\right)^2$. In PyTorch, this can be done with the command `torch.var`. Note that Bessel's correction is not important since it only rescales the loss. We theoretically introduce the log-variance divergence (and the corresponding log-variance loss) in Definition 2.4 and Equation (8) in the paper.  

    We also added further empirical evaluation of the variance-reducing effect of the log-variance loss in Figure 5 in the appendix. This supports our analysis in Proposition 2.5 (robustness at the solution) and Remark A.1 in the appendix (control variate interpretation). Moreover, we refer to [1, 2] for further analysis.

Reviewer comment (Presentation):
I found the paper quite dense, although I acknowledge that this is partly due to the nature of the contribution which is to propose a unifying framework. In particular, Section 3 assumes precise knowledge of the various samplers recently proposed in the literature. While it is not a problem in itself, I would suggest in case of lack of space to move part of the derivations of this Section to the Appendix.

Author response:
**Dense presentation** 

We agree that our presentation is rather dense, which is indeed also caused by the page limit. We tried to incorporate proper storytelling and moved multiple technical details to the appendix. Section 2 focuses on our main contribution (the variational path space perspective). Section 3 then draws connections to existing methods (Schrödinger bridge, DIS, PIS, DDS). Additionally, we added a high-level explanation of the sampling methods in Section A.3 and pseudocode in Algorithm 1 in the appendix (see the PDF). Moreover, Table 1 helps to understand the technical connections between the different methods. Note that we wrote Section 3 to be self-contained, i.e., so that one can, in principle, understand the precise connections even without knowing the various other methods. This is why we need to rely on some formulas.   

If you have more suggestions, we would be very happy to improve our presentation further.

Reviewer comment (Theory):
Regarding the choice of \( w \), it is stated that this choice allows a balance between exploitation and exploration. I do not understand what this precisely means. I believe this statement needs to be backed by either a theoretical analysis or a numerical experiment (or be removed). It is also stated without justification that choosing \( w \) different from \( u \) might be beneficial to avoid mode collapse. At the same time, in the experiments, \( w \) is chosen equal to \( u \) (according to Appendix A.4), and yet mode collapse is already avoided. Furthermore, I did not see where it is specified in Section 3 how \( w \) should be chosen for the connection with the various approaches to hold.

Author response:
*Role of the forward process $X^w$*. 

Thank you for bringing up this point. 

The connections in Section 3 hold for any forward process $X^w$, but the exact algorithms from previous works (i.e., DIS, PIS, DDS) rely on the KL divergence, which requires $w = u$ by design. In fact, it is a contribution of our work that one can now also approach the methods DIS, PIS, and DDS with the log-variance loss and an arbitrary forward control $X^w$.

We note that the improved performance of the log-variance loss does not only come from making choices of $w$ different from $u$. As you correctly observe, also the standard choice $w=u$ yields better results (and counteracts mode collapse). First, recall that the log-variance loss yields a provably zero variance of its gradients for any $w$ *at the optimum*. 
However, for the choice $w=u$, the derivative w.r.t. $u$ can generally be interpreted as a control variate version of the derivative of the KL divergence (see Remark A.1 in the appendix).
We added a remark on this important fact below Proposition 2.5. 

We also mention that, in the case of the log-variance loss, one has not only freedom in picking the control $w$ but also the initial condition of the process $X^w$. Motivated by recent work (Zhang et al., 2023), we use this fact to devise an extension of our method, which uses trajectories starting at random points on the space-time domain, thus actively promoting exploration (in particular, including low probability regions of $X^u$). We achieve this by optimizing over subtrajectories of $X^w$ on $[t,T]$ for uniformly chosen $t$ and uniformly chosen starting point $X^w_t$, where we use the partially learned density as new prior. We refer to the new section in Appendix A.7 for further details.

By construction, this approach reduces simulation time. Moreover, we rely only on zero-th-order information of the target and can still (slightly) improve our best results for the DIS method. In comparison, the KL-loss is not even amendable to such a derivative-free setting, and we show in Table 3 that omitting the target in the parametrization negatively affects performance (as also observed by Zhang & Chen (2022)).

Reviewer comment (Evaluation):
Regarding the computational benefit, the proposed method still requires to compute gradients with respect to \( u \) and \( v \) even though gradients wrt \( w \) are not required. It is not clear why this entails a significant improvement in computational efficiency. The naive implementation with stop_gradient still requires backpropagating through the SDE, although the computational graph is lighter. Experimentally, the time improvement shown in Figure 6 is marginal. A more efficient ad hoc implementation that requires only a forward pass through the SDE might be possible, but this is not explained in the paper. Furthermore, unless I am mistaken, simulating the SDE is still required in any case. This should be mentioned to avoid inducing the reader into thinking that it is not required (like in score matching).

Author response:
*Backpropagation through SDE*. 

You mentioned that you think "a more efficient ad hoc implementation that requires only a forward pass through the SDE is possible." The log-variance loss indeed only requires a forward pass (and no backward pass) through the SDE. When we simulate the process $X^w$, we detach $w$ from the computational graph such that no gradient of the stochastic process $X^w$ has to be computed. The SDE must still be simulated to evaluate respective quantities and, e.g., compute the derivatives w.r.t. $u$ and $v$. However, unlike the KL case, no gradients of the stochastic process have to be computed.
In particular, we can use any black-box SDE solver and do not need to track the derivatives "through the SDE solver" for automatic differentiation.

To give an example, the losses usually contain the term $\log \rho(X_T^w)$ to incorporate the target information—for the KL-based losses, one needs to compute derivatives of this quantity, whereas, for the log-variance loss, one does not. This *derivative-free setting* is crucial for problems where the target is only available as a black box. 

We consider a corresponding experiment in Appendix A.7 and added further explanations below Definition 2.4.

Reviewer comment (Experiments):
Proposition 2.5 is not uninteresting, but it is difficult without further explanations or experiments to understand if it plays a role in the improvement observed in practice in the paper.

Author response:
*Further explanations of the role of Proposition 2.5.*

This is a great suggestion! We added Figure 5, which compares the variance of the loss and its average gradient for the KL and LV divergence. As predicted by our theory, we observe a substantial variance reduction when converging to the optimum for the log-variance loss ("sticking-the-landing") and a rather high, non-vanishing variance for the KL-based loss.  

To answer your question, the "variance of the Monte-Carlo estimator of the variance" (for the log-variance loss) is thus actually *lower* than the variance of the Monte-Carlo estimator of the mean (for the KL-based loss). This motivates why the improvements of the log-variance should also hold for lower batch sizes. As suggested by you, we added an ablation for a lower batch size in Table 4, which verifies this claim.

Reviewer comment (Writing):
Proposition 2.3: I think \( \rho \) is undefined.

Author response:
**Minor remarks** 

*Proposition 2.3: I think $\rho$ is undefined.* We added the definition of $\rho$ below Proposition 2.3. We think that the unnormalized density $\rho$ was actually defined in the first equation in the introduction. But we agree that the reader should again be reminded of its definition.

Reviewer comment (Presentation):
The title of the paper is not very informative. I think it would be interesting to consider a title more specifically linked to the contributions of the paper.

Author response:
**Minor remarks** 

*The title of the paper is not very informative.* Thank you for the suggestion. Upon acceptance, we will discuss the possibility of changing the title with the AC.

Reviewer comment (Experiments):
Limited validation. The generalizations proposed are compared on rather simplistic examples. Though the authors appear to target scenarios where data is not available, data-based modeling is clearly an important application of a diffusion-based sampler. Is there a reason the models are not compared to other non-diffusion methods, e.g., MCMC, normalizing flow, autoregressive, or GAN models?

Author response:
**Limited validation**

* *Rather simplistic examples are considered*. We agree that, building upon our principled framework for diffusion-based samplers, future work should tackle even more complicated problems. Our empirical validation shows a clear advantage of our proposed methods and, compared to previous diffusion-based sampling methods (e.g., DIS, DDS, PIS), we already made significant improvements. See also our new ablation studies and extensions to partial trajectory optimization in the appendix. 

    We add that, from a sampling point of view, our examples are already quite challenging. Sampling from multimodal (unnormalized) densities with isolated modes is known to be hard, even in moderate dimensions. In fact, the Funnel and DW problems, as well as complex GMMs, are considered challenging benchmark problems for recently developed methods. This is underlined by the performance of state-of-the-art methods on our considered benchmarks (see Table 6 in the appendix).   

    Let us also refer to Figure 8 in the appendix, which considers a 1000-dimensional DW example, showing that the effect of the log-variance loss does translate even to very high dimensions. We note that the DW problems resemble typical and relevant problems in molecular dynamics, where potentials contain multiple energy barriers and resulting local minima correspond to meaningful configurations of a molecule.

* *Data-based modeling is also important*. We agree that the setting of having data samples available, but no density, is also an interesting problem. Although some advancements of one setting can be translated to the respective other, there are crucial differences. In particular, note that a direct computation of divergences between path space measures seems to be not possible in data-based modeling since the analytical expression of the target density (which, in this setting, is usually not available) appears in the related Radon-Nikodym derivative (see Proposition 2.3). However, we agree that future work can now explore the benefit of additionally using our objectives for the (arguably rare) scenario where both data and an unnormalized density are available.

* *Comparison to other sampling methods*. Note that, additionally to other diffusion-based samplers, we have compared our results to a number of alternative sampling methods in Table 6 in the appendix. We observe that we can outperform state-of-the-art methods based on MCMC, normalizing flows, and SMC. To the best of our knowledge, GANs are not applicable in the scenario where one only has access to an unnormalized density, since the discriminator relies on data samples. Standard autoregressive models seem to be used for density estimation tasks, and typical models for continuous state spaces can not be used for sampling. However, we are very open to comparing to further baselines suggested by you.

Reviewer comment (Theory):
For example, to introduce the derivation of Equation (10), the paper uses “defined in (2) with $u$ replaced by $r \in \mathcal{U}$”. I find this description a little confusing, as Equation 10 assumes three SDEs.

Author response:
Dear reviewer, 

We thank you for your extensive review and your very good and precise summary of our work. We are happy that you value both our contributions and our presentation. We particularly appreciate that you agree with the significance of the log-variance divergence and that you conclude that our novel (and more general) framework will be helpful in the further understanding and development of diffusion-based sampling methods.

Thank you for commenting on potential minor notational issues. In particular, in the example you pointed out, we now refer to the appendix (due to space restrictions), where we added further clarifications.
More general, we improved the descriptions in the main part and included more explanations and additional material (extensions of our method and ablation studies) in the appendix, see the PDF. 

We thank you again for your comment which helped us to improve our manuscript! Feel free to hint at further notational improvements, which we would gladly like to incorporate.

Reviewer comment (Writing):
Clarifying such descriptions would be helpful to potential readers who are not familiar with SDEs and relevant backgrounds.

Author response:
Dear reviewer, 

We thank you for your extensive review and your very good and precise summary of our work. We are happy that you value both our contributions and our presentation. We particularly appreciate that you agree with the significance of the log-variance divergence and that you conclude that our novel (and more general) framework will be helpful in the further understanding and development of diffusion-based sampling methods.

Thank you for commenting on potential minor notational issues. In particular, in the example you pointed out, we now refer to the appendix (due to space restrictions), where we added further clarifications.
More general, we improved the descriptions in the main part and included more explanations and additional material (extensions of our method and ablation studies) in the appendix, see the PDF. 

We thank you again for your comment which helped us to improve our manuscript! Feel free to hint at further notational improvements, which we would gladly like to incorporate.