Improved sampling via learned diffusions
Authors: Lorenz Richter
Julius Berner
Venue: ICLR 2024 poster (2024)
Keywords: Schrödinger bridge
sampling from densities
stochastic optimal control
diffusion-based generative modeling
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: We propose a unifying framework to improve diffusion-based samplers.

Abstract: Recently, a series of papers proposed deep learning-based approaches to sample from target distributions using controlled diffusion processes, being trained only on the unnormalized target densities without access to samples. Building on previous work, we identify these approaches as special cases of a generalized Schrödinger bridge problem, seeking a stochastic evolution between a given prior distribution and the specified target. We further generalize this framework by introducing a variational formulation based on divergences between path space measures of time-reversed diffusion processes. This abstract perspective leads to practical losses that can be optimized by gradient-based algorithms and includes previous objectives as special cases. At the same time, it allows us to consider divergences other than the reverse Kullback-Leibler divergence that is known to suffer from mode collapse. In particular, we propose the so-called \textit{log-variance loss}, which exhibits favorable numerical properties and leads to significantly improved performance across all considered approaches.

Published as a conference paper at ICLR 2024
IMPROVED SAMPLING VIA LEARNED DIFFUSIONS
Lorenz Richter* Zuse Institute Berlin dida Datenschmiede GmbH
Julius Berner
Caltech
jberner@caltech..
ri eichter@zike
ABSTRACT
Recently, series of papers proposed deep learning- approaches to sample from target distributions using controlled diffusion processes, being trained only on the unnormalized target densities without access to samples. Building previous work, we identify these approaches as special cases of a generalized Schrödinger bridge problem, seeking a stochastic evolution between given prior distribution and the specified target. We further generalize this framework by introducing a variational formulation based on divergences between path space measures of time- reversed diffusion processes. This abstract perspective leads to practical losses that can be optimized by gradient-basee algorithms and includes previous objectives as special cases. At the same time, it allows to consider divergences other than the reverse Kullback-Leite In particular, we propose the so-called log-variane loss, which exhibits favorable numerical properties and leads to significantly improved performance across all considered approaches.
divergence that known to suffer from mode collapse.
INTRODUCTION
Given function p: Rdd 10, oc), consider the task of sampling from the density
Ptarget
P Z
with Z p(x) dr,
where the normalizing constant Z is typically intractable. This represents a crucial and challenging problem in various scientific fields, such as Bayesian statistics, computational physics, chemistry, or biology, see, c.g., Liu Liu (2001); Stoltz et al. (2010). Fueled by the success of denoising diffusion probabilistic modeling (Song & Ermon, 2020; Ho et al., 2020; Kingma et al., 2021; Vahdat et al., 2021; Nichol & Dhariwal, 2021) and deep learning approaches to the Schrödinger bridge (SB) problem (De Bortoli al., 2021: Chen et al., 2021a: Koshizuka & Sato, 2022), there is significant interest in tackling the sampling problem by using stochastic differential equations (SDEs) which are controlled with learned neural networks to transport given prior density Pprior the target Ptarget.
Recent works include the Path Integral Sampler (PIS) and variations thereof (Tzen & Raginsky, 2019; Richter, 2021: Zhang& & Chen, 2022; Vargas et al., 2023b), the Time-Reversed Diffusion Sampler (DIS) (Berner et al., 2024), as well as the Denoising Diffusion Sampler (DDS) (Vargas et al., 2023a). While the ideas for such sampling approaches based on controlled diffusion processes date back to earlier work, see, e.g., Dai Pra (1991); Pavon (1989), the development of corresponding numerical methods based on deep learning has become popular in the last few years.
However, up to now, more focus has been put on generative modeling, where samples from Ptarget are available. As a consequence, seems that for the classical sampling problem, i.e., having only an analytical expression for p Ptarget. but samples, diffusion-basee methods cannot reach state-of-the-an performance yet. Potential drawbacks might be stability issues during training, the need differentiate through SDE solvers, mode collapse due to the usage of objectives based reverse Kullback-Leible (KL) divergences, see, e.g., Zhang & Chen (2022); Vargas et al. (2023a).
In this work, overcome these issues and advance the potential of sampling via learned diffusion processes toward more challenging problems. Our contributions can be summarized as follows:
"qqual contribution (the author order determined by numpy. random. rand (1)).

---

Published as a conference paper at ICLR 2024
We provide a unifying framework for sampling based on learned diffusions from the perspective of measures on path space and time-reversals of controlled diffusions, which for the first time connects methods such SB, DIS, DDS, and PIS.
20.0
17.
log vortance loss
15.0
This path space perspective, in consequence, allows us to consider arbitrary divergences for the optimization objective, whereas existing methods solely rely on minimizing reverse KI. divergence, which prone to mode collapse.
10.0
75-
30000
G0000
gradient
In particular, we propose the variaaiae divergence that avoids differentiation through the SDE solver and allows to balance exploration and exploitation, resulting significantly improved numerical stability and performance, see Figure 1.
Figure 1: Improved convergence of our proposed log-variance loss for a double well problem, see Section 4 for further details.
1.1 RELATED WORK
There exists a myriad of Monte Carlo-based methods for sampling from unnormalized densities, e.g. Annealed Importance Sampling (AIS) (Neal, 2001), Sequentil Monte Carlo (Del Moral et al., 2006; Doucet et al., 2009) (SMC), or Markov chain Monte Carlo (MCMC) (Kass et al., 1998). Note, however, that MCMC methods are usually only guaranteed reach the target density asymptotically, and the convergence speed might be too slow in many practical settings (Robert et al., 1999). Variational methods such as mean-field approximations (Wainwright et al.. 2008) and normalizing flows (Papamakarios et al., 2021; Wu et al., 2020; Midgley et al., 2023; Vaitl et al., 2024) provide an alternative. Similar to our setting, the problem of density estimation cast into an optimization problem by fitting parametric family of tractable distributions to the target density.
We build our theoretical foundation on the variational formulation of bridge problems proposed by Chen et al. (2021a). We recall that the underlying ideas were established decades ago (Nelson, 1967; Anderson, 1982; Haussmann & Pardoux, 1986; Föllmer, 1988), however, only recently applied to diffusion models (Song et al., 2020) and SBs (Vargas, 2021; Liu et al., 2022). While the numerical treatment of SB problems has classically been approached via iterative nested schemes, the approach in Chen et al. (2021a) uses backward SDEs (BSDEs) to arrive at single objective based KL divergence. This objective includes the (continuous-time) ELBO of diffusion models (Huang et al., 2021) as special case, which can also be approached from the perspective of optimal control (Berner et al., 2024). For additional previous work on optimal control in the context of generative modeling, we refer to De Bortoli et al. (2021); Tzen & Raginsky (2019); Pavon (2022); Holdijk al. (2022).
Crucially, we note that our path space perspective on the variational formulation of bridges has not been known before. Our novel derivation only relies on time-reversals of diffusion processes and shows that, in general, corresponding losses (in particular the one in Chen et al. (2021a)) do not have a unique solution as they lack the entropy constraint of classical SB problems. However, in special cases, we recover unique objectives corresponding to recently developed sampling methods, e.g., DIS, DDS, and PIS. Moreover, the path space perspective allows us to extend the variational formulation of bridges different divergences, in particular to the log-variance divergence that has originally been introduced in Nüsken & Richter (2021). Variants of this loss have previously only been analyzed in the context of variational inference (Richter et al., 2020) and neural solvers for partial differential equations (PDEs) (Richter Berner, 2022). Extending these works, we prove that the beneficial properties of the log-variance loss also hold for the general bridge objective, which incorporates two instead of only one controlled stochastic process. Finally, we refer to Vargas et al. (2024) for concurrent work on the path space perspective on diffusion-basee sampling.
1.2 OUTLINE OF THE ARTICLE
The rest of the article is organized follows. In Section 2, we provide an introduction to diffusion- based sampling from the perspective of path space measures and time-reversed SDEs. This can be understood as a unifying framework allowing for generalizations to divergences other than the KL divergence. We propose the log-variance divergence and prove that it exhibits superior properties. In Section 3, we will subsequently outline multiple novel connections to known methods, such as SBs in Section 3.1, diffussin-basee generative modeling (i.e.. DIS) in Section 3.2, and approaches based on reference processes (i.e., PIS and DDS) in Section 3.3. For all considered methods, we can find compelling numerical evidence for the superiority of the log- varvinre divergence, see Section 4.
2

---

