Rethinking Losses for Diffusion Bridge Samplers

We appreciate the reviewer's thoughtful feedback and are grateful for the acknowledgment of our "strong theoretical analysis" regarding the relationship between the gradients of the Log Variance (LV) loss and the reverse KL divergence (rKL) loss. We are also pleased that our experiments demonstrate the consistent superiority of our proposed loss over the LV and rKL-R losses.

Response to Comment W1: Counterexample is not convincing

For diffusion bridge samplers in the discrete domain, the learnable prior distribution $q(Y)$ could indeed becomes a $\delta$ distribution. So we believe that this example is a valid counterexample. However, to further illustrate our point of the DPI violation, we provide an additional counterexample based on continuous distributions:

Consider the following distributions:

• $q(x, y)$ is uniform on the triangle defined by the vertices $(0,0)$, $(1,0)$, and $(0,1)$.
• $p(x, y)$ is uniform on the entire unit square $[0,1]^2$. Thus, we have:
• $q(x) = 2 (1 - x)$
• $p(x) = 1$

In this example $q(y|x) = 1/(1-x)$ for $y \in [0,1-x]$ and $0$ otherwise. Thus $y$ does contain information on $x$, like in a diffusion process.

This example yields:

• $\mathrm{var}_{x,y \sim q(x,y)} \log \frac{q(x,y)}{p(x,y)} = 0$ for the joint distributions.
• $\mathrm{var}_{x \sim q(x)} \left( \log \frac{q(x)}{p(x)} \right) \approx 0.25$ for the marginals $q(x)$ and $p(x)$.

Hence, this simple example again demonstrates that the LV violates the DPI. The distribution $q(x, y)$ can be regarded as a Markov process and is thus relatable to discretized diffusion processes on continuous domains. Please also note that in this example, the LV for the joints vanishes, i.e., it is at a global minimum. However, the marginals are still different, which shows that the LV loss can be a misleading loss function for the goal of approximating the target distribution. We will include this example in the appendix.

Response to Comment W2: Theoretical convergence of LV loss

The statement that "Theoretically, at least LV loss can converge to the optimal only when the two diffusion processes are identical" is not necessarily true. This holds only when the support of the proposal distribution includes the support of the variational and the target distribution, which cannot be guaranteed during training. Our newly provided counterexample shows that the LV loss of the joint distributions can be zero even though $p(x, y)$ and $q(x, y)$ are not identical. This is particularly relevant for the On-Policy LV loss (OP-LV), where the assumption is certainly not satisfied, and thus the OP-LV loss might possess suboptimal solutions as global minima. For other divergences commonly known in information geometry, such failure cases do not exist.

Response to Comment W3: Verification on larger datasets

We would like to emphasize that our experiments on the Log-Gaussian Cox Process (LGCP) involve 1600-dimensional data, where we observe that the rKL-LD loss outperforms the LV loss in all four cases. To our knowledge, there are no other popular benchmarks in dimensions of 1000 or 10000 in recent literature (see, e.g., [1], [2], [3]). We would like to highlight the comprehensiveness of our experimental evaluation. Specifically, we tested 5 different loss functions across 5 Bayesian learning tasks and 4 multimodal benchmarks, each assessed using 2 distinct diffusion bridge samplers. For each configuration, we performed a grid search followed by final run evaluations. In total, this resulted in over 1,350 individual experiments (calculated as 5 × 9 × 2 × 12 + 5 × 9 × 2 × 3). This level of experimental rigor aligns with the standard protocol commonly adopted in published work within this domain [2]. Evaluating multimodal distributions in such high dimensions is memory-intensive. Computing Maximum Mean Discrepancy (MMD) and Sinkhorn distances requires many samples to achieve sufficient precision, which incurs higher memory costs. Our current code runs out of memory, and we would need more time and several code changes to evaluate MMDs and Sinkhorn distances in such high dimensions. Nevertheless, we have conducted experiments on multimodal target distributions in higher dimensions using GMM-100D and MoS-100D (see below). Our results indicate that the rKL-LD loss performs comparably to the LV loss in terms of Sinkhorn distance and Entropic Mode Coverage (EMC), but outperforms the LV loss in terms of Evidence Lower Bound (ELBO), log Z, and Maximum Mean Discrepancy (MMD). On GMM-100D, the LV loss combined with CMCD exhibits divergent behavior in a few seed runs, which significantly degrades its average performance metrics. This problem also occured in MoS-100D for DBS trained with the LV loss. Notably, we again observe that learnable diffusion coefficients improve performance when trained with rKL-LD but worsen performance when trained with LV.

Response to Limitations: Clarify the advantage of this method.

As stated in several sections of our paper and as the reviewer has stated in the “Summary” and “Strength” sections of their review, our proposed loss offers multiple advantages. It is less sensitive to hyperparameters, more stable during training, and achieves superior results across a wide range of benchmarks. Unlike the LV loss, our method remains stable when diffusion coefficients are trained, and its performance can be further enhanced by learning these coefficients. We thank you once again for the time and effort you have dedicated to reviewing our work. We hope that our responses have adequately addressed your questions.

[1] Sendera, Marcin, et al. "Improved off-policy training of diffusion samplers." Advances in Neural Information Processing Systems 37 (2024): 81016-81045.

[2] Chen, Junhua, et al. "Sequential Controlled Langevin Diffusions." The Thirteenth International Conference on Learning Representations.

[3] Blessing, Denis, et al. "Underdamped Diffusion Bridges with Applications to Sampling." The Thirteenth International Conference on Learning Representations.

We sincerely appreciate your positive evaluation. It’s encouraging to know that our work is perceived as a good quality paper with significant findings and good coverage of the problem background. Please find our detailed responses below.

Weaknesses:

On Clarity:

Thank you for your suggestion, we improved the writing style in our experiment section as you suggested, i.e. shorter paragraphs that are dedicated to specific findings.

On Originality:

The fact the log-derivative trick was employed in prior work for discrete diffusion samplers is regarded as the main weakness of our work. We challenge the associated argumentation by the reviewer since in the mentioned works on discrete diffusions this method is applied to diffusion methods that do not require any backpropagation through the forward diffusion process, i.e. these are not diffusion bridge samplers. For diffusion bridge samplers the application of the log-derivative is thus novel, regardless of whether the considered domain is continuous or discrete. Apart from this conceptual difference these discrete samplers are studied on entirely different benchmarks. Importantly, within the significantly larger field of continuous domain samplers the log-derivative is of course well-known but as pointed out in our manuscript it was dismissed for questionable reasons based on the variance of the gradient estimators that do ultimately not have any predictive power over its actual performance compared to the reparametrization trick. One of our novel insights is that the log-derivative trick is indeed very competitive and often even superior, refuting the argumentation in prior works against it. While this is just one of our contributions, we still think that it is a genuinely new and highly important insight that challenges the predominant choice of gradient estimators and losses in this field.

Questions:

On Results in Table 2:

We would like to note that these performance differences w.r.t. Sinkhorn are not statistically significant (hence multiple results are bold) whereas significant performance improvements can be achieved by rKL-LD with learned diffusion parameters w.r.t. ELBO and MMD. Hence, in the metrics that yields a single best method on GMM40 with statistical significance, we always find that rKL-LD performs best.

Impact of DPI violation

It is beyond the scope of this paper to empirically investigate practical drawbacks of the LV loss that stem from its DPI violation. However, in our response to Weakness 1 of reviewer iCJN we provide a more compelling counterexample that shows that the LV loss violates the DPI. What this new example illustrates is that unless the learned distribution is absolutely continuous with respect to the target distribution, it is possible that the LV loss on the joints reaches its global minimum, i.e. it vanishes, while the corresponding marginal distributions are not equal. In this situation the LV loss on the joint distribution cannot provide any insights into the divergence between the marginals which are the ultimate target of these samplers. In practice, this situation cannot be excluded due to finite numerical precision. While this is still a theoretical argument, we hope that it supports our scepticism that the LV loss is a reliable loss for diffusion bridge samplers.

Limitations

On More Difficult Tasks:

Due to the limited time of the rebuttal period we had to make a compromise on which additional benchmarks to address to satisfy all the reviewers' requests (see reviewer iCJN's request). We chose to expand our experimental evaluation with GMM 100-D and MoS 100-D below. The results support once more the strength of our rKL-LD-based methods.

We thank you again for your time and effort used to review our work. If you are satisfied with our answers and extensions, we kindly ask you to take this into account in your final decision.

We thank the reviewer for this very detailed and thoughtful critique of our manuscript.

W1 Limitations of rKL:

First, we would like to stress that the potential mode collapse issue applies analogously to the LV loss which yields - in the absence of shared parameters in the forward and backward process, identical gradients as the rKL-LD loss. We agree that this is a central problem and hence we follow other works and evaluate our method with metrics like the sinkhorn distance, MMD, and EMC which provide insights into mode coverage (Blessing et al., 2024). Whether or not mode collapse occurs depends on many factors - including the choice of the gradient estimator. As our results in Tab. 2 and 3 show rKL-LD yields overall the best results the aforementioned metrics which probe mode coverage. The second critique in this point is that the log-derivative trick can have a higher variance than the reparametrization trick. Frequently, it is assumed that this statement holds true in general but this is not the case. Mohamed et al., 2019, section 5.3.2 discuss that it depends on the Lipschitz constant of the loss function. More generally, even if the variance of the log-derivativative trick gradient estimator is larger it is not obvious whether this additional source of exploration is detrimental or beneficial. For example, low-variance gradients as in large batch SGD are known to yield worse generalization (Keskar et al. 2016), Tab. 2 and 3 suggest that rKL-LD is less prone to mode collapse.

W2: Apparent Contradiction in Gradient Equivalence:

We would like to clarify that there is no contradiction in our formulation. As explicitly stated in the Abstract (line 4 ff.), Introduction (line 46 ff.), Discussion (line 340 ff.), and the paragraph beginning on line 204 ff., the gradient equivalence between rKL-LD and LV (Equations 12–13) is established under the condition that \phi = v = empty, i.e. there are no learnable forward process parameters and no learnable shared parameters between the forward and reverse process. Therefore, we explicitly state in our paper that for diffusion bridge samplers such as CMCD, where the forward and backward processes share the same network parameters, this equivalence does not hold. This observation is precisely what motivates our investigation into the performance of the rKL-LD loss. We hope that the references provided in the submitted version of the paper clarify this point and address the reviewer's concerns.

W3 Lack of Clarity in Objectives:

We appreciate the reviewer's feedback and acknowledge that implementing a loss function to achieve the desired gradient can indeed be challenging. To address this, we will include pseudocode for the rKL-R, rKL-LD, and LV loss functions in the final version of the manuscript. Below, we provide the pseudocode for implementing the rKL-LD loss.

Pseudocode:

Computation of the rKL-LD Loss

Averages are always computed over the batch dimension.

---

Given:

A batch of diffusion paths $X_{0:T} \sim q_{\alpha, \nu}(X_{0:T})$ computed via Euler-Maruyama integration (see Eq. 4).

---

1. Detach gradients:

$$X^\prime_{0:T} \gets \mathrm{stop\_grad}(X_{0:T})$$

2. Compute log-densities:

$$\log q_{\alpha, \nu}(X^\prime_{0:T}) \quad \text{(Eq. 5)}$$ $$\log p_{\phi, \nu}(X^\prime_{0:T}) \quad \text{(Eq. 6)}$$

3. Compute control variate:

$$b = \mathrm{mean} \left[ \log \frac{q_{\alpha, \nu}(X^\prime_{0:T})}{p_{\phi, \nu}(X^\prime_{0:T})} \right]$$

4. Compute advantages (detached):

$$A^* = \mathrm{stop\_grad} \left[ \log \frac{q_{\alpha, \nu}(X^\prime_{0:T})}{p_{\phi, \nu}(X^\prime_{0:T})} - b \cdot \vec{1} \right]$$

5. Compute loss:

$$\mathcal{L}(\alpha, \phi, \nu) = \mathrm{mean} \left[ A^* \odot \log q_{\alpha, \nu}(X^\prime_{0:T}) \right] - \mathrm{mean} \left[ \log p_{\phi, \nu}(X^\prime_{0:T}) \right]$$

6. Backpropagate and update parameters:

$$(\alpha, \phi, \nu) \leftarrow \mathrm{Adam} \left( \nabla_{(\alpha, \phi, \nu)} \mathcal{L}(\alpha, \phi, \nu) \right)$$

W4 Limited Scope of Evaluation:

Training diffusion samplers on this benchmark involves the usage of memory expensive GNN architectures, which is why losses have to be used that allow for subsampling over diffusion times steps to improve memory efficiency. Due to that to the best of our knowledge, diffusion bridges have not yet been successfully scaled to the benchmarks you suggested (see, e.g., Rissanen et al., who report that they were unable to make CMCD work on this benchmark). Nevertheless, we have extended our experimental evaluation to include multimodal benchmarks in higher dimensions such as GMM-100D and MoS-100D. Our results indicate that the rKL-LD loss performs comparably to the LV loss in terms of Sinkhorn distance and Entropic Mode Coverage (EMC), but outperforms the LV loss in terms of Evidence Lower Bound (ELBO), log Z, and Maximum Mean Discrepancy (MMD). On GMM-100D, the LV loss combined with CMCD exhibits divergent behavior in a few seed runs, which significantly degrades its average performance metrics. This problem also occured in MoS-100D for DBS trained with the LV loss. Notably, we again observe that learnable diffusion coefficients improve performance when trained with rKL-LD but worsen performance when trained with LV.

We would like to highlight the comprehensiveness of our experimental evaluation. Specifically, we tested 5 different loss functions across 5 Bayesian learning tasks and 4 multimodal benchmarks, each assessed using 2 distinct diffusion bridge samplers. For each configuration, we performed grid search followed by final run evaluations. In total, this resulted in over 1,350 individual experiments (calculated as 5 × 9 × 2 × 12 + 5 × 9 × 2 × 3). This level of experimental rigor aligns with the standard protocol commonly adopted in published work within this domain (Chen et al.).

We thank you again for your time and effort used to review our work and hope that all of our answers properly address your questions.

Rissanen, S., OuYang, R., He, J., Chen, W., Heinonen, M., Solin, A., & Hernández-Lobato, J. M. (2025). Progressive Tempering Sampler with Diffusion. arXiv preprint arXiv:2506.05231.

Junhua Chen, Lorenz Richter, Julius Berner, Denis Blessing, Gerhard Neumann, and Anima Anandkumar. Sequential controlled langevin diffusions. arXiv preprint arXiv:2412.07081, 2024a.

Blessing, Denis, et al. "Beyond ELBOs: A large-scale evaluation of variational methods for sampling." arXiv preprint arXiv:2406.07423 (2024).

Mohamed, Shakir, et al. "Monte Carlo gradient estimation in machine learning. arXiv e-prints, page." arXiv preprint arXiv:1906.10652 (2019).

Keskar, Nitish Shirish, et al. "On large-batch training for deep learning: Generalization gap and sharp minima." arXiv preprint arXiv:1609.04836 (2016).

We appreciate the reviewer's thoughtful and constructive feedback. We are also grateful for the acknowledgment of our paper's clarity, the significance of our finding that the LV loss violates the Data Processing Inequality, and the recognition of our experimental results.

Weaknesses:

W1: We agree that both the log-derivative trick and the reparametrization trick play such a central role in this paper that they deserve a self-contained discussion. However, due to the space limitation of 9 pages, we cannot add any further material to the main paper. Hence, we decided to add clear references to sections in the appendix that discuss both methods and provide intuitive explanations of the corresponding computational graphs in the context of the vanishing/exploding gradient problems.

W2: Thank you for your comment, we changed this in the updated version and stick to notation style $\mathbb{E}_{X \sim p} [O(X)]$.

W3: Regarding the problems of the rKL-R we mention three times in the paper the vanishing / exploding gradient problem. However, the reviewers statement clarified to us that we did not explicitly point out that this problem does not arise to the same extent with the log-derivative trick. We will explain the underlying reason for this in detail in the newly added appendix section that we mentioned in our response to W1. Here we will try to give a brief outline of this reason. For the reparametrization trick, backpropagation propagates through the consecutively generated samples $X_t$ along the diffusion path from the final sample ($X_0$) all the way back to the initial sample from the prior $X_T$. The corresponding paths taken by the backpropagation pass through the computational graph of the entire diffusion process with all of its steps, e.g. T = 128 in our case. These long paths can result in the vanishing / exploding gradients problem. Indeed, as mentioned in our manuscript, prior work points out this issue specifically in the context of diffusion samplers. See in particular the discussion in Vargas, et al., 2023 appendix C. 8.1 where it is shown that adding additional stop_gradient operators helps in alleviating this problem, at the cost of introducing biased gradient estimates. This problem does not arise for the log-derivative trick since it does not involve gradient flow through the consecutive sampling processes in diffusion samplers and is thus not subject to this problem. Details on the involved computational graphs will be provided in the updated appendix. We would like to stress, however, that this difference between the reparametrization trick and log-derivative trick or considerations of the variance of these gradient estimators (see first weakness of reviewer KBTa) are not sufficient to argue that one is better than the other. In general, the success of non-convex optimization depends on trade-offs in terms of exploration and exploitation. Evidently it is hard to argue theoretically whether methods that boost exploration or exploitation are generally preferable. Consequently, empirical investigations play a central role in addressing these questions. One of our contributions is that we are the first to empirically investigate the log-derivative trick for diffusion bridge samplers in continuous domains and we find that it does outperform popular methods.

Questions:

Q1: See our response to W2 above.

Q2: Thanks for pointing this out; it is a typo. It will be corrected to refer to the “reparametrization trick”. Here we refer to the issue addressed in W3 above, which is associated with long paths of the backpropagation algorithm through the computational graph that is generated by repeated Euler-Maruyama steps, i.e., by the consecutive diffusion steps. Also, here we will point the reader to the aforementioned new appendix on the issue with the reparametrization trick.

Q3: Frequently, the reparametrization trick is argued in the literature to have a lower variance in the gradient estimation than the log-derivative trick. However, this is only true if the Lipschitz constant of the loss is sufficiently small (Mohamed et al, 2019 section 5.3.2). Arguably, it is, however, not necessarily the case that a low-variance gradient estimator outperforms a high-variance estimator. The reason is that in non-convex optimization the trade-off between exploration and exploitation is key. Hence, higher-variance gradient estimators could contribute to more exploration and hence be beneficial. Thus investigating the gradient variances cannot provide definitive insights into why one gradient estimator is better than another. For an illustration of the potentially detrimental effects of lower variance gradient estimators, see the influential work by Keskar et al., 2016.

We appreciate your time and effort in reviewing our work. If our responses and revisions meet your expectations, we kindly ask that you consider them in your final decision.

Vargas, Francisco, Will Grathwohl, and Arnaud Doucet. "Denoising diffusion samplers." arXiv preprint arXiv:2302.13834 (2023).

Mohamed, Shakir, et al. "Monte Carlo gradient estimation in machine learning. arXiv e-prints, page." arXiv preprint arXiv:1906.10652 (2019).

Keskar, Nitish Shirish, et al. "On large-batch training for deep learning: Generalization gap and sharp minima." arXiv preprint arXiv:1609.04836 (2016).