Rethinking Losses for Diffusion Bridge Samplers
Authors: Sebastian Sanokowski
Lukas Gruber
Christoph Bartmann
Sepp Hochreiter
Sebastian Lehner
Venue: NeurIPS 2025 poster ()
Keywords: Diffusion Samplers
Diffusion Bridges
Bayesian Learning
Diffusion Models
Primary area: probabilistic_methods
TL;DR: We demonstrate that using reverse Kullback-Leibler loss with the log-derivative trick for training diffusion bridges outperforms the commonly used Log Variance loss, providing better results and more stable training.

Abstract: Diffusion bridges are a promising class of deep-learning methods for sampling from unnormalized distributions. Recent works show that the Log Variance (LV) loss consistently outperforms the reverse Kullback-Leibler (rKL) loss when using the reparametrization trick to compute rKL-gradients. 
While the on-policy LV loss yields identical gradients to the rKL loss when combined with the log-derivative trick for diffusion samplers with non-learnable forward processes, this equivalence does not hold for diffusion bridges or when diffusion coefficients are learned.
Based on this insight we argue that for diffusion bridges the LV loss does not represent an optimization objective that can be motivated like the rKL loss via the data processing inequality. Our analysis shows that employing the rKL loss with the log-derivative trick (rKL-LD) does not only avoid these conceptual problems but also consistently outperforms the LV loss. Experimental results with different types of diffusion bridges on challenging benchmarks show that samplers trained with the rKL-LD loss achieve better performance. From a practical perspective we find that rKL-LD requires significantly less hyperparameter optimization and yields more stable training behavior.

Rethinking Losses for Diffusion Bridge Samplers
Sebastian Sanokowski! 1.2
Lukas Gruber
Christoph Bartmann
Sepp Hochreiter
Sebastian Lehner
Technical University Munich, Chair of Robotics, Artificial Intelligence and Embedded Systems ELLIS Unit Linz, LIT AI Lab, Johannes Kepler University Linz, Austria
NXAI Lab & NXAI GmbH, Linz, Austria
sebastian, www.ssnnnokkssseo. [at] tum. de
Abstract
Diffusion bridges are a promising class of deep-learning methods for sampling from unnormalized distributions. Recent works show that the Log Variance (LV) loss consistently outperforms the reverse Kullback-Leibler (rKL) loss when using the reparametrization trick compute rKL-gradients. While the on-policy LV loss yields identical gradients to the rKI. loss when combined with the log-derivative trick for diffusion samplers with non-learnable forward processes, this equivalence does not hold for diffusion bridges or when diffusion coefficients are learned. Based on this insight we argue that for diffusion bridges the LV loss does not represent an optimization objective that can be motivated like the rKL loss via the data processing inequality. Our analysis shows that employing the rKL loss with the log-derivative trick (rKL-LD) does not only avoid these conceptual problems but also consistently outperforms the LV loss. Experimental results with different types of diffusion bridges on challenging benchmarks show that samplers trained with the rKL--DD loss achieve better performance. From practical perspective we find that rKL-LDD requires significantly less hyperparameter optimization and yields more stable training behavior
1 Introduction
We consider the task of learning to generate samples X E R.N from an unnormalized target distribu- tion:
To exp (-E(X)) with 2 jp
f exp exp((x(X))
(1)
Z
where denotes the partition funetion and & RR. R is the energy function. this setting, assumed that the energy function E of the target distribution can be evaluated and 2 is unknown and computationally intractable. Importantly, there are no available samples from To- Sampling prob- lems of this kind represent fundamental challenges in computational physics and chemistry and in Bayesian inference [Noé and Wu 2018 Wu et al., 2019 Shih and Ermon 2020). Recent approaches have focused training probability distributions parameterized by neural networks to approximate target distributions. Early deep barming-based methods explored exact likelihood models such as normalizing flows (Noé and Wu! 2018) and autoregressive models (Wu al.) 2019 Nicoli et al.)
2020). while more recent work has turned to approximate likelihood models like diffusion models
continuous [Zhang and Chen 2022 Berner et al. 2022 and discrete domains (Sanokowski al.
Our code at
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

2024 2025]. In the continuous domain, initial research explored these so-called diffusion samplers that aim at reversing the forward time evolution processes that are given by the variance- preserving
and variance-exploding stochastic differential equations (SDEs) [Zhang and Chen 2022 Vargas
al. 2023 Berner al,, 2022]. More recently, diffusion samplers based on diffusion bridges em-
ploy learable forward processes and have established a new state-of-the-art in the field (Richter and Berner 2024 Vargas et al., 2024 Blessing et al., 2025]. These models are typically trained using either (i the reverse Kullback-Leibler divergence loss with gradients from the reparameterization trick (rKL-R) or (ii) the Log Variance (LV) loss where backpropagation through an expectation is not necessary [Richter and Berner, 2024). However, diffusion bridge samplers face significant prac- tical challenges. The rKL-R loss is susceptible to vanishing and exploding gradient problems when employing numerous diffusion steps (Zhang and Chen 2022 Vargas et al. [023) and has been em- pirically shown to underperform compared to the LV loss. Conversely, diffusion bridge samplers trained with the LV loss are prone training instabilities and thus need extensive hyperparameter
tuning.
In this work, make three key contributions to address these limitations of diffusion bridge sam- plers. (i) We identify crucial problems in the application of the popular Log Variance (LV) loss in diffusion bridges. While the gradients of the LV loss and the gradients of the reverse KL divergence, when optimized with the log-derivative trick (rKL-LD), are identical when only the reverse diffusion control is learned, this equivalence does not hold in the context of diffusion bridge samplers when diffusion coefficients are Iearned. We demonstrate, in this case, the LV loss suffers from training instabilities and is based divergence that does not satisfy the data processing inequality which represents an important theoretical criterion for losses of latent variable models. (ii) Instead of using the LV loss, we advocate for training diffusion bridge samplers with the rKL-LD loss and show that it consistently outperforms diffusion samplers trained with LV and rKL-R losses. Additionally, we find that it is significantly less sensitive hyperparameter tuning. (iii) In addition the usual learned drift terms of SDEs we motivate learnable diffusion terms that enable dynamic adaptation of the exploration-exploiatiii significantly improved performance when applied to hitherto state-of of-the-art samplers that build on diffusion bridges.
trade-off and show that combined with the rKL-LD loss they yield
2
Problem Description
2.1 Diffusion Bridges
Diffusion models are generative models that learn to transport samples Xy e IR- from a simple prior distribution Ty to samples that are distributed according to a target distribution X₂ ~ Tu- This transport map is defined by the reverse diffusion process, whose parameters are learned with the objective of inverting a forward diffusion process. The parameters of the corresponding for- ward SDE are either learned or predefined. Examples for later case are the variance-exploding or variance-preserving SDEs [Song et al.) [202T]. For these predefined forward SDEs the selection of parameters--prtiiulle diffusion coefficients must be carefully adjusted since they determine the stochastic process which the diffusion model needs to revert in order to transport from тт to To- In diffusion bridges, however, also the forward SDE learnable. These models are thus equipped with the flexibility to freely learn which transport path to use between to wg- Due to this flexibility diffusion bridges have recently attracted increased research interest and represent the state-of sute-of-the- the-hhan art in wide range of applications, including the data generation and pairing problem [De Bortoli et al.) 2021 2024) or the problem of sampling from unnormalized distributions, which is the focus of this paper.
In diffusion bridges, the forward diffusion process is defined by the following SDE:
dx,- fo(,,, di dWW, where t [0, T], Xo no,
where W. e N is a Brownian motion, i.c. dW, = Wi+ W. ~ NO IN dt) which is multiplied with diffusion coefficients o, e R.N and & are the learnable parameters of the forward diffusion control Sx( () I The reverse process is defined as:
dx, - r. (X, T) o, dW, where TE [0,
X,-0 ~ TT,
(3)
with learnable reverse control r. (X,,7) e R-N. Here a denotes learnable parameters of the reverse diffusion process, which evolves in the opposite time direction t= TT T'-7
2

---

In practice the reverse process is often simulated via Euler-Maruyama integration with step size AT Åt for steps so that the reverse SDE generation process given by:
go (r, 7) X X, TT,, AT o, V ÂT E,.
(4)
The conditional probability for a step the reverse direction then given by: (x
(5)
and for a step in the forward direction by:
Ps (x( NX(x; + fa(X--1,t1) At, (1
(6)
2.2 Parameterization of Diffusion Bridge Samplers
In the following, we will introduce two common parameterizations of diffusion bridge samplers with a focus on the distinction between SDE parameters that are shared or non-shared between the forward and reverse processes. These considerations will reveal the conceptual problem of the L.V loss and motivate the introduction of diffusion sampler training via rKL-LD. For this purpose we denote with & and a the parameters of the forward and reverse process, respectively, that are non- shared between the two processes. The parameters that are shared are denoted via p. We denote the entirety of the parameters of diffusion bridge as f (a, (., This formulation encompasses two popular diffusion bridge samplers: Denoising Bridge Samplers (DBS) and Controlled Monte Carlo Diffusion (CMCD).
Denoising Bridge Samplers: DBS, introduced by Richter and Berner 2024 uses separate neural networks for reverse drift ra and forward drift fe with non-shared parameters. In the underdamped DBS variant proposed in Blessing et al. [2255] additional SDE-related parameters including diffu- sion coefficients are learned, resulting
# Ø.
2024 parameterizes both reverse and forward diffusion processes with the same parameterized control U, with shared parameters y. In CMCD, all learnable parameters are shared, i.e. ø= a = Ø. The control terms the forward and reverse diffusion processes are guided by the score Vx log log (X) of interpolation between the forward and reverse process: Te (X,) #o (o(..))) x,))(t) Tr (X.)) 1-(((), where 7)()) [0, 1] a learned monotonically decreasing function with (7)) = and (T) = 0 (see App. D.3 for more details). The control of the reverse diffusion processes in Eq. is then defined by:
Controlled Monte Carlo Diffusion:
CMCD, introduced by Vargas et al.
Tr (X,,
Vx, log n. (X,) + (X,,T)
and the control of the forward diffusion process in Eq by:
o2 V x, log m. (X.) - U (X,,
(x,,,,
2
CMCD is initialized such that u. (X,,t) = O, which ensures that before training, the simulation of the reverse diffusion process Eq. is equivalent to the Unadjusted Langevin Annealing. Intuitively, CMCD learns to modify unadjusted Langevin dynamics with an additive control term so that the resulting process efficiently transports samples from the prior the target distribution within limited number of diffusion steps T. In CMCD, the parameter structure can be represented by (7, 01) (.).......) when also diffusion coefficients are learned.
2.3 Training of Diffusion Samplers
The loss to train diffusion models typically based on divergence between the underlying marginal distribution (Xo) induced by the reverse diffusion process and the target distribution To (Xo). In this context, f-divergences are a popular choice [Csiszár 1967]:
D,(((X)))((X) f Q())s
P(X)
dx,
Q(X)
where f is convex function satisfying f(1) 0 and P and Q are two probability distributions satisfying P << Q, i.e. P is absolutely continuous with respect to Q. The choice of I significantly influences the learning dynamics in terms of mode-secking vs. mass-covering properties.
3

---

Reverse KL Divergence: The rKL corresponds to f(t) = - -lo((t) and is given by:
Jo (X)
DeL (X), [x)) = Ex-)x. log
To (X)
The rKL divergence is convenient choice for learning to sample from unnormalized target distri- butions as the expectations are calculated over model samples X qa., However, the rKL, exhibits mode-seeking behavior, i.e. In.v tends to focus on the high-density regions of the target distribution тo, potentially ignoring low-density regions or modes with smaller probability mass which leads to bias between the target and variational distribution that cannot be corrected with Neural Importance Sampling or Neural Markor-Chain-Moote-Carrr (Nicoli et al. 2020].
Forward KL Divergence: The IKL corresponds to f(t) +log(t) and is given by:
(x( DEL D (o., = Ex- log (x,,(()
(8)
In contrast to the rKL it is mass mass-covering, i.e. qa, tends to cover the entire support of To, even at the cost of placing high probability in low-density regions. While this mass covering property useful for sampling applications, because it enables asymptotically unbiased estimates of observables it comes at the cost of sample efficiency, and the fKI. divergence cannot straightforwardly used as a loss function as samples from X ~T (x() are not available. This prohlem can, however, be mitigated by either rewriting the expectation with respect to samples from xo (X) to an expectation over samples from ao (X) by incorporating self-normalized importance weights as done in Müller itt al. (211)) Jing et al. 22022), Midgley et al. [2023]. Sanokowski et al. 2025].
Jeffrey Distance: Another possible choice is the Jeffrey distance which corresponds to f(t)
(t 1) log(t). This yields:
Dieft (no X), (x)) Dx (To ga. (X)) DKL. ((a, mo
(9)
This divergence is sum the rKL and iKL and is therefore a trade-off between the mode-seeking and mass-covering behavior of the rKL and fKL divergence. In the context of sampling this loss is for example used in Noé and Wu (2018).
Data Processing Inequality: For expressive latent variable models like diffusion samplers, the marginal variational distribution (x.) is typically intractable. Therefore, losses for diffusion models are in practice often motivated by the data processing inequality for f-divergences that sat-
isfy the following monotonicity relation [Zhang and Chen 2022
Vargas al. 2023 Sanokowski
et al.
2024 Blessing et al., 2025]:
D (T(((())) |I (X.o) s Dy (Po. ) ll (Xxr)).
(10)
The right-hand side of this inequality is (up to normalization constant) tractable for diffusion bridges as is based on joint distributions of the diffusion paths Nor- Using Eq. 3]and Eq the corresponding joint probability distributions can be efficiently calculated with:
Po. (X.(X--)) (a.r(Xor))= = TI (XT) I1 (x 1-1
(o.T(XaT) = (Xo)
[-1
The parameters (a, o, v) of the diffusion bridge (Sec. 2.2 are optimized with the objective of min- imizing the right-hame side of Eq. [T] This training objective corresponds to maximizing the Ev- idence Lower Bound (ELBO) in latent variable models [Blessing et al.) 2024]. Eq. T0 states ex- plicitly that any reduction of the right-hand side leads tighter bound on the divergence between the intractable marginals on the left-hand side. In fact it is well known in generative modeling that diffusion losses that are based on properly bounded objectives yield better sample likelihoods (Song et al. [222]. We therefore argue that the data processing inequality is an essential property for losses in sampling setting and that losses that do not provide such a conceptual footing might be practically problematic (see Sec. 2.
2.3.1 Log Variance Loss
Richter et al. 20200) propose the LV loss as a convenient replacement of the rbsed loss func- tions in Bayesian variational inference and it has since then been frequently used to train diffusion

---

samplers |Rchher and Berner 2024 Vargas et al., 2024, Sendera et al. 2024 Chen et al. 2024a
Blessing 20.4|0221111 He He al. 2025 In the context of GilowNets Bengio al. [2021] this loss
also known as the Trajectory Balance loss. For diffusion bridges, i.e. when both the reverse process la and the forward process Po.v involve learnable parameters the L. takes the following form:
1 -E xxr-- Div (Xxr) I| (Xor))
Jo [[1os La (XaT) Po (Xor)
)
(11)
2
where is a suitable proposal distribution and U = Exar-w log Por Pau (Nor) the support of Za.v and Po.r ( the L.V loss is zero if and only ga.0 (Xor) PP ( In the simplest case Log Variance (OP-LV) loss. For diffusion samplers, the general LV loss and the OP-LV loss were
If w contains
P
stop_gt cooppgrrddennt(e... == Lo for which we will refer to as the on-policy
recently proposed in Richter and Berner (2024), however, also noted in Malkin et al. [2022] the LV loss is not an f-divergence and we show via the following simple counter-example that the LV loss does not satisfy the data processing inequality.
Violation of the Data Processing Inequality by the Log Variance Loss
Consider the following distributions on the unit square Ω [0,1]
p(z, v) (uniform distribution over ?
if x + y
Un (n v)
m- if
x 1,,
where n Rz0 determines the ratio between the amount of probability mass lo-
cated below and above the diagonal. For this choice, both Varz-n. log (x) p(a)
and
Var(s.y)-e. log gn(2.z) p(x.y) can be computed analytically (see App. A.5
As illustrated below, these calculations show that for n < 10-2 the Log Variance loss does not satisfy the data processing inequality (DPI):
Varzonn log In (n(xx) p(x)
$ Var(z.)) log On (I, v) p(æ, y)
where the joint distribution (xx,)) 4) is absolutely continuous with respect to p(2, y). and vice versa.
Gulx Joint: Varg, pix.y
0.4
Marginal: Varg,[log 9.0k)
DPI violated region
Variance
DPI satisfied region
DPI Satisfied
DPI Violated
0.0
10-6
10-5
10
10-3
10-2
10-1
100
n
Consequently, we argue that its application to latent variable models potentially problematic since it conflicts with the rationale of diffusion sampler training based on divergences of joint probabilities motivated by the data processing inequality (see See. 2.3). We provide additional counter example
in discrete state spaces is proved in App. A.5
5

---

3 Method
3.1 Reverse KL Loss with Log-derivative Trick and Control Variate
One way to compute the gradient of the rk1. loss is to use the reparametrization trick [Kingma and Welling 2014] as explained in more detail in App. A.1 In diffusion samplers, the repeated appli- cation of the reparametrization trick is, however, likely to give rise to the vanishing or exploding gradient problem, which might explain the suboptimal behavior of rKL loss when minimized with the usage of the reparametrization trick (Zhang and Chen 2022 Vargas et al., 2023]. Several re- cent works demonstrate that the LV loss outperforms the rLL loss with the reparametrization trick [Richter and Berner 2024 Vargas et al. 2024 Chen et al.) 2024a It is argued that this is due to the mode-seeking tendency associated with the rKL objective which results in mode collapse. While the aforemennioned works applied the reparametrization trick in conjunction with the rKL objective we investigate the rKL objective with the log-derivative gradient estimator. The gradients corresponding to the rKL.-ID with respect to the parameters (a, 1, v) read (App. A.4):
log (XoT) D (Nor)
V LD DKL (Ma,v Il Exar- -4r,
Va log 9a (Xo:r)
V Dxt. (yo, Il )= -Exor [V log Po.e (Xor)]
(12)
[[0g log Ya (Xor))
V go DKL PKL(qo PPov)) )== Exar
Vu log qa, (Xorr)
Páiu (Xo.)
a,
L Naar
VV log Par ( (XoT),]
where V I.D.ui is an operator that computes the gradient of parameters with respect the expectation over distribution wy by applying the log-derivative trick and by reducing the variance with a control variate 0g as defined in App. A.2
3.2 Gradient of the Log Variance Loss
The corresponding gradients of the LV loss with respect to bridge parameters (ar, 6, are given by
(see App. A
Exor-w
log Ya (Xerr)
Va log ao (XaT)
a V Div (Ia,r.Dov)=
Po. (Xor)
Exxorww
log Lo (Xarr
ba
V log u (Xor)
(13)
Po.c (xarr)
a,
V
Experww
log 2o.u (NoT)
b
V, log Ca. (Xur)
Po.u (Xer)
Po. (X.r)
The gradient of the LV loss is related to well-known f-divergences in the following way. With the choice w stop-stee gradieet(qa... =: 4a.v one obtains:
V. Tai Dix (fo,r. = LDD.a. DKL (ga. Pl...,
Hence, for non-shared parameters of the reverse diffusion process a the gradient estimator of the LV loss coincides with the gradient of the rKL divergence computed with the log-derivative trick.
When the off-policy distribution is stop slop-gradienn((p..)) po., the gradient with respect of given by:
VoDVy (ga.v.Po.d) V. DKLL(PNN Ia,
and therefore exactly the same as the gradient of the forward KI. divergence with log-derivative trick. For the OP-LV loss, however, we have at initialization Ip - To.vl >> and therefore
Va.y VoD (ga.v, p).v) can be seen biased estimate of LDD.De. ....... DKI(Ps.v 1l (o..)) ) as it can al-
ternatively be estimated with the usage of Neural Importance Sampling in the following way:
wr (e((Nor) loog (Xo.)) Pa. (Xor) a.thi
L.D. DEL (Po. ga.v) =-Bxar
V log (urr)
where (with abuse of notation) wi Xur) N (X.r(XoT)) (Xurr) with Ω = ΣN,, gn (Xor.s)) are self normal- ized importance weights. However, as ya.v approximates Po. and therefore w (XoT) 1, this bias is continually reduced.
(NoT.))

---

For parameters that are shared between the forward and backward diffusion process & we have for OP-LV at optimality when P Vo.
V V. Vo. (9a.v.P...) A. P... (KL
(p (l D + V L.D.aa. DDL. DKL (l
where the right-hand side is related to the gradient of the Jeffrey distance with a stop-gradient operation in the second argument of each KL divergence. Therefore, at initialization when Ip Io. >> 0 the gradient is biased as w (xXrr) 1 and additionally there is a difference to the gradient of the Jeffrey distance due to the stop gradient operations.
A comparison of these gradients with the gradients of the rKL.-DD loss in Eq. 12 shows that when Ø the gradients of the OP-LV loss and rKL.-ID loss are identical. Conversely, when w # stop--racee crop--gaaeeet(( go) these two losses yield, in general, different gradients. Similarly, when Ø + 2, 1/ + Ø, the gradients with respect to these parameters are different, in which case an evaluation of the rKL-LD loss has not been considered in recent literature. These considerations and the fact that the LV loss violates the data processing inequality put the validity of the LV loss for diffusion bridge samplers in question and in fact we observe in our experiments that the LV loss often exhibits unstable behavior and requires hyperparameters to be carefully tuned (see Sec. [5}. Additionally, we provide in Sec. 5 a detailed comparison between the LV and rKL-LD loss for DBS and CMCD, where we observe that rKI.-I. yields better diffusion bridge samplers than the LV loss.
3.3 Learning of SDE parameters
SDEs in diffusion samplers are governed by diffusion coefficients that significantly influence their sampling performance. While previous work has predominantly focused on non-learned diffusion coefficients, the choice of these parameters involves nuanced considerations that can strongly affect the model's ability to capture complex distributions. this section, we motivate why is beneficial to learn these coefficients.
The diffusion coefficients control the amount of stochasticity injected throughout the sampling pro- cess and thereby regulate how broadly the model explores the data space. Larger coefficients in- crease randomness, allowing the sampler better cover the support of high-entropy or multimodal target distributions. However, excessive noise reduces the topal- ratio and can obscure im- portant structural information in the data, leading to poorer reconstructions and slower convergence. Conversely, too little noise restricts exploration and may cause the sampler to miss significant re- gions of the target distribution.
This interplay between coverage and fidelity reveals a fundamental trade-off in selecting appropri- ate diffusion coefficients. Learning these coefficients rather than fixing them a priori allows the model to adaptively tune the level of stochasticity across time and dimensions, achieving a better balance between exploration and precision. This trade-off motivates the introduction of learnable diffusion coefficients. The optimal magnitude of noise injection varies across different problems and even across different dimensions of the same problem, suggesting that adaptively learning these parameters can lead to more efficient approximations [Blessing et al. 2025). In our experiments in Sec. [ we demonstrate that diffusion coefficient learning yields significantly better performance in combination with the rKL.-DD
-L. loss. While we observe that diffusion coefficient learning consis- tently improves results with rKL-LD it tends to worsen performance and increases hyperparameter sensitivity when using the LV loss, highlighting the benefit of rKL.-ID.
4 Related Work
For continuous diffusion samplers, the rKL loss is typically used with the reparametrization trick VVargas et al. 2023 Berner al. 2022). More recently, Richter and Berner [2024] proposed the LV loss as an alternative and showed that it outperforms the rKL loss with the reparametrization trick. The corresponding experiments are performed with the Path Integral Sampler (Zhang and Chen 2022] and the Time-Reversed Diffusion Sampler [Berner et al. 2022], in which the forward diffusion process is not learned. However, they report that the application of LV to diffusion bridges results in poor performance and numerical instabilities. In Vargas et al. [2024). Chen et al. [2024a] the L.V loss is employed in conjunction with diffusion bridges and both works report that outperforms the rKL. loss with the parametrization trick. Frequently, the mode-collapse tendency of rKI. is given as
[
7

---

an explanation for its inferior performance in the context of diffusion samplers [Richter and Berner, 2024]. The rKL with log-derivative trick has recently been successfully applied to diffusion samplers in discrete domains. Examples of such problems arise in combinatorial optimization and statistical
physics of spin lattices [Sanokowski al. 2024 2025 Some of the considerations in Sec. 3.2 are
also made in the context of Gflow Nets in Malkin et al. 2022].
0.0
CNCO. LV
0.030
(RLLD
CRCDEEY
0.225
CRCD: IRU-LD
00250
-0.4
CHCD: FRLL
DELAD)
0.0000
dan
0 0175
dar
Car
-0.0
0 01150
0.125
1002 1500 2000 2500
4000
15
500 1000
2000 2002 3500 4000
Epochs
Dimension
Epochs
Figure 1: Models with fixed aur are marked with * Left: Training curves on the Brownian task of CMCD trained with LV loss and rKL-LD loss. Middle: Plot of the learned dur ascending order of the CMCD-rKL-LD run from the left figure. Right: Training curves on the Seeds task, where CMCD-rKL-LD Odiff init compared CMCD-1KL-LD diiff,init at different initializations of Odiff
ELBO (1)
Seeds (26d)
Sonar (61d)
Credit (25d)
Brownian (32d)
L.GPP (1600d)
CMCD: r
74.37-0.01
109 6.0...0.ooo c.s
504. 194.99...ou
0.30.0.0s
471.91=0.291
CMCD: LV
74.13.0.0
109.53... 23..53...0zz 0.0
504. .0.002
0.05.o.002
460. 04.000
CMCD: LV
73
109.66.0.01
32
6.05 f
447.7420.0
CMCD: rKL-LD
74.10_o.0
109. 109.25.c.cor
504. 114.88.u 88.0.008
0.3620.001
466,73 166.73.0.027
CMCD: rKL-LD
-33.4520.01
-108.83 030.00
504.58.0.01
1.06.0.0
100.800..00
DBS: rKL-R
75.09
-119.58 anss
-528.15.o. 15
-3. 13
461.09.m.11a
DBS: LV .
-4.112.000
-110.662o.o1 ona
-506. 2006:21-a0 10s
9.39- t
460.48.1.00
DBS: LV
o7..14.004
-111.14..com
-573. 23
-9.39 f
455.49...m f
DES: rKL.LD
-74.03
-109.41
100944.0..13
-534. w00.
-0.16
469.73-n..d
DBS: rKL.IDD
-73.50.a.0
-108.88+0.00s
-500.71.0.017
0.85x0.00s
469.89~0.0se
Table 1: Results Bayesian learning benchmarks. The ELBO (the higher the better) is reported for various methods and tasks. denotes that daift is not learned during training. Divergent runs are denoted with t.
5 Experiments
Task
GMM-40 (50d)
MrS- (0
Metric
Sinkharn (L)
ELRO (t)
EMC (t)
MMD (1)
Sinkhorn (L)
ELBO (f)
EMC (t)
MMD (L)
Ground Truth
875.21
0.07:0.01
449.06.
0.
1.
0.07.0.00
CMCD: rKL-R
00830.16.aies..
-37.37
0.400.0.00
1.77.000
1605. 1605.24.sa 24, 19
19.888...4
0.628.n
0.58.a.oo
CMCD: LV
2559.201 211
-37.37.0.10
0.996.0.01
0.12
1263.78.5..11
552.52
0.971...00:
0.35.s.000
CMCD: LV
2627. 96.suus
-15.85 15..85ea.r
0.9
0.13waa
1181
s0a
13.63
0.99-4ae.om
0.36.a.000
CMCD: rKL-LD
2362. 40000
-26. 28..0s
0.997 00.0
0.09.0000
915.91.Fea
-34.93
0
0.29.o.000
2301.16.a..us
-21.94aa.10
0.997.0.000
0.08.0.
915.52444.0
-31.93.o.s
4.981
0.29.0.000
DBS: rKL-R
2569 611. ose
-80. 29.
0. w0.00
0.12.0.00
1908 32
-33,
833.874
0. 0.124.c.0zz
0. 0.13.0.
DBS: LV
2073.09.ea.s4
-35. 45.00o
0.998.0..00
0.12.aa
1220.20
200.27.47
0.980aa.on -07-494n.m
0.36.n.00
DBS: IN
2067.56 56 aus
42. os.
0.998 ose
0.12,aa
1284.96.sa
-58 -58.22 28.
0.96 an ous
0.36 nm
DBS: rKL-LD
2128 99.
-33. 97
998
0.12amm
1052.31.aus
-43.67
0. 0.983ue.um
0.33
DBS: rKL-LD
2133.500mmon.
-30.4420007
0.9 0.906.. oee
0.11,00
1051
1051.342.0ss
-43.6Gun.44 448
0.903
0.33 noe
Table 2: Results on GMM-40 (50d) and MoS-10 (50d). Arrow 1 denotes that lower values of the
metric are hetter, and denotes that the higher values are the better. * denotes that diff is not learned during training. Runs that diverge after reaching the reported value marked with
Following Chen et al.] [2024a] we evaluate our model on two types of tasks: In Tab. [ we evaluate on Bayesian learning problems, where we report the ELBO due to the absence of data samples (see App. C. In Tab. 2 and Tab. 3 we evaluate Synthetic targets, where we report the Sinkhorn distance, Entropic Mode Coverage (EMC) Blessing et al.) [024] and Maximum Mean Discrepancy (MMD), which are all based samples from the diffusion sampler and samples from
Benchmarks:
8

---

the target distribution (see App. C). On multimodal tasks, a combination of high ELBO and EMC values and low Sinkhorn and MMD distances indicate good performance. Detailed descriptions of all problem types are provided in App. C.T To obtain comprehensive evaluation of the compared loss functions, we assess their performance using both diffusion bridge sampling methods CMCD and DBS. Due to different parameter counts between the two methods, we do not aim for a fair com- parison between CMCD and DBS. Our experimental setup mirrors Chen et al. [2024a]. i.c. training all methods for 40.000 training iterations with batch size of 2.000 on all tasks besides LGCP with batch size of 300. Models are trained using 128 diffusion steps. We use a commonly used diffusion sampler architecture for all diffusion- based methods described in App. D.3 In our experiments, we compare different variations of CMCD and DBS trained with different loss functions, denoted as rKL-R (Eq. 15), rKL-LD (Eq. 12), and LV (Eq. [ We provide a pseudoalgorithm for each loss in App. D.4 For each loss, we report the results of a variation, where the diffusion coefficients diff of the underlying SDE are not learned. Whenever diffusion coefficients are not learned we denote it with a * For rKL.-LD and LV we also report the results of the method when daiff is learned. For each setting. we perform grid search over the learning rates, the initial diff and the initial variance of a prior distribution prioo as described in App. D
Ablation on Learnable Diffusion Coefficients and Divergent Behavior of Log Variance Loss:
Task
Funnel (10d)
Many Well (5d)
Metric
Sinkhorn (1)
BLBO (1)
Sinkhorn (4)
ELBO (t)
We first investigate
Ground Truth
64.77_o.m1
0.
0.12,0.0
-0.54
the impact of learnable Oai R when combined
CMCD: rKL-R
113.38
-2.46
2.05x0.009
-4. 4.1420.00
CMCD: LV
95.90x2.s
-0.6710.00
0.1320.
-1.0210.01
CMCD: LV
102.941a.04 t
-0.461c.01 t
0.1320.0 t
- -1.40... t
with the rKL-LD loss, as shown in Fig. | (left and right). For simplicity we chose aaif to be constant across time and leave time dependent odiff up for future work. In Fig. [ (left) we evaluate CMCD under several configura- tions on the Brownian Bayesian learning task (see App. C.I.. Our baseline comparison is with LV loss, where daiff are not updated during training. The LV loss is highly sensitive to the initial choices of aprior and oaiff and often exhibits unstable behavior. When training oui with the LV loss, we observe divergent behavior across all hyperparameter choices. For comparison, we study our proposed loss function in two variants, with trainable and frozen Jaiff. For the frozen version we performed hyperparameter tuning the initial value of Jain- The experimental results demonstrate that the rKL.-LD loss consistently outperforms the LV loss across all tested configurations and that incorporating learnable diffusion coefficients further enhances model performance when using the rKL-LD loss. In Fig. [ (middle), we analyze the learned Jaiff across dimensions, showing their values in ascending order along with standard deviations computed from three independent seeds. The results reveal that aai systematically adopts different scales across dimensions while maintaining consistency between seeds. To test the hypothesis, whether different values of daie in each dimension are indeed beneficial, we additionally train a diffusion sampler with the rKL-LD loss with frozen sigma parameters initialized at the average Juift of the best previously trained run (CMCD: rKL-LD * Jdiff.avr in Fig. I (left)). The results show that this uniform choice of daifr across dimensions yields inferior results.
CMCD: rKL.IDD
24.04
-0.54.
0.12.0.0
-0.75.0.003
CMCD: rKL-LD
94.1622.5
0.232ooo
0.1220.0
-0.7420.003
DBS: rKL.R
111.12 sow
-0.00
0.42
-7.21.0.005
DBS: LV
335. 1257
- 2.66 asa
0.15 t
-6.27.1.200 t
DBS: LV
151.83.s..r 647
2.32.o.111
0.20.0.0
-2.92.0.012
DBS: rKL.LDD
105.97+2.377
-0500-0
0.12.0.0
-0.65.0.052
DBS: rKL-LD
108.88 88.3.8.34
-0.
0.
-0.6520.002
Table 3: Results on Funnel (10d) and Many Well (5d). Arrow de- notes that lower values of the metric are better, and denotes that the higher values are the better. denotes that daitt not learned during training. Divergent runs marked with a t.
Fig. I (right) shows how the initial value of daff affects performance by comparing CMCD: rKL- LD with and without learned diffusion coefficients on the Seeds dataset. We track the convergence using AELBO, defined as JELBOopt ELBOJ, where we estimate ELBO = -73 nats to enable visualization on logarithmic scale. The results demonstrate that when CMCD learns daift achieves much lower ELBO values regardless of the initial parameter choice. Without learned diffusion parameters, the model struggles to achieve low ELBO values.
Results on Bayesian and Synthetic Targets: Overall, the results in Tab. and Tab. Lub.3iine indicate that the divergent behavior when using the LV loss is present on most investigated benchmarks.

---

Our results further show that on Bayesian tasks in Tab. 1 when using the rKL-LD loss CMCD with frozen diff significantly outperforms its counterpart trained with LV loss on all tasks, and the version trained with rKI.-R on 4 out of 5 tasks. Similarly, DBS with frozen daift in combination with rKI.-I.D significantly outperforms DBS variants trained with other losses in 4 out of 5 cases. If we additionally train diff, we observe that CMCD improves on 4 out of 5 tasks and DBS on all tasks when trained with rKL-LD. In contrast, we observe that learning diff in combination with the LV loss deteriorates the performance of CMCD and of DBS in 4 out of 5 cases. In fact, learning diff with the LV loss often leads to unstable learning dynamics as the runs diverge in 3 out of cases for CMCD and in 4 out of 5 cases for DBS. On synthetic tasks in Tab. 2 and 3 CMCD with rKL-LD achieves the best Sinkhorn distance on MoS- 10 (50d) and GMM-40 (50d) by a significant margin, and insignificantly better Sinkhorn distance than CMCD: LV on Funnel. In terms of MMD, using the rKL.-DD loss for CMCD significantly outperforms training with the other losses in all cases. In terms of ELBO, training CMCD with rKL.-ID achieves better results than with LV on all 4 synthetic tasks. For DBS we observe that rKL.-LD * outperforms all variants of LV and rKL.RR in terms of Sinkhorn in 3 out of cases, in terms of ELBO in all 4 cases and in terms of MMD in out of 2 cases. On synthetic targets training diff with DBS improves the method only significantly on GMM-40 (50d) in terms of ELBO and MMD, and on Funnel in terms of ELBO. However, learning aain never detoriates the performance of DBS when trained with rKL-LD which is not the case for the LV loss. All methods except samplers trained with rKL-R achieve on MoS-10 (50d) and GMM-40 (50d) an EMC value close to 1., indicating that all modes are covered. Furthermore, we extend experiments on GMM-40 and MoS- 10 to a higher-dimensions setting and results in Tables and 5 support our previous findings.
6 Conclusion and Limitations
In this study, we advocate for the training of diffusion bridge-based samplers using gradients of the reverse Kullback-Leible divergence, estimated with the log-derivative trick (rKL-LD). Our analysis reveals an important insight: while the Log Variance (LV) loss and reverse KL loss are equivalent when training solely the reverse diffusion process, this equivalence does not hold when dealing with diffusion bridge samplers or learning diffusion coefficients. Furthermore, we show that the LV loss does not, in general, satisfy the data processing inequality, raising questions about its suitability for diffusion bridge samplers. Empirical results highlight the superiority of the proposed rKL-LD loss over the LV loss. Notably, employing the rKL-LD loss allows for further improvement of diffusion bridges by learning the diffusion coefficients, which also diminishes sensitivity to hyperparameter choices. We find that in contrast the rKL.IDD loss the LV loss frequently yields unstable training in particular when diffusion coefficients are Iearned. While the rKL-LD loss outperforms the LV loss and rKL loss with parametrization trick on wide range of challenging sampling benchmarks, it remains susceptible to mode collapse when hyperparameters, particularly learning rates, are not sufficiently well tuned. Future research aimed reducing the mode-collapsing tendency of rKL- based losses presents an interesting direction for further investigation. The LV loss has recently been combined with an off- off-policy sample buffer that incorporates MCMC updates and resampling strategies, which improves its stability [Sendera et al.) 2024, Chen et al.) [202a]]. A comprehensive comparison between the LV loss and rKL.LDD loss in such setting remains an important direction for future work.
7 Acknowledgements
The ELLIS Unit Linz, the LIT AI Lab, the Institute for Machine Learning, are supported by the Federal State Upper Austria. We thank the projects FWF AIRI FG 9-N (10.55776/FG9) AlAGrrenHeatingrriiss (FFG- 899943), Stars4Waters FWF Bilateral Artificial Intelligence (10.55776/OEE1). We thank NXAI GmbH, Audi AG, Sil- icon Austria Labs (SAL), Merck Healthcare KGaA, GLS (Uniy. Waterloo), TÜV Holding GmbH, Software Competence Center Hagenberg GmbH, dSPACE GmbH, TRUMPF SE Co. KG.
CITZON-CL.-2021-CCLINN
10

---

References
Michael Arbel, Alex Matthews, and Arnaud Doucet. Annealed flow transport monte carlo. In International Conference Machine Learning, pages 318-330. PMLR, 2021.
Yoshua Bengio, Tristan Deleu, Edward J. Hu, Salem Lahlou, Mo Tiwari, and Emmanuel Bengio. Gflownet foundations. CoRR, abs/2111.09266, 2021. URL htttts://aaxiiiii org/ 2b21111. 09266
Julius Berner, Lorenz Richter, and Karen Ullrich. An optimal control perspective on diffusion-based generative modeling. arXiv preprint arXiv:2211.01364, 2022.
Denis Blessing, Xiaogang Jia, Johannes Esslinger, Francisco Vargas, and Gerhard Neumann. Be- yond elbos: A large-scale evaluation of variational methods for sampling. arXiv preprint orXiv:2406.07423. 2024.
Denis Blessing, Julius Berner, Lorenz Richter, and Gerhard Neumann. Underdamped diffusion bridges with applications sampling. In The Thirteenth International Conference on Learning Representations, 2025.
Junhua Chen, Lorenz Richter, Julius Berner, Denis Blessing, Gerhard Neumann, and Anima Anand- kumar. Sequential controlled langevin diffusions. arXiv preprint arXiv:2412.0708) 2024a.
Wenlin Chen, Mingtian Zhang, Brooks Paige, José Miguel Hemánndez-Lbatto and David Barber. Diffusive gibbs sampling. In Proceedings of the s1 International Conference on Machine Learn- ing, pages 7731-7747, 2024b.
Imre Csiszár. On information-type measure of difference of probability distributions and indirect observations. Studia Sci. Math. Hungar. 2:299-318, 1967.
Marco Cuturi, I Laetitia Pang-aapaxan,hos, Yingtao Tian, Charlotte Bunne, Geoff Davis, and Olivier Teboul. Optimal transport tools (ott): A jax toolbox for all things wasserstein. arXiv preprint arXiv:2201.12324, 2022.
Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems, 34:17695-17709, 2021.
Valentin De Bortoli, Iryna Korshunova, Andriy Mnih, and Arnaud Doucet. Schrodinger bridge flow for unpaired data translation. Advances Neural Information Processing Systems, 37:103384- 103441, 2024.
Tomas Geffner and Justin Domke. Langevin diffusion variational inference. In International Con- ference on Artificial Intelligence and Statistics, pages
PMLR, 2023.
Glasserman. Monte carlo methods in financial engincering, 2004.
Jiajun He, Yuanqi Du, Francisco Vargas. Dinghuai Zhang, Shreyas Padhy, RuiKang Ou Carla Gomes, and José Miguel Hernández-Lobate No trick, no treat: Pursuits and challenges towards simulation-fire free training of neural samplers. arXiv preprint arXiv:2502.066: 06668, 2025.
Jeremy Heng, Adrian N Bishop, George Deligiannidis, and Amaud Doucet. Controlled sequential monte carlo. The Annals of Statistics, 48(5):2904-290 2929, 2020.
Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi Jaakkola. Torsional dif- fusion for molecular conformer generation. Advances Neural Information Processing Systems, 35:24240-2423 2022.
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In International Conference on Learning Representations (ICLR), 2014. URL URL https://arxiv org/ 1312. 6114
Liyun Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenRe- niew.net, 2020. htt.::ttpp://oootet http:/rrwwweeeeeeeotu/r/a/aaee net./torumiddrrgg/aaaaae

---

Nikolay Malkin, Salem Lahlou, Tristan Deleu, Xu Ji, Edward Hu, Katie Everett, Dinghuai Zhang, and Yoshua Bengio. Gflownets and variational inference. arXiv preprint arXiv:2210.00580, 2022.
Laurence Illing Midgley, Vincent Stimper, Gregor N. C. Simm, Bernhard Schölkopf, and José Miguel Hernández-Lobato. Flow annealed importance sampling bootstrap. In The Eleventh International Conference on Learning Representations, 2023. URL
net/forum?ic /ttp
Shakir Mohamed, Mihaela Rosca, Michael Figurnov, and Andriy Mnih. Monte carlo gradient esti- mation in machine learning. Journal of Machine Learning Research, 21(132):1-62, 2020.
Jesper Møller, Anne Randi Syversveen, and Rasmus Plenge Waagepetersen. Log gaussian cox processes. Scandinavian journal of statistics, 25(3):451-445 1998.
Thomas Müller, Brian McWilliams, Fabrice Rousselle, Markus Gross, and Jan Novák. Neural im- portance sampling. ACM Trans. Graph. 3855)::5:::111445:5199 2019. doi: 10.1 10.1/43334415 URL
https://doi. org/ 10. 1145/3341156
Radford M Neal. Slice sampling. The annals of statistics, 31(3):705-767, 2003.
Kim A. Nicoli, Shinichi Nakajima, Nils Strodthoff, Wojeiech Samek, Klaus-Rober Müller, and Pan
Kessel. Asymptotically unbiased estimation of physical observables with neural samplers. Phys. Rev. E, 101:023304, Feb 2020. doi: 10.1 www.100/hhysee...00.0233355 org/doi/100 www.:/1/////y.ee..100.... 101. 023304
aps.
Frank Noé and Hao Wu. Boltzmann generators sampling equilibrium states of many-body systems with deep learning. CoRR, abs/182..017222,, 2018. UtLhttp::/aasiiv.orr
Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trends@ in Machine Learning. 11(5-6):3556677, 2019.
D. J. Rezende, S. Mohamed, and D. Wierstra. Stochastic backpropagation and approximate in- ference in deep generative models. In International Conference on Machine Learning (ICML), 2014.
Lorenz Richter and Julius Berner. Improved sampling via learned diffusions. In International Conference on Learning Representations, 2024. URL htttpp://ppenrevvee..et///oo ImPpbbPjQb/
Lorenz Richter, Ayman Boustati, Nikolas Nüsken, Francisco Ruiz, and Omer Deniz Akyildiz. Var- grad: a low-variance gradient estimator for variational inference. Advances Neural Information Processing Systems, 33:13481 2020.
Sebastian Sanokowski, Sepp Hochreiter, and Sebastian Lehner. A diffusion model framework for unsupervised neural combinatorial optimization. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 43346 43367. PMLR, 21-27 Jul 2024. [tt. mlr press/v235. saaooooski12a html
Sebastian Sanokowski, Wilhelm Berghammer, Martin Ennemoser, Haoyu Peter Wang, Sepp Hochre- iter, and Sebastian Lehner. Scalable discrete diffusion samplers: Combinatorial optimization and statistical physics. In The Thirteenth International Conference on Learning Representattonss 2025. URL
Marein Sendera, Minsu Kim, Sarthak Mittal, Pablo Lemos, Luca Scimeca, Jarrid Rector-Brooks Alexandre Adam, Yoshua Bengio, and Nikolay Malkin. Improved off-policy training of diffusion samplers. In The -iigth-Eighth Annual Conference Neural Information Processing Systems, pages 1-30. ACM, 2024.
Andy Shih and Stefano Ermon. Probabilistic circuits for variational inference in discrete graphical models. Advances in neural information processing systems, 33:4635-4666 2020.
Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances neural information processing systems, 34:1415- 1428, 2021.
12

---

Francisco Vargas, Will Grathwohl, and Arnaud Doucet. Denoising diffusion samplers. arXiv preprint wrXiv:2302.18882 2023.
Francisco Vargas, Shreyas Padhy, Denis Blessing, and Nüsken. Transport meets variational
ference: Controlled monte carlo diffusions. In The Twelfth International Conference Learning Representations, 2024.
Dian Wu, Lei Wang, and Pan Zhang. Solving statistical mechanics using variational autoregressive networks. Phys. Rev. Lett. 122:080602, Feb 2019. doi: 10.1 10.1103/PhysRevLett 122.00002..0000 URL
poctps://likk: aps
http://.....00/////eteeettttttt. 122. 080602
Qinsheng Zhang and Yongxin Chen. Path integral sampler: A stochastic control approach for sam- pling. In Proceedings of the International Conference on Learning Representations (ICLR), 2022. URL http:://ppomorrrrre....fffrorrrp
13

---

A Derivations and Proofs
A.1 Reverse KL Divergence with Reparametrization Trick
A popular choice for training diffusion bridges is the rKL objective Richter and Bemer 2024
Vargas al. 2024 Chen et al., 2024a Blessing et 2024 He al. 2025 Blessing et 2025 which is given by:
log Ya (XTT) Po.e (Xorr) (Xor()
DKL DKi Po.e =Ex go
(14)
This loss involves expectation over the variational distribution qa.e. This is a typical use case for the reparameterization trick [Glasserman, 2004). This method gained popularity as a gradient estimator for training generative models [Kingma and Welling 2014 Rezende al.) 2014 This technique frequently yields lower variance than the log dog-derivative trick Mohammed et al. 2020 a discussion). the context of diffusion samplers, the reparameterization trick was introduced in Zhang and Chen 2022]. In this method the trajectories of the reverse process Xo.r Ta (X are reparameterized as a function (a.r(<0:TT) : [9m,..0, -T. of independent Gaussian noise ~ (O,T) T) and the model parameters vo, Samples X, at time step f. are successively gen- crated via the Euler-Maruyama update function Ja.rt := Xx-1 (X,, e.), which is used to numerically integrate the reverse process Eq. Substituting this reparameterization into the rKL loss yields, with slight abuse notation:
DKL DKI(Ia.(EETT) (o.o(E0:T) Bear-N((..)) log (f Pe.u (fa (Cor)) (fa,((EUT))
(15)
This expression highlights that the gradients of the loss with respect to (a, v) can propagate into the expectation of the deterministic function fa.v. enabling direct gradient calculation. In diffu- sion samplers the frequent iterative application of Ja.u. is likely to contribute to the vanishing or exploding gradient problem, which might explain the suboptimal behavior of rKL loss when min- imized with usage of the reparametrization trick Zhang and Chen [2022]. Moreover, Vargas et al. 2023) demonstrates that applying stop-gradient operation to the Langevin parametrization term in the PISgradNet architecture (see App. enhances the stability of the reparametrization trick in diffusion samplers. Without this modification, the direct application of the reparametrization trick through the gradient of the energy function is unstable, though the stop-gradient approach introduces an additional bias into the gradient.
A.2 Definition of v LDww Operator
Let (o(()) (X) be a variational probability distribution parameterized by 0 and let Oo be function that depends these parameters. We want to define the L.D.a operator based on the gradient of the expectation of the observable with respect to samples X we (X) given by (O(X)) X-)x(X)-
The L.Dw is then defined in the following way:
V DD.00 (Os(X)) (x) := Ex (X) [(Oo(X)-
- by Ve log (oogo(())] Ex [V)Og(X)],
where 1 Ex [log wo(X) When applying the V.D. operator we often use that in the case of O(X) log wo (X) we have E Ex- o log (/) [X) = 0.
A.3 Gradient of the Log Variance Loss
In the following, we derive the gradient of the log variance loss, where we consider a reverse and forward diffusion process with shared parameters 0:
14

---

(1OS go (o(oorr) po (
E Exorrw log (o(r)) po (X(T))
27
VoDir((o.P) Volxor
log 90 (Xu.r)
log go (o(x)r) Pa (Xor))
1).
go (X(xr))
Vo log (X:T) pe Pa(Nos (XoT)
-xxor
Exorraw
V log
Exorma
po po (Xar)
po (Xorr)
2 log 40 2o ((oor) po (X(TT)
{10R
log (o(xuur'
go (Xur)
-xxorr
Exor
Va log
pe (o(Xo:))
po (Xo.r)
Where we have used that
Exoz=w
2 log (Xorr) po (XXor)
E Exorrw
log
9o (X(T)))) pe (Xo:))
))
Exor
Vo log po (Xoo))
go (Xor))
since E Exrrow tion.
Ve log
Qa( Nor) po (Xor)
does not depend on Xo.r and can be pulled out of the first expecta-
We can recover the gradients in Eq. 13 with respect to a by setting a, (x(x))) 9a (Xa:) and po (Xar) - p( ((oor)). Likewise we can recover gradient with respect to with qo (XaT) (( (XoT) and Po (Xor)) Po( (ar) and with respect to with (XoT) qo (XoT) and po (x(o)) (X(T))..
A.4 Gradient of the Reverse Kullback-Leibler Divergence Loss
In the following the gradient of the rKL divergence is derived, when it is optimized with the usage of the dervvvvive trick:
log go Xo.r) po (Xar)
VoDKu FoDki(o((aar)) po V VoLxxrru
log 9 (Xa:))
Ve log (o(Xr)) po (Xu:r)
= Exur-go
Vo log qo (Xorr)
Exx
Po (Xurr)
log go (((x))) Po (o(oT)
log 2o (Xorr)
V log go (x:r) po (Xor)
Exar=go
E Ex
Vo logog go (Xorr)
Exar--o
Po (XXrr)
[og y0( (o(Xnr))
log Lo (Xorr)
= Exxr-ao
Exar
Vo log go (Xar)
E EXor-go Ve log pe (Xor))
po (Xu:r)
po (XuT)
where we use lines two to three the fact that 6 X [V log ge (Xo:T) and that [Xarr] is baseline that leads to gradient updates with lower variance. In line
Exur-go log [o())r)
three to four we have used that Exær~qo [Vo log (Xorr) Exar=ga (X(r)) Vogo (xgo() I VVag (o = Ve. go (Xaa) d'xor=
We can recover the gradients in Eq. 12 with respect to by setting 0 + qo (o(o))) qa (Xo:T) and (Xor)) + p(Xorr) Likewise we can recover gradient with respect to with qo (X(x)) ((oTT) and po (XoT) (X())) and with respect to with (x(T)) (Xo.)) and (((T))) + (XaT)..
A.5 Counterexamples to Data Processing Inequality for the Log Variance Loss
A.5.1 Counterexample in the Continuous Domain
In the following we provide a counter example where
Var(z)- log (xx) p(x)
& Var -g log g(2 y) p(a, y)
(16)
15

---

where the distributions g(2,3) 3) is absolutely continuous with respect to p(x,3) 3) and the other way around. By definition, absolute continuity q << holds if every measurable set [0. 1]2 for which y( A) = 0 also satisfies ({4) O.
Consider the following distributions:
2n a n 2
x+ <
p(x,y) 1,
(a,
(r. e [0,1]
b
x+y21,
1++1
Joint variance. Since the lower and upper triangles A {x + y < 1} and B {o + 2 13 have equal area I. under g the probabilities of the two constant values are (r,(A) and Pr, (B) 2" 2
With p(2,1) y) 1 we have log ((..)) log ((,, y) taking only values log a, log Hence
p(a.x)
a
log + log b
2
2
b log b, 2
Ey log" a(X, Y)]
log 2
and therefore
Var log (X, Y) p(X, Y)
log* a log b
(
a
log
2 logb)
b
Marginal variance.
First obtain the marginal g(x):
(x)) = So' a( dy So 1-x 2n
n+ dy [r]
2
1 n dy n+
(n x0 n)).
Set
C= n
== 1 n,
ma (u [n, 1]).
Then, with the change of variables
(so dx du the moments under X g
Eg log (X)] S/¹ (ogg()))))))))))))))) dr m
C J wlog(Cu) du,
logg (x)] = C f. ulog" du.
m
Use log(Cu)) = log C + log u and the elementary integrals
1 n
du
2
S. ulog du
1
n
n²
log
4
4
S( u logg du
1
n
n²
n²
log
log
n
4
2
Let
1- n²
n
2
1
n
n
n²
A ==
B
log
+
D :=
log
+
log
2
2
4
Then
C. Eg [log (X)] ((log C)A+ B).
C'
log x(X]]
m
and hence
(xX) pxX)
Varx-
log
Log {(X)] q( ( (x))))²
with the expectations given above (substitute C²= 2 n-T m n., and A, B, D defined).
16

---

Numeric evaluation for = 10 yields
Substituting = 10 (so a a 0.001998002, x 1.998001998)
Va log pl
s 0.04762179,
Varx-q log p
0.24995228 = 0.25,
showing the claimed DPI violation.
A.5.2 Counterexample in the Díscrete Domain
Let p,g be probability distributions, similar as in Eq. 1 but without learnable parameters, with re-
spective marginal and conditional distributions and (X, ~ q. Then the data processing inequality for the LV loss is
Var(x)- log ((X) p(X)
log ((,))
S Var(x,))~
(17)
p(X,Y)
which can be decomposed such that
Var log gX,,Y)
log ((X)
log x((x)
p(X,)) = -Var(x.Y)
Var(x,Y)-g
p(X)
x
x( .log g(xx
g(Ylx)
+ Co
log
p(X)
Therefore, if we find distributions p,q such that
Va log ( 2 Cov(X
log 2 log g(X) p ((
50
p(Y[X)
p(X)
we disprove inequality Eq. 17 For this, let X (0, P.g be defined on X' x² with marginal probabilities
p(X X = 0) 0.1, pX 1) 0.9, ((X=0)=00 0.9, (x= = 1) 0.1, =)= 1, 1(V = 1)=00
and conditional probabilities
p(Y 0| 0) = 0.5, p(Y 0| 0.1.
From :(Y 1) we get ((X, Y 1) 0. By standard arguments e.g., used for KL-divergence we can interpret
q(X,Y = 1) (Y 1|x) = (Y = 1[X) IY 1[X)
as being zero since lim x xlogx= 0 which results the following simplifications
log (((1x) x()
log ((vlz)
Σ g(x',v) ./ log (1011²')
Var(x,Y)-g
Σ
((x,1))
p(y(x)
p(2)(x'))
(x.,))ex22
log 2(0)x) p(0)x))
log (0px²)) p(0lx))
Σ (xx, 0) zet
Σ
qa
#ee
2
Σ (()) [1oRM(0/0)- zex
Σ (((x) (ogg'))))())))))
by using ((x,0)= ((x). Analogously we get for
(((x) llog (x) (V(X)) p(x)
Cov(X,Y)~
log
- [logp(0/))- Σ ((x') log ((<))log logp(0)z²)) x'EX
(log ((x) p(x)
Σ ((x') log ((x')
p(x*)
zed
Z'E
By inserting the corresponding probability values we have
(YIX)
log gx (X)
g( ((YX))
Co log
-0.6365, Var(x.y)
log
0.2331
px (X)
p(Y
which demonstrates that the data processing inequality does not always hold for the LV loss.
17

---

B Additional Experiments
We extend our experimental evaluation include multimodal benchmarks in higher dimensions such as GMM-40 (100d) and MoS-10 (100d). The results in Table 4 and S demonstrate that the proposed rKL-LD loss achieves performance comparable to the LV loss with respect to Sinkhorn distance and Entropic Mode Coverage (EMC), while yielding superior results in terms of the Evi- dence Lower Bound (ELBO), log Z, and Maximum Mean Discrepancy (MMD). On GMM-100D, we observe that combining the LV loss with CMCD leads to divergent behavior in several random seed runs, substantially degrading the average performance metrics. A similar issue arises on MoS- 100D when training DBS with the LV loss. Notably, learnable diffusion coefficients consistently enhance performance when trained with the rKL-LD loss, but degrade performance when used with the LV loss.
Task
GMM-40 (100d)
Metric
Sinkhorn (L)
EIBBO (1)
log (t)
EMC (1)
MMD (4)
CMCD: rKL.R
7152.88.1mm.5
-795.09.6.ap
0.99520
0.11620:001
CMCD: LV
22237. 10827.67..
(s1) x 10
-774.91 647
0. 0.866.0.084
0.268.o.101
CMCD: LV
12420.83.cm 6008..11
(-2.14.2.s5) x 10
0.951.0.043
0.138
CMCD: rKL-LD
a079.38=206. 248
-45. 1020.091
-18.09 20.099...40 341
0.99620.000
0. 0.084:o.001
CMCD: rKL-LD
(043.26z16..4
-5..44...00
-18.88a0.300 3944
0.99620.000
0.085x0.001
DBS: rKL.RR
5637.36 953
270.61.0.rr
133.60=0.850
0.997z0.001
0. 1.0000000
DBS: LV
5024.582171.00
230.83.0.10
-05.53
0.99820.oo1
0.1 7..000000
DBS:
4938.60.10s.cc 103.66001844.
252.56.0.20
-116. 166.76.. 7..1
0.998
0. 0.184.0.00.
DBS: r
5163.31 280
225.99.0.132
-101. -101.80.a..o 8..0.955
0.997
0.1 78.00000
DBS: rKL-LLD
5132.655124.06s
198.53.0.127
-87.59 8188
0.997
0..00
Table 4: Results on GMM-100D tasks. Arrow denotes that lower values of the metric are better, and denotes that higher values are better. * denotes that daiff is not learned during training. Best results are highlighted in bold.
Task
MoS- 10 (100d)
Metric
Sinkhorn (4)
ELBO (t)
log Z (t)
EMC (1)
MMD
CMCD: rKL-R-
1809.48..00 512
114.862ao02
188.25a0.a
0.977
0.265=0.00
CMCD: LV
27.6.16..0... ass
-9..44...222
-6.377..400
0.981.a.00s
0.286.0.00.
CMCD: rKL-LD -
2852.13.2
-68.21 -68.21.0.22
-31.71a0.173
0.957=0.007
0.260a0.001
DBS: rKL-R
62.32.2..761...mm
-183.53
-113.88
0.990
0. 0.298-0.001
DBS: LV
(1.9610.7e)
10
(-1.68.0.67) x 10
484009 484409.64. 189e2.
0.303.c.201
0.298aa.com
DBS: rKL-LD
3626.16x1s57.oom
-121.77a0.920
-69.53.0.347
0.996
0.29520.001
Table 5: Results on MoSD tasks. Arrow denotes that lower values of the metric are better, and denotes that higher values are better. denotes that daif is not learned during training. Best results are highlighted in bold.
C Metrics
Evidence Lower Bound The evidence lower bound is lower bound on log and computed with:
Po (XxT)
ELBO E Exerrga (Xer)
$ log Z
(18)
Jo (Xorr)
Sinkhorn Distance The Sinkhorn distance an entropic regularization of the 2- Wasserstein (W)) optimal transport (OT) distance Peyré et al. [2099) providing a principled alternative to ELBO for evaluating sample quality. Unlike ELBO, which is often insensitive to mode collapse Blessing et al. [2024), the Sinkhorn distance measures the discrepancy between generated and target distributions, offering better insights into sample diversity and multimodal coverage. As it requires access to ground-truth samples, its use is limited to synthetic tasks Chen et al.) 2024a Following Blessing
18

---

et al. 2024 Chen et al. [024a], we compute the Sinkhorn distance using the ott package Cuturi al. 2022 and report it as a primary metric for the uppropriate benchmarks. Sinkhorn distances are computed on Funnel using 2000 samples and on GMM and MoS using 16000 samples. On May Well Sinkhorn distances are computed with 100000 samples using code from Richter and Berner 2024].
Maximum Mean Discrepancy Given samples X p(X) and samples Y q(Y the Maximum Mean Discrepancy (MMD) computed in the following way:
MMD(p,9) (ani))V-Vxxxeeeenn xxxxx- [A(X, X(] Ey,y"=0(r) (N(Y,Y') 2I --x (k(X,
2025 and compute the MMD with an average bandwidth of 100 and using K = 10 kemelss. The code is based on an implementation of Chen al. [2024b].
where k(z,y)= Di=1 exp (- | y|l²) where Ky is the bandwith. We follow He
Entropic Mode Coverage
The Entropic Mode Coverage Blessing et al. [0244] is given by
1 N
M Σ mm
EMC == Ex Xx)))
Σ Xo
pC, Xo) log p(g, Xo)
where i {1....,M) M) and p(E., Xo) is the probability corresponding to the mixture component with the highest likelihood at Xo- The optimal value of EMC is 1, i.e. every mode from the target distribution is covered and it is O. in the worst case.
C.1 Benchmarks
C.1.1
Bayesian Learning tasks
These tasks involve probabilistic inference where the true underlying parameter distributions are unknown, requiring Bayesian approaches for estimation.
Bayesian Logistic Regression (Sonar and Credit). We consider Bayesian logistic regression for binary classification two well-established benchmark datasets, frequently used for evaluating variational inference and Markov Chain Monte Carlo (MCMC) methods. The model's posterior distribution is given by:
Ptarget (x) p(x) II 1-1
Bernoulli (y; (x uu))
where the dataset consists of standardized input-output pairs ((u, y1 ))) Our evaluation includes the Sonar dataset (d 61, n = 208) and the German Credit dataset (d 25, n = 1000). The prior distribution is chosen as standard Gaussian p NOO, I) for Sonar, whereas for German Credit, follow the implementation of Blessing et al. [002], which omits an explicit prior by setting p
Random Effect Regression (Seeds): The Seeds dataset (d = 26) is modeled using hierarchical random effects regression framework, which captures both fixed and random effects to account for variability across different experimental conditions. The generative process is specified as:
T Gamma(O. (0.01, 0.01)
ao, 01, 02 a12 (0,
1
by N 0, /F/
i= 1,...,21,
logits, ag + 01ι; 023i 012 by i= 1=..,,,2 .,,.,,,,, Binomial (logits,, N.). i= 1. 21.
The inference task involves estimating the posterior distributions т. ag. d1, 02, d12, and the random effects b, given observed values of æ;, Dt. and N,. This model is particularly relevant for analyzing seed germination proportions, where the inclusion of random effects accounts for heterogeneity in experimental conditions; see Geffner and Domke 2023 for further details.
19

---

Time Series Models (Brownian): The Brownian motion model (d 32) represents a discretized stochastic process commonly used in time series analysis, with Gaussian observation noise. The generative model follows:
Ginn LogNormal 2),
Globs LogNormal 2),
(00 Amma),
2, ([ Oinn), i 2,...,3..,, Vi ~ (( Gobs),
i = 1,...,.3 30.
The inference objective is estimate Oins: Oobs: and the latent states { 1 based on the available 10 and {3/ i-20- with the middle observations missing. This missing-date struc-
observations {y: i-|
ture increases the difficulty of inference, making it useful benchmark for probabilistic time series modeling; see Geffner and Domke [223]
Spatial Statisties (LGCP): The Log- Gaessaassin Cox Process (LGCP) is a widely used spatial model in statistics [Moller et al. as the locations of tree saplings. The target density defined over a discretized spatial grid of size
1998). which describes spatially distributed point processes such
d 40 40 1600, and follows:
1.Yi exp (I;) (1t))
Prarget N(: Σ)
II exp i=1
where y represents observed dataset, and μ and Σ define the mean and covariance of the prior distribution. This formulation leads to a complex spatial dependency structure. We focus on the more challenging unwhitened variant of the model, which retains the full covariance structure and thus introduces stronger dependencies between grid locations, as described in |Heng et al. 2000)), Arbel et al. 2021].
C.1.2 Synthetic targets:
For these tasks, ground-truth samples are available, allowing for direct evaluation of inference accu- racy.
Mixture distributions (GMM and MoS): We consider mixture models where the target distribu- tion consists of m mixture components, defined as:
Prarget
1=1
The Gaussian Mixture Model (GMM), adapted from Blessing al. 2024 is constructed with
m = 40 Gaussian components:
p: = N
[4 W4-ua(--44 (-40, 40),
where (a u) denotes a uniform distribution over [4, u We set the dimensionality to d 50 in the experiments in Tab.
The Mixture of Student's 's t-distributions (MoS) follows a similar construction but uses Student's t-distributions with two degrees of freedom (t2) as the mixture components:
pj = 12 + 14g #4 (-10
where μ4 denotes the translation of each component. We set the dimensionality to d 50 in the experiments in Tab. 2
For both the GMM and MoS tasks, the component locations μ4 remain fixed across experiments using a predefined random seed to ensure reproducibility.
20

---

Funnel: The Funnel distribution, originally introduced in Neal (2003). serves as challenging
benchmark due to its highly anisotropic shape. is defined as:
Ptarges (x) = (2 02) o")) Zi0
exp(xi))),
(19)
where x² 9 for any number of dimensions d 2 In our main experiments, we consider the case
d = 10. To maintain consistency with prior benchmarks [Blessing et al. [0024), we apply hard constraint by clipping all sampled values to the interval [-30,30].
Many Well: A standard problem in molecular dynamics is sampling from the stationary distribu- tion of a Langevin dynamic. For our experiments, the resulting d-dimensional many-well potential corresponds to the following unnormalized density
cxp(
( m
Ptarget (x)) = exp
1 -
-8)
x
with separation parameter å E (0, 00) and m N combined double wells. Following Chen et al.
[024aa]. Berner et al. [0022 we set d m = 5 and 8 4 which results 2n 32 well-separated modes. Since Ptarget factorizes in the dimensions ground truth statisties and samples can be obtained by numerical integration and rejection sampling.
D Experimental Details
D.1 Evaluation
In the Bayesian learning task, we compute the average of the ELBO over three independent runs, each estimated using 2000 samples. For the LGCP task, the evaluation is performed using 300 samples. The ELBO values reported in Tab. | represent the best ELBO achieved during training. For synthetic tasks in Tab. 2 and Tab. 3 we report all metrics based on the checkpoint the end of training. When the run diverges on synthetic benchmarks, we report the result based on the checkpoint with the best ELBO. We use 16000 samples on MoS-40 50D and GMM-40 50D, 2000 samples on Funnel, and 100000 samples on Many Well to estimate the Sinkhorn distance. In Tab.[]] Tab. 2 and Tab. 3 we report the average metric value together with the standard error over three seeds. For MoS-40 50D and GMM-40 50D we calculate the standard error over 7 independent seeds, where for each seed the metrics are averaged over 30 repetitions due to large standard errors of the sinkhorn distance. MMD values are computed by using 4000 samples. Ground truth Sinkhorn distances and MMD values are computed by calculating these distances between two independent set of samples from the target distribution. The ground truth values are then averaged over three independent seeds.
D.2 Hyperparameter tuning
Benchmarks In benchmark experiments in Sec. 5 we perform for each method a grid search over ddiff, prior. the learning rate of the model. The learning rate of the diffusion parame- ters such as Opior and daiff is always chosen to be equal to the model learning rate. On all Bayesian learning tasks, we perform for CMCD and DBS a grid search over duiffinn = 0.31, priior,inis init (0.5. 1.0} and the learning rate Amodel.SDE e (0.005, 0.002, 0.001 On Brownian and German Credit, we found that if oaif is not learned a finer grid-search over dan is necessary. There- fore on Brownian, we additionally add dai 0.05 and on German Credit dai 0.01 to the grid search.
On MoS 50D and GMM 50D, we follow Chen et al. [024a] and fix ppio to a high initial value. We found that o prio 80 yielded the best results. We found that small model learning rates and compared to that large learning rates of the interpolation parameters between the prior and the target distribution work well. Therefore we adapt the grid search to o jnit Gfiftiini Ainterpol = {0.01, 0.001} and the learning rate Amodel,SDE {0.0001, 0.00005, 0.00001) } for CMCD. For DBS, the interpolation between the prior and the target distribution is not learned. Therefore, we did not additionally search over Ainterpol but increased the size of the grid search by searching over iprior,init = (60, 80). On Many Well we conduct grid search over aiiffinn = (0.05, 0.2). init = (0.5, 1.0, 2.0), Amodel,SDE e (0.001, 0.0001, 0.00001).
1.5),
21

---

Grid searches are performed over 8000 training iterations on all targets except MoS 50d and GMM 50d, where 12000 training iterations are performed. The best run is chosen according to the best ELBO value at the end of training on Bayesian tasks and on Synthetic targets according to the run with the best Sinkhorn distance at the end of training. The best hyperparameters are then run for 40000 training iterations. On MoS 50d and GMM 50d in Tab. 2 the best Sinkhorn distance is sometimes achieved initialization. In this case, the checkpoint is excluded as it has only slightly better Sinkhorn values but much worse ELBOs than the runs at the end of training.
All scripts for running grid searches can be found in the code in /Configs/Sweep// and the final selected hyperparameters can be found in
Ablations In ablation experiments in Sec. 5 we iteratively tuned hyperparameters such as prior,init and inif and learning rates for each method. For CMCD-LV and CMCD-LV * we found it hard to find a good-performing diffusion coefficient. Therefore, we used the learned av- erage diffusion coefficient of CMCD-rKL-LD daiff as starting point for iterative hyperparameter tuning which resulted in decent performance of CMCD-LV and CMCD-LV
D.3 Architecture
Score parametrization way:
The learned score with parameters parameterized the following
so sX(), clip(55(X,,t)) + (t clip
log Ix, 10²). -00",
(20)
where &g t)) and (t) are two separate neural networks which are parameterized with the PIS- gradnet architecture from Vargas al. [2024) with 64 hidden neurons and 2 layers. The score is related to the control by ug (X,) se (X,). The usage of the Vx, T term in the neural
network parametrization is often refered langevin parametrization He al. 2025].
Parametrization of prior distribution: Similarly to Chen et al. 2024a and Blessing et al. 2024 and parameterize the prior distribution गт the following way:
mr (Mo, diag(exp((o))
where /io E IRd and logarithmic standard deviations lo IRd are learnable parameters. In contrast to Chen et al. [2024a] and Bllsssigg et al. [2024] we do not update μi and lo via the reparameterization trick training progresses but also with the usage of the log- log-derivative trick.
Parametrization of interpolation parameters: For each diffusion time step E (0, (O..., T 1) we parameterize the interpolation parameter (t) in the following way:
softplus(O.)
(o(t) ΣnD softpluus())
(21)
where 0 E T are learnable parameters. Each variable of 0 E R is initialized to zero.
Parametrization diffusion coefficient: We keep diffusion coefficients constant across time steps and parameterize it as o; expy, where y = log Piitt In principle, one could parameterize similarly as the interpolation parameters, which would allow for time-adaptive schedule. However, we leave this for future work.
Training All parameters trained with the usage of the RAdam Liu et al. (0200 optimizer. We use gradient clipping by norm at the value of 1. The learning rate decays with a cosine learning rate schedule from Astart to >start/ 100.
D.4 Pseudoalgorithms
In the following we provide pseudoalgorithms for the computation of the LV loss, rKI.-LD loss and the rKL-R Loss:
22

---

Algorithm Computation of the LV Loss
1: Given: Batch of diffusion paths Xor Yo.v computed with the Buler-Manuyama integration
(Eq. 4)
2: Xo:r stop-grad (grrr(Xor))
detach the gradient from the Fuler-Maruyama integration
3: compute log I (Xgrr 9.2 with Eq.
4: compute Po. (X D:T.) with Eq. 6
5: compute Loss L(a, V) = Var (log Ia. X
(
6: Backpropagate through the loss and update (a, Φ, 1)) using Adam optimizer
Algorithm 2 Computation of the rKL-R Loss
Given: Batch of diffusion paths Xor qa.v computed with the Huler-Maruyama integration (Eq. 4)
2: compute log 9o. (Xor) with Eq.
3: compute log Po. (XaT) with Eq. 6
4: compute Loss L(a, o, v) = mean log Po. (Xar)
(o.r))r)
5: Backpropagate through the loss and update (a, Ø, v) using Adam optimizer
Algorithm 3 Computation of the rKL-LD Loss:
Averages are always computed over the batch dimension
1: Given: Batch of diffusion paths Xor (Xur) computed with the Buler-Maruyane inte-
gration (Eq. 4)
2: X6:r stop grad Xo.T)
3: compute log I X1
9. with Eq.
4: compute log X
(.7)) with Eq. 6
5: compute control variate 590,0
a.0.0 mean log Yo (n..(
Po. (o.o(Xor)
6: compute advantages A stop atto-gradient log Rowul X
fai 0,
po Po N2
7: compute Loss L(a, db v) mean [A* log (a xbr) mean Plog Po. Xbr)
8: Backpropagate through the loss and update (a, Φ, v) using Adam optimizer
23

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Claims in the abstract are supported in the method section and in the experi- ment section.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
There is Conclusion and Limitations section the end of the paper.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specificattin, asymptotic approximations only holding locally). The au- thors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, c.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the ap- proach. For example, facial recognition algorithm may perform poorly when image resolution is low images taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to ad- dress problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't i acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [es]]

---

Justification:
Full proofs are given in the Appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theo- rems.
The proof' can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be comple- mented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Code is available and all experimental details are either described in the main text or in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken make their results reproducible verifiable.
Depending on the contribution, reproducibility can accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or the contribution is a specific model and empirical evaluation, it may be necessary either make it possible for others replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of model checkpoint, other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all sub- missions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or a way to re- produce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case au- thors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models. it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
25

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
Code is provided together with a description to run the code.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines lhttps://nips: cc/
public/guides/CdeeSuuaaaa sor.eeeeSsbbmisiiooooooo..... for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: Cubm ssi onPolicy) for more details.
"niipe.cc/pulic//u...oo Sutpa/////uudee//eeeububmmbbbb
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: All important experimental details are properly described.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question:
Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: All results are reported together with error bars.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
26

---

The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It OK to report I-sigma error bars, but one should state it. The authors should prefer- ably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of eπrоrs is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
Guidelines:
The answer means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics htttp::///eurpp:: http:///rrrr..cccoooac//ccaaaaaa
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This is mainly basic research and does not have negative or positive societal impact at the current stage.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact spe- cific groups), privacy considerations, and security considerations.
27

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitiga- tion strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility MI).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Misuse is not possible here.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by re- quiring that users adhere usage guidelines or restrictions access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Citations are given.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name the license (c.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use the pack- age should be provided. For popular datasets, wwwppppwpiiiiaeee thoode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the li- cense of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
28

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documenta- tion provided alongside the assets?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
For crowdsourcing experiments and research with human subjects, does the pa- per include the full text instructions given to participants and sereenshots, applicable, well as details about compensation (if any)?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contri- bution of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, cura- tion, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equiva- lent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
29

---

Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy pattps://neurl/e http:e://heerrppsscc///ooeeeerrp /20.ppps...////teeeeeee///////// for what should or should not be described.
30


Reviews:
[{"review_id": "FePu4IjyH2", "reviewer": "Reviewer_Ce3h", "summary": "In this paper the problem setup considered is learning to sample from a data distribution $\\pi_0$ with no available samples/observations from $\\pi_0$. For this problem setup, different loss functions are contrasted and compared: namely, the reverse Kullback-Leibler (rKL) loss or the log variance (LV) loss.  In prior works, the LV loss has performed better than the rKL loss but the paper argues this is due to the application of the reparametrisation trick when computing the rKL loss and instead motivates using the rKL loss with the log derivative trick.", "questions": "1. See Weakness 2. above.\n2. Can you explain what they mean by “frequent iteration of the parametrization trick” in line 169? Do you mean due to the successive application of Euler-Maruyama to get the values $X_t$ of one trajectory for different values of $t$?\n3. Related to Weakness 3 above, can you say anything about the gradients obtained via the two different tricks and why the log derivative trick is expected to outperform the reparametrisation. Or for example, can you quantify the error induced via the reparametrisation trick?\n4. Currently, I’m not completely sure what I should take away from the gradients of the rKL loss with the log derivative, other than in specific circumstances (as stated in section 3.2) it coincides with the LV loss. Are there any more conclusions to be drawn for example in relation to the performance of the loss?", "limitations": "Yes.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n1. The paper is well-written and structured and I found it easy to follow with illuminating descriptions of losses through the lens of $f$-divergences. The paper presents concise comparisons between these commonly used losses as well as the log variance loss(es) which are not $f$-divergences.\n\n2. The submission shows that the log variance loss does not satisfy the data processing inequality. I find this to be an important observation when choosing such a loss for diffusion models.\n\n3. The submission shows through experiments that the rKL loss with the log derivative trick performs well in comparison to both the rKL loss with the reparametrisation trick as well as the LV loss. The experiments conducted are the same as those in Chen et al. (2024a). \n\n\nWeaknesses:\n\n1. One of the main points made in the paper is to advocate for the use of the rKL loss with the log derivative trick. However, the definition of the log derivative trick is relegated to the appendix, alongside the reparametrisation trick. I would find it beneficial to have this in the main paper. Moreover, I would find it beneficial to have more of an explanation of the log derivative trick alongside the definition. \n\n2. Related to this, in A.2 is there a reason for switching to the notation $\\langle O(X) \\rangle_{X \\sim p}$ as opposed to $\\mathbb{E}_{X\\sim p} [O(X)]$ for the expectation as used in the rest of the paper?\n\n3. The paper motivates using the rKL objective with the log derivative trick instead of the LV loss well. However, I do not understand why it should be expected that the rKL objective with the log derivative trick is expected to do better than the rKL loss with the reparametrisation trick, or how these methods relate to each other. A more concrete comparison would improve the clarity of this work.\n\n\nTo conclude, I think this is a well-written paper and I find the contributions of comparing the loss functions compelling. However, I think there are some weaknesses to be addressed, mostly in terms of a comparison between the reparametrisation trick and the log derivative trick for the gradients. For this reason I currently lean towards a weak accept, but I’m open to revising my score should the weaknesses be easy to resolve.", "was_revised": true, "final_justification": "My main weakness was comparing the log derivative and reparametrisation trick which has been addressed in the rebuttals.", "rebuttal": "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.\n\n**Weaknesses:**\n\n**W1:** \nWe 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.\n\n**W2:** \nThank you for your comment, we changed this in the updated version and stick to notation style $\\mathbb{E}_{X \\sim p} [O(X)]$.\n\n**W3:** \nRegarding 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.\nWe 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.\nFor 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.\nWe 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.\n\n**Questions:**\n\n**Q1:**\n See our response to W2 above.\n\n**Q2:** \nThanks 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.\n\n**Q3:** \nFrequently, 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.\n\nWe 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.\n\n**Vargas, Francisco, Will Grathwohl, and Arnaud Doucet. \"Denoising diffusion samplers.\" arXiv preprint arXiv:2302.13834 (2023).**\n\n**Mohamed, Shakir, et al. \"Monte Carlo gradient estimation in machine learning. arXiv e-prints, page.\" arXiv preprint arXiv:1906.10652 (2019).**\n\n**Keskar, Nitish Shirish, et al. \"On large-batch training for deep learning: Generalization gap and sharp minima.\" arXiv preprint arXiv:1609.04836 (2016).**"}, {"review_id": "lokylixc70", "reviewer": "Reviewer_ET35", "summary": "- Examines the losses used to train diffusion bridges and controlled monte carlo diffusion where both the forwards and reverse process are parameterized. \n- Examine potential issues with the log variance (LV) loss typically used, focusing on stability and adherence to the data processing inequality \n- Argue for the application of the rKL loss with the log-derivative trick\n- Demonstrates empirical gains with their proposed rKL-LD loss when applied to both CMCD and DBS training when compared to LV and rKL-R loss\n- Additionally, show that rKL-LD allows for further improvements when combined with learnable diffusion parameters, which tend to destabilize training with LV loss", "questions": "In Table 2, it seems that rKL-LD provides significant benefits over LV when applied to CMCD for the GMM-40 task. Furthermore, the ability to learn the diffusion parameters also provides additional gains. However, when applied to DBS, it seems that the opposite holds: rKL-LD decreases performance relative to the LV loss, and LV seems to further improve performance with learnable diffusion parameters whereas rKL-LD performance deteriorates. Why does this occur? \n\nWhat are the practical effects of the LV loss not satisfying the data processing inequality? While I understand this as a theoretical issue, I am not sure if this results in tangible effects. It would be nice if the authors could connect this back to observable phenomena that occur during training. \n\nOverall, I am leaning weakly towards acceptance. I think if the above questions are answered, I will be happy to increase my score.", "limitations": "I think that it would be nice if there were empirical results on more difficult tasks, such as NICE [1], or alanine dipeptide [2], but the results included are sufficient already. \n\n[1] Score-Based Diffusion meets Annealed Importance Sampling. Doucet et al. NeurIPS 2022 \n[2] Flow Annealed Importance Sampling Bootstrap. Midgley et al. ICLR 2023.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "*Quality*\n \nOverall, this is a good quality paper. While the experiments are not very high dimensional, this is typical for the field of diffusion sampling (from what I know), and the authors do include results on LGCP, which is 1600 dimensional. The main point — that LV loss is not suitable for diffusion sampling when both the forward and reverse process is learned — seems fairly clear. The analytical comparison of the gradients for the different losses seems accurate. \n\n*Clarity*\n\nOne significant strength of this paper is its coverage of the problem background and prior work. I appreciate the thoroughness of this section.\nOverall, the paper is well organized and the writing style is very easy to follow. I do think that the writing style could be improved for the experiments section. In particular, the paragraphs were very long and difficult to follow — it would be better to break them up into smaller paragraphs that focus on specific findings. But I do not find this to be that significant of an issue \n\n*Significance*\n\nThe findings of this paper are fairly significant. I agree with the discussion in section 3.3, which discusses the importance of learning diffusion parameters, as different target distributions may require different diffusions to adequately capture the entropy of the target distribution while balancing the signal to noise ratio to ensure stable training.  \n\n*Originality*\n\nThe biggest weakness of this paper is the originality: as noted by the authors in the related works, the application of the log-derivative trick to the rKL loss has been explored before in the context of discrete diffusion samplers. \n\nWhile the proposed loss is not new, I think the analysis and the application towards CMCD / DBS with learnable diffusion parameters is novel. Thus I find the overall originality to be fair.", "was_revised": true, "final_justification": "I initially had concerns regarding the novelty, and the importance of the DPI violation. The authors have thoroughly addressed each part, and thus I feel that there are no other issues with this paper as far as I can see. I think this would be fairly impactful paper for the field of diffusion bridges / diffusion samplers, so I think a 5 is appropriate.", "rebuttal": "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.\n\n**Weaknesses:**\n\n**On Clarity:** \n\nThank 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.\n\n**On Originality:** \n\nThe 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.\nApart 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.\n\n**Questions:**\n\n**On Results in Table 2:** \n\nWe 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. \n\n**Impact of DPI violation**\n\nIt 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.\n \n**Limitations**\n\n**On More Difficult Tasks:**\n\nDue 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.\n\n| GMM 100-D | Sinkhorn | ELBO | log Z | EMC | MMD |\n|----------|----------|----------|----------|----------|----------|\n| CMCD: rKL-R $\\Large \\star$ | $7152.88\\pm 177.312$    | $-795.09\\pm 6.491$    | $-528.09\\pm 2.995$     | $0.995\\pm 0.000$     | $0.116\\pm 0.001$     |\n| CMCD: LV $\\Large \\star$   | $22237.03\\pm 10627.664$    | $-1957894528.00\\pm 1354142383.491$    | $-774.91\\pm 466.647$   | $0.866\\pm 0.084$    | $0.268\\pm 0.101$    |\n| CMCD: LV  | $12420.83\\pm 6668.17$   | $-2141430400.00\\pm 2031539132.986$    | $-1037.09\\pm 931.118$    | $0.951\\pm 0.043$    | $0.138\\pm 0.015$    |\n| CMCD: rKL-LD $\\Large \\star$   | $6079.38\\pm 205.248$     | $-45.10 \\pm 0.091$    | $-18.09\\pm 0.341$    |  $\\mathbf{0.996 \\pm 0.000}$   | $0.084\\pm 0.001$    |\n| CMCD: rKL-LD    | $\\mathbf{6043.26\\pm 166.427}$     | $\\mathbf{-45.84\\pm 0.08}$    | $\\mathbf{-18.88\\pm 0.394}$   | $\\mathbf{0.996\\pm 0.000}$     | $\\mathbf{0.085\\pm 0.001}$    |\n\n| GMM 100-D | Sinkhorn | ELBO | log Z | EMC | MMD |\n|----------|----------|----------|----------|----------|----------|\n| DBS: rKL-R $\\Large \\star$    | $5637.36 \\pm 155.953$     | $-270.61 \\pm 0.275$    | $-138.60\\pm 0.85$    | $\\mathbf{0.997\\pm 0.001}$ | $0.179\\pm 0.000$   |\n| DBS: LV $\\Large \\star$  | $\\mathbf{5024.58 \\pm 171.079}$     | $-230.83\\pm 0.151$    |$-105.53\\pm 0.972$    | $\\mathbf{0.998\\pm 0.001}$    | $0.178\\pm 0.000$    |\n| DBS: LV   | $\\mathbf{4938.60\\pm 103.434}$   | $-252.56\\pm 0.299$    | $-116.76\\pm 1.061$   | $\\mathbf{0.998\\pm 0.001}$    |  $0.184\\pm 0.001$    |\n| DBS: rKL-LD $\\Large \\star$   | $\\mathbf{5163.31\\pm 206.48}$     | $-225.99 \\pm 0.132$     | $-101.80\\pm 0.953$    | $\\mathbf{0.997\\pm 0.001}$     | $0.178 \\pm 0.000$    |\n| DBS: rKL-LD   | $\\mathbf{5132.65\\pm 124.068}$   | $\\mathbf{-198.53\\pm 0.127}$    | $\\mathbf{-87.59\\pm 0.818}$    | $\\mathbf{0.997\\pm 0.001}$     | $\\mathbf{0.171\\pm 0.000}$     |\n\n\n| MoS 100 D | Sinkhorn | ELBO | log Z | EMC | MMD |\n|----------|----------|----------|----------|----------|----------|\n| CMCD: rKL-R $\\Large \\star$ | $1809.48\\pm 298.512$\t| $-114.86\\pm 0.021$\t| $-68.25\\pm 0.48$ \t| $0.977\\pm 0.004$ \t| $0.265\\pm 0.001$ \t|\n| CMCD: LV $\\Large \\star$\t| $2716.16\\pm 808.633$\t| $-97.44\\pm 0.622$\t| $-66.37\\pm 0.404$   | $\\mathbf{0.981\\pm 0.005}$\t| $0.286\\pm 0.001$\t|\n| CMCD: rKL-LD $\\Large \\star$  | $\\mathbf{1852.13 \\pm 271.333}$ \t| $\\mathbf{-68.21\\pm 0.322}$\t| $\\mathbf{-31.71\\pm 0.173}$\t|  $0.957\\pm 0.007$   |$\\mathbf{0.260\\pm 0.001}$\t|\n\n| MoS 100 D | Sinkhorn | ELBO | log Z | EMC | MMD |\n|----------|----------|----------|----------|----------|----------|\n| DBS: rKL-R $\\Large \\star$ | $9732.62\\pm 3761.867$\t| $-183.53\\pm 0.867$\t| $-113.88\\pm 1.51$ \t| $0.990\\pm 0.003$ \t| $0.298\\pm 0.001$ \t|\n| DBS: LV $\\Large \\star$  | $1964053495808.00 \\pm 783156518610.415$ \t| $-168435.12 \\pm 66865.384 $\t|$-48409.64 \\pm 18962.519 $\t| $0.393\\pm 0.201$\t| $0.298\\pm 0.001$\t|\n| DBS: rKL-LD $\\Large \\star$   | $\\mathbf{3626.16\\pm 1557.097}$ \t| $\\mathbf{-121.77 \\pm 0.926}$ \t| $\\mathbf{-69.53 \\pm 0.347}$\t| $\\mathbf{0.996 \\pm 0.001}$ \t| $\\mathbf{0.295 \\pm 0.001}$\t|\n\nWe 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."}, {"review_id": "rqvz6tA0y4", "reviewer": "Reviewer_KBTa", "summary": "This paper advocates for training diffusion samplers with the reverse Kullback-Leibler (rKL) loss using the log-derivative trick (rKL-LD), arguing it is more effective than the commonly used Log Variance (LV) loss. The authors demonstrate that the equivalence between LV and rKL losses is applicable with non-learnable forward processes, while it does not apply to diffusion bridges, and the LV loss fails to satisfy the data processing inequality in this context. Experimental results show that the rKL-LD loss consistently outperforms the LV loss on several benchmarks, offering more stable training and reduced sensitivity to hyperparameters, especially when diffusion coefficients are learned. While rKL-LD shows superior performance, the paper notes its potential for mode collapse if not well-tuned and suggests future work could explore its combination with off-policy sample buffers.", "questions": "When comparing the performance of rKL-LD to the LV loss, did your implementation of the LV baseline include a replay buffer? Recent work has shown that off-policy sample buffers significantly improve the stability and performance of LV-based training, making this a critical detail for a fair comparison.", "limitations": "The primary limitation is a potential contradiction in the paper's core theoretical argument, as the claimed gradient equivalence between rKL-LD and LV loss does not actually hold for practical implementations like CMCD where parameters are shared. Furthermore, the method still inherits the known drawbacks of reverse KL, such as a tendency for mode-collapse, and its overall robustness could be better demonstrated on more complex benchmark tasks.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\n- The paper provides a new perspective to understand the connection between reverse KL and log variance.\n- The rKL-LD loss provides more stable training behaviour than log-variance and rKL with the reparametrisation trick (rKL-R).\n\n**Weaknesses**\n\n- **Unaddressed Limitations of Reverse KL:** The paper rightly focuses on the benefits of rKL-LD, but it doesn't fully address the known drawbacks. The reverse KL divergence is inherently mode-seeking and prone to mode collapse, a problem the authors acknowledge rKL-LD is still susceptible to. Similarly, the log-derivative trick typically produces higher variance gradients than the reparameterization trick. A deeper analysis of why rKL-LD manages to outperform other rKL methods, despite these issues, would strengthen the paper\n- **Apparent Contradiction in Gradient Equivalence:** A central weakness lies in the claimed equivalence between the gradients of rKL-LD and LV loss (Equations 12-13). This equivalence holds only when the forward process parameters are not learned. However, in models like CMCD, the forward and backward processes share the same network parameters. This means the gradient of rKL-LD is *not* equivalent to the gradient of LV in the practical CMCD and DBS setup. The authors should clarify this in the paper to avoid misunderstanding.\n- based on the equation 12 and 13, the gradient of rKL-LD is equivalent to LV when \\phi = v = empty. However, this is generally not true for CMCD. Note that CMCD parametrise the dirft term of the forward and backward using the same network parameters, which means both q and p depends on the same parameters and the gradiet of rKL-LD is not equivalent to the gradient of LV\n- **Lack of Clarity in Final Objectives:** To address the ambiguity mentioned above, the paper would be significantly improved by explicitly stating the final training objective functions for each method tested (e.g., CMCD-rKL-LD, DBS-rKL-LD). Including a clear algorithm or pseudo-code detailing the training procedure would further enhance reproducibility and clarify the exact implementation differences\n- **Limited Scope of Evaluation:** The experimental validation could be made more compelling by including more challenging benchmark tasks. Evaluating the proposed method on complex energy landscapes, such as the Lennard-Jones potential, would provide a more robust assessment of its capabilities and limits.", "was_revised": true, "final_justification": "I increased my score from 3 to 4 since most of my concerns have been addressed, such as the clarity in final objectives and the gradient equivalence. Although there still exists the limitation of empirical study, I think it is okay due to the scalability of the diffusion bridge sampler. I believe this paper still introduces good insights to re-understanding the connection between log-variance and KL-divergence.", "rebuttal": "We thank the reviewer for this very detailed and thoughtful critique of our manuscript.\n\n**W1 Limitations of rKL:**\n\nFirst, 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.\nThe 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.\n\n\n**W2: Apparent Contradiction in Gradient Equivalence:**\n\nWe 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.\nTherefore, 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.\n\n**W3 Lack of Clarity in Objectives:**\n\nWe 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.\n\n**Pseudocode:**\n\n### **Computation of the rKL-LD Loss**\n### *Averages are always computed over the batch dimension.*\n---\n\n**Given:**\n\nA batch of diffusion paths $X_{0:T} \\sim q_{\\alpha, \\nu}(X_{0:T})$\ncomputed via Euler-Maruyama integration (see Eq. 4).\n\n---\n1. **Detach gradients:**\n$$\nX^\\prime_{0:T} \\gets \\mathrm{stop\\_grad}(X_{0:T})\n$$\n2. **Compute log-densities:**\n$$\n\\log q_{\\alpha, \\nu}(X^\\prime_{0:T}) \\quad \\text{(Eq. 5)}\n$$\n$$\n\\log p_{\\phi, \\nu}(X^\\prime_{0:T}) \\quad \\text{(Eq. 6)}\n$$\n3. **Compute control variate:**  \n$$\nb = \\mathrm{mean} \\left[ \\log \\frac{q_{\\alpha, \\nu}(X^\\prime_{0:T})}{p_{\\phi, \\nu}(X^\\prime_{0:T})} \\right]\n$$\n4. **Compute advantages (detached):**\n$$\nA^* = \\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]\n$$\n5. **Compute loss:**\n$$\n\\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]\n$$\n6. **Backpropagate and update parameters:**\n$$\n(\\alpha, \\phi, \\nu) \\leftarrow \\mathrm{Adam} \\left( \\nabla_{(\\alpha, \\phi, \\nu)} \\mathcal{L}(\\alpha, \\phi, \\nu) \\right)\n$$\n\n**W4 Limited Scope of Evaluation:**\n\nTraining 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.\n\nWe 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.).\n\nWe thank you again for your time and effort used to review our work and hope that all of our answers properly address your questions.\n\n**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.**\n\n**Junhua Chen, Lorenz Richter, Julius Berner, Denis Blessing, Gerhard Neumann, and Anima Anandkumar. Sequential controlled langevin diffusions. arXiv preprint arXiv:2412.07081, 2024a.**\n\n**Blessing, Denis, et al. \"Beyond ELBOs: A large-scale evaluation of variational methods for sampling.\" arXiv preprint arXiv:2406.07423 (2024).**\n\n**Mohamed, Shakir, et al. \"Monte Carlo gradient estimation in machine learning. arXiv e-prints, page.\" arXiv preprint arXiv:1906.10652 (2019).**\n\n**Keskar, Nitish Shirish, et al. \"On large-batch training for deep learning: Generalization gap and sharp minima.\" arXiv preprint arXiv:1609.04836 (2016).**\n\n| GMM 100-D | Sinkhorn | ELBO | log Z | EMC | MMD |\n|----------|----------|----------|----------|----------|----------|\n| CMCD: rKL-R $\\Large \\star$ | $7152.88\\pm 177.312$    | $-795.09\\pm 6.491$    | $-528.09\\pm 2.995$     | $0.995\\pm 0.000$     | $0.116\\pm 0.001$     |\n| CMCD: LV $\\Large \\star$   | $22237.03\\pm 10627.664$    | $-1957894528.00\\pm 1354142383.491$    | $-774.91\\pm 466.647$   | $0.866\\pm 0.084$    | $0.268\\pm 0.101$    |\n| CMCD: LV  | $12420.83\\pm 6668.17$   | $-2141430400.00\\pm 2031539132.986$    | $-1037.09\\pm 931.118$    | $0.951\\pm 0.043$    | $0.138\\pm 0.015$    |\n| CMCD: rKL-LD $\\Large \\star$   | $6079.38\\pm 205.248$     | $-45.10 \\pm 0.091$    | $-18.09\\pm 0.341$    |  $\\mathbf{0.996 \\pm 0.000}$   | $0.084\\pm 0.001$    |\n| CMCD: rKL-LD    | $\\mathbf{6043.26\\pm 166.427}$     | $\\mathbf{-45.84\\pm 0.08}$    | $\\mathbf{-18.88\\pm 0.394}$   | $\\mathbf{0.996\\pm 0.000}$     | $\\mathbf{0.085\\pm 0.001}$    |\n\n| GMM 100-D | Sinkhorn | ELBO | log Z | EMC | MMD |\n|----------|----------|----------|----------|----------|----------|\n| DBS: rKL-R $\\Large \\star$    | $5637.36 \\pm 155.953$     | $-270.61 \\pm 0.275$    | $-138.60\\pm 0.85$    | $\\mathbf{0.997\\pm 0.001}$ | $0.179\\pm 0.000$   |\n| DBS: LV $\\Large \\star$  | $\\mathbf{5024.58 \\pm 171.079}$     | $-230.83\\pm 0.151$    |$-105.53\\pm 0.972$    | $\\mathbf{0.998\\pm 0.001}$    | $0.178\\pm 0.000$    |\n| DBS: LV   | $\\mathbf{4938.60\\pm 103.434}$   | $-252.56\\pm 0.299$    | $-116.76\\pm 1.061$   | $\\mathbf{0.998\\pm 0.001}$    |  $0.184\\pm 0.001$    |\n| DBS: rKL-LD $\\Large \\star$   | $\\mathbf{5163.31\\pm 206.48}$     | $-225.99 \\pm 0.132$     | $-101.80\\pm 0.953$    | $\\mathbf{0.997\\pm 0.001}$     | $0.178 \\pm 0.000$    |\n| DBS: rKL-LD   | $\\mathbf{5132.65\\pm 124.068}$   | $\\mathbf{-198.53\\pm 0.127}$    | $\\mathbf{-87.59\\pm 0.818}$    | $\\mathbf{0.997\\pm 0.001}$     | $\\mathbf{0.171\\pm 0.000}$     |\n\n\n| MoS 100 D | Sinkhorn | ELBO | log Z | EMC | MMD |\n|----------|----------|----------|----------|----------|----------|\n| CMCD: rKL-R $\\Large \\star$ | $1809.48\\pm 298.512$\t| $-114.86\\pm 0.021$\t| $-68.25\\pm 0.48$ \t| $0.977\\pm 0.004$ \t| $0.265\\pm 0.001$ \t|\n| CMCD: LV $\\Large \\star$\t| $2716.16\\pm 808.633$\t| $-97.44\\pm 0.622$\t| $-66.37\\pm 0.404$   | $\\mathbf{0.981\\pm 0.005}$\t| $0.286\\pm 0.001$\t|\n| CMCD: rKL-LD $\\Large \\star$  | $\\mathbf{1852.13 \\pm 271.333}$ \t| $\\mathbf{-68.21\\pm 0.322}$\t| $\\mathbf{-31.71\\pm 0.173}$\t|  $0.957\\pm 0.007$   |$\\mathbf{0.260\\pm 0.001}$\t|\n\n| MoS 100 D | Sinkhorn | ELBO | log Z | EMC | MMD |\n|----------|----------|----------|----------|----------|----------|\n| DBS: rKL-R $\\Large \\star$ | $9732.62\\pm 3761.867$\t| $-183.53\\pm 0.867$\t| $-113.88\\pm 1.51$ \t| $0.990\\pm 0.003$ \t| $0.298\\pm 0.001$ \t|\n| DBS: LV $\\Large \\star$  | $1964053495808.00 \\pm 783156518610.415$ \t| $-168435.12 \\pm 66865.384 $\t|$-48409.64 \\pm 18962.519 $\t| $0.393\\pm 0.201$\t| $0.298\\pm 0.001$\t|\n| DBS: rKL-LD $\\Large \\star$   | $\\mathbf{3626.16\\pm 1557.097}$ \t| $\\mathbf{-121.77 \\pm 0.926}$ \t| $\\mathbf{-69.53 \\pm 0.347}$\t| $\\mathbf{0.996 \\pm 0.001}$ \t| $\\mathbf{0.295 \\pm 0.001}$\t|"}, {"review_id": "9Mu1zdddCq", "reviewer": "Reviewer_iCJN", "summary": "This paper reveals the relationship between the gradients of Log Variance(LV) loss and the reverse Kullback-Leibler (rKL) loss for various cases. It also observes that for diffusion bridges, the LV loss does not satisfy the data processing inequality. Therefore, the authors argue that LV loss is not a good choice for diffusion bridges and advocate for training diffusion bridge samplers with the rKL loss with the log-derivative trick.\nExperiments show that it consistently outperforms diffusion samplers trained with LV and original rKL-R losses.", "questions": "1. Larger datasets should also be verified.\n2. Clarify the problems written above about the data processing inequality and the advantage of this method.", "limitations": "Yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths:\n1. This paper gives a strong theoretical analysis of the relationship between the gradients of Log Variance(LV) loss and the reverse Kullback-Leibler (rKL) loss, and finds that their gradients do not hold in the context of diffusion bridge samplers or when diffusion coefficients are learned. This analysis provides more insights for understanding these two common losses.\n\n2. This paper demonstrates that using rKL loss with the log-derivative trick can achieve better performance and is significantly less sensitive to hyperparameter tuning. Combined with learnable diffusion terms, the performance can be improved further.\n\nWeakness:\n1. The paper shows a counter-example to prove LV loss does not satisfy the data processing inequality, but for their counter-example, the distribution p(X, Y) is a discrete distribution and p(Y) is a $\\delta$-distribution, which seems not rational for its assumption, Y is a diffused distribution from X. So I don't think this counter-example is convincing to support their points.\n\n2. Just for my understanding, no matter that the LV loss does not satisfy the data processing inequality or the non-equivalence gradients between LV loss and rKL loss, they can't strongly support that the LV loss is not better than rKL loss from a theoretical view.  Theoretically, at least LV loss can converge to the optimal only when the two diffusion processes are identical. Therefore, the advantage of the proposed method is mostly from experimental observation rather than a convincing theoretical analysis.\n3. For experiments, the proposed method is only verified in some small-dimensional datasets. Can it also be verified on a larger dataset(larger than 1000 or 10000)?", "was_revised": true, "final_justification": "Since most of my concerns have been addressed, I raise my score to 4.", "rebuttal": "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.\n\n**Response to Comment W1: Counterexample is not convincing**\n\nFor 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.\nHowever, to further illustrate our point of the DPI violation, we provide an additional counterexample based on continuous distributions:\n\nConsider the following distributions:\n- $ q(x, y) $ is uniform on the triangle defined by the vertices $ (0,0) $, $ (1,0) $, and $ (0,1) $.\n- $ p(x, y) $ is uniform on the entire unit square $ [0,1]^2 $.\nThus, we have:\n- $ q(x) = 2  (1 - x) $\n- $ p(x) = 1 $ \n\nIn 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.\n\nThis example yields:\n- $ \\mathrm{var}_{x,y \\sim q(x,y)} \\log \\frac{q(x,y)}{p(x,y)} = 0 $ for the joint distributions.\n- $ \\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) $.\n\nHence, 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.\nWe will include this example in the appendix.\n\n**Response to Comment W2: Theoretical convergence of LV loss**\n\nThe 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.\n\n**Response to Comment W3: Verification on larger datasets**\n\nWe 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]).\nWe 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].\nEvaluating 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.\nNevertheless, 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.\n\n**Response to Limitations: Clarify the advantage of this method.**\n\nAs 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.\nWe 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.\n\n\n[1] Sendera, Marcin, et al. \"Improved off-policy training of diffusion samplers.\" Advances in Neural Information Processing Systems 37 (2024): 81016-81045.\n\n[2] Chen, Junhua, et al. \"Sequential Controlled Langevin Diffusions.\" The Thirteenth International Conference on Learning Representations.\n\n[3] Blessing, Denis, et al. \"Underdamped Diffusion Bridges with Applications to Sampling.\" The Thirteenth International Conference on Learning Representations.\n\n\n| GMM 100-D | Sinkhorn | ELBO | log Z | EMC | MMD |\n|----------|----------|----------|----------|----------|----------|\n| CMCD: rKL-R $\\Large \\star$ | $7152.88\\pm 177.312$    | $-795.09\\pm 6.491$    | $-528.09\\pm 2.995$     | $0.995\\pm 0.000$     | $0.116\\pm 0.001$     |\n| CMCD: LV $\\Large \\star$   | $22237.03\\pm 10627.664$    | $-1957894528.00\\pm 1354142383.491$    | $-774.91\\pm 466.647$   | $0.866\\pm 0.084$    | $0.268\\pm 0.101$    |\n| CMCD: LV  | $12420.83\\pm 6668.17$   | $-2141430400.00\\pm 2031539132.986$    | $-1037.09\\pm 931.118$    | $0.951\\pm 0.043$    | $0.138\\pm 0.015$    |\n| CMCD: rKL-LD $\\Large \\star$   | $6079.38\\pm 205.248$     | $-45.10 \\pm 0.091$    | $-18.09\\pm 0.341$    |  $\\mathbf{0.996 \\pm 0.000}$   | $0.084\\pm 0.001$    |\n| CMCD: rKL-LD    | $\\mathbf{6043.26\\pm 166.427}$     | $\\mathbf{-45.84\\pm 0.08}$    | $\\mathbf{-18.88\\pm 0.394}$   | $\\mathbf{0.996\\pm 0.000}$     | $\\mathbf{0.085\\pm 0.001}$    |\n\n| GMM 100-D | Sinkhorn | ELBO | log Z | EMC | MMD |\n|----------|----------|----------|----------|----------|----------|\n| DBS: rKL-R $\\Large \\star$    | $5637.36 \\pm 155.953$     | $-270.61 \\pm 0.275$    | $-138.60\\pm 0.85$    | $\\mathbf{0.997\\pm 0.001}$ | $0.179\\pm 0.000$   |\n| DBS: LV $\\Large \\star$  | $\\mathbf{5024.58 \\pm 171.079}$     | $-230.83\\pm 0.151$    |$-105.53\\pm 0.972$    | $\\mathbf{0.998\\pm 0.001}$    | $0.178\\pm 0.000$    |\n| DBS: LV   | $\\mathbf{4938.60\\pm 103.434}$   | $-252.56\\pm 0.299$    | $-116.76\\pm 1.061$   | $\\mathbf{0.998\\pm 0.001}$    |  $0.184\\pm 0.001$    |\n| DBS: rKL-LD $\\Large \\star$   | $\\mathbf{5163.31\\pm 206.48}$     | $-225.99 \\pm 0.132$     | $-101.80\\pm 0.953$    | $\\mathbf{0.997\\pm 0.001}$     | $0.178 \\pm 0.000$    |\n| DBS: rKL-LD   | $\\mathbf{5132.65\\pm 124.068}$   | $\\mathbf{-198.53\\pm 0.127}$    | $\\mathbf{-87.59\\pm 0.818}$    | $\\mathbf{0.997\\pm 0.001}$     | $\\mathbf{0.171\\pm 0.000}$     |\n\n\n| MoS 100 D | Sinkhorn | ELBO | log Z | EMC | MMD |\n|----------|----------|----------|----------|----------|----------|\n| CMCD: rKL-R $\\Large \\star$ | $1809.48\\pm 298.512$\t| $-114.86\\pm 0.021$\t| $-68.25\\pm 0.48$ \t| $0.977\\pm 0.004$ \t| $0.265\\pm 0.001$ \t|\n| CMCD: LV $\\Large \\star$\t| $2716.16\\pm 808.633$\t| $-97.44\\pm 0.622$\t| $-66.37\\pm 0.404$   | $\\mathbf{0.981\\pm 0.005}$\t| $0.286\\pm 0.001$\t|\n| CMCD: rKL-LD $\\Large \\star$  | $\\mathbf{1852.13 \\pm 271.333}$ \t| $\\mathbf{-68.21\\pm 0.322}$\t| $\\mathbf{-31.71\\pm 0.173}$\t|  $0.957\\pm 0.007$   |$\\mathbf{0.260\\pm 0.001}$\t|\n\n| MoS 100 D | Sinkhorn | ELBO | log Z | EMC | MMD |\n|----------|----------|----------|----------|----------|----------|\n| DBS: rKL-R $\\Large \\star$ | $9732.62\\pm 3761.867$\t| $-183.53\\pm 0.867$\t| $-113.88\\pm 1.51$ \t| $0.990\\pm 0.003$ \t| $0.298\\pm 0.001$ \t|\n| DBS: LV $\\Large \\star$  | $1964053495808.00 \\pm 783156518610.415$ \t| $-168435.12 \\pm 66865.384 $\t|$-48409.64 \\pm 18962.519 $\t| $0.393\\pm 0.201$\t| $0.298\\pm 0.001$\t|\n| DBS: rKL-LD $\\Large \\star$   | $\\mathbf{3626.16\\pm 1557.097}$ \t| $\\mathbf{-121.77 \\pm 0.926}$ \t| $\\mathbf{-69.53 \\pm 0.347}$\t| $\\mathbf{0.996 \\pm 0.001}$ \t| $\\mathbf{0.295 \\pm 0.001}$\t|"}]

Author rebuttal:

Meta-review: This paper studied two loss functions used in literature for training diffusion samplers: log variance and reverse KL losses. It advocates for using the reverse KL loss with the log-derivative trick (rKL-LD), particularly for diffusion bridges where both forward and reverse processes are learned. The conclusion goes against an earlier work that shows log variance loss outperforms rKL loss - which the author pointed out to be an artifact of using the reparametrization trick for optimization. 

All reviewers vote for acceptance, highlighting the significance of the finding, the paper's excellent clarity, and the empirical gains. Several concerns regarding equivalence results and counter examples was sufficiently addressed during rebuttal. Therefore, I recommend acceptance.
Decision: Accept (poster)