Published as a conference paper at ICLR 2024
2 DIFFUSION-BASEL SAMPLING
In this section, we will reformulate our sampling problem time-reversal of diffusion processes from the perspective of measures on path space. Let us first define our setting and notation.
2.1 NOTATION AND SETTING
We denote the density a random variable X by px- For a suitable R -vllue stochastic process X (X.)) define its density Px w.r.t to the Lebesgue measure by px (x(;,t) px, for t t [0, T). For suitable functions f E C(RR [0. 7), IR) and E C(TR [0. 7), RR,,, we further define the deterministic and stochastic integrals
Rx(X)= = S[= s(),,)) ds and
Sx) = l0" w((,, dW.,
(1)
where W a standard d-dimensional Brownian motion. We denote by P the set of all probability measures the space of continuous functions C([O, Rd) and define the path space measure Px E P as the law of X. For time-dependent function μr, we denote by j the time-reversal given by (i(()= u(T 1). We refer to Appendix A.I for technical assumptions.
2.2 SAMPLING AS TIME-REVERSAL PROBLEM
The goal of diffusion- based sampling to sample from the density Ptarget = & by transporting a prior density Pprior via controlled stochastic processes. We consider two processes given by the SDEs
axx (n + ou) s) + (s))
X
Pprior:
(2)
dY = ---+ ds + ((s) dW.,
Yo ~ Ptarget
(3)
where we aim to identify control functions u, E U in a suitable space of admissible controls I C C(RR x [0, 7], (R:)) in order to achieve Xy Ptarget and Y Pprior- Specifically, seek
controls satisfying
Pprior y
Ptarget
in the sense that Y is the time-reversed process of X and vice versa, i.e., Px- = py=. In this context, we recall the following well-known result on the time-reversal of stochastic processes (Nelson, 1967; Anderson, 1982; Haussmann Pardoux, 1986; Föllmer, 1988).
Lemma 2.1 (Time-reversed SDEs). The rime-reversed process Y", given by the SDE dÝ (1+ ov + ao V log Fy.) (V" 8) a(s) satisfies that Pp" ..
V YP,
(4)
Proof. The result can be derived by comparing the Fokker-Planck equations governing Pp" and py see, e.g., Chen et al. (2021a); Huang al. (2021); Berner al. (2024).
Let us now define the problem of identifying the desired control functions u and from the perspective of path space measures on the space of trajectories C( 7), Rd), as detailed in the sequel.
Problem 2.2 (Time-reversal). Let P X be the path space measure of the process X defined in (2)
and let IP ge be the path space measure of y" the time-reversal of y", given in (4). Further, let
D: P x P + R>0 be a divergence, i.c., a non-negative function satisfying that ((P, Q) = Of and only if P Q. We aim to find optimal controls u², ,,"" s.t.
u D" E arg min D xm PPx)) u.eEuxu
We note that Problem 2.2 aims to the processes X" and Y with respect to each other while obeying the respective initial values X²' ~ Pprior and Y ~ Ptarget- For the actual computation of suitable divergences, we derive the following fundamental identity.
3

---

Published as a conference paper at ICLR 2024
Proposition 2.3 (Likelihood of path measures). Let x be process as defined in (2) with being replaced by w E U and let and be as in (1). We can compute the Radon-Nine MadoonNikodym derivative as
dP T
dPp (x
Z exp (Riggg 11,.000
+ +Su0 + (X
(6)
with B(X") = log Pprior (Xo)
and
Surid = (u+
w V - + V. (ou
p( X
Proof. The proof combines Girsanov's 's theorem, Ito's lemma, the HJB equation governing log py", and the fact that p Peargets Appendix A.2.
Note that can remove the divergence V (ov - 11) in (6) by resorting backward stochastic integrals, see Remark A. in the appendix. Using the path space perspective and the representation of the Madon-Nikodym derivative Proposition 2.3. we may now, in principle, choose any suitable divergence D in order to approach Problem 2.2. Using our path space formulation, we are, to the best of our knowledge, the first to study this problem such generality. In the following, we demonstrate that this general framework unifies previous approaches and allows us to derive new methods easily,
2.3 COMPARISON OF THE KL AND LOG- VARIANCE DIVERGENCE
Most works in the context of diffusion-based sampling rely on the KI. divergence. Choosing D DKL, which implies u in (6), we can readily compute
Bridge + B) + log DKL (II x IP x-) E' [R
with fuiridee
+ (are - ui), where we used the fact that the stochastic integral Supo has vanishing expectation. Note that in practice we minimize the objective
||++=l? 2
(KI(u, == (Px
- log Z.
This objective is analogous to the one derived in Chen et al. (2021a) for the bridge problem, see also Section 3.1 and Appendix A.4. Unfortunately, however, the KL divergence is known have some evident baaks such as mode collapse (Minka et al., 2005; Midgley et al., 2023) or potentially high variance of Monte Carlo estimators (Roeder et al., 2017). To address those issues, propose another divergence that has been originally suggested in Nüsken & Richter (2021) and extend it to the setting of two controlled stochastic processes.
Definition 2.4 (Log-variance divergence). Let be a reference measure. The log-variance divergence
between the measures IP and Q w.r.t. the reference IP is defined
dPP
(P, = Vp log dQQ
Note that the log-variance divergence is symmetric in IP and Q and actually corresponds to family of divergences, parametrized by the reference measure P. Obvious choices in our setting are P Px., PP == Px,,, and Q == Pee, resulting in the log- variance loss
C v) = Dpx xw (Tx', Pgo) = [R Bridge + S + ()
(8)
Since the variance is shift-invarii inft-innvrriant, we can omit log Z in the above objective.
Compared to the Ka-based loss (7), the log-variances loss (8) exhibits the following beneficial properties. First, by the choice of the reference measure Px", one can balance exploitation and exploration. To exploit the current control u, one can set w u, but one can also choose another control or another initial condition X We can leverage this to counteract mode collapse by optimizing the loss in (8) along (saj-)traeectories where IP x'x has low probability, see Appendix A.7. Next, note that the log- varianne loss in (8) does not require the derivative of the process X" w.r.t. the control (which, for the Appendix A.6). While we still need to simulate the process X",, we can rely on any (black-box) box) SDE solver and do not need track the computation of X for automatic differentiation. This
u, implemented by detaching stopping the gradient,
4

---

Published as a conference paper at ICLR 2024
implies that the log- variance loss does not require derivatives' of the unnormalized target density p. which is crucial for problems where the target is only available as a black box. In contrast, the KL-based loss in (7) demands to differentiate X w.r.t. u, requiring to differentiate through the SDE solver and resulting in higher computational costs. Particularly interesting is the following property, sometimes referred to ficcing--he-landing (Roeder al., 2017). states that the gradients of the log-variance loss have zero variance at the optimal solution. This property does, in general, not hold for the KL-based loss, such that variants of gradient descent might oscillate around the optimum.
Proposition 2.5 (Robustness at the solution). Let LLV be the Monte Carlo estimator of the log- variance loss in (8) and let the controls variances of the respective derivatives vanish at the optimal solution (u" ,²") = (ug Voy i.e.
and v Vty be parametrized by 0 and y. The
W VoLiv(upr May =0 and
W
V by-) 0,
for all w E U. For the estimator ÊKL of the KL-based loss (7) the variances do not vanish.
Proof. The proof is based on technical calculation and Proposition 2.3, see Appendix A.2.
For the case u, can further interpret the log-variance loss control variate version of the KL-based loss, see Remark A.2 in the appendix. We can empirically observe the variance reduction for the loss and its gradient in Figure in the appendix.
3 CONNECTIONS AND EQUIVALENCES OF DIFFUSION-BASED SAMPLING APPROACHES
In general, there are infinitely many solutions Problem 2.2 and, in particular, our objectives in (7) and (8). In fact, Girsanov's is theorem shows that the objectives only enforce Nelson's 's identity (Nelson, 1967), i.e.,
" + o" V Voogxx o V log
see also the proof of Proposition 2.3. In this section, we show how our setting generalizes existing diffusion- sampling approaches, which in turn ensure unique solutions to Problem 2.2. Moreover, with our framework, we can readily derive the corresponding versions of the log-variance loss (8). We refer to Appendix A.3 for high-level overview of previous diffusion-basse sampling methods.
3.1 SCHRÖDINGER BRIDGE PROBLEM (SB)
Out of all solutions u" fulfilling (9), the Schrödinger bridge problem considers the solution u" that minimizes the KL divergence DKL(PP x [Pxx) ) to given reference process X", defined in (2) with u replaced by E U, see Appendix A.4 for further details. Traditionally, r = 0, i.e., the uncontrolled process X is chosen. Defining
+ ))
fref = (u -r). W
(10)
Girsanov's theorem shows that dpx
dP (Xw) = exp (R + (X"-) which implies that
DKL. IP xx-) E Roont (x4)
(11)
see, e.g., Nüsken & Richter (2021, Lemma A.1) and the proof Proposition 2.3. The SB objective can thus be written as
min EE well
R R (X") X
Ptarget
(12)
De Bortoli al. (2021); Caluya & Halder (2021); Pavon & Wakolbinger (1991); Benamou & Brenier (2000); Chen et al. (2021b); Bernton et al. (2019). We note that the above can also be interpreted as entropy-regularized optimal transport problem (Léonard, 2014). The entropy constraint in (11) could now be combined with our objective in (5) by considering, for instance,
(x*) + AD (P (Px))p)))
min
!hhle, by default, the samplers presented in the following use V log p their parametrization of the control u, we present promising results for the derivative-free regime in Appendix A.7.
5

---

Published as a conference paper at ICLR 2024
where À E (0, oo) is sufficiently large Lagrange multiplier. In Appendix A.4 show how the SB problem (12) can reformulated as a system of coupled PDEs or BSDEs, which can alternatively be used to regularize Problem 2.2, also Liu al. (2022); Koshizuka & Sato (2022). Interestingly, the BSDE system recovers our KI.-based objective in (7), as originally derived in Chen et al. (2021a).
Note that via Nelson's identity (9), optimal solution u" the SB problem uniquely defines an optimal control v* and vice versa. For special cases of SBs, we can calculate such v² or an approximation
v². Fixing e in (5) and only optimizing for u appearing in the generative process (2) then allows us to attain unique solutions to (an approximation of) Problem 2.2. We note that the approximation v* incurs an irreducible loss given by
dPP (X")
dl
(X"),
(13)
dP
AP
thus requiring an informed choice of and Ppriors such that y y We will consider two such settings in the following sections.
3.2 DIFFUSION-BASEE GENERATIVE MODELING (DIS)
We may set 0 == 0, which can be interpreted as a SB with u² =r=r=o loggyo and Pprior = Pys such that the entropy constraint (1I) can be minimized to zero. Note, though, that this only leads to feasible sampling approaches if the functions and in the SDEs chosen such that the distribution of Pyp is (approximately) known and such that we can easily sample from it. practice, one chooses functions and such that Pyp R Pprior = NN I²),, see Appendix A.6. Related approaches are often called celfiusion-basete generative modeling or denoising diffiusion probabilistic modeling since the (optimally controlled) generative process X can be understood as the time-reversal of the process y that moves samples from the target density Gaussian noise (Ho al., 2020; Pavon, 1989; Huang al., 2021; Song et al., 2020). Let us recall the notation from Proposition 2.3 and define fi = şBridge =u-
[nll V μ. Setting v = 0 in (7), we readily get the loss
2
LKL (u) E [R + B)
which corresponds to the Reversdd Diffusion Sampler (DIS) derived in Berner et al. (2024). Analogously, our path space perspective and (8) yield the corresponding log-variance loss
(u) W [ [RI + Sa + B) (X*)
(14)
3.3 TIME-REVERSAS OF REFERENCE PROCESSES (PIS & DDS)
In general, we may also set := a logpx- - r. Via Lemma 2.1 this implies that IP x IP Ppiirtt where e.ef is the process Y as in (3), however, with Ptarget replaced by Pref pxp, i...
dy (-ji + 5T) (Y s) ds + ((s) dW,,
ye
Pref-
In other words, B, is the time-rewe reversal of the reference process X", Using (6) with Prof instead of Prarget = 4. we thus obtain that
dPx
(X Pprior (X
exp (Riie Bridge + Sr+e)
(15)
Prof (x²)
This identity leads to the following alternative representation of Proposition 2.3.
Lemma 3.1 (Likelihood w.r.t. reference process). Assuming Px- = IP sonee, it holds that
dIP x (x = ZZ exp RRyse + Su-r + Bree) (x"),
dI o
where frot is defined as in (10) and Bref (Xx") == log Penf (XX)).
Proof. The result follows from dividing aP to Remark A.3 for alternative derivation that does not rely the concept of time-reversals. D
dP dPxx (X") (6) by aP aP
dpx
(X') in (15). We also refer

---

Published as a conference paper at ICLR 2024
Table I: Comparison of the objectives with Ry, Sus B, Bret fBridge fref u.a
as defined the text.
Lxt.
Ciy (ours)
Pprior
Bridge
I RR Bridge B) Pe
W (Rgggg Susa (X()))
arbitrary
learned
DIS
E R
+ B))
W
BBggg S. {B)(Xm)
pxf
+B*))(x)))
Brrs + S. B (X
8no
VV Pxg
DDS
(Ratt.
+ (x7))
(Rggs... + Su-r - (x")]
Pyg.
T V log Pryon
Note that computing the Radon-Nikodym derivative in Lemma 3.1 requires to choose r, Pprior μ., and such that Pref pxy. is tractable². For suitable choices of (see below), one can, for instance, the SDEs with tractable densities stated in Appendix A.5 with Pprior 5x Pprior N(,, r²), mixture of such distributions. Recalling (13) and the choice = T log
px loggx- need guarantee that Y = Y Let us outline two such cases in the following.
т, also
PIS: We first consider the case = Lemma 3.1 and taking D = DKL in Problem 2.2 then yields
CKL (u) DKL P P
Pgs) - log Z = (R prof + Brel) (x*)
This objective has previously been considered by Tzen & Raginsky (2019): Dai Pra (1991) and corresponding numerical algorithms, referred to Path Integral Sampler (PIS) in Zhang & Chen (2022), have been independently presented in Richter (2021); Zhang & Chen (2022); Vargas et al. (2023b). Choosing D DLv, we get the corresponding log-variance loss
(Tv W [ + S. (x
which has already been stated by Richter (2021, Example 7.1). Typically, the objectives used with Pprior = 82,, since Doob's h-transform guarantees that 전 = 2", i.e., we can solve the SB exactly, see Rogers& Williams (2000) and also Appendix 1.4.1. In this special case, the SB often referred to as Schrödinger half-bridge.
DDS: Next, we consider the choices r a V log Pyone, 8 0, and Pprior == Pyg.rt which in turn yields special case of the setting from Section 3.2. Using Lemma 3.1, we obtain the objective
(ul E (R + Biet) (xx)
This corresponds to the Denoising Diffusion Sampler (DDS) objective stated by Vargas et al. (2023a) when choosing j and such that y is a VP SDE, see Appendix A.5. Choosing the invariant distribution Pref == (0, of the VP SDE, see (26) in the appendix, we have that px-(-, f)
Dy (., (- Pref Pprior for [o, TJ. and, particular, r(a, t) - - Finally, with our general framework, the corresponding log- borviariance loss can now readily be computed as
(u) W [R +S. Brd
We refer to Table I for a comparison of all our different objectives.
4 NUMERICAL EXPERIMENTS
In this section, compare the KL-based loss with the log-variance loss the three different approaches, i.e., the general bridge, PIS, and DIS, introduced in Sections 3.2, and 3.3. As DDS can be seen as a special case DIS (both with V see also Berner al., 2024, Appendix A.10.I). we do not consider it separately. We can demonstrate that the appealing properties of the log-variance
In general, it suffices be able compute pxę up its normalizing constant.

---

Published as a conference paper at ICLR 2024
Groundtruth
KL-PIS
LV-PIS lours)
KL_
LV-DIS (ours)
5
X1
Figure 2: KDE plots of (1) samples from the groundtruth distribution, (2 & 3) PIS with KL divergence and log-variance loss, and (4 & 5) DIS with KL divergence and log-variance loss for the GMM problem (from left to right). One can see that the log-variance loss does not suffer from mode collapse such as the reverse KL divergence, which only recovers the mode of Perior (0.
KL-MS
UADIS
Density
Marginal
Hatogram of
0
-2
Figure 3: Marginals of the first coordinate of samples from PIS and DIS (left and right) for the DW problem with d = 5, = 5, 4. Again, one observes the mode coverage of the log- lor-veriance loss as compared to the reverse KL divergence. Similar behavior can also be observed for the other marginals (see Figure 6) and higher-dimensional settings (see Figure for an example in d 1000).
loss can indeed lead to remarkable performance improvements for all considered approaches. Note that we always compare the same settings, in particular, the same number of target evaluations, for both the log-va log-variance and KL-based losses and use sufficiently many gradient steps to reach convergence. See Appendix A.6 and Algorithm I for computational details Still, we observe that qualitative differences between the two losses are consistent across various hyperparameter settings. We refer Appendix A.8 for additional experiments.
4.1 BENCHMARK PROBLEMS
We evaluate the different methods the following three numerical benchmark examples.
Gaussian mixture model (GMM): We consider p(x) m W (4; E,) and choose m = 9, Σ, = 0.3I, (4 = (-5,0,5) x xf 0,,05) C R² to obtain well-separated modes, see Figure 2.
Funnel: The 10-dimensional Funnel distribution (Neal, 2003) is challenging example of- ten used to test MCMC methods. It is given by the density p(x) Ptarget (æ)
M(11; ??) (e,,, for z (x E R with 17 =
Double well (DW): A typical problem in molecular dynamics considers sampling from the stationary distribution of a Langevin dynamics. In example consider d-dimensional double well poten- tial, corresponding to the (unnormalized) density p(z) = exp Em (2) - 8)2 - E 23
with m E N combined double wells (i.e., 2" modes) and a separation parameter E (0, oo), see also Wu al. (2020) and Figure 3. We choose large value of J to make sampling particularly challenging due to high energy barriers. Since factorizes in the dimensions, we obtain reference solutions by numerical integration and ground truth samples using rejection sampling with Gaussian mixture proposal distribution, see also Midgley et al. (2023).
4.2 RESULTS
Let us start with the bridge approach and the general losses in (7) and (8). Table 5 (in the appendix) shows that the lar-neri loss can improve considered metrics significantly. However, the general bridge framework still suffers from reduced efficiency and numerical instabilities. For high-dimensional problems, it can be prohibitive to compute the divergence of using automatic
"The repository can be found httpt::://// http://gttthubcoo../.. uuberner ssdeenrr//dde_pppp er.
8

---

Published as a conference paper at ICLR 2024
Table 2:
PIS and DIS metries for the benchmark problems of various dimensions d. We report the median over five independent runs, see Figure 9 for corresponding boxplot. Specifically, we report errors for estimating the log-normaiizing constant A log Z as well the standard deviations A std of the marginals. Furthermore, we report the normalized effective sample size ESS and the Sinkhorn distance W (Cuturi, 2013), see Appendix A.6 for detailsMax- -oseaa--oospp-hhulle The arrows 1 and + indicate whether we want to maximize or minimize a given metric.
Problem
Method
Loss
A log Z 4 W L
ESS 1
A std
GMM (d (d=22 2)
PIS
KL (Zhang & Chen, 2022)
1.094
0.467
0.0051
1.937
LV (ours)
0.046
0.020
0.9093
0.023
DIS
KL (Bemnr et al., 2024)
1.551
0.064
0.0226
2.522
LV (ours)
0.056
0.020
0.8660
0.004
Funnel (d = 10)
PIS
KL (Zhang & Chen, 2022)
0.288
5.639
0.1333
6.921
LV (ours)
0.277
5.593
0.0746
6.850
DIS
KL (Bemer et al., 2024)
0,433
5.120
0.1383
5,254
LV (ours)
0.430
5.062
0.2261
5,220
DW (d (dd=5,m0 5,m 5,6 5,5=44 4)
PIS
KL (Zhang & Chen, 2022)
3.567
1.699
0.0004
1,409
LV (ours)
0.214
0.121
0.6744
0.001
DIS
KL (Berner et al., 2024)
1,462
1.175
0.0012
0.431
LV (ours)
0.375
0.120
0.4519
0.001
DW (d = 50, m = 5,0 2)
PIS
KL (Zhang & Chen, 2022)
0.101
6.821
0.8172
0.001
LV (ours)
0.087
6.823
0.8453
0.000
DIS
KL (Berner et al., 2024)
1.785
6.854
0.0225
0.009
LV (ours)
1.783
6.855
0.0227
0.009
differentiation, and relying on Hutchinson's 's trace estimator introduces additional variance. We refer to Remark A.1 for further discussion. The instabilities might be rooted in the non-uniqueness of the optimal control (which follows from our analysis, ef. Section 3). Furthermore, such issues are also commonly observed in the context of SBs (De Bortoli et al., 2021; Chen et al., 202 Fernandes al., 2021), where two controls need to be optimized. Therefore, for the more challenging problems, we focus on DIS and PIS, which do not incur the described pathologies.
We observe that the log-variance loss significantly improves both DIS and PIS across our considered benchmark problems and metrics, see Table 2. The improvements are especially remarkable consider- ing that we only replaced the KL-based loss LKL. by the log- variance loss LLV without tuning the hyperparameter for the latter loss. In the few cases where the KI. divergence performs better, the difference seems rather insignificant. In particular, Figures and show that the log-variance loss successfully counteracts mode collapse, leading to quite substantial improvements. The benefit of the log-variance loss can also be observed for the benchmark posed in Wu et al. (2020), which aims to sample a target distribution resembling a picture of Labrador, see Figure 8 in the appendix. In Ap- pendix A.8, we present results for further thigh-dimensiona)) targets, showing that diffusion-based samplers with -orrvanieece loss are competitive with other state-of-the-art sampling methods.
5 CONCLUSION
In this work, we provide a unifying perspective on diffusion-based generative modeling that is based on path space measures of time-reversed diffusion processes and that, for the first time, connects methods such as SB, DIS. PIS, and DDS. Our novel framework also allows us to consider arbitrary divergences between path measures as objectives for the corresponding task of interest. While the KL. divergence yields known methods, we find that choosing the log-variance divergence leads to novel algorithms that are particularly useful for the task of sampling from (unnormalized) densities. Specif- ically, this divergence exhibits beneficial properties, such as lower variance, computational efficiency, trade-offs. We can demonstrate in multiple numerical examples that the log-var variance loss greatly improves sampling quality across range of metrics. We believe that problem and approach-speciffe finetuning might further enhance the performance of the log variance loss, thereby paving the way for competitive cesttsiooo-bsedd sampling approaches.
and inttorattio-xxlloiattti

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
We thank Guan-Horng Liu for many useful discussions. The research of was funded by Deutsche Forschungsgemminschatt (DFG) through the grant CRC 1114 "Scaling Cascades in Complex Systems" (project A05, project number 235221301). J.B. acknowledges support from the Wally Baer and Jeri Weiss Postdoctoral Fellowship.
REFERENCES
Brian DO Anderson. Reverse-time diffusion equation models. Applications, 12(3):313-326, 1982.
Stochastic Processes and their
Michael Arbel, Alex Matthews, and Arnaud Doucet. Annealed flow transport Monte Carlo. In International Conference Machine Learning, pp. 318-330. PMLR, 2021.
Ludwig Amold. Stochastic Differential Equations: Theory and Applications. A Wiley-Interscienee publication. Wiley, 1974.
Baldi. Stochastic Calculus: An Introduction Through Theory and Exercises. Universitext. Springer International Publishing, 2017.
Jean-David Benamou and Yann Brenier. computational fluid mechanics solution the Monge- Kantorovich mass transfer problem. Numerische Mathematik, 84(3):375-23 -2775,---33-- 2000.
Julius Berner, Lorenz Richter, and Karen Ulirich. An optimal control perspective on diffusion-based generative modeling. Transactions on Machine Learning Research, 2024.
Espen Bernton, Jeremy Heng, Arnaud Doucet, and Pierre E Jacob. Schrödinger Bridge Samplers. 2019.
arXiv preprint
Kenneth F Caluya and Abhishek Halder. Wasserstein proximal algorithms for the Schrödinger bridge problem: Density control with nonlinear drift. IEEE Transactions on Automatic Control, 67(3): 1163-1178, 2021.
Tianrong Chen, Guan Moan-Hoong Liu, and Evangelos Theodorou. Likelihood training of Schrödinger Bridge using Forward-Backward SDEs theory. arXiv preprint arXiv:2110.11291, 2021a.
Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. Stochastic control liaisons: Richard Sinkhorn meets Gaspard Monge on a Schrödinger bridge. SIAM Review, 63(2):249-313, 2021b.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, pp. 2292-2300, 2013.
Paolo Dai Pra. A stochastic control approach to reciprocal diffusion processes. Applied mathematics and Optimization, 23(1)311-3329 1991.
Valentin De Bortoli, James Thornton, Jeremy Heng, and Araud Doucet. Diffusion Schrödinger bridge with applications to score-based generative modeling. Advances Neural Information Processing Systems, 14:17695-77709 2021.
Pierre Del Moral, Arnaud Doucet, and Ajay Jasra. Sequential Monte Carlo samplers. Journal of the
Royal Statistical Society: Series (Statistical Methodology). 68(3):411 (133:411-436, 2006.
Arnaud Doucet, Adam M Johansen, et al. A tutorial on particle filtering and smoothing: Fifteen years later. Handbboo of nonlinear fintinearfiltering, 12(656-704):3, 2009,
David Lopes Fernandes, Francisco Vargas, Carl Henrik Ek, and Neill DF Campbell. Shooting Schrödinger's cat. In Fourth Symposium on Advances in Approximate Bayesian Inference, 2021.
Wendell Fleming and Halil Mete Soner. Controlled Markov processes and viscosity solutions, volume 25. Springer Science & Business Media, 2006.
Hans Föllmer. Random fields and diffusion processes. In École d' Été de Probabilités de Saint-Flour XV-XVII, 1985-87, pp. 101-203. Springer, 1988.
10

---

Published as a conference paper at ICLR 2024
Carsten Hartmann, Lorenz Richter, Christof Schütte, and Wei Zhang. Variational characterization of free energy: Theory and algorithms. Entropy, 19(11):626, 2017.
Ulrich G Haussmann and Etienne Pardoux. Time reversal diffusions. The Annals of Probability, pp. 1188 1188-1205, 1986.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6888 -851, 2020.
Lars Holdijk, Yuanqi Du, Ferry Hooft, Priyank Jaini, Bernd Ensing, and Max Welling, Path integral stochastic optimal control for sampling transition paths. arXiv preprint arXiv:2207.02149, 2022.
Chin-Wa Huang, Jae Hyun Lim, and Aaron C Courville. A variational perspective diffusion-basee generative models and score matching. Advances in Neural Information Processing Systems, 34, 2021.
Robert E Kass, Bradley P Carlin, Andrew Gelman, and Radford M Neal. Markov chain Monte Carlo in practice: a roundtable discussion. The American Statistician, 52(2):93-100, 1998.
Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances in Neural Information Processing Systems, 34:21696-21707, 2021.
Takeshi Koshizuka and Issei Sato. Neural Lagrangian Schrödinger bridge. arXiv preprint D.Xiv:2204.048535 2022.
Hiroshi Kunita. Stochastic flows and jump-diffusions. Springer, 2019.
Christian Léonard. Some properties of path measures. In Séminaire de Probabilités XLVI, pp. 207-230. Springer, 2014.
Guan Guan-Horng Liu, Tianrong Chen, Oswin So, and Evangelos A Theodorou. Deep generalized Schrödinger bridge. arXiv preprint (tt:::200909999) 2022.
Jun S Liu and Jun S Liu. Monte Carlo strategies in scientific computing, volume 10. Springer, 2001.
Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio. Trajectory balance: Improved credit assignment in gflownets. Advances Neural Information Processing Systems, 35: 5955-5967, 2022a.
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.005580, 2022b.
Laurence Illing Midgley, Vincent Stimper, Gregor Simm, Bernhard Schölkopf, and José Miguel Hemández-LLobato Flow annealed importance sampling bootstrap. In The Eleventh International Conference on Learning Representations, 2023.
Tom Minka et al. Divergence measures and message passing. Technical report, Citeseer, 2005.
Radford M Neal. Annealed importance sampling. Statistics and computing, 11(2):125-133) 2001. 2001.
Radford M Neal. Slice sampling. The Annals of Statistics, 31(3)7005-767, 2003.
Nelson. Dynamical theories Brownian motion. Press, Princeton, NJ, 1967.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, pp.
PMLR, 2021.
Nikolas Nüsken and Lorenz Richter. Solving high-dimensionall Hamilon- PDEs using neural networks: perspectives from the theory of controlled diffusions and measures on path space. Partial Differential Equations and Applications, 2(4): 2021.
Bernt Øksendal and Bernt Øksendal. Stochastic differential equations. Springer, 2003.
George Papamakarios. Eric T Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. Mach. Learn. Res., 22(57): 2021.
11

---

Published as a conference paper at ICLR 2024
Michele Pavon.
Stochastic control and nonequilibrium thermodynamical systems. Applied Mathe- matics and Optimization, 1911:187-20,, 1989.
Michele Pavon. On local entropy, stochastic control and deep neural networks. arXiv preprint arXiv:2204. 13049, 2022.
Michele Pavon and Anton Wakolbinger. On free energy, stochastic control, and Schrödinger processes. In Modeling, Estimation and Control of Systems with Uncertainty, pp. 334-348. Springer, 1991.
H. Pham. Continuous-tion Stochastic Control and Optimization with Financial Applications. Stochas- tic Modelling and Applied Probability. Springer Berlin Heidelberg, 2009.
Lorenz Richter. Solving high-dimensional PDEs, approximation of path space measures and impor- tance sampling of diffusions. PhD thesis, BTU
2021.
Lorenz Richter and Julius Berner. Robust SDE-based variational formulations for solving linear PDEs via deep learning. In International Conference Machine Learning, pp. 18649-18666.
PMLR, 2022.
Lorenz Richter, Ayman Boustati, Nikolas Nüsken, Francisco Ruiz, and Omer Deniz Akyildiz. VarGrad: low-variance gradient estimator for variational inference, Advances in Neural Information Processing Systems, 33:13481-13499 2020.
Christian P Robert, George Casella, and George Casella, Monte Carlo statistical methods, volume 2. Springer, 1999.
Geoffrey Roeder, Yuhuai Wu, and David K Duvenaud. Sticking the landing: Simple, lower-variance gradient estimators for variational inference. Advances in Neural Information Processing Systems, 30, 2017.
L Chris G Rogers and David Williams. Diffiusions, Markov processes and martingales: Volume Itô calculus, volume 2. Cambridge university press, 2000.
Abul Hasan Siddiqi and Sudarsan Nanda. Functional analysis with applications. Springer, 1986.
Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. Advances in neural information processing systems, 05:12233-12444 2020.
Yang Song, Jascha
Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. International Conference on Learning Representations, 2020.
Vincent Stimper, Bernhard Schölkopf, and José Miguel Hermández-Lobato, Resampling base distri- butions of normalizing flows. In International Conference on Artificial Intelligence and Statistics, pp. 4915-4936, 2022.
Gabriel Stoltz, Mathias Rousset, et Free energy computations: mathematical perspective. World Scientific, 2010.
Belinda Tzen and Maxim Raginsky. Theoretical guarantees for sampling and inference in generative models with latent diffusions. In Conference on Learning Theory, pp. 3084-31 14. PMLR, 2019.
Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. Advances in Neural Information Processing Systems,
2021.
Lorenz Vaitl, Ludwig Winkler, Lorenz Richter, and Pan Kessel. Fast and unified path gradient estima- tors for normalizing flows. In The Twelfth International Conference Learning Representations, 2024.
Francisco Vargas. Machine-learning approaches for the empirical Schrödinger bridge problem. Technical report, University of Cambridge, Computer Laboratory, 2021.
Francisco Vargas, Will Grathwohl, and Arnaud Doucet. Denoising diffusion samplers. In International Conference Learning Representations, 2023a.
12

---

Published as a conference paper at ICLR 2024
Francisco Vargas, Andrius Ovsianas, David Fernandes, Mark Girolami, Neil D Lawrence, and Nikolas Nüsken, Bayesian learning via neural Schrödiger-Föllmee flows. Statistics and Computing, 33 (1):1-22, 2023b.
Francisco Vargas, Shreyas Padhy, Denis Blessing, and N Nüsken. Transport meets variational inference: Controlled Monte Carlo diffusions. In The Twelfth International Conference on Learning Representations, 2024.
Martin Wainwright, Michael Jordan, al. Graphical models, exponential families, and variational inference. Foundations and Trends in Machine Learning, 1(1-2):1-305 2008.
Hao Wu, Jonas Köhler, and Frank Noé. Stochastic normalizing flows. Advances in Neural Information Processing Systems, 33:5933-59944, 2020.
Dinghuai Zhang, Ricky Tian Qi Chen, Cheng-Hao Liu, Aaron Courville, and Yoshua Bengio. Diffu- sion generative flow samplers: Improving learning signals through partial trajectory optimization. arXiv preprint arXiv:2310.02679, 2023.
Qinsheng Zhang and Yongxin Chen. Path Integral Sampler: a stochastic control approach for sampling. In International Conference on Learning Representations, 2022.
13

---

Published as a conference paper at ICLR 2024
A APPENDIX
A.1 ASSUMPTIONS
In our proofs, we assume that the coefficient functions of all appearing SDEs are sufficiently regular such that Novikov's 's condition is satisfied and such that the SDEs admit unique strong solutions with smooth and strictly positive densities px, for E (0, T), see, for instance, Arnold (1974); Øksendal & Øksendal (2003); Baldi (2017).
A.2 PROOFS
Proof of Proposition 2.3. Let us define the path space measures P' x and Py as the measures of X and Y" conditioned on X x and yo with Rdd respectively. We can then compute
dP x (X = log dP or
dP X (X + log
dTP dl (X1"))
log
dPP pn
dIP 9:
(16)
dP x
(X") + log Perior (X))
log
dPp
pu.
Pp" (X 0)
We follow Liu al. (2022) and first note that the time-reversal of the process Y defined in (3) is given by
d (a + aa VVg ov)(Y" 8) ds + ds) where we abbreviate = Py, see Lemma 2.1. Let us further define the short-hand notations = + o Vg and μ + (u h). Then, we can write the SDEs in (2) and (3)
dx = (b+ oh) 8) ds + (s)s -r. = b(v",
,, s)dsss ds + (s)
We can now apply Girsanov's theorem (see, e.g.. Nüsken & Richter, 2021, Lemma A.1) to rewrite the logarithm of the Radon-Nikodym derivative R log xP urr. (X") in (16)
(a -h) (X",,) dxw - 1/² (o- h) (x.", 8) ds
1 ("
R-
2 s)l
(w u) h + 1 [11111] (x." s) ss (
(17)
/-¹
T (w-ww) u) (u + -utu-oTve Vg) + vgll²) (X.") 5) ds + (X"
[-' (ov 1)
R RBaddeee
Biidde
+ (u + (ou
Vgg- - T (x)
(X ds (x")..
Further, may apply Itô's lemma to the function g get
T) ( 0) 0)= - l' (a.9+ V Vg Vg (u+ow)+ Tr (oo ( (X." 8) ds S, (
Noting that g = log pv- fulfills the Maaiitoon-acobi-BBelllnn equation (see, e.g., Berner al., 2024)
1 Tr (oo TV'g) + (av- - 4) Vg+ + V (ov 4) - gl.
1 o T
0.g 2
get
Vg g(X", T) -) (v Tv } oloT g/ (X , 3) + S vX
112)(x
Finally, combining this with (16) and (17) and noting that
g(XF. = log pp (XT, T) log py (X7", 0) Ptarget (XT).
yields the desired expression.
14

---

differentiable at The Gâteaux derivative of L w.r.t. in direction is then defined as Published as a conference paper at ICLR 2024
Remark A.1 (Divergence- free objectives). One remove the divergence from the Radon-Niie Radon-Nikodym derivative (6) and thus from corresponding losses by noting the identity
I v.
T (av- (Xx s)ds L'" (u-o (X) -dW.)
where for a suitable function E C(R2 x [0, 12) the backward integration w.r.t. Brownian motion defined as
L'
(xx aw, == lim Ar-0
Σ (X bnt1) (Wint11 - Ws,))
which, in contrast to the definition of the usual Itō integral,
L' ((x aW, lim Σ (XX tn) (W Wen).
considers the right endpoint when discretizing the integral on a time grid 0 = to < h < .... < iN T with step size At = fa+1 fu. The above definitions via refined partitions readily bring implementation schemes for both integrals when choosing a fixed step size At > O. Divergence-free objectives might be partieularly beneficial in higher dimensions, where is typically expensive to compute the divergence using automatic differentiation. For further details, we refer to Vargas et al. (2024) and Kunita (2019).
Proof of Proposition 2.5. Let us first recall the notion of Gâteaux derivatives, Siddiqi & Nanda (1986, Section 5.2). We say that C: x u + 22 Gâteaux differentiable U for all v. U the mapping
- L(u + EΦ, v)
8u L(u, 0z = de L(u + E6,
The derivative of L w.r.t. v is defined analogously. Let now = ug and Uy be parametrized' by E IRP and 7 E IP. Relating the Gâteaux derivatives to partial derivatives w.r.t. and 7, respectively, let us note that we are particularly interested in the directions ø Og, uy and Φ = 0, Vy for E (1.....p) p). This choice is motivated by the chain rule of the Gâteaux derivative, which, under suitable assumptions, states that
6
8
0g, C(ug, 22) 8u oumu uu
C [u,, Uy; 0o. and 0. Cuus, v.,)
L (ug, ii 0r Un)).
Analogous the computations in Nüsken & Richter (2021), the Gâteaux derivatives of the Monte Carlo estimator fty of the log-variance loss Liv in (8) with K N samples is given by
8
CH Civ
2
Σ
(b) u.u.w
(b)
B
1 K
Σ
B
(18)
fu
u.an,
where the superseript (k) denotes the index of the k-th i.i.d. sample in the Monte Carlo estimator Fry and we define the short-hand notations
A. = (R Bridge
(((k) + B) (Xxm.(k)) + log Z
and
(()) 1 (Rymn + s(()) (x)) with Pum (w u) Φ
Now, note that the definition of the log-variance loss and Proposition 2.3 imply that for the optimal choices
u²,v v- o² holds that
We only assume that and y in the same space RP for notational simplicity.
15

---

Published as a conference paper at ICLR 2024
almost surely for every l {l...., K} and w u. This readily implies the statement for the derivative w.r.t. the control My. The analogous statement holds true for the derivative w.r.t Ey, as we can compute
2 (u, u; 0;
2
A (k) u,0,00
(i) o,u.o
8u
K
K
where
(R) u,uu
(Rging. ayy
+ (k)
(Xx:
with
Su = (v + w) + V. (oo).
For the derivative of the Monte Carlo version of the loss LKL as defined in (7) w.r.t to we may compute
! ((u u) v). x V (ao)) (X.:(k) ,s) ds.
(u, (u, v)
Σ
8v
k=1
We note that even for Monte Carlo estimator to be zero. For the computation of the derivative w.r.t. we refer to Nüsken & Richter (2021, Proposition 5.3).
and 8 we can usually not expect the variance of the corresponding
Remark A.2 (Control variate interpretation). For the gradient of the loss LKI. w.r.f. to u we may compute
8 5u
(u, 0, E L' ((u + v) Φ) 3) +( (R (x") B(X*) Saa())))
IF
Au, A (
where we used Girsanov's s theorem and the Itō isometry. Comparing with (18), we realize that the derivative of LLV w.r.t u for the choice u can be interpreted as control variate version of the derivative of LKL. thereby promising reduced variance of the corresponding Monte Carlo estimators, cf. Nüsken Richter (2021); Richter al. (2020).
In the context of reinforcement learning, such a control variate is also known as local baseline. As an alternative, global baselines have been proposed, where the batch-dependent scaling of the local baseline is replaced by an exponentially moving average. This corresponds to replacing the variance the loss with the second moment and additionally optimizing an approximation of the log-normalizing constant (with specific learning rate), see Malkin et al. (2022b). The resulting loss then known as (second) moment loss (Nüsken & Richter, 2021; Richter et al., 2020) or trajectory balance objective (Malkin et al., 2022a).
Remark A.3 (Alternative derivations of Lemma 3.1). The expression in Lemma 3.1 can also be derived via
dP (X") dP 8
dP dIP x
dP y aot dIP
dIP (Xw) PxF (XT) )
(X
(x")=
dlP xr
Ptarget
where the first factor can computed via recalling
dl'x
(X")= = exp (Ryet... (X-.))
(19)
dPxx
Yet another viewpoint based on importance sampling in path space, see, e.g., Hartmann et al. (2017). Since our goal is to find an optimal control u² such that we get samples Xy." Ptarget, we may dP xuu
define our target path space measure IP x via dP xr
(Xm)) Prarget (X*). We can then compute
pXT
dP PM (X") JP X
dIP
dP x
(X"),
dP
dPP X
which, together with (19), is equivalent to the expression in Lemma 3.1. Note that in the importance sampling perspective do not need the concept of time-reversals.
16

---

Published as a conference paper at ICLR 2024
A.3 SAMPLING VIA LEARNED DIFFUSIONS
In the following, we provide high-level overview of sampling methods based on controlled diffusion processes. We base our explanation on the general KL-based loss stated in (7) since most previous methods are special cases of this formulation, see Table 1.
Let us recall that we want learn a control u such that X Ptarget- We first observe that the terminal costs (Xx" or Bref (X.)) contain the term - log p - Pearget log Z, which penalizes X7 for ending up at regions with low probability w.r.t. the target density. The other terms of the terminal cost, together with the running costs R RBrride,, are enforcing additional constraints on the trajectories of our process X", In our formulation, they generally enforce X to be the time-reversal of Y For special choices of u, this yields the following settings:
For the PIS method, we minimize the reverse KL. divergence of the controlled process X" to the uncontrolled process X promoting to be as close to zero possible. This corresponds to a classical Schrödinger bridge problem, see Appendix A.4, which, for the simple initial condition X 8m can be solved without sequential optimization routines, see also Section 1.1 and Appendix A.4.1.
The DIS and DDS methods motivated by diffusion-basee generative modeling (Ho et al., 2020; Kingma et al., 2021; Nichol & Dhariwal, 2021; Vahdat et al., 2021; Song & Ermon, 2020). In particular, they minimize the reverse KL divergence the time-reversed "noising" process Y In other words, X" is enforced to denoise the samples Y, in order to yield samples from 1
Ptarget-
While we base our unifying framework in Section 2 on the perspective of path measures, the respective methods for the KL divergence can also be derived from the underlying PDEs or BSDE systems, see Appendix A.4 and Berner et al. (2024).
A.4 THE SCHRÖDINGER BRIDGE PROBLEM
In this section, provide some background information the classical Schrödinger bridge problem. Recall from Section 3.1 that out of all solutions u" fulfilling the general bridge problem stated in Problem 2.2, which can be characterized by Nelson's identity in (9), the Schrödinger bridge problem considers the solution w" that minimizes the KI. divergence DKI. (IP x=* [Pxx+) to a given reference process X", defined in (2) with u replaced by r E U,
-x = (μ + (X,, 8) ds + (s)
X ~ Pprior-
Traditionally, the uncontrolled process X with r = 0 is chosen, i.e., dxo x(XP,,) ds (s))
X ~ Pprrior
In the following, we will formulate optimality conditions for the Schrüdinger bridge problem defined in (12) for this standard case 0. Moreover, we outline how the associated BSDE system leads to the same losses as given in (7) and (8), respectively. The ideas are based on Chen et al. (2021a); Vargas (2021): Liu al. (2022); Caluya & Halder (2021).
First, we can define the
1 2
((x, min IF wEll
[ }) ds X x. X4
Ptarget
By the dynamic programming principle it holds that & solves the Hamiitonn-accobi-eellee (HJB) equation
aΦ - Vo (with unknown boundary conditions) and that the optimal control satisfies
1 Tr (oo V
2 1 v
(20)
a = - Vg.
Together with the corresponding Fokker-Plande equation for X" this yields necessary and sufficient conditions for the solution to (11). I. Now, we can transform the Fokker-Plang equation and the HJB equation (20) into a system of linear equations, using the exponential transform
/ == exp(p and + = Px= exp(0) = Px=
(21)
1
17

---

Published as a conference paper at ICLR 2024
often referred to the Hopf-Cole transform. This yields the following well-known optimality conditions of the Schrödinger bridge problem defined (12).
Theorem A.4 (Optimality PDEs). The solution u to the Schrödinger bridge problem (12) is equivalently given by
1. 4 = Vo, where px=* and & are the unique solutions to the coupled PDEs
O.px'' - (px (μ - oo T0 Vo)) + Tr (oo V'xx")
O.b - Ve
Tr (oo V22) + T
with boundary conditions
Px (,,0)
Pprior
Px (:,T) = Ptarget
2. = TV log th, where v and Hb are the unique solutions to the PDEs
{ 0 - P
I Tr [o V
( Vp + Tr 00 V2
(22)
with coupled boundary conditions
(1. (., 0) Ppriors
(23)
(.)()) = Ptarget-
The optimal control v" is given by Nelson's 's identity (9), i.e.,
= o Voogxx- u² o T Vog
(24)
Using Itô's lemma, we now derive a BSDE system corresponding to the PDE system in (22).
Proposition A.5 (BSDEs for the SB problem). Let us assume &p and & fulfill the PDEs (22) with boundary conditions (23) and let us define the processes
v log v (X s).
5 log (X ,8,
z = o V log (X 5) = =====((A (X.",,
3 V log X v² vx(x"", 8),
where the process X is given by
dx;" (μ (X," ds ds + (8) dw,
with w E U being an arbitrary control function. We then get the BSDE system
dy IZ. w(Xx s) - 1)) ds + - dW.,
d 11:2112 V (z m(x,", s)) + 2 w(x" s)) ds + 2g dW..
Furthermore, it holds
y + T: = Px (X. s) log py (X",
(25)
Proof. The proof similar to the one Chen et (2021a). For brevity, we define D 100 We can apply Itô's lemma to the stochastic process y log (X",, s) and get
dy (a, log v + V log e (μ ow) + Tr [ (X 5) ds+a T V log ())
Further, via (22) it holds
1 (-V p TTV(()
log t Tr DV'
0 log d
2b
18

---

Published as a conference paper at ICLR 2024
and we note the identity
V
(V
V
log
6
Combining the previous three equations, we get
Vih Vi (V(b)
))
d
V log 1o - TT D
(X 8) ds o T log s) dW,
222
2 w (X 8) 2 |1 112
1
ds 2g" dW..
Similarly, we may apply Itô's lemma to DTH = log 15 (x 8) and get
V Th 00 log T + V log T. (a + ow) + Ir (VV² log x) (x,", 8) ds + 2 dW.,
Now, via (22) it holds that
0, log & = (-V v. (vR) + (DV2T)) = V log 2 v. Tr DV
Combining the previous two equations, we get
- 6 +
Tr
D
DV
1o2 $
V u o V log w) (x s) ds + 2. dWW.
Now, noting the identity
V
D Vzy
Tr D
+ DV2 log $
2T
2 T V log 2212
log 2 + v (oo Too log 1)
can get the relation
. = 2
2 |o T v log क 2 V. (oo Vollologo T0 - 14) + T
1
Voog x. w (X.. s) ds 2 dW,
= - 2 12 + V (a (1) + 2,r w) s) -.r dW..
1
which concludes the proof.
Note that the BSDE system is slightly more general than the one introduced in Chen et al. (2021a), which can be recovered with the choice wr( w(X"", ,)) n Also, the roles of Pprior and Ptarget. are interchanged in Chen et al. (2021a) since they consider generative modeling instead of sampling from densities.
A valid loss can now be derived by adding the two BSDEs and recalling relation (25), which yields
Ptarget (x
B(x log( /)) = llog
(2 + 5t) (D/
Pprior (X11)
Rridee " w.
+ -w. (XX)
almost surely. Analogous to Berner et al. (2024); Huang et al. (2021) in generative modeling, the above equality suggests a parameterized lower bound of the log-likelihood log Pprior when replacing the optimal controls in 2 = u (X",, s) and 2. = 2 (X"", s) with their approximations u and v, Chen et al. (2021a). This lower bound exactly recovers the loss given in (7). Further, note that the variance of the left-hand minus the hand side is zero, which readily yields our log- variance loss as defined in (8).
19

---

Published as a conference paper at ICLR 2024
A.4.1 SCHRÖDINGER HALF- BRIDGES (PIS)
For the Schrödinger half-bridge, also referred as PIS, introduced in Section 3.3, we can find an alternative derivation, motivated by the PDE perspective outlined in Appendix A.4. For this derivation is crucial that assume the prior density to be concentrated at single point, i.e. Pprior 62o for some 20 E R (typically Fo = 0), see Tzen & Raginsky (2019); Dai Pra (1991). We can recover the corresponding objectives by noting that, in the case Pprior = dz.- the system of PDEs in (22) can be decoupled. More precisely, we observe that the second equation in (22) is the Fokker-Plance equation of X and we have that
Pxu* exp(o) = Px and (=, 0) = Pxy = 620-
In view of (24), we note that this defines 2- o² V log pxo. By (21), we observe that &b Px=* which yields the boundary condition
Px
pxt
Zpxp P
Tl:T) = - log (h, = log
log
Ptarget
to the HJB equation in (20). By the verificution theorem (Dai Pra, 1991; Pavon, 1989; Nüsken & Richter, 2021; Fleming & Soner, 2006; Pham, 2009), we thus obtain the PIS objective
(u) = E 2 L
Pxp (X7)
[u((X, s)ll ds + log
= E (R + B (xx))
((x4)
Moreover, the optimal control given by u" -a T a T log v. We can also derive this objective from the BSDE system Proposition A.S. Since x.) 0) = S0n9 we may focus on the process DTHH log (X.", only, and get
VT Vi ['
- w ;) 1 ds [ z dw
The PIS objective now follows by choosing w/ (X 8) z and noting that
VT = log (XT, T) = log Prarget (XT)).
Px
Recalling our notation in (I), this also shows that the log- variance loss can be written as
(u) = W (R + Sy + Brert
A.5
TRACTABLE SDES
Let us present some commonly used SDEs of the form
dx u(X", s) ds (s) aW,
with affine drifts that have tractable marginals conditioned on their initial value, see Song et al. (2020). For notational convenience, let us define
a(t) S B(s)ds
with suitable 3 e C((0, TJ, (0, oo)).
Variance--reeserving (VP) SDE: This Ornstein- Wrnstein-Uhlenneck process is given by
(t) = 222)(t) and u(xxt)= = -B(t)xr
with E (0, oo). Then, we have that
NNeeNaaaaa (o)l v2 (1-e-2010) 1). X,IX, ~ N
This shows that for (T') sufficiently large it holds that Xx 2 N (0, 12).. For Xo Nm(, E), further have that
X ~ N (e (e-al) -2((() (I- v²1) + +121)
(26)
20

---

Published as a conference paper at ICLR 2024
Algorithm Training of generalized time-reversed diffusion sampler
Input:
neural networks ug, Vyy with initial parameters 0, 7, optimizer method step for updating the parameters, number of steps K, batch size m Output: optimized parameters 0 for k t 1...., K do +- choose KL-based loss LKI. (7) or log-variance loss Liv in (8)
......K.
Setup
if L LKI. then
w t Ug
p Pprior
else
choose (detached) control for the forward process
p + choose initial distribution for the forward process
end if
for
1...., 772 do
Approximate cost (batched in practice)
x. f- sample from p
(W, x f- simulate discretizations of Brownian motion W and SDE X with X x R Sridge B) &- compute approximations of the running and terminal costs using X Ju
rnd, x R BBridge +
if C LIv then
Supta, + compute approximation of the stochastic integral using W rnd; rnd, + Suote,
end if
end for
mean
f- 1 Ei rnd;
P Compute loss
m
if L LKL then
C + mean
else
C + I-I Ei=1 (ndd, mean)²
end
if
4- step
0, VoL.
D Gradient descent
- step
7: V.
end for
Variance-explodingg (VE) SDE/ scaled Brownian motion: This SDE is given by scaled Brown- ian motion, i.c., = 0 and a as defined above. It holds that
X, ~ N (Xo, (()))
For Xo ~ (m,, E), we thus have that
~ (m, Σ).
A.6 COMPUTATIONAL DETAILS
For convenience, we first outline our method in Algorithm 1. Recall that the methods DIS, PIS, and DDS can be recovered when making particular choices for v, r, and Ppriors see Table 1. We specify the corresponding setting and further computational details in the following.
General setting: Every experiment executed on a single GPU and, in our PyTorch implementa- tion, we generally follow the settings and hyperparameters of DIS and PIS presented in Berner et al. (2024), which itself is based on the implementation of Zhang & Chen (2022). In particular, we the Fourier MLPs of Zhang & Chen (2022), batch size of 2048, and the Adam optimizer. To facilitate the comparisons. we use a fixed number of 200 steps for the Buler-Maruyamt scheme. A difference Berner et al. (2024) is that observed better performance (for all considered methods
21

---

Published as a conference paper at ICLR 2024
and losses) by using exponentially decaying learning rate starting 0.005 and decaying every 100 steps to a final learning rate of 10-4, We use 60000 gradient steps for the experiments with d < 10 and 120000 gradient steps otherwise approximately achieve convergence. However, we observed that the differences between the losses are already visible before convergence, see, e.g., Figure 1.
PIS: We follow Zhang & Chen (2022) and use Brownian motion starting at 6g for the uncontrolled SDE X Furthermore, we also leverage the score of the target density V log (typically given in closed-form or evaluated via automatic differentiation) for the parametrization of the control see Zhang & Chen (2022): Berner et al. (2024).
DIS: We use the VP-SDE in Song et al. (2020) for the SDE y Specifically, we use 1 and
B(t) (1 -) t) Smin + t.3max E 1],
with 8min 0.05 and Smax see Appendix A.5. Moreover, we employ linear interpolation of v log and v log Pprior for the parametrization of the control u, see Berner et al. (2024).
Bridge: For the general bridge, we consider the loss (7), which corresponds to the setting Chen et al. (2021a) adapted to unnormalized densities. We use an analogous setting to DIS; however, we additionally employ Fourier MIP to control the process y Since Y is already close to Pprior by construction of the VP-SDE, we use a lower initial learning rate of 10-4 for the control e. While these choices already provide better results for the KL divergence, see Table 5, we note that more sophisticated, potentially problem-specific choices might be investigated in future studies. In particular, for the general bridge, we would be free to choose the prior density Pprior as well as the drift function μ in the SDEs (2) and (3).
Log-variance loss: For the log- log-variance loss, we only change the objective from CKL to Liv- where used the default choice of w u, i.e., X" X". We emphasize that do not need to differentiate w.r.t. w, which results in reduced training times, see Figure 9. In practice, we can thus detach X= from the computational graph without introducing any bias. This be achieved by the detach and stop-gradis ent operations in PyTorch and TensorFlow, respectively. We leave other choices of future research and anticipate that choosing noisy versions of u in the initial phase of training might lead to even better exploration and performance. Furthermore, we use the same hyperparameters for the log-variance loss as for the Kl--based loss. these settings originate form Berner et al. (2024) and have been tuned for the KL-bbased loss, we suspect that optimizing the hyperparameters for the log- variance loss can lead to further improvements.
Evaluation: To evaluate our metries, we consider 100 samples (a(1)) and use the ELBO as an approximation to the log normmmnalgggng constant log Z, see Appendix A. We further compute the (normalized) effective sample size
(Ei w(())
ESS
11 E (w(()))
where (z)())²z are the importance weights the samples (1))) path space. Finally, we estimate the Sinkhorn distance W (Cuturi, 2013) and report the error for estimating the average standard deviation across the marginals, i.e.,
std
W[GH], where G
Ptarget
A. COMPUTATION OF LOG-NORMALIZING CONSTANTS
For the computation of the log normalzzng constant log Z in the general bridge setting, we note that for any u, W it holds that
dIP dlP X
E
(X")
1
Our implementation of the Sinkhorn distance is based on https://githhb. com/ fwi w amss scal aable peteee-yytcchh--- a inhhorn with the default parameters.
22

---

Published as a conference paper at ICLR 2024
Together with Proposition 2.3, this shows that
log Z - log E eep (- Bridge + Sutp + B) (x*))]-
(27)
If u² and v = v*, the expression in the expectation is almost surely constant, which implies
log Z = - -RR "=""""." + Su + B) (x*)).
(28)
If we only have approximations of ut and v², Jensen's 's inequality shows that the right-hand side in (28) yields lower bound to log Z. For PIS and DIS, the log-normalizing constants can computed analogously, see Zhang & Chen (2022); Berner et al. (2024). If not further specified, we use the lower bound as an estimator for log Z in our experiments.
A.7 PARTIAL TRAJECTORY OPTIMIZATION
In this section, present method that does not need the simulation of entire trajectories but rely on subtrajectories only. On the one hand, this promises faster computations and, on the other hand, it can be used for exploration strategies since subtrajectories can be started arbitrary prescribed locations, independent of the control u. Crucially, this strategy only works for the log-variance loss and not for the KL.-based loss.
Let us recall that the log- reg-variance loss in (8) is defined for any process X In particular, in addition to the control w, we can also freely choose the initial condition of X see the proof of Proposition 2.3. Motivated by Zhang et al. (2023), we can leverage this fact to train the DIS or DDS methods on smaller time-intervals [t, T). Specifically, recall that the log-variance loss for the DIS method in (14) is given by
(u) = W [ [RIDI + Sa + (X
(29)
where the optimal control is defined by
ut* = o TV log Pyo,
(30)
see Section 3.2. Also, recall that
B(X log Pprior (X0)
(31)
p(X4)
where Pprior RJ pyp- Now, assuming an approximation ((,()) & log for t E [0, 7), we can replace log Pprior in (31) by ( t). and consider the corresponding sub-problems time intervals [t, T.. For the log- variancee loss, we can then sample ~ (n0 Ti), choose an arbitrary initial condition X and minimize the loss (29) on the time interval T).
In order to obtain an approximation ( f) Rs logPye, we could train a separate network, e.g., by using the underlying optimality PDEs in Appendix A.4. However, for the DIS method, this is not needed if we parametrize the control
V
(32)
where & is a neural network. Based on (30) we can then use p f) as an approximation of log PY; during training. Therefore, we can therefore optimize the loss
CtV sub (4) = V (R
[[(
+ SovV + B ant)
w.r.t. the function db In the above, we can pick 1 ~ Unif((O, TJ), X with v being suitable probability measure on Rdd and
Bo,su((X)) = 4(X," t) - log p(XT),
where (with slight abuse of notation) the integrals and 5 defined in (1) now run from f to T.
The subtrajectory training procedure has three potential benefits. First, training may be accelerated since we consider smaller time-intervals [t. T]. leading to faster simulation the SDEs. Second,
choose X in suitable way to prevent mode collapse, e.g., can sample X from
23

---

Published as a conference paper at ICLR 2024
Table 3: We compare the performance DIS without using the target information V logg in the parametrization. In this case, the performance of the KL-based loss is generally decreasing, as also observed in Zhang & Chen (2022). For the log-variance loss, we can counteract this decrease by relying on sub-trajectories starting at random / ~ Uni TJ) and sufficiently large a (0, oo)) order to facilitate exploration, see Appendix A.7. This allows to obtain competitive results without using any gradient information of the target.
Unif([-a, a] a!() (for
Problem
Loss
A log Z
W
ESS f
A Astd
GMM (d (d=2)
KIL-DIS (Berner et al., 2024)
2.291
3.661
0.8089
3.566
LV-DIS-Subtraj. (ours)
0.059
0.020
0.8613
0.008
DW (d 5, m, =,
KL-DIS (Berner al., 2024)
3.983
5.517
0.3430
1.795
LV-DIS- LV-DIS-Subtraj. (ours)
0.394
0.121
0.4378
0.002
Porer
(eurs)
Su-Dis:
-40
-40
-20
0
20
40
-40
-20
0
20
40 -40
-20
0
20
40
440
-20
0
20
X2
Figure 4: Contour plots of a Gaussian mixture model Pharget with 40 modes analogous to the problem proposed Midgley et al. (2023). We plot samples the prior Pprior = NO I) (left) and the DIS method trained with the KL-based loss, the log- variance loss, and partial trajectory optimization (from left to right), see Appendix A.7. For all methods, we use T = 2 to guarantee Рyş ~ Pprior- Using the setting from Table 3, subtrajectory training can recover all modes without gradient information from the target, whereas other methods suffer from mode collapse despite making Midgley et al. (2023) report mode collapse on this benchmark for several state-of- the art methods. We remark that LV-DIS (unlike KL-DIS) recovers all modes when slightly increasing the prior variance.
of V log p.
distribution y with sufficiently large variance. Third, (32), we consider a parametrization of u that does not rely V log p. which may not be available expensive compute.
In addition to these benefits, we show in Table that subtrajcctory training can achieve competitive performance compared to the results in Table 2. On the contrary, show that the performance of DIS with the KL-based loss gets worse when not using parametrizations that contain the term V log p. Note that subtrajectory training cannot be used with the KL-based loss since, for this loss, need to sample the initial condition according to X."
In Figure 4, compare the performance the benchmark proposed in Midgley al. (2023). We show that partial trajectory optimization can identify all 40 modes of the Gaussian mixture model, significantly outperforming the DIS method, even when using the log-variance loss and the derivative of the target in the parametrization. We note that Midgley et al. (2023) report mode collapse on this benchmark for other state-of-the-ars methods, such Stochastic Normalizing Flows (Wu et al., 2020), Continual Repeated Annealed Flow Transport Monte Carlo (Arbel et al., 2021), and flows with Resampled Base Distributions (Stimper et al., 2022).
A.8 FURTHER EXPERIMENTS AND COMPARISONS
In this section, present further results and ablation studies. In Table 4, we show that the log- variance loss also leads to improvements for smaller batch sizes. This can be motivated by its variance-reducing effect, Proposition 2.5, Remark A.2, and Figure 5. In Figures 6 and 7,
24

---

Published as a conference paper at ICLR 2024
Table 4: The same setting as in Table 2 is considered, however, with smaller batch size of 512 instead of 2048. We again observe that the log-variance loss consistently yields better performance. This can be motivated by the inherent variance reduction of its gradient estimators (particularly helpful for smaller batch sizes), see Proposition 2.5, Remark A.2, and Figure 5,
Problem
Method
Loss
A log Z 4
W
ESS 1
std 4
GMM (d = 2)
PIS
KI. (Zhang & Chen, 2022)
2.201
2,708
0.0002
3.576
LV (ours)
2.200
2.629
0.0002
3,564
DIS
KL (Berner et al., 2024)
1.725
0.088
0.0045
2.711
LV (ours)
0.063
0.020
0.8517
0,004
DW (d = 5, m = 5, 4)
PIS
KL (Zhang & Chen, 2022)
3.693
4.949
0.0001
1.793
LV (ours)
0.285
0.124
0.5957
0,008
DIS
KL (Berner et al., 2024)
4.047
5.068
0.0015
1.797
LV (ours)
0.447
0.121
0.3917
0,002
avg. gradlent stddev.
100
Jow womplllll yo
LV-DIS (ours)
stolde..
10
10-3
10-2
Hum
10-1
10-1
Mageet Mmmann
20000
40000
60000
20000 40000
60000
gradient steps
gradient
Figure 5: We compare the standard deviations of the loss and (average) gradient estimators using either the KL-based loss the log-variance loss. Each standard deviation is computed over 40 simulations of the loss without updating the parameters. We show results for the DIS method on the 5-dimensional DW target. As predicted by our theory, the log-variane, loss exhibits significantly smaller standard deviations for both the loss and its gradient.
show that the log-variance loss can counteract mode collapse in both moderate as well as very high dimensions. Moreover, we present the results for the general bridge approach in Table 5, and consider problem from Wu et al. (2020) in Figure 8. In Figure 9, we present boxplots to show that our results from Table are robust w.r.t. different seeds.
Finally, in Table 6, we show that our methods are competitive to other state-of the-hhh-rtt sampling baselines. However, we want to emphasize that the focus of our work is not to extensively compare against MCMC methods or normalizing flows. Our goal is to show that recently developed methods, such as PIS, DDS, and DIS, can be unified under a common framework, which enables the usage of different divergences. We then propose the log-variance divergence, which makes diffusion-based samplers even more competitive and mitigates potential downsides compared to other methods. The fact that there general trade-off's between the considered diffusion bastud samplers and variants of MCMC has already been discussed and numerically analyzed in the papers introducing the respective methods, see Zhang& & Chen (2022); Vargas et al. (2023a).
25

---

Published as a conference paper at ICLR 2024
KL-PIS
LV-P15 (ours)
KL-D15
LV-DIS (ours)
Density
Histagram of
X1
X1
X;
Density
2-
X2
X2
X2
X2
Densey
Xa
Xa
X
Density
2
Xa
X
X
X1
Dersity
%
Figure 6: Marginals of samples from PIS and DIS (left and right) for the DW problem with d
5, and d = 4. The mode coverage of the log-variance loss is superior to the KL-based loss for all marginals (from top to bottom).
KL-DIS
DVDIS (ours)
Density
Marginal
Histogram xt
X1
Dousity
X2
Density
X
Xs
Demnity
1-
X4
Figure 7: Marginals of the first four coordinates of samples from DIS for a high-dimensional shifted double well problem in dimension d = 1000 with m and 5 = 2 (see Section 4.1), using the KL or the log-variance loss, respectively. Again one observes the better mode coverage of the log-variance loss compared to the reverse KL divergence.
26

---

Published as a conference paper at ICLR 2024
Table Metrics the general bridge approach in Section 2.3 for selected benchmark problems. We observe a clear improvement using the log- variance loss. Moreover, for the KL divergence, we note that the general bridge framework can obtain better results than DIS PIS, see Table 2. As in Table we report the median over five independent runs. We show emrors for estimating the log-normalizing constant A log Z as well the standard deviations std of the marginals. Furthermore, we report normalized effective sample size ESS and the Sinkhorn distance W (Cuturi, 2013). The arrows and | indicate whether we want to maximize or minimize a given metric.
Problem
Method Loss
log Z
W
ESS f
std
GMM (d = 2)
Bridge
KL (Chen et al., 2021a)
0.328
0.393
0.0180
0.698
LV (ours)
0.084
0.020
0.9669
0.010
DW (d 5, m 8,8 4)
Bridge
KL (Chen al., 2021a)
0.872
0.132
0.0561
0.099
LV (ours)
0.215
0.119
0.5940
0.002
Groundtruth
KL-PIS
LV-PIS (ours)
KL-DIS
LV-DIS (ours)
Figure 8: Comparison samples for the target in et al. (2020). For the KL-based losses, large fraction of the samples (PIS: 86%, DIS: 67%) lies outside of the domain despite the low density values. On the other hand, the log-variance loss significantly improves performance, yielding competitive performance compared to stochastic normalizing flows presented in Wu et al. (2020).
Table 6: We compare our methods to Continual Repeated Annealed Flow Transport Monte Carlo (CRAFT), see Arbel et al. (2021). We adapt the proposed configurations to use the same batch size and number of iterations as our methods and evaluate all methods using 10a samples. We see that diffusion- sampling, in combination with the log-variance loss, can provide competitive performance across range of metrics. We report the median over five independent runs and compare the normaaing constant (using the reweighted estimator in (27)), the Sinkhorn distance ground truth samples, and the error in estimating the average standard deviation.
Problem
Method
A log Z (rw) $
W J
std +
GMM (d (d=3) 2)
CRAFT (Arbel et al., 2021)
0.012
0.020
0.019
LV-PIS (ours)
0.001
0.020
0.023
LV-DIS (ours)
0.017
0.020
0.004
Funnel (d 10)
CRAFT (Arbel al., 2021)
0.123
5.517
6.139
LV-PIS (ours)
0.097
5.593
6.852
LV-DIS (ours)
0.028
5.075
5.224
DW (d 5, m m 5,6=4)
CRAFT (Arbel et al., 2021)
0.001
0.118
0.000
LV-PIS (ours)
0.000
0.121
0.001
LV-DIS (ours)
0.043
0.120
0.001
DW (d 50
CRAFT (Arbel et al., 2021)
0.000
6.821
0.001
LV-PIS (ours)
0.001
6.823
0.000
IV-DIS (ours)
0.422
6.855
0.009
'The configuration files can be found at transport/bllc//asttee
27

---

Published as a conference paper at ICLR 2024
GMM (d 2)
Punnel
DW 5,m= 5, 4.
DW (d = 50, m 5, 6=
0.0
0.0
-0.5-
-0.2
-1.0-
42.50
Target
-0.4
00.1
-15--
KL.
42.25
(ourst
0.6
-20-
-0.2
42.00
-0.8
255
41.75
-0.3-
30
-122
41.50
-35
-144
-04-
-1.0-
41.25
-1.0-
41.00
0.8
5
0.7
5.-
0.6
4-
5
0.5
3-
0.4
3-
0.3
2-
2-
0.2
0.1
1--
1-
1-
0.0
1.0
1.0
0.8
0.8
0.6
0.6
0.6
0.6
0.4
0.4
0.4
04-
0.2
0.2
0.2
0.0
0.0
0.0
2.00
1.044
1.75
3.5
7-
1.50
1.042
3.0
1.25
1.040
5-
1.00
2.5
1.038
0.75
2.0
3-
0.50
1.036
0.25
1.5
1.031
0.36
0.52
time par gradient stop [2]
0.56
0.54
0.48-
0.55
0.50
0.46
0.54
0.48
0.50
0.53
0.40
0.44
0.46
0.52
0.46
0.42
0.64- -
0.51
0.44
0.42-
0.50
0.40
0.42
PIS
DIS
PIS
105
P15
D15
PIS
DIS
Figure 9: Boxplots for five independent runs for each problem and method (KL-PIS, LV-PIS (ours), KI.-DIS, LV-DIS (ours) from left to right in each plot) in the settings of Table 2 and corresponding ground truth or optimal values (dashed lines). It can be seen that the performance improvements of the log- log-variance loss are robust across different seeds. At the same time, the log-variance loss reduces the average time per gradient step by cireumventing differentiation through the SDE solver.
28


Reviews:
[{"review_id": "b0Ofls4O5B", "reviewer": "Reviewer_y7Cc", "summary": "The paper tackles sampling problems, such as sampling from unnormalized densities (i.e., target distributions), by using stochastic differential equations (SDEs).\n\nTo do that, the paper first proposes a novel framework to unify previous SDEs-based sampling methods: (1) variational formulation of the Schrödinger bridge (SB) problem, (2) path integral sampler (PIS), (3) time-reversed diffusion sampler (DIS), and (4) denoising diffusion sampler (DDS). In the framework, the authors show that we can obtain the analytic form of the Radon-Nikodym (RN) derivative of one SDE wrt another SDE, evaluated at any path generated by the third SDE, as long as all SDEs have the same noise coefficient term (while time direction can vary). Then, the authors show that we can define proper divergences between two SDEs using the RN-derivative (by taking expectation with the RN-derivative). Furthermore, the authors show that the previous SDE-based sampling problems can be defined as the minimization of the reverse KL divergence (or its variant) wrt a reference measure; more specifically, the choice of the reference measure and forward will determine the corresponding sampling problems.\n\nSecond, based on the framework, the authors propose to use the log-variance divergence instead of the reverse KL divergence thanks to several benefits of the log-variance against the reverse KL; for example, the log-variance divergence uses a reference measure and thus does not require to differentiate through the SDE solver. The paper shows all four variants based on the log-variance divergence corresponding to SB, PIS, DIS, and DDS.\n\nFinally, the paper demonstrates the performance of the proposed methods against the previous works by using some benchmark datasets.", "questions": "N/A", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "In my understanding, the paper's contributions are clear, and I also consider that the results are essential for several reasons:\n\n1. The paper proposes a novel framework that provides better understanding of previous literature on sampling problems using SDEs.\n2. The paper well motivates the log-variance divergence-based methods so that readers can understand how each step contributes to the merits of the proposed method.\n3. I found that the paper has a well-organized structure that makes it clear to understand the proposed method. Significantly, the paper shares sufficient information to understand the new divergence’s behavior, which will help readers understand relevant backgrounds and the proposed method.\n\nIn general, the paper's contributions wrt the novelty are clear, and the proposed methods are well-defined. In addition, I found that the paper has a well-organized structure that makes it clear to understand the proposed methods.", "weaknesses": "In general, I find that the paper is well-written. However, descriptions of some derivations can be improved for clarification. For example, to introduce the derivation of Equation (10), the paper uses “defined in (2) with $u$ replaced by $r \\in \\mathcal{U}$”. I find this description a little confusing, as Equation 10 assumes three SDEs. Clarifying such descriptions would be helpful to potential readers who are not familiar with SDEs and relevant backgrounds.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HqtsvFh8DM", "reviewer": "Reviewer_i9dA", "summary": "The authors propose a unifying framework for various existing diffusion-based samplers. A log-variance based (rather than KL-based) loss is proposed. The proposed work is evaluated against competing methods on the problem of sampling without data --- i.e., with respect to an unnormalized density function.", "questions": "Please see weaknesses.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- There is an extensive theoretical discussion of the provided method and its analytical properties.\n- The authors provide a useful guide to organize recent works on diffusion-based approaches to density-based sampling.\n- The work appears to be mathematically sophisticated and relatively rigorous.", "weaknesses": "- Limited validation. The generalizations proposed are compared on rather simplistic examples. Though the authors appear to target scenarios where data is not available, data-based modeling is clearly an important application of a diffusion-based sampler. Is there a reason the models are not compared to other non-diffusion methods, e.g., MCMC, normalizing flow, autoregressive, or GAN models?\n\n- Unclear abstract: First, the abstract appears to claim that the authors \"identify [diffusion models] as special cases of the Schrödinger bridge problem\". This feels far too bold of a claim to me, as many papers have connected diffusion models to Schrodinger Bridges (and the authors are clearly aware of this). After all, nearly all modern generative diffusion models sample from a target density. Second, the authors focus entirely on the problem of sampling without data -- i.e., w.r.t. a density function. This is in stark contrast to what a general ICLR reader may consider as sampling in diffusion models --- i.e., sampling from a model trained on data. Although it appears that the proposed framework can perhaps be formulated to include this case, the goal of the article is clearly the former and not the latter. In fact, there do not appear to be any meaningful experiments on the latter case. This should be clarified in the abstract, as this fact was not at all clear except for a sentence in the third paragraph of the introduction.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OmUY0lizTi", "reviewer": "Reviewer_29Lh", "summary": "This paper proposes a framework for learning to sample from unnormalized target densities. Specific cases of this framework include several recently proposed algorithms (Schrodinger Bridge, Time-Reversed Diffusion Sampler, Path Integral Sampler, Denoising Diffusion Sampler). The framework involves a general divergence between path measures of diffusion processes. This encompasses the most common case of the reverse KL divergence, but other divergences can also be used. The paper introduces another divergence named the log-variance divergence, which is shown empirically to outperform the KL by avoiding mode collapse.", "questions": "Do the authors think the proposed methodology could be adapted to the case of generative modeling?\n\nOne could imagine that the variance of the Monte-Carlo estimator of the variance is high, which could explain why a large batch size (2048 in the experiments) is required. Did the authors try comparing their method with the KL divergence for smaller batch sizes? Do the authors still observe an improvement wrt the KL divergence?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper proposes a theoretical framework that encompasses several recently proposed methods for sampling via diffusions. This is a very interesting contribution, in particular from a theoretical point of view, since it allows a unified conceptualization of these methods. The proposed divergence seems to outperform the KL divergence, making it also an interesting contribution to the community.", "weaknesses": "[Update on Nov 22: authors have answered in details the various points raised below.]\n\nI found the paper quite dense, although I acknowledge that this is partly due to the nature of the contribution which is to propose a unifying framework. In particular, Section 3 assumes precise knowledge of the various samplers recently proposed in the literature. While it is not a problem in itself, I would suggest in case of lack of space to move part of the derivations of this Section to the Appendix.\n\n**More importantly, I am not convinced by the explanations on the advantages of the proposed log-variance divergence in Section 2.3**, although I agree that the experiments show that it outperforms the KL divergence in the presented settings. This is an important issue in my opinion, especially since these advantages are part of the narrative of the paper (see the introduction and the conclusion).\n\nFirst, regarding the choice of $w$, it is stated that this choice allows a balance between exploitation and exploration. I do not understand what this precisely means. I believe this statement needs to be backed by either a theoretical analysis or a numerical experiment (or be removed). It is also stated without justification that choosing $w$ different from $u$ might be beneficial to avoid mode collapse. At the same time, in the experiments, $w$ is chosen equal to $u$ (according to Appendix A.4), and yet mode collapse is already avoided. Furthermore, I did not see where it is specified in Section 3 how $w$ should be chosen for the connection with the various approaches to hold.\n\nRegarding the computational benefit, the proposed method still requires to compute gradients with respect to $u$ and $v$ even though gradients wrt $w$ are not required. It is not clear for me why this entails a significant improvement in computational efficiency. I think the naive implementation with stop_gradient still requires backpropagating through the SDE, although I agree that the computational graph is lighter. Experimentally, the time improvement shown in Figure 6 is marginal. Nevertheless, I believe that a more efficient ad hoc implementation that requires only a forward pass through the SDE is perhaps possible. However, this is not explained in the paper. Furthermore, unless I am mistaken, simulating the SDE is still required in any case. This should be mentioned to avoid inducing the reader into thinking that it is not required (like in score matching).\n\nFinally, Proposition 2.5 is not uninteresting, but it is difficult without further explanations or experiments to understand if it plays a role in the improvement observed in practice in the paper.\n\nAll in all, I think that either more precise explanations must be given on these various points, or the claims regarding the proposed log-variance divergence must be adapted.\n\n**Minor remarks that do not influence the rating**\n+ Proposition 2.3: I think $\\rho$ is undefined.\n+ The title of the paper is not very informative. I think it would be interesting to consider a title more specifically linked to the contributions of the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DsdexjwMk1", "reviewer": "Reviewer_kAh6", "summary": "The work proposes a unified framework that clearly explains recent sampling algorithms based on generative models. The improvement is led by replacing KL divergence with log variance loss. Experiments show that with new loss, existing generative sampling can achieve better performance.", "questions": "See above", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. Authors present the work and its theoretical analysis clearly, good writing.\n2. The usage of the new divergence metric leads to much more robust training/optimization, which is a critical issue for existing learning-based sampling algorithms.", "weaknesses": "1. Although the author claims the framework is new. The underlying math theory and contributions have existed long before.\n2. To further improve the readability, authors should clearly define the log-variance loss and outline the algorithm pseudocode in the appendix. \n3. PIS work shows the network that has information of target distribution gradients can perform much better. Is it also true with new log-variance loss?\n4. I am a bit confused discussion of stop gradient, does it lead to a biased gradient estimation? Is it the same trick used in Chen 2021a?\n5. Can authors provide more descriptions and analysis on log-variance and its computations? How do authors estimate the mean, which is required to calculate the variance?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)