Improved sampling via learned diffusions

Dear reviewer,

Thank you for your review. We are happy that you rate the soundness and presentation of our paper as excellent and value our contribution of significantly improving sampling performance via our novel framework. Let us address your concerns and questions in the following.

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

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

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

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

---

[1] Nikolas Nüsken and Lorenz Richter. 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):1–48, 2021.

[2] Lorenz Richter, Ayman Boustati, Nikolas Nusken, Francisco Ruiz, and Omer Deniz Akyildiz. VarGrad: low-variance gradient estimator for variational inference. Advances in Neural Information Processing Systems, 33:13481–13492, 2020.

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

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.

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.

---

We are grateful for your constructive suggestions, which significantly improved our manuscript. We hope that we have addressed all your concerns. If so, we would be very happy if you consider updating your score. If not, please let us know what else we could improve. Thank you again!

Dear reviewer,

We thank you very much for your detailed review and are happy that you value our contribution as being very interesting, allowing for a unified conceptualization of sampling based on diffusions. We also appreciate that you consider our log-variance loss valuable for the ICLR community. Let us respond to your concerns and questions in the sequel.

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.

Clarifications regarding the log-variance divergence

Thanks for bringing up this topic. We also think that the advantages of the log-variance loss are part of the paper's narrative—therefore, we were keen on making them clearer and are grateful for any further suggestions. In particular, we added more discussion after Definition 2.4, an ablation with a smaller batch size in Table 4, an empirical validation of the variance-reducing effect of the log-variance loss in Figure 5, and an extension of our method based on partial trajectory optimization (actively promoting exploration) in Appendix A.7 (see the PDF). Let us answer your specific questions in more detail in the following:

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.

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.

Dear reviewer,

We thank you for your review and are happy that you value our theoretical as well as practical contributions. Let us respond to your comments and questions in the sequel.

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.

Unclear abstract

We thank the reviewer for pointing out this aspect and agree that we can improve the clarity of our abstract. We adjusted the abstract accordingly to emphasize that we are interested in sampling from unnormalized densities rather than from a distribution that is (only) specified by data samples (as in the classical diffusion-based generative modeling case, where one does not have access to densities).

We also agree that we are definitely not the first to connect diffusion-based sampling to Schrödinger bridges. We have mentioned several related references in the text and also the fact that our work was inspired by [1]—however, with our path space measure perspective, one can see that the problem formulation is non-unique, which is not discussed in [1]. By this argument, we admit that the abstract was not exact in the sense that we are considering a generalized Schrödinger bridge that only demands some bridge between prior and target, not necessarily minimizing $\operatorname{KL}(\mathbb{P}\_{X^u}| \mathbb{P}\_{X^0})$—thank you for the good catch.

We have adapted the abstract accordingly (see the updated PDF). Let us know if you think it is more suitable now or if you would like to see further changes.

---

Thanks again for improving the overall readability and clarity with your suggestions! If we addressed all your concerns, we would be happy if you consider potentially updating your score. Otherwise, we would be happy if you let us know of any further improvements.

---

[1] Chen, T., Liu, G. H., & Theodorou, E. A. (2021). Likelihood training of Schrödinger bridge using forward-backward SDEs theory. ICLR 2022.

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.