An optimal control perspective on diffusion-based generative modeling
Authors: Julius Berner
Lorenz Richter
Karen Ullrich
Venue: Accepted by TMLR ()

Abstract: We establish a connection between stochastic optimal control and generative models based on stochastic differential equations (SDEs), such as recently developed diffusion probabilistic models. In particular, we derive a Hamilton--Jacobi--Bellman equation that governs the evolution of the log-densities of the underlying SDE marginals. This perspective allows to transfer methods from optimal control theory to generative modeling. First, we show that the evidence lower bound is a direct consequence of the well-known verification theorem from control theory. Further, we can formulate diffusion-based generative modeling as a minimization of the Kullback--Leibler divergence between suitable measures in path space. Finally, we develop a novel diffusion-based method for sampling from unnormalized densities -- a problem frequently occurring in statistics and computational sciences. We demonstrate that our time-reversed diffusion sampler (DIS) can outperform other diffusion-based sampling approaches on multiple numerical examples.

Published in Transactions on Machine Learning Research (02/2024)
An optimal control perspective on diffusion-based generative modeling
Julius Berner*i Caltech
jbernerecattech. edu
Lorenz Richter*
richter@aib.cc de
Zuse Institute Berlin dida Dulenschmiede GmbH
Karen Ullrich
karenusmmetaa. com
Meta AI
Reviewed on OpenReview: https: // openreview. net/ forum? id- oYIju37P7P
Abstract
We establish a connection between stochastic optimal control and generative models based on stochastic differential equations (SDEs), such as recently developed diffusion probabilistic models. In particular, we derive a Hamilton-Jacoe evolution of the log-densities of the underlying SDE marginals. This perspective allows to transfer methods from optimal control theory to generative modeling. First, we show that the evidence lower bound is direct consequence of the well-known verification theorem from control theory. Further, we can formulate diffusion-based generative modeling as minimization of the Kullback Leibler divergence between suitable measures in path space. Finally, we develop novel diffusion-barte method for sampling from unnormalized densities a problem frequently occurring statistics and computational sciences. We demonstrate that our time-reversod diffusion sampler (DIS) can outperform other diffusion- basod sampling approaches multiple numerical examples.
equation that governs the
1 Introduction
Diffusion (probabilistic) models (DPMs) have established themselves as state-of-the-art in generative modeling and likelihood estimation of high-dimensional image data (Ho al., 2020; Kingma et al., 2021; Nichol & Dhariwal, 2021). In order to better understand their inner mechanisms, multiple connections adjacent. fields have already been proposed: For instance, in discrete-time setting, one can interpret DPMs types of variational autoencoders (VAEs) for which the evidence lower bound (ELBO) corresponds to multi-scale denoising score matching objective (Ho et al., 2020). In continuous time, DPMs can be interpreted in terms of stochastic differentiaal equations (SDFs) (Song et al., 2021; Huang et al., 2021; Vahdat et al., 2021) or as infinitely deep VAEs (Kingma et al., 2021). Both interpretations allow for the derivation of a contimuous-timo ELBO, and encapsulate various methods, such as denoising score matching with Langevin dynamics (SMLD), denoising diffusion probabilistic models (DDPM), and continuous -ietiie normalizing flows as special cases (Song et al., 2021; Huang et al., 2021).
In this work, we suggest another perspective. We show that the SDE framework naturally connects diffusion models to partial differential equations (PDEs) typically appearing in stochastic optimal control and reinforcement learning. The underlying insight is that the Hopf- Cole transformation serves as a means of
"Equal contribution (the author order was determined by numpy (andom.rrand())
Work done during internship Meta

---

Published in Transactions on Machine Learning Research (02/2024)
Yo X
-
-
Y
Diffusion models
Generative SDE (Xx)sejor with densely px(+
density
(x(x,) f) solvos Fokker-Plande equation
ELBO
drift of drift of
(x I - f) solova (genoratized)
Path space measures
Kolmagorov equation
Optimal control
Sampling from
(unnoraaized) densitiea
2153
2135 2265 2XC
Figure The connection between diffusion models and optimal control is outlined in Section 2. Three possible implications of this connection are the deriva- tion of the ELBO via the verification theorem (Sec- tion 2.2), a path space measure interpretation of dif- fusion models (Section 2.3), and novel approach for sampling from (unnormalized)) densities (Section 3).
Figure 2: reversed process X, we show the density (r(,,)= (x ---)) (top), the drift f(-,t) of Y (middle), and the drift u(- () of X (bottom) for E [0, 12.3.11 T], see (2) and (3).
For stochastic process Y and its time-
27
showing that the time-reversed log-density of the diffusion process satisfies Hamitton-Jacobi-eellm (HJB) equation (Section 2.1). This relation allows to connect diffusion models to control problem that minimizes specific control costs with respect to a given controlled dynamics. We show that this readily yields the ELBO of the generative model (Section 2.2), which can subsequently be interpreted in terms of Kullback-Leibler (KL) divergences between measures on the path space of contimous time stochastic processes (Section 2.3).
While one of our contributions lies in the formal connection between stochastic optimal control and diffusion models, described in Section 2.4, we moreover demonstrate the practical relevance of our analysis by transferring methods from control theory to generative modeling. More specifically, in Section 3, we design novel algorithm for sampling from (unnormalized) densities a problem which frequently occurs in Bayesian statistics, computational physics, chemistry, and biology (Liu & Liu, 2001; Stoltz et al., 2010). As opposed to related approaches (Dai Pra, 1991; Richter, 2021; Zhang & Chen, 2022a), our method allows for more flexibility in choosing the initial distribution and reference SDE, offering the possibility to incorporate specific prior knowledge. Finally, in Scction 4, we show that our sampling strategy can significantly outperform related approaches across number of relevant numerical examples.
In summary, work establishes new connections between generative modeling, sampling, and optimal control, which brings the following theoretical insights and practical algorithms (sce Figure for an overview):
PDE and optimal control perspectives: We identify the HJB equation as the governing PDE of the score function, which rigorously establishes the connection of generative modeling to optimal control (Section 2.1). The HJB equation can further be used for theoretic analyses as well for practical algorithms for the numerical approximation of the score, e.g., based neural PDE solvers.
Using the HJB equation, show that the objectives in diffusion- generative modeling can be solely derived from the principles of control theory (Section 2.2).
ELBO:
Path space perspective: This perspective diffusion models offers an intuitive representation of the objective (and the variational gap) in terms of KL divergences between measures on path space (Section 2.3). Crucially, it also allows to consider alternative divergences, such as, e.g., the log-variance divergence, leading to improved loss functions and algorithms.
Sampling: Our control perspective leads to a novel, biffusiin-bbsed method to sample from unnormalized densities (Section 3). This method can already outperform previous diffusion-basod hate-of-the-ar samplers. Moreover, the connection diffusion models allows transfer noise schedules, integrators, and further popular techniques from generative modeling to sampling.
2

---

Published in Transactions on Machine Learning Research (02/2024)
1.1 Related work
Starting from DPMs (Sohl-Dickstein et al., 2015), number of works have contributed to the success of diffusion-based generative modeling, see, e.g., Ho et al. (2020); Kingma et al. (2021); Nichol & Dhariwal (2021); Vahdat et al. (2021); Song & Ermon (2020). We are building upon the SDE-based formulation developed by Song et al. (2021) which connects diffusion models to score matching (Hyvärinen & Dayan, 2005). The underlying idea of time-reversing a stochastic process dates back to work by Nelson (1967); Anderson (1982); Haussmann & Pardoux (1986); Föllmer (1988). Notably, Pavon (1989) connects the log-density of such a reverse-lime SDE to an HJB equation. In this setting, we exlend the results of Huang et al. (2021) on the (contimuous-time) ELBO of diffusion models and provide further insights from the perspective of optimal control and path space mcasures.
Diffusion models:
For further previous work on optimal control in the context of generative modeling, we refer the reader to Tzen & Raginsky (2019); De Bortoli et al. (2021); Pavon (2022); Holdijk et al. (2022). Connections between optimal control and time-reversals of stochastic processes have also been analyzed by Maoutsa & Opper (2022); Maoutsa (2023). Finally, we refer to Richter & Berner (2024) for further generalizations based on the path space perspective and to Zhang& & Katsoulakis (2023) for recent work on interpreting generative modeling via mean-field games.
Sampling from (unnormalized) densities: Monte Carlo (MC) techniques are arguably the most common methods to sample from unnormalized densities and compute normalizing constants. Specifically, variations of Annealed Importance Sampling (AIS) (Neal, 2001) and Sequential Monte Carlo (Del Moral et al., 2006; Doucet et al., 2009) (SMC) are often referred to as the "gold standard" in the literature. They apply a combination of Markov chain Monte Carlo (MCMC) methods and importance sampling to a sequence of distributions interpolating between a tractable inltial distribution and the larget density. Even though MCMC methods are guaranteed to converge to the target distribution under mild assumptions, the convergence speed might be too slow in many practical scttings (Robert et al., 1999). Variational methods such as mean-field approximations (Wainwright et al., 2008) and normalizing flows provide an alternative" (Papamakarios et al., 2021). By fitting a parametric family of tractable distributions to the target density, the problem of density estimation is cast into an optimization problem. In Salimans et al. (2015) it is noted that stochastic Markov chains can be interpreted as the variational approximation in an extended state space, thus bridging the two techniques and allowing to apply them jointly.
In this context, diffusion models have been employed to approximate the extended target distribution needed in the importance sampling step of AIS methods (Doucet et al., 2022). Moreover, diffusion models have been trained on densities by simulating samples using importance sampling with the likelihood of the partially-traneed model (computed via the probability flow ODE) as proposal distribution (Jing et al., 2022). In this work, we propose a novel variational method based on diffusion models, which we call lime-reversed diffusion sampler (DIS). It is based on minimizing the reverse KL divergence between a controlled SDE and the reverse-time SDE. This is intimately connected to methods developed in the field of stochastic optimal control based on minimizing the reverse KL divergence to a reference process starting at a Dirac delta distribution (Dai Pra, 1991; Richter, 2021; Zhang& & Chen, 2022a; Tzen & Raginsky, 2019; Vargas el al., 2023b). Finally, we mention that concurrently to our work, Denoising Diffusion Samplers (Vargas et al., 2023a) have been proposed, which can be viewed as special case of our approach, see Scction A.10.
Schrödinger bridges: Schrödinger bridge (SB) problems (Schrödinger, 1931) aim to find the minimizer of the KL divergence to a reference process, typically a Brownian motion, subject to the constraint of satisfying given marginal distributions at the initial and terminal time. Such problems include diffusion-based methods as a special case where the reference process is given by the uncontrolled inference process (De Bortoli et al., 2021; Chen et al., 2021; Koshizuka & Sato, 2022). Since in this case only the reverse-time process is controlled, its initial distribution must correspond to the terminal distribution of the inference process. In practice, this constraint is usually only fulfilled in an approximate sense, resulting in prior loss, see Section 2.3. Note that the previously mentioned sampling methods in the field of optimal control rely on a Dirac delta as the
I Note that diffusion models give rise continuous- time normalizing flow, the so-called probability flow onlinary differential equation (ODE), having the same marginals as the reverse-time diffusion process, see Song al. (2021) and Section A.3.
3

---

Published in Transactions on Machine Learning Research (02/2024)
initial distribution and solve the so-called Schrädinger balf-bridge problem, constraining the KL minimization problem only to the terminal distribution (Dai Pra, 1991).
1.2 Notation
We denote the density of random variable Z by pz. For an Rd-valued stochastic process Z = (Ze) define the function pz by pz(:,t) Pz, for every t [0, We further denote by Pz the law of Z on the space of continuous functions CI,0,T.RR) For a time-dependent function f, we denote by f the time-reversal given by ((t) == J(T - t). Finally, we define the divergence of matrix-valued functions row-wise. More details on our notation can be found in Section A.1.
2 SDE-based generative modeling as an optimal control problem
Diffussin models can naturally be interpreted through the lens of continuous- time stochastic processes (Song el al., 2021). To this end, let us formalize our selting in the general context of SDE-based generative models. We define our model as the stochastic process X (X.))ee(or1) characterized by the SDE
dx,- () 8) ds + ((s)dB,
(1)
with suitable? drift and diffusion coefficients μ:: Rd [0, I&d and [0. + Rddd Learning the model in (1) now corresponds to solving the following problem.
Problem 2.1 (SDE- based generative modeling). Learn an initial condition Xo as well as coefficient functions I and a such that the distribution of XT approximates given data distribution D.
While, in general, the initial condition Xo as well as both the coefficient functions and a can be learned, typical applications often focus on learning only the drift (Ho et al., 2020; Kong et al., 2021; Corso et al., 2022). The following remark justifies that this is sufficient to represent an arbitrary distribution D. Remark 2.2 (Reverse-time SDE). Naively, one can achieve X7 ~ D by setting
Xo Yr and ;== oo Vlogpy f,
(2)
where f:Rd x [0,T] + Rd and Y is a solution to the SDE
dY, = f(Y,, ss + dB,, Yo ~ D.
(3)
This well-known result dates back to Nelson (1967); Anderson (1982); Haussmann & Pardoux (1986); Föllmer (1988). More generally, it states that X can be interpreted as the time-reversal of Y, which we will denote by Y, in the sense that px = Py almost everywhere, see Figure and Section A.3. Even though the reverse-time SDE provides a valid solution to Problem 2.1, apparent practical challenges are to sample from the terminal distribution of the generative SDE, X, Yr, and to compute the so-called score V log.yy In most scenarios, we either only have access to samples from the true data distribution D or access to its density. If samples from D are available, diffusion models provide working solution to solving Problem 2.1 (more details in Section 2.4). When, instead, having access to the density of D, general time-reversed diffusions have, to the best of our knowledge, not yet been considered. In Section 3, we thus propose a novel strategy for this scenario.
As already apparent from the optimal drift j in the previous remark, and noting again that py = ix in this case, the reverse-time log-density logãx of the process X will play a prominent role in deriving a suitable objective for solving Problem 2.1. Hence, in the next section, we derive the HJB equation governing the evolution of log Px, which then provides the bridge to the fields of optimal control and reinforcement learning.
Motivated by Remark 2.2, we start with time-reversed drift and diffusion coefficients ja and 8. Further, we assume certain regularity on the coefficient functions all appearing SDFs, see Section A.1
4

---

Published in Transactions on Machine Learning Research (02/2024)
2.1 PDE perspective: HJB equation for log- log-density
We start with the well-known Pokker-Planeke equalion, which describes the evolution of the density of the solution X to the SDE in (1) via the PDE
0px = div (div (Dpx) -mpx):
(4)
where set D = zoo for notational convenience. This implies that the Lime-reversed density Fx satisfies (generalized) Kolmogorov backward equation given by
0,Px = div div (Dpx) - μpx) - I μ Vpx + div(p)px
(5)
The second equality follows from the identities for divergences in Section A.2 and the fact that a does not depend on the spatial variable x. Now we use the Hopf-Col Cole transformation to convert the linear PDE in (5) to an HJB equation which is prominent control theory, see Pavon (1989); Fleming Rishel (2012), and Section A.5.
Lemma 2.3 (HJB equation for log-density). Let us define -logpx Then V solution to the HJB equalion
0.V = - TT (DVVVV) + p VV - div(u)+ l11 VVI2
(6)
with terminal condition V(,T)) lxg
For approaches to directly solve Kolmogorov backward or HJB equations via deep learning rather high dimensions, we refer to, e.g., Richter (2021); Berner et al. (2020); Zhou al. (2021); Nüsken & Richter (2023; 2021); Richter & Berner (2022), see also Sections A.4 and A.11. In the following, we leverage the HJB equation and tools from stochastic control theory to derive suitable objective for Problem 2.1.
2.2 Optimal control perspective: ELBO derivation
We will derive the ELBO for our generative model (1) using the following fundamental result from control theory, which shows that the solution to an HJB equation, such the stated in Lemma 2.3, related to optimal control problem, Dai Pra. (1991); Pavon (1989); Nüsken & Richter (2021); Fleming & Soner (2006); Pham (2009) and Section A.7. For general introduction to stochastic optimal control theory we refer to Section A.6.
Theorem 2.4 (Verification theorem). Let V be solution to the HJB equation in Lemma 2.3. Further, let u C C (TR. x [0, be a suitable set of admissible controls and for every control u E U let Y be the solution to the controlled SDE
dY = = 104-41 u) (Y. 8) ((8)
Then il holds almost surely that
V(Y"", 0) min EE [Y log (V7) well
(8)
where - -x,
Px, (XF) constitutes the terminal costs, and the running costs are defined as
R² S'' (div(w) + + [ull²) (Y." s) ds.
(9)
Moreover, the unique minimum, attained by = TVV.
"In Section A.3, provide proof for the reverso-time SDE in Remark 2.2 for general depending and However, for simplicity, restrict ourselves to only depending the time variable in the following. usually done, assume that the initial condition 1/" of solution y controlled SDE does depend the control
5

---

Published in Transactions on Machine Learning Research (02/2024)
Plugging in the definition of V from Lemma 2.3, this readily yields the following ELBO of our generative model in (1). The variational gap can be found in Proposition 2.6 and Remark A.6. Corollary 2.5 (Evidence lower bound). For every U it holds almost surely that
Px, (8)) 2 E [ (8
-
(10)
where equality is obtained for u o V log px
Comparing (8) and (10), we see that the ELBO equals the negative control costs. With the initial condition YO D, it represents lower bound on the negative log-likelihood of our generative model. In practice, one can now parametrize u with, for instance, neural network and rely on gradient-based optimization to maximize the ELBO using samples from D. The optimality condition in Corollary 2.5 guarantees that px- = Py almost everywhere if we ensure that Xo ~ YI see Section A.3. In particular, this implies that XT ~ D, ie., our generative model solves Problem 2.1. However, we still face the problem of sampling Xo ~ YT since the distribution of Y depends on the initial distribution D. In Sections 2.4 and we will demonstrate ways circunvent this problem.
2.3 Path space perspective: KL divergence in continuous time
In this section, we show that the variational gap corresponding to Corollary 2.5 can be interpreted as a KL divergence between measures on the space of continuous trajectories, also known as path space (Üstünel & Zakai, 2013). Later, we can use this result to develop our sampling method. To this end, let us define the path space measure Py as the distribution of the trajectories associated with the controlled process Y" defined in (7), also Section A.1. Consequently, we denote by Py the path space measure associated with the uncontrolled process y with the choice = 0.
We can now state path space measure perspective on Problem 2.1 by identifying formula for the target measure IP Py which corresponds to the process in (7) with the optimal control u = u² = o logāx- We further show that, with the correct initial condition Y ~ Xr, this target measure corresponds to the mcasure of the Lime-reversed process, iee, Py = =Px, see Remark 2.2 and Section A.3. For the proof and further details, see Proposition A.9 in the appendix.
Proposition 2.6 (Optimal path space measure). The optimal path space measure can be defined via the work functional W: C(10, Rd) - R and the Radon-Nikodle derivative
dP y dl yo
Pxo (Y%)
exp (-W) with W(Y = R,,( log
(11)
pxr (YO))
where (Y9) is as in (9) with u = 0. Moreover, for any u E Ud the expected variational gap
G(u)= = E [logxx, (Y,') [logpx, (Y7) (Y
(12)
of the ELBO in Covollary 2.5 satisfies that
G(u)= =
DKL (Pr DKL (Py (Px))-
(13)
Note that the optimal change of measure (11) can be seen as version of Doob's h-transform, see Dai Pra & Pavon (1990). Furthermore, it can be interpreted as Bayes' rule for conditional probabilities, with the target IP, y denoting the posterior, P' the prior measure, and exp (-W) being the likelihood. Formula (13) emphasizes again that we can solve Problem 2.1 by approximating the optimal control u* and having matching initial conditions, see also Remark 2.2. The next section provides an objective for this goal that can be used in practice.
IIn case one has acoess to samples from the data distribution D. one could use these as initial data Y in order to simulate No yy In doing so, however, one expect to recover the entire distribution D. but only the empirical distribution of the samples.

---

Published in Transactions on Machine Learning Research (02/2024)
2.4 Connection to denoising score matching objective
This section outlines that, under reparametrization of the generative model in (1), the ELBO in Corollary 2.5 corresponds to the objective typically used for the training of continuous-time diffusion models. We note that the ELBO in Corollary 2.5 in fact equals the one derived in Huang et al. (2021, Theorem 3). Following
the arguments therein and motivated by Remark 2.2, we can now use the reparametrization = ow to arrive at an uncontrolled inference SDE and a controlled generative SDE
dY, = f(Y,, ss (ss) You D, and ax = (ōū - ds
(14)
In practice, the coefficients f and a are usually constructed such that Y an process with Yr being approximately distributed according to standard normal distribution, also (139)
(OU)
in Section A.13. This is why the process Y is said to diffuse the data. Setting XØ- (0, thus satisfies that Pxg = py, and allows to easily sample Xp-
The corresponding ELBO in Corollary 2.5 now takes the form
logpxy (Y)) 2 E [logPxy (Yr) Ro (Y)o)
(15)
where, in analogy to (13), the expected variational gap given by the forward KT. divergence G(u) (II(PPyPPx")) DDL( (xPx;))) see Proposition A.10. We note that the process in the ELBO does not depend on the control anymore. Under suitable assumptions, this allows to rewrite the expected negative ELBO (up to constant not depending on u) a denoising score matching objective (Vincent, 2011), i.c.,
CDSM (u) = T 2
[(Yr,))) - (T) (X (Y+1X))]
(16)
where F ~~ ((00 T1) and Py, denotes the conditional density of Y, given y,, see Section A.9. We emphasize that the conditional density can be explicitly computed for the OU process, see also (137) in Section A.13. Due to its simplicity, variants of the objective (16) are typically used in implementations. Note, however, that the setting in this section requires that one has access to samples of D in order to simulate the process Y. In the next section, we consider a different scenario, where instead we only have access to the (unnormalized) density of D.
3 Sampling from unnormalized densities
In many practical settings, for instance, in Bayesian statistics or computational physics, the data distribution 'D admits the density p/2, where p is known, but computing the normalizing constant 2 == Sge p(ar) da is intractable and samples from D are not easily available. In this section, we propose novel method based on diffusion models, called time-reversed difffusion sampler (DIS), which allows to sample from D, see Figure 3. To this end, we interchange the roles of X and yu in our derivation in Section 2, i.e., consider
dY s) ds (s)dB, Yo D. and dx (लū - (xx s) ds d()
(17)
where we also renamed μ to f to stay consistent with (14). Aualogously to Theorem 2.4, the following corollary specifies the control objective, see Corollary A.12 in the appendix for details and the proof. Corollary 3.1 (Reverse KL divergence). Let X and Y be defined by (17). Then it holds
(RI Pp)) DRL (P Px 1xx x ) (P Pxx PYr)
(18a)
PYz (XT)
E
R (x*) + log
+ log Z + DKL, (P xy (P (Pyy)
(18b)
p(X))
LDIS(U) log Z,
(18c)
For coefficients typically used in practice (leading, for instance, contimnous-time analogs of SMLD and DDPM) we refer Song et (2021) and Section A.13.
7

---

Published in Transactions on Machine Learning Research (02/2024)
0.2 initial density
0.6 0.8
1.0
0.0
0.5
1.0
-2
2
terminal density
X1
Figure Illustration of our DIS algorithm for the double well example in Section with d = 20, == = 3. The process X starts from Gaussian (spproximately distributed Yr) and the control u is trained such that the distribution terminal time X approximates the target density p/Z. The plot displays some trajectories as well as histograms at initial and terminal times. Tn the right panel, we show a KDE density estimation of a 2d marginal of the corresponding double well.
where
PXP (XD))
(DIS(W) = EE (x() + log
(19)
p(Xp)
particular, this implies that
DKL (Px PYr) (Py loog Z min Cpis (u)
(20)
and, assuming that X ~ Yr, the minimizing control u² = o V logpy guarantees that Xi ~ D.
As in the previous chapter, XØ ~ Y7 can be approximately achieved by choosing X ~ (N(OTT I) and f and a such that py, = (O, which incurs an irreducible prior loss given by DKL (Py,). In practice one can now minimize the control objective (19) using gradient-baser optimizatiin techniques. Alternatively, one can solve the corresponding HJB equation obtaín an approximation to u², see Section A.11. Note that, in contrast to (15) or (16), we do not need access samples Yo D Y does not need to be simulated for the objective in (19). However, since the controlled process X now appears in the objective, the whole trajectory of X and cannot resort to a Monte Carlo approximation, such as in the denoising score matching objective in (16). In Section A.14, we show that our path space perspective allows us to introduce divergences different from the reverse KL-divergence in (18a), which allows for off- off-policy training and can result in improved numerical performance.
need to simulate
3.1 Comparison to Schrödinger half-bridges
The idea to rely on controlled diffusions in order to sample from prescribed target densities is closely related to Schrödinger bridges (Dai Pra, 1991). This problem has been particularly well studied in the case where the initial density follows Dirac distribution, i.e., the stochastic process starts pre-specified point-ofter referred to Schrödinger half-bridge. Corresponding numerical algorithms based deep learning, referred to as Path Integral Sampler (PIS) in Zhang & Chen (2022a), have been independently presented in Richter (2021); Zhang & Chen (2022a); Vargas et al. (2023b). Combining ideas from Dai Pra (1991) and Fleming & Soner (2006), it can be shown that corresponding control problem can be formulated as
Pxę (X7) p(X
log min LrIs(u), with (PIS(u) = RØ((")+ + log well
(21)
where the controlled diffusion X defined as in (17) with a fixed initial value X²' 2g E Rd. In the above, Pxo denotes the density of the uncontrolled process X at time T. Equivalently,
show that P
x..

---

Published in Transactions on Machine Learning Research (02/2024)
satisfies the optimal change of measure, given by
dP x (X) = 2xx
(Xr):
(22)
dP x2
for all suitable stochastic processes X on the path space. Using the Girsanov theorem (Theorem A.7), we see that
(R = Cris(u)+ log Z,
(23)
see also Zhang& & Chen (2022a); Nüsken & Richter (2021). Similar to Corollary 3.1, one can thus show that the optimally controlled process satisfies X7 "D, see also Tzen & Raginsky (2019) and Pavon (2022).
Comparing the two attempts, i.e., the objectives (19) and (21), we can identify multiple differences:
Different initial distributions, running costs, and terminal costs are considered. In the Schrödinger half-bridge, the terminal costs consist of p and Pxp: whereas for the diffusion model they only consist of p.
For the Schrödinger half-bridge, f needs to be chosen such that Pxp. is known analytically (up to a constant). For the time-reversod diffusion sampler, f needs to be chosen such that pY, pxg in order to have a small prior loss DKI. (Pxx xg (ryy).
In the Schrödinger half-bridge, XØ starts from arbitrary, but fixed, point, whereas for the diffusion-based generative modeling attempt Xø must be (spproximalely) distributed as Yr.
In Appendix A.10, we show that the optimal control u' of our objective in (19) can be numerically more stable than the one in (21). In the next section, we demonstrates that this can also lead to better sample quality and more accurate estimates of normalizing constants.
4 Numerical examples
The numerical experiments displayed Figures 4 and 5 show that our time-reversed diffusion sampler (DIS) succeeds in sampling from high-dimensional multimodal distributions. We compare our method against the Poth Integral Sampler (PIS) introduced in Zhang & Chen (2022a), which uses the objective from Section 3.1. As shown in Zhang& & Chen (2022a), the latter method can already outperform various state-of-the-art sampling methods. This includes gradient-guided MCMC methods without the annealing trick, such as Hamiltonian Monte Carlo (HMC) (MacKay, 2003) and No-U-Turn Sampler (NUTS) (Hoffman & Gelman, 2014), SMC with annealing trick such as Annealed Flow Transport Monte Carlo (AFT) (Arbel et al., 2021), and variational normalizing flows (VINF) (Rezende Mohamod, 2015).
Let us summarize our setting in the following and note that further details, well as the exact hyperparameters, can be found in Section A.13 and Table 2. Our PyTorch implementation" is based on Zhang & Chen (2022a) with the following main differences: We train with larger batch sizes and more gradient steps to guarantee belter convergence. We further use a clipping schedule for the neural network outputs, and an exponential moving average of the parameters for evaluation (as is often done for diffusion models (Nichol & Dhariwal, 2021)). Moreover, we train with a varying number step sizes for simulating the SDEs using the Euler- Maruyama (EM) scheme in order to amortize the cost of training separate models for different evaluation settings. To have fair comparison, we trained PTS using the same setup. We also present comparison to the original code in Figure 11, showing that our setup leads to increased performance and stability. The optimization routine is summarized in Algorithm 1.
Similar to PTS, we also use the score of the density V logg P (typically given in closed-form or evaluated via automatic differentiation) for our parametrization of the control uo, given by
uo xx, t) == 4 (z. + D (t),())
(24)
Here, 8 is a lincar interpolation between the initial and terminal scores V Vlog p and V loggxx, thus approxi- mating the optimal score V log at the boundary values, see Corollary 3.1. This parametrization yielded
TTee associated repasitory can be found at

---

Published in Transactions on Machine Learning Research (02/2024)
Algorithm 1 Time-reversed diffusion sampler (DIS): Training the control in Section via deep learning.
input neural network ug with initial parameters ((0), optimizer method slep for updating the parameters, number of steps K, batch size output parameters (0( k=1 for h + ...... K - do (1(()) - sample from (1, (DIS(ug(s) + estimate the cost in (19) using the EM scheme with Xg - =x(0), i= 1,...,m ((++1) + step (otk), ( VLDIs(ug())
end for
the best results in practice, see also the comparison other paramctrizations in Section A.14. For the drift and diffusion coefficients, we make of successful noise schedules for diffusion models and choose the suriance-preserving SDE from Song et al. (2021).
We evaluate DIS and PIS on number of examples that we outline in the sequel. Specifically, we compare the approximation of the log-normalizing constant log Z using (approximately) unbiased estimators given by importance sampling in path space, see Section A.12 for DIS and Zhang & Chen (2022a) for PIS. To evaluate the sample quality, we compare Monte Carlo approximations expectations as well as estimates of standard deviations. We further refer Figure 7 in the appendix for preliminary experiments using Physics informed neural networks (PINNs) for solving the corresponding HJB equation, well as to Figure where employ a divergence different from the KL divergence, which is possible due to our path space perspective.
4.1 Examples
Let us present the mumerical examples which evaluate method.
Gaussian mixture model (GMM): We consider
M
( = Σ
ComN( Nom Emn);
(25)
m=1
with Σm= Om 1. Specifically, we choose m = 9, E. = 0.3I, and (um))m=1 (- [-5,0,5) The density and the optimal drift are depicted in Figure
(-5,000 c R
Funnel:
The 10-dimensional Funnel distribution (Neal, 2003) is challenging example often used to test MCMC methods. Tt. is given by
d II 2=2
=(x)= = N(t; (2) ((x;0,e0))
(26)
for x = (x))11 E 10 with 3.
Double well (DW): A typical problem in molccular dynamics considers sampling from the stationary distribution of a Langevin dynamics, where the drift of the SDE is given by the negative gradient of a potential v namely
dx, dx, + (s)
(27)
see, e.g., (Leimkuhler & Matthews, 2015). Given certain assumptions on the potential, the stationary density of the process can be shown to be px.. = e- Potentials often contain multiple energy barriers, and resulting local minima correspond to meaningful configurations of a molecule. However, the convergence speed of px, can be very slow in particular for large energy barriers. Our time-reversed diffusion sampler,
Due typo the code, the results presented by Zhang Chen (2022a) seem consider 3 for the baselines, but the favorable choice of for evaluating their method.
10

---

Published in Transactions on Machine Learning Research (02/2024)
GMM
Funnel
DW (d=20,
DIS
0.2
PIS
0.0
ground-truth
0.1
-0.2
12
logi22)
-0.4
0.0
11
-0.6
-0.1
-0.8
200
integrator steps
inlegraior steps
integrator steps
Figure 4: We compare our DIS method against PIS on the ability to compute the log-normalizing constant log Z (median and interquartile range over 10 training seeds) when using N E {100, 200, the Euler-Maruyama scheme. Our method outperforms PIS clearly for the GMM and Funnel examples and offers a slight improvement for the DW example. Each model has been trained with each N for 1/4 of the total gradient steps (starting with 100 and ending with 800 steps). See also Figure 10 for a comparison of models trained on a single step size.
steps of
the other hand, can (at least in principle) sample from densities already fixed time T. In our example, shall consider a d-dimensional double well potential, corresponding to (unnormalized) density e given by
d Σ
p(x) = exp
2 - 5)2
x
(28)
2 w-
with w E N combined double wells and a separation parameter å E (0, 20), see also Wu et al. (2020). Note that, duc to the double well structure of the potential, the density contains 2n modes. For these multinodal examples, we can compule a reference solution of the log-normalizing constant log Z and other statistics of interest by numerical integration since p factorizes in the dimensions.
4.2 Results
Our experiments show that DIS can offer improvements over PIS for all the tasks we considered, i.e., estimation of normalizing constants, expectations, and standard deviations. Figures 4 and 5 display direct comparisons for the examples in Section 4.1, noting that we use the same training techniques but different objectives for the respective methods. Moreover, we recall that our setup of the PIS method already outperforms the default implementation by Zhang & Chen (2022a), see Figure 11 in the appendix. We also emphasize that the reference process of PIS satisfies that X ~ (),, which is already the correct distribution for d - w dimensions of the double well example. Nevertheless, our proposed DIS method converges to a better approximation and provides strong results even for multimodal, high-dimensional distributions, such as d = 50 with 32 well-separated modes, see Figure 5. Finally, we provide further results for two promising future directions: We show that replacing the KL divergence in Corollary 3.1 with the log-variance divergence (Nüsken & Richter, 2021: Richter et al., 2020) can further improve performance, see Figure 6 and Section A.14 for details. In Figure 7, we show that solving the HJB equation in Lemma 2.3 using cryssics-infomml neural networks (PINNs) can provide competitive results, see Scction A.11 for further explanations. We present more numerical results and comparisons in Sections A.13 and A.14.
4.3 Limitations
In this section we shall discuss limitations of the optimal control perspective and our resulting sampling algorithm. First, note that the control perspective so far only yields new algorithms if the larget densily is known (up to the normalization constant). In particular, we note that the direct minimization of divergences, such as the KL. or log-variance divergence, nceds access to the (unnormalized) target density. This is typically not the case for problems in generative modeling where one has only access to data samples and where one resorts to optimizing the EL,.O,, as outlined in Section 2.4. However, we emphasize that classical sampling
11

---

Published in Transactions on Machine Learning Research (02/2024)
DW (d=10,m=3,6-2)
DW (d=20, w 5,6
DW (d= 5,,66
DIS
12.60
PIS
54.25
ground-truth
Ix?
29
12.50
54.00
28
53.75
12:40
53.50
12.35
27
53.25
20.30
42.6
9.56
expeclation Mit
20.25
42.5
42.4
20.00
9.52
42.3
9.50
19.75
19.50
42.1
9.46
19.25
1.18
standard deviation
1.1100
1.0375
1.16
1.1075
1.1050
1.14
1.0325
1.1025
1.12
1.0300
1.1000
1.0275
1.10
100
200
400
100
200
800
inlegrator steps
integrator steps
integralor steps
Figure 5: We compare our DIS method against PTS on the ability to estimate the expectations IF. [[1%]]]] and E [olla] and the average standard deviation 1/ Σi= [[ see (148). Each model has been trained with each N E (100, 200, 400, 800} for 1/4 of the total gradient steps (starting with 100 and ending with 800 steps). We compare the methods using N steps of the Euler-Maruyama scheme when computing our estimates (median and interquartile range over 10 training seeds). Our method outperforms PIS in all considered settings in terms of accuracy.
problems (without any samples from the data distribution) can arguably be more challenging, e.g., due to exploration exxploottiioooooiiiiionnnn trade-offs and potential mode collapse.
Compared to classical sampling methods such as, c.g., MCMC and SMC methods, the sampling time of our algorithu DIS typically much laster. However, this comes at the cost of training a model first, which might only amortize when a larger number of samples is needed. While well-trained model enjoys good sampling guarantees (cf. De Bortoli (2022); Chen et al. (2022); Lee al. (2023)), approximating the optimal control (along the controlled dynamics) can be challenging and sensitive to the chosen hyperparameters. We have observed that clever choices of alternative divergences can improve convergence and counteract mode collapse, however, it is generally difficult to provide convergence guarantees.
5 Conclusion and outlook
We propose connection of diffusion models to the fields of optimal control and reinforcement learning, which provides valuable new insights and allows to transfer established tools from one field to the respective other. As first steps, we have shown how to readily derive the ELBO for continuous-time diffusion models solely from control arguments, we provided an interpretation of diffusion models via measures on path space (eventually leading to improved losses), and we extended the diffusion modeling framework to sampling from unnormalized densitics. We could demonstrate that our framework offers significant numerical advantages over existing diffusion-based sampling approaches across series of challenging, high-dimensional problems.
12

---

Published in Transactions on Machine Learning Research (02/2024)
GMM
Punnel
DW (d=3,w=3,8w4)
DIS (log (logvariiance)
DIS (K1.)
0.5
0.1
-1.0
ground-truth
0.4
0.0
1.5
-0.6
-0.1
-0.8
2.5
200
200
400
Integrator steps
Integrator steps
Integrator steps
Figure 6: We use the settings in Figure 4 and only change the reverse KI. divergence in Corollary 3.1 to the log-var oogvaariacc divergence (Nüsken & Richter, 2021). Training with the latter objective can significantly improve the estimate of the log-normalizing constant log 2 for the GMM, the Funnel, and a DW example.
With the new connections at hand, we anticipate further fruitful theoretical as well as practical insights. For instance, the PDE perspective offers to rely on PDE techniques and solvers that have been developed in the past, see Section A.11 for details. For high-dimensional settings, tensor-based methods (Richter et al., 2023), as well as SDE-based solvers (Nüsken & Richter, 2021), seem to be well suited for numerical approximation. Alternatively, one can readily apply numerical methods from control theory to approximate the score function, such as, for instance, methods based on tensor trains (Oster al., 2022) policy iteration, which has recently been combined with deep learning (Zhou et al., 2021).
Note that one usually controls linear, Ornstein-Uhlenbee type processes. This motivates to approximate the target density by Gaussian density (which corresponds to quadratic costs in the control problem). Then, one can pretrain the score model by solving a linear-quadrraiicconttto problem, for which very efficient. numerical methods have been developed extensively. On the other hand, we may add additional running costs the control objective in order to promote or prevent certain regions in the domain, which can incorporate domain knowledge and improve sampling for respective applications. Interesting future perspectives also include the usage of other divergences on path space (motivated by our promising results with the log-variance divergence), well as the extension to Schrödinger bridges, for instance based on the work by Chen et al. (2021). These methods are particularly well-suited for the task of sampling from unnormalized densities, where we cannot leverage the efficient denoising score matching objective.
Acknowledgments
We would like to thank Nikolas Nüsken and Ricky T. Q. Chen for many useful discussions. The research of Lorenz Richter has been partially funded by Deutsche Forschungseemiischaft (DFG) through the grant CRC 1114 "Scaling Cascades in Complex Systems" (project A05, project number 235221301). Julius Berner is grateful to G-Rescarch for the travel grant.
13

---

Published in Transactions on Machine Learning Research (02/2024)
References
Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12
(3):313-326, 1982.
Michael Arbel, Alex Matthews, and Arnaud Doucet. Annealed flow transport Monte Carlo. In International Conference on Machine Learning, pp. 318- 330... PMLR, 2021.
Ludwig Arnold. Stochastic Differential Equalions: Theory and Applicalions. A Wiley-Interscience publication. Wiley, 1974.
P. Baldi. Stochastic Calculus: An Introduction Through Theory and Exercises. Universitext. Springer International Publishing, 2017.
Christian Beck, Sebastian Becker, Philipp Grohs, Nor Jaafari, and Armulf Jentzen. Solving the Kolmogorov PDE by means of deep learning. Journal of Scientific Computing, 88:1-28, 2021.
Richard Bellman. Dynamic programming. Princcton University Press, 1957.
Nils Berglund. Kramers' law: Validity, derivations and generalisations. arXiv preprint arXiv:1106.5799 2011.
Julius Berner, Markus Dablander, and Philipp Grohs. Numerically solving parametric families of high- dimensional Kolmogorov partial differential equations via deep learning. Advances in Neural Information Processing Systems, volume 33, pp. 16615-16627, 2020.
Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim, and Anru R Zhang. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. ar preprint arXiv:2209.11 11215, 2022.
Tianrong Chen, Guan-Horng Liu, and Evangelos Theodorou. Likelihood training of Schrödinger Bridge using forward-backward SDEs theory. arXiv preprinl urXiv:2110.1129) 2021.
Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi Jaakkola. Diffdock: Diffusion steps, twists, and turns for molecular docking. arXiv preprint arXim:2210.0111 01776, 2022.
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 diffusson processes. Applied mathematics and Optimization, 23(1):313-32) 329, 1991.
Paolo Dai Pra and Michele Pavon. On the Markov processes of Schrödinger, the Feynman-Kac formula and stochastic control. In Realization and Modelling in System Theory, pp. 197-504. Springer, 1990.
Valentin De Bortoli. Convergence of denoising diffusion models under the manifold hypothesis. ar preprint WwXiv:2208.053/1 2022.
Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion Schrödinger bridge with applications to basee generative modeling. Advances in Neural Information Processing Systems, 34: 17695-17709, 2021.
Pierre Del Moral, Arnaud Doucet, and Ajay Jasra. Sequential Monte Carlo samplers. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 68(3):4111433 2006.
Arnaud Doucet, Adam M Johansen, et al. A tutorial on particle filtering and smoothing: Fifteen years later. Handbook of nonlinear filtering, 12(656-704):3, 2009.
Arnaud Doucet, Will Sussman Grathwohl, Alexander GDG Matthews, and Heiko Strathmann. Score-basod diffusion meets annealed importance sampling. In Advances in Neural Information Processing Systems, 2022.
14

---

Published in Transactions on Machine Learning Research (02/2024)
Lawrence C. Evans. Partial differential equations. American Mathematical Society, Providence, R.I., 2010.
Wendell H. Fleming and Raymond W Rishel. Deterministic and stochastic optimal control, volume 1. Springer Science & Business Media, 2012.
Wendell H 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' d'Été de Probabilités de Saint-Floud Flour XV-XVII, 1985-87, pp. 101-203. Springer, 1988.
Carsten Hartmann and Lorenz Richter. Nonasymptotic bounds for suboptimal importance sampling. arXiv preprint trXiv:2102.0900 09600,, 2021.
Carsten Hartmann, Lorenz Richter, Christof Schütte, and Wei Zhang. Variational characterization free energy: Theory and algorithms. Entropy, 19(11):626, 2017.
Ulrich G Haussmann and Etienne Pardoux. Time reversal of diffusions. The Annals of Probability, pp- 1188-1205, 1986.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (GELUs). arXio preprint arXi::1666.08 08415, 2016.
Jonathan Ho, Ajay Jain, and Picter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840 6851, 2020.
Matthew D Hoffman and Andrew Gelman. The No-U-Turn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo. Journal of Machine Learning Research, 15(1):1593-100 1623, 2014.
Lars Holdijk, Yuanqi Du, Ferry Hooft, Priyank Jaini, Bernd Ensing, and Max Welling, Path integral stochastic optimal control for sampling transition paths. arXiv preprint arXi::2277.001 2022.
Chin- Wei Huang, Jae Hyun Lim, and Aaron C Courville. A variational perspective on diffusion- based generative models and score matching. Advances Neural Information Processing Systems, 34, 2021.
Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
Bowen Jing, Gabricle Corso, Jeffrey Chang, Regina Barzilay, and Tommi Jaakkola. Torsional diffusion for molecular conformer generation. arXiv preprint arXiv:2206. 01729, 2022.
Patrick Kidger, James Foster, Xuechen Chen Li, and Terry Lyons. Efficient and accurate gradients for neural SDEs. Advances in Neural Information Processing Systems, 34:18747-18761, 2021.
Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances in Neural Information Processing Systems, 02216666-27777 2021.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint (alo:441.69980 2014.
Peter Kloeden and Eckhard Platen. Stochastic differential equations. Tn Numerical Solution of Stochastic Differential Equations, pp. 103 160. Springer, 1992.
Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: versatile diffusion model for audio synthesis. In International Conference on Learning Representations, 2021.
Takeshi Koshizuka and Issei Sato. Neural Lagrangian Schrödinger bridge. arXiv preprint arXiv:2204.04853, 2022.
Isaac E Lagaris, Aristidis Likas, and Dimitrios Fotiadis. Artificial neural networks for solving ordinary and partial differential equations. IEFF transactions on neural networks, 9(5):987- 1000, 1998.
15

---

Published in Transactions on Machine Learning Research (02/2024)
Holden Lee, Jianfeng Lu, and Yixin Tan. Convergence of score-based generative modeling for general data distributions. In International Conference on Algorithmic Learning Theory, pp. 946-983 PMTR, 2023.
Flavien Léger and Wuchen Li. Hopf Cole transformation via generalized Schrödinger bridge problem. Journal of Differential Equations, 274:788 -2778 2021.
Ben Leimkuhler and Charles Matthews. Molecular dynamics. Interdisciplinary applied mathemalics, 39:443, 2015.
Christian Léonard. Some properties of path measures. In Séminaire de Probabilités XLVI, pp. 207-230. Springer, 2014.
Xuechen Li, Ting-Kam Leonard Wong, Ricky TQ Chen, and David Duvenaud. Scalable gradients for stochastic differential equations. In International Conference Artificial Intelligence and Statistics, pp. PMLR, 2020.
Pierre-Louis Lions.
Optimal control diffusion processes and Hamillon-Jacobi Bellman equations part 2: viscosity solutions and uniqueness. Communications in partial differential equations, 8(11):1229 1276, 1983.
Jun S Liu and Jun $ Liu. Monte Carlo strategies in scientific compuling, volume 10. Springer, 2001.
David MacKay. Information theory, pattern recognition and neural networks. In Proceedings of the 1sl
International Conference Evolutionory Computation, Cambridge University Press, 2003.
Dimitra Maoutsa. Geometrie constraints improve inference of sparsely observed stochastic dynamics. arXiv preprint ar arXi::3044004225 2023.
Dimitra Maoutsa and Manfred Opper. Deterministic particle flows for constraining stochastic nonlinear systems. Physical Review Research, 4(4):043335, 2022.
Laurence Illing Midgley, Vincent Stimper, Gregor NC Simm, Bernhard Schölkopf, and José Miguel Hernández- Lobato. Flow annealed importance sampling bootstrap. In NewrIPS 2022 Al for Science: Progress and Promises, 2022.
Tom Minka al. Divergence measures and message passing. Technical report, Citeseer, 2005.
Radford M Neal. Annealed importance sampling. Statisties and computing, 11(2):125 139, 2001.
Radford M Neal. Slice sampling. The Annals of Statistics, 31(3):705-767, 2003.
E Nelson. Dynamical theories of Brownian motion. Press, Princeton, NJ, 1967.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, pp. 8162-8171. PMLR, 2021.
Nikolas Nüsken and Lorenz Richter. Solving high-dimensional Hamiton-Jacohi Bellman PDEs using neural networks: perspectives from the theory of controlled diffusions and measures on path space. Partial Differential Equations and Applications, 2(4):1-48, 2021.
Nikolas Nüsken and Lorenz Richter. Interpolating between BSDEs and PINNs: deep learning for elliptic and parabolic boundary value problems. Journal of Machine Learning, 2023.
Bernt Øksendal and Bernt Oksendal. Stochastic differential equations. Springer, 2003.
Mathias Oster, Leon Sallandt, and Reinhold Schneider. Approximating optimal feedback controllers of finite horizon control problems using hierarchical tensor formats. SIAM Journal Scientific Computing, 44(3): B746 B770, 2022.
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshmi- narayanan. Normalizing flows for probabilistic modeling and inference. Mach. Learn. Res., 22(57):1-64. 2021.
16

---

Published in Transactions on Machine Learning Research (02/2024)
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et. al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
Michele Pavon. Stochastic control and nonequilibrium thermodynamical systems. Applied Mathematics and Optimization, 19(1):187-202, 1989.
Michele Pavon. On local entropy, stochastic control and deep noural networks. arXiv preprint arXiw:2204.13040, 2022.
H. Pham. Continuous-lime Stochastic Control and Optimization with Financial Applications. Stochastic Modelling and Applied Probability. Springer Berlin Heidelberg, 2009.
Maziar Raissi, Paris Perdikaris, and George Em Karniadakis. Physics informed deep learning (part I): Data-driven solutions of nonlinear partial differential equations. arXiv preprint arXto:171111061,, 2017.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International conference on machine learning, pp. 1530-1538. PMLR, 2015.
Lorenz Richter. Solving high-dimensional PDEs, approzimation of path space measures and importance sampling of diffusions. PhD thesis, BTU Cottbus-Senftenberr,, 2021.
Lorenz Richter and Julius Berner. Robust SDE-based variational formulations for solving linear PDEs via deep learning. In International Conference on Machine Learning, pp. 18649 18666. PMLR, 2022.
Lorenz Richter and Julius Berner. Improved sampling via learned diffusions. In International Conference Learning Representations, 2024.
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 13492, 2020.
Lorenz Richter, Leon Sallandt, and Nikolas Nüsken. From continuous-time formulations to discretization schemes: lensor trains and robust regression for bsdes and parabolic PDEs. arXio preprind trXiw:2307.1544 2023.
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.
Tim Salimans, Diederik Kingma, and Max Welling. Markov chain Monte Carlo and variational inference: Bridging the gap. International conference on machine learning, pp. 1218-1226. PMLR, 2015.
E. Schrödinger. Über die Umkehrung der Naturyesetze. Verlag der Akademie der Wissenschaften in Kommission bei Walter De Gruyter u. Company, 1931.
Justin Sirignano and Konstantinos Spiliopoulos. DGM: deep learning algorithm for solving partial differential equations. Journal of compulational physics, 375:1339 1364, 2018.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pp. 2256 2265. PMLR, 2015.
Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. Advances in neural information processing systems, 33:12438 -24488 2020.
Yang Song, Sahaj Garg, Jiaxin Shi, and Stefano Ermon. Sliced score matching: A scalable approach density and score estimation. In Uncertainty in Artificial Intelligence, pp. 574-584. PMLR, 2020.
17

---

Published in Transactions on Machine Learning Research (02/2024)
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishck Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference Learning Representations, 2021.
Gabriel Stoltz, Mathias Rousset, al. Free energy computations: A mathematical perspective. World Scientific, 2010.
Matthew Tancik, Pratul Srinivasan, Ben Mildenhall, Sara Fridlvich-Keil, Nithin Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan Barron, and Ren Ng. Fourier features let networks learn high frequency functions in low dimensional domains. Advances in Newral Information Processing Systems, 33:7537- 7547, 2020.
Sep Thijssen and HJ Kappen. Path integral control and state-dependent feedback. Physical Review E, 91(3): 032104, 2015.
Belinda Tzen and Maxim Raginsky. Theoretical guarantees for sampling and inference in generative models with latent diffusions. In Conference on Learning Theory, pp. 3084 3114. PMLR, 2019.
A Sülcyman Ustünel and Moshe Zakai. Transformation of measure Wiener space. Springer Science & Business Media, 2013.
Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. Advances in Neural Information Processing Systems, 14:11287-11302, 2021.
Ramon Van Handel. Stochastic calculus, filtering, and stochastic control. Leclure Noles, 2007.
Francisco Vargas and Nikolas Nüsken. Transport, variational inference and diffusions: with applications to annealed flows and Schrödinger bridges. orXiv preprint arXiv:2307.01050 2023.
Francisco Vargas, Will Grathwohl, and Arnaud Doucet. Denoising diffusion samplers. In International Conference on Learning Representations, 2023a.
Francisco Vargas, Andrius Ovsianas, David Fernandes, Mark Girolami, Neil D Lawrence, and Nikolas Nüsken. Bayesian learning via neural ceriiddinng-Fööllle flows. Statistics and Computing, 33(1):1-22, 2023b.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation, 23 (7):1661-167,, 2011.
Martin J Wainwright, Michael Jordan, et 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 5944, 2020.
Benjamin J Zhang and Markos A Katsoulakis. A mean-field games laboratory for generative modeling, arXiv preprint XrXi::2304. 18534, 2023.
Qinsheng Zhang and Yongxin Chen. Path integral sampler: stochastic control approach for sampling, In International Conference on Learning Representations, 2022a.
Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator. arXie preprint ar arx:::2204.1 13902, 2022b.
Mo Zhou, Jiequn Han, and Jianfeng Lu. Actor-critic method for high dimensional static Hamilton-Jccobi Bellman partial differential equations based on neural networks. SIAM Journal on Scientific Computing, 43(6):A4043 A4666, 2021.
18

---

Published in Transactions on Machine Learning Research (02/2024)
A
Appendix
Contents
A.1 Setting
A.2 Identities for divergences
20
A.3 Reverso-time SDEs
20
A.4 Further details on the HJB equation
21
A.5 Hopf Cole transformation
21
A.6 Brief introduction to stochastic optimal control
22
A.7 Verification theorem
23
A.8 Measures on path space
................................
24
A.9 ELBO formulations...
27
A.10 Diffusion-based sampling from unnormalized densities
28
A.10.1 DDS as a special case DIS
................................
29
A.10.2 Optimal drifts of DIS and PIS
30
A.11 PDE-based methods for sampling from unnormalized densities
30
A.11.1 Physics informed neural networks
31
A.11.2 BSDE-based methods and feymman-Kat formula
32
A.12 Importance sampling in path space
................................
33
A.13 Details on implementation
34
A.14 Further numerical results
37
A.1 Setting
Let d, E N and T E (0, oo). For random variable X, which is absolutely continuous w.r.l. the Lebesgue measure, we write px for its density. We denote by B a standard d-dimensional Brownian motion. We say that continuous va-uule stochastic process Y (Y)) has density py: Rd [0, TJ + [0, oc) if for all [0, 7 the random variable Y, has density py py(:, (-,t) w.r.t. to the d-dimensional Lebesgue mcasure, i.c., for all / E [0, 7) and all measurable A C Rd it holds that.
P P f (r, J (x:) dx.
(29)
We denote by Py the law of Y the space of continuous functions C((O, Ti, Rd) equipped with the Borel measure. We assume that the coefficient functions and initial conditions of all appearing SDEs are sufficiently regular such that the SDEs admit unique strong solutions and Novikov's condition is satisfied, see, c.g., Øksendal Oksendal (2003, Section 8.6). Furthermore, we assume that the SDE solutions have densities that are smooth and strictly positive for E (0, T) and that can be written as unique solutions corresponding Fokker-Planck equations, see, for instance, Arnold (1974, Section 2.6) and Baldi (2017, Section 10.5) for the details. For a function f Rdd x [0, T + R² we write F for the time-reversed function given by
f(,,t) fx,T - t), (x,t) [0.T].
(30)
For scalar-valuee function g: 2d 10, T R, we denote by Vg. V'g, and div(g) its gradient, Hessian matrix, and divergence w.r.t. to the spatial variable For matrix-valued functions A: Rd - Rdxd we denote
19

---

Published in Transactions on Machine Learning Research (02/2024)
by
(r(A)
Σ
Aa
(31)
i=1
the trace its output. Finally, we define the divergence of matrix-valued functions row-wise, see Section A.2.
A.2 Identities for divergences
Let A: Rd + Raxd Rd + Rd, and Rd - R. We define the divergence of A row-wise, i.c.,
div(A) = (div(A.)) Σ 02, A.
(32)
where Ag. and A, denote the i-th row and j-th column, respectively. Then the following identitics hold truc:
div(div(A)) Eij=1 2z Oz,
2. div(rg) = div(v)g + Vg
div(Ag) div(A)g AVg
4. div(Av) div(AT) oV.
A.3 Reverse-time SDEs
The next theorem shows that the marginals of a time-reversed Itô process can be represented marginals of another Itô process, Huang et al. (2021); Song al. (2021); Nelson (1967); Anderson (1982); Haussmann & Pardoux (1986); Föllmer (1988). We present formulation from Huang al. (2021, Appendix G) which derives whole family of processes (parametrized by function A). The relations stated in equation (2)
follow from the choice À = 0 and the fact that div(D) = 0 if a does not depend on the spatial variable æ.
Theorem A.1 (Reverse-time SDE). Let f E R2d x (0, T] TRd and a: Rd x [0, T) - Rdddd let Y (Ys)se(o,T)) be the solution to the SDE
dY, = J(Y,, ds +
(33)
and assume that Y has density py, which satisfies the Pokker-Planch equation given by
0,py div (div (Dpy) - fpy),
(34)
where D = To For every X (3, 1]) the solution y (Y se)se(o,71 to the reverse-time SDE
d² x1 ( ., ds ds (8,) 8) dB,, 80~ Yr,
(35)
with
( (2 (2- div (D) + (2 A)DV logpy f.
(36)
and
(A) = -
(37)
has densily py given by
p+(+t)= = by
(38)
almost everywhere for every E [0, T). other words, for every E (0, T] holds that Y V..
Proof.
Using the Fokker- Planck equation in (34), we observe that
0,py = div (divv (Dix) + Jix))
(39)
20

---

Published in Transactions on Machine Learning Research (02/2024)
The negative divergence, originating from the chain rule, prohibits us from directly viewing the above equation as Fokker-Planck equation. We can, however, use the identities in Section A.2, to show that
div (Dpy)= - div (D) By + Dvp, = (div (D) + DV log py pY.
=div This implies that we can rewrite (39) as
(40)
O:py div (1L (1- (1= div div (Diy)
(2 A) div (Dōr) fFx) = div (div (5th pr) - x pr):
(4la)
(41b)
where D(A) = (o( As the PDE in (41b) defines a valid Fokker-Planck equation associated to the reverse-time SDE given by (35), this proves the claim.
A.4 Further details on the HJB equation
In order to solve Problem 2.1, one might be tempted to rely on classical methods to approximate the solution of the HJB equation from Lemma 2.3 directly. However, in the setting of Remark 2.2, one should note that the optimal drift,
== oo Vlogyy - -oo T VV f.
(42)
contains the solution itself. Plugging it into the HJB equation Lemina 2.3, we get the equation
a,v = Tr (DV"V) - f vV + div(f) - fllo T VVI2 V(.TT) = logPx.
(43)
Likewise, when applying the Hopf Cole transformation from Section A.5 to py directly, ie., considering V := - loggyy py, where Y is the solution SDE (3), we get the same PDE. We note that the signs in (43) do not match with typical HJB equations from control theory. In order to obtain an HJB equation, consider the line-reversse function V, which satisfies
--=-Th Tr (DV*V) -.
Vē - div(F) 1 2 1 V V(,T) log Px-
(44)
Unfortunately, the terminal conditions in both (43) and (44) are typically not available in the context of generative modeling, Specifically, for the optimal drift in (42), they correspond to the intractable marginal densities of the inference process Y, since it holds that px = py, see Remark 2.2. However, the situation is different in the case of sampling from (unnormalized) densities, see Section 3.
A.5 Hopf-Cole transformation
Recall that the Kolmogorov backward equation follows from the Fokker-Planck equation by using the divergence identities from Section A.2, i.e.
0,5x div (- div (D5x) μpx) div (-DVPx) + - Vix + div( = -r (DV*Fx) + μ Vpx + div(H)
(45a)
(45b)
The following lemma details the relation of the HJB equation in Lemma 2.3 and the linear Kolmogorov backward equation" in (5). A proof of the Hopf-Cole transformation can, e.g., be found in Evans (2010, Section 4.4.1) and Richter & Berner (2022, Appendix G). Lemma 2.3 follows with the choices b == - and h = div(u).
Lemma A.2 (Hopf Cole transformation). Let h: Rdd x [0, + R and let p E CRd x [0, IR) solve the linear PDE
1
0.p Tr(oo T V2p) b Vp + hp.
(46)
Then V = 1--log logp satisfies the HJB equation
a V.= = Tr vLV)
1
(vlooTVvV:-- - VV h -V+ loo
(47)
For div(r) this be viewed the adjoint the Fokker-Planck equation.
21

---

Published in Transactions on Machine Learning Research (02/2024)
For further applications of the Hopf- Cole transformation we refer, for instance, to Fleming & Soner (2006); Hartmann et al. (2017); Léger & Li (2021).
A.6 Brief introduction to stochastic optimal control
In this section, we shall provide a brief introduction to stochastic optimal control. For details and further reading, refer the interested reader to the monographs by Fleming & Rishel (2012); Fleming & Soner (2006); Pham (2009); Van Handel (2007). Loosely speaking, stochastic control theory deals with identifying optimal strategies in noisy environments, in our case, contimuuu--time stochastic processes defined by the SDE
dx - X Us) + 8, dB..
(48)
where and ā are suitable functions, B d-dimensional Brownian motion, and is progressively measurable, Rd-valued random control process. For ease of presentation, we focus on the frequent case where U is Markor control, which means that there exists a deterministic function such that U, u(X!.s). s). In other words, the randomness of the process U is only coming from the stochastic process X The function class U then defines the set of admissible controls. Very often, one considers the special cases
u C C(R* x [0, R*),
(zz, s, u) = p(z, s) + (s))(((, s) and u.) = os)),
(49)
where μ and a might correspond the choices taken in Scction 2.4, and one might think of as a steering force as to reach certain larget.
The goal is now to minimize specified control costs with respect to the control To this end, we can define the cost functional
[. M(xy s8, 8, u(xy ss) ds + 9(X4) x Tiuit
J Tinit Tinit 0) E.
(50)
where h: R x J0, Rd + R specifies runming costs and g Rd + R represents terminal costs. Furthermore we can define the cost-lo-go as
J(ux,t) f'xx s, u(x! $)) ds + s(x4) x =
(51)
now depending on respective initial values (z, E Rd [). The objective in optimal control is now to minimize this quantity over all admissible controls E U and we, therefore, introduce the so-called value function
7).
Viz, = inf Ju; well
(52)
as the optimal costs conditioned on being in position x at time f.
Motivated by the dynamic programming principle (Bellman, 1957), one can then derive the main result from control theory, namely that the function V defined in (52) fulfills a nonlinear PDE, which can thus be interpreted as the determining equation for optimality10.
Theorem A.3 (Verification theorem for general HJB equation). Let V E CR. x [0,2], fulfill the PDE
1 0,V =- inf (((,aa)+ + + VV + Tr (GoT)(...aVV41) V(.T)) = g.
(53)
NERd
2
such that
[I(2, x [0,2] 1+ |||| 00
sup
(54)
I0 practice, solutions optimal control problems may posses cnough regularity order formally fulfill the HJB equation, such that a complete theory optimal control needs introduce an appropriate concept of weak solutions, leading to so-called viscosity solutions that have been studied, for instance, in Fleming & Soner (2006); Lions (1983).
22

---

Published in Transactions on Machine Learning Research (02/2024)
and ossume that there erists a measurable function U u² :R' x T0 Rd that attains the above infimum for all (x,t) E Rd x [). Further, let the corvespondingly controlled SDE in (48) with U = u² (X." ,8) have a strong solution X U Then V coincides with the value function as defined in (52) and u² is an optimal Morkovian control.
Let us approciate the fact that the infirumm in the HJB equation in (53) is merely over the set Rd and not over the function space U as in (52), the minimization reduces to pointwise operation. A proof of Theorem A.3 can, for instance, be found in Pham (2009, Theorem 3.5.2).
In many applications, in addition to the choices (49), one considers the special form of running costs
Ma,, 8, u(z, 8)) = h(r, 8) + 2 [u((, s)1??
(55)
where h R& x [0, T] RR. In this setting, the minimization appearing in the general HJB equation (53) can be solved explicitly, therefore leading to closed-form PDE, as made precise with the following Corollary. Corollary A.4 (HJB equation with quadratic running costs). If the diffusion coefficient ō does not depend on the control, the control enters additively in the drift as in (49), and the running costs take the form
Mx,, 8, u(r, 8)) h(x, s) + |u({,
(56)
then the general HJB equation in (53) can be stated in closed form as the HJB equation in (47).
Proof. We formally compute
((,,,)) vV) h VV + inf aERa
{ 2
[all| + σa vr).
inf
(57)
OER"
and realize that the infiruum is attained when choosing a - -V Rd T), resulting in the optimal control u² -o VV. Plugging this into the general HJB equation (53), we readily get the PDE in (47).
(z, t) for each corresponding (xx t) E
A.7 Verification theorem
The verification theorem is a classical result in optimal control, and the proof can, for instance, be found in Nüsken & Richter (2021, Theorem 2.2), Fleming & Soner (2006, Theorem IV.4.4), and Pham (2009, Theorem 3.5.2), see also Section A.6. For the interested reader, we provide the theorem and a self-contained proof using Itô's lemma in the following. Theorem 2.4 follows with the choices t == 0, b = -4, h = div(p), and g == Pog
Theorem A.5 (Verification theorem). Let V be a solution to the HJB equation in (47) with terminal condition V(, g. Further, let E [0, T) and define the set of admissible controls by
|u[,, s)ll 1+ ||]|||
u u CR" (T,T),
sup (1.77
(58)
For every control u E U let 7 ( be the solution to the controlled SDE
dz = (ou + b) (2:h 8) + (8))
(59)
and let the cost of the control u be defined by
(u: J(::))= = [ T
(1 h + (111)))))) (Z) ds ds + g(27) Z
(60)
Then for every u E 1 it holds almost surely that
1 T
vizi, 1) +
TV
+ uul (2,, 8) ds Zi = ((;
(61)
2
23

---

Published in Transactions on Machine Learning Research (02/2024)
In particular, this implies that . min (u; almost surely, where the unique minimum is attained by u² = -a VV.
Proof. Let us derive the verification theorem directly from Ito's lemma, which, under suitable assumptions, states that
v(Zz-7)- S. ( VV (TV 8) +
(62)
almost surely, where
5 s= f= (a (o (oTV) dB.;
(63)
sec, e.g., Theorem 8.3 Baldi (2017). Combining this with the fact that solves the HJB equation in (47)
and the simple calculation
0
VV + u/l
0 VV VV
+1) (o VV +u)
(64a)
2
-
+ 111
(64b)
shows that
(++ [.((++]] 1lul 2 vo + u1(") s) ds g(*7) 5
T
1
V(Z,, t) -
(65)
almost surely. Under mild regularity assumptions, the stochastic integral S has zero expectation conditioned ZH", which proves the claim.
Remark A.6 (Variational gap). We interpret the term
T LV + u|| (2+, ds Zi
E
(66)
in (61) as the variational gap specifying the misfit of the current and the optimal control objective. In the
setting of Corollary 2.5 it takes the form
1 I
E
Ilu-o o T Vlogxx 1| (YT 8) ds Yo
(67)
2
and can be compared to Huang et al. (2021, Theorem 4), where, however, the factor 1/2 seems to be missing.
A.8 Measures on path space
In this section, elaborate on the path space measure perspective diffusion-based generative modeling, introduced in Scction 2.3. Recalling our setting in Scction A.1, first state the Girsanov theorem that will be helpful in the following, see, for instance, Üstünel & Zakai (2013, Proposition 2.2.1 and Theorem 2.1.1) for a proof.
Theorem A.7 (Girsanov theorem). For every control u E let Y (Y,")s((o))) be the solution to the controlled SDE
dy (ou - 11) (Y.", 8) (s)) dB,
(68)
For u, v E U, it then holds that
dlP d'y dP
(Y") RS (Y") + /
T (u v)(,,, s). dB.
log
(69)
and, in particular, that
d.P T (Y")
(Px P
= F. log
= F. [RY""("")
(70)
dl y
24

---

Published in Transactions on Machine Learning Research (02/2024)
Note that the expression for the KL divergence in (70) follows from the fact that, under mild regularity assumptions, the stochastic integral in (69) is martingale and has vanishing expectation. Now, we present a lemma that specifies how the KL divergence behaves when we change the initial value of path space measure.
Lemma A.8 (KL divergence and disintegration). Let X, Y. Z be diffusion processes. Further, let Z be such that Zo ~ Yo and dZ = dx. Then it holds that
(P PPy + DKL Px,)
(71)
Proof. Let Py= be the path space measure of the process Y with initial condition Y, - Rd let Pr, be the marginal at time 0, and similarly define the corresponding quantities for the processes X and Z. Since our assumptions guarantee that Px= P2= and Pz. Py: the disintegration theorem, see, e.g., Léonard (2014), shows that
dPy
dmy= Y= dPy. Yo
dF y= dP Yo dy dPy dPy Yo.
(72)
dPxx dlP
dPP Xa
dIP Z= dP Zo dPx. dIP dP
which implics the claim.
Using the previous two results, we can now provide a path space perspective on the optimal control problem
in Theorem 2.4. Note that Proposition 2.6 follows directly from this result. For further connections between optimal control and path space measures, we refer to Hartmann al. (2017); Thijssen & Kappen (2015).
Proposition A.9 (Optimal path space measure). We define the work functional W: C(O,T), Rd) - RR for all suitable stochastic processes by
log Pxo (Yr)
W(Y) = RY)
(73)
pxr (Y))*
where (Y") as in (9) with u = 0. Further, let the path space measure Q defined via the Radon-Niiee Maddon-Nikoyym derivative
dQ
exp (-W).
(74)
dP y
Then it holds that Q TP yut where u² o T logix in Corollary 2.5, and for every l have that
dlP
dly T
= + /'" w(Y.") dB, + log pxo (Yr)
Pxr (Yi1)
log
(75)
aIP Y
In particular, the expected variational gap
G(u) = EE [logxx, (Y))) -E - RY)()))
(76)
of the ELBO in Corollary 2.5 satisfies that
G(u) DKL (Py PP DKL (Py (Pxx) DKL (P' Px
(77)
Proof. Similar to computations in Nüsken & Richter (2021), we may compute
dPy= dQQ
ddyy dlyo dP dQ
log
(Yn) log
(Y")
(78a)
RC L u(Y.) dB, + RY - llog
Px, Px, (Y7))
(78b)
px,, (Y/')
RX" [' u(Y.") dB, + log Xx (Y
(78c)
px. (Y
1iThis means that and X are governed by the same SDE, however, their initial conditions could be different.
25

---

Published in Transactions on Machine Learning Research (02/2024)
where used the Girsanov theorem (Theorem A.7) and (74) in (78b). This implics that
dIP Y dQ
DeL (PPuQQ = E
log
(Ym) = E (Y log (XT) + IB [xo (66')
(79)
Comparing to Theorem 2.4, realize that the above KL divergence is equivalent to the control costs up [X"", 0)] (X6") Using (8), we thus conclude that
(R [Q) 0
(80)
for u² = - VV VV in (75).
= o log which implies that Q = P yyu. Together with (78), this proves the claim
Now, we can express the expected variational gap G(u) (76) in terms of KL divergences. We can prove the first equality in (77) by combining the identity Q Py with (79). The second equality follows from Lemma A.8 and the observations that Y)" ~ Yo and dy =dX, see Theorem A.1. This concludes the proof.
Note that the Girsanov theorem (Theorem A.7) shows that the expected variational gap G(u) = (Py-Pyy|PPxr as in (77) equals the quantity derived in Remark A.6. The following proposition states the path space measure perspective on diffusion-based generative modeling.
Proposition A.10 (Forward KL divergence). Let consider the inference SDE and the controlled generative SDE as in (14), i.c.,
aY, ((Y,, dB) Yo ~ D. - + ((s) dBs,
(81)
and the associaled puth space measures Py and Px.. Then the expected varialional gap is given by
G(u) = IE Pog (Y)) - EE llogxx (Yr) - R Ra (IP )) DKL (Py, PPx))
(82a)
(82b)
Proof. Let us consider the SDE
dY (aw (Y 8) (B) You D.
(83)
With the choice μ f, then have fu.au Y, which actually does not depend anymore, and dy" ,a- dx" where u' o TV logEx.. Together with the fact that Yo Yo oou- Lemma A.8 implies that
DKL (P (PyneP-- on-)) DKI. x*) DKL. (Pr (P.
(84)
and Proposition A.9 (with μ ou f) yields the desired expression.
In practice, the expected variational gap G(u) in (82a) cannot be minimized dircctly since the evidence logpxy (Y) is rntraclable and one resorts maximizing the ELBO instead. The same problem occurs when considering divergences other than the (forward) KL divergence DEL in (82b). The situation is different in the context of sampling from densities, where we directly minimize the (reverse) KL divergence, Corollary 3.1, allowing us to use other divergences, such as the log-variance divergence, see Section A.14.
12in view Lemma A.8, the variational gap can also be written as C(u) (I' [1'2") where d2 dx and 2 Yo. i.c., the process z governed by the same SDE X", however, has the initial condition of Y. One could, however, the probability flow ODE, corresponding to the choice in Theorem A.I, which then resembles training a normalizing flow.
26

---

Published in Transactions on Machine Learning Research (02/2024)
A.9 ELBO formulations
Here provide details on the connection of the ELBO to the denoising score matching objective. Using the reparametrization in (14), Le. taking μ = ou f. Corollary 2.5 yields that
1xg (Y)) 2 IE logpx; (Yr) - Rrn Rou- ( xm/
(85)
The next lemma shows that the expected negative ELBO in (85) equals the denoising score matching objective (Vincent, 2011) up a constant, which does not depend on the control u.
Lemma A.11 (Connection to denoising score matching). Let us define the denoising score matching objective by
(u) = I [|(((=.7) T) - - - (x))logps, (XF1))))
(86)
2
Then it holds that
(DSM(u) -E'' plo (YT) - R Rou-,(()) +C,
(87)
with
C = E log (Vr) + divv() + (T)V ( (Y>1
(88)
where u(1, T)) and py,|%, denotes the conditional density of Y given Yo.
Proof. The proof closely follows the one in Huang et al. (2021, Appendix A). For notational convenience, let first define the abbreviations
p(z,s) s) PY. Yo) and r(a,s) s) (oT
(89)
for every x E Id and s [0, T}. Note that we have
2 || - o Vloggll = 2
|ul| - r + z Vooolll'''
(90)
which shows that
(D) IE T (T()) |ul| -r) (Y-: +
T l|o (+)oooo
(91)
Focusing the first term the right-hand side, Fubini's theorem and Monte Carlo approximation establish that, under mild regularity conditions, it holds that
1 |1ul -r (x.. L {
- S. [Br((., ds,
IE T
(92)
where the expectation on the left-hand side is over the random variable (Yr, ~ PY-|>Pr, where ps is the density 7 ~ 4(10, T)) and lr is the conditional density of Y, given т.
Together with (91), this implics that
LDSM (u) = -E [log [xo (Yr) - R²) - [ ()d C.
(93)
Focusing on the term 8(),, for fixed a [0, T), remains to show that
BEr(.,, s)] -E(div(ou)((,...
(94)
Using the identities for divergences in Section A.2, show that
div(oup) - div(ou)p = (au) Vp (au) (p V logp) rp
(95)
27

---

Published in Transactions on Machine Learning Research (02/2024)
Further, Stokes' theorem guarantees that under suitable assumptions it holds that
J (ar 0.
(96)
Thus, using (95) and (96), we have that
E (dvv(ow)(Y, = J/e div(ou)(x; s)p(s, da
(97a)
Raa s)dz E s)
(97b)
Combining this with (93) finishes the proof.
Note that one can also establish equivalences to explicit, implicit, and sliced score matching (Hyvärinen & Dayan, 2005; Song et al., 2020), see Huang et al. (2021, Appendix A). Using the interpretation of the ELBO in terms of KL divergences, see Huang et al. (2021, Theorem 5) and also Proposition A.10, one can further derive the ELBO for discrete-time diffusion models as presented Ho et al. (2020); Kingma et al. (2021).
A.10 Diffusion-based sampling from unnormalized densities
In the following, we formulate an optimization problem for sampling from (unnormalized) densitics. The proof and statement are similar to Theorem 2.4 and Proposition A.9 with the roles of the generative and inference SDEs interchanged. Note that Corollary 3.1 is direct consequence of this result.
Corollary A.12 (Reverse KL divergence). Let X and Y be defined by (17). Then, for every a it holds that
dP xM = + f ((xx)) dB, log (X,) log 2, p(X*)
log
(98)
dP x
where u² = o logpr. In particular, it holds that
(Px- DKI, x x DKL. (KI xp (Py)) (
CDIS(U) log Z,
(99)
where
pxy (Xx)
LDIS(W) = E (x²) + +og
(100)
p(X{))
This further implies that
DKI. x PYr) log 2 = min (Dis(u),
(101)
well
and, assuming that X Yr, the minimizing control u* guarantees that Xy D.
Proof. Analogously the proof of Lemma 2.3 and noting the interchanged roles of X and Y, we see that
logTy satisfies the HJB equation
8 0p - -T Vp div(Ap, p(,T) log PYo z
(102)
We can now proceed analogously to the proofs of Theorem 2.4 and Proposition 2.6. Using the identity
log px: (X:)
DKL. Xx PY,))
IF.
(103)
py, Xx)
this proves the statements in (99) and (101). The last statement follows from Theorem A.1 and the fact that the minimizer u" == a log py is the scaled score function.
28

---

Published in Transactions on Machine Learning Research (02/2024)
Table 1:
We compare the error of DDS (Vargas et al., 2023a) and DIS estimating the log nog-noomalizing constant and the average standard deviation on two problems. We report the average crror over three independent. runs. For fair comparison, we the same prior (0, I), batch size m 2048, and number of gradient steps K = 20000, and a fixed number of integrator steps N 256, see Table 2. For DDS, we present the best results for different cosine schedules with 0uma E (0.25, 0.5, 1.5, 2, 2.5} (sce Vargas et al. (2023a) for the precise definition). For comparison, we present results in their training setup as well as ours, i.e., with exponentially moving average of the parameters, clipping, and learning rate scheduler. For DIS, we also present results for the log- varvariaece divergence, see Section A.14.
Problem
Method
A log 2 (rw) +
Astd +
GMM
DDS (a 0.5)
0.176
2.48
DDS (a 0.5, our setup)
0.090
2.52
DIS (KL)
0.046
2.83
DIS (log-variance)
0.013
0.01
Funnel
DDS (Amax 0.5)
0.121
7.20
DDS (Omax 1, sctup)
0.046
6.49
DIS (KL)
0.017
5.50
DIS (log-variance)
0.021
5.46
A.10.1 DDS as special case of DIS
In this section, we compare the denoising difflusion sampler (DDS) suggested in Vargas et al. (2023a) to our lime-reversed diffusion sampler (DIS). We present numerical comparison in Table 1 and refer to Richter & Berner (2024, Section 3) and Vargas & Nüsken (2023) for further details. In particular, these works show that DDS can be derived using reference process with known time-reversal. In the following, outline how DDS can be considered a special case of DIS. To this end, we note that DDS considers the process
X. dx,- = (33(s)x,+ + 22 (5(X,,)) + xx2 ) ds ds+ 2s+ dB. 1/2
(104a)
= (3())x,+ - 22) ds + n (3)
(104b)
with Xg ~ NO. 11), see also (136). To case notation, let us define (t) == 7/ (t) f(z,)) = -3(t)x and u(x, 1) = (1) (ar 1) to get
dx = (u)), s) s) +(s) dB.
(105)
in (17). We added the superscript u to the process in order to indicate the dependence the control For DDS, the loss
(DDS(1)= = IE I/"' T(X* , 9) ( ds log
:(X;:0.172) p(X+)
(106)
2
is considered. In order to show the mathematical equivalence of DDS and DIS for the special choice of a and f, we may apply Ttô's lemma to the function V(z, t) 1 ||r. This yields that
E
-
2 |
|
'/" (-J+ हū (X X2 + + 2d)3(s ds
(107)
where X is defined in (105). Noting that
1 = 2np ||:0||2 log (2 ?)
(108)
and
V = v (-3(())) - dz(1).
(109)
29

---

Published in Transactions on Machine Learning Research (02/2024)
can rewrite (107) as
%
TT'
[ 5ū 1/2
E log
E
(X X²" + div i (Xx x d8
(110)
Now, we observe that the special choice of a and f implics that
172 - (x s) X [|ū (ixx,,))] TT(X [Xxx 8)
s)(l
(111)
s7²
which, together with (110), shows that the DDS objective in (106) can be written
CDDS(u) E (x + log W(Xx:0,nTT))
(112)
p(x7)
Recalling that pxg- NN0,, 17),, we obtain that CDDs(u) = Lpis(u), where the DIS objective is given as in (19).
We emphasize, however, that the derivation of the DDS objective in (106) relies a specific choice of the drift and diffusion coefficients f and μ in (106). In contrast, the DIS objective in (19) is, in principle, valid for an arbitrary diffusion process.
A.10.2 Optimal drifts of DIS and PIS
In this section, we analyze the optimal drift for DIS and PIS in the tractable case where the target p (mm, v²I) is multidimensional Gaussian. This reveals that the drift of DIS can exhibit preferable numerical properties in practice. We refer to Figures 14 and 15 for visualizations.
Using the VP SDE in Section A.13 with (as in our experiments), the optimal drift for DIS is given by
(17) 3(t)a" =(t)= = (²) 1 28(1) (v2 (v2-1) 1)
e-((t) m +- g((x)))
1
(I,,
(113)
v Fl
where is as in (138), (137). Note that for 1, rewrite the score
V (a, a(t) = (1 a(t) V log (1) te a(0) v log p(z),
(114)
which reminiscent of our initialization given by linear interpolation and, in fact, the same for the (admittedly exotic) choice o² - see Section A.13.
In comparison, the optimal drift of PIS (using f = 0, i.c., a pinned Brownian motion scaled by E (0, oo) as in the original paper) can be calculated via the Feynman-Ko Kac formula (Zhang & Chen, 2022a; Tzen & Raginsky, 2019) and is given by
2T driftris (I, tv² + oTTTT -t)
Tx t
I +
m
(115)
t)
see Vargas et al. (2023a, Appendix A.2). Due to the initial delta distribution, the drift of PIS can become unbounded for - 0, which causes instabilities practice. Note that this not the case for DIS.
A.11 PDE-based methods for sampling from unnormalized densities
In principle, the different perspectives on diffusion- based generative modeling introduced Section 2 allow for different numerical methods. While our suggested method for sampling from (unnormalized) densities Section 3 directly follows from the optimal control viewpoint, an alternative route is motivated by the PDE perspective outlined in Section 2.1. While classical numerical methods for approximating PDFs suffer from
30

---

Published in Transactions on Machine Learning Research (02/2024)
the curse of dimensionality,, we shall briefly discuss methods that may be computationally feasible even in
higher dimensions.
Let us first rewrite the PDEs from Section 2.1 in a way that makes their numerical approximation feasible. To this end, we recall that the Fokker-Planck equation is a linear PDE and its time-reversal can be written as a generalized Kolmogorov backward equation. Specifically, in the setting of Section 3, holds for the time-reversed density By that
0.Py = div (-div (Dir) + JTr) -
-T - + +J. Vyr + div(()
(116)
analogously to (5). We note that, due to the linearity of the PDE in (116), the scaled density Zpy also satisfics Kolmogorov backward equation given by
Opp = Tr (DV*p) + Vp+ div(fl, p(,T) =
(117)
The corresponding HJB equation following from the Hopf Cole transformation V = - -oog (Zpy) (as outlined in Section A.5) equals
a,v = TDDVVV) + VV - div(F) + flbo T V(,T) = log
(118)
in analogy to Lemma 2.3, however, with a modified terminal condition since we omitted the normalizing constant. Now, as p is known, viable strategies to learn u* would be to solve the PDFs for or V via deep learning.
The general idea for the approximation of PDE solutions via deep learning is to define loss functionals C such that
((p) is minimal
p is solution to (117);
(119)
or, respectively,
C(V) is minimal
V is a solution (118).
(120)
This variational perspective then allows to parametrize the solution or V by neural network and to minimize suitable estimator versions of the loss functional L via gradient-based methods. We can then use automatic differentiation to compute
u
V loggy = o Vogg = o TV
(121)
as the score is invariant to rescaling. This adds an alternative directly optimizing the control costs in (19) in order approximately learn u u².
A.11.1 Physics informed neural networks
A general approach to obtain suitable loss functionals L for learning the PDEs, dating back to the 1990s (sce, e.g., Lagaris et al. (1998)), often referred to Physics informed neural networks (PINNs) or deep Galerkin method (DGM). These approaches have recently become popular for approximating solutions to PDEs via combination of Monte Carlo simulation and deep learning (Nüsken & Richter, 2023; Raissi et al., 2017; Sirignano& & Spiliopoulos, 2018). The idca is to minimize the squared residual of the PDE in (118) for an approximation V ~ V, i.e.,
=E =E [ +IH (on=+=(DVeV--- - vv + div(F) z (V (6.7).
(PDE(
(122)
as well as the squared residual of the terminal condition, i.e.;
Lhoundary V IE (V(()))+ + +lg p(E) )l.
(123)
where (E,7) suitable random variable distributed on Rd x [O. This results in loss
LVINN(V) LYDE(V) Lboundarr (V),
(124)
31

---

Published in Transactions on Machine Learning Research (02/2024)
3,,3,,,))
12.60
7.50
9.56
PIS
12.55
PINN
7.45
expectaaion
12.50
expectation
9.54
ground truth
logiz)
9.52
7740
12.45
9.50
7.35
12.40
9.18
7.30
9.46
100
200
400
800
100
200
400
800
100
200
400
800
integrator steps
integrator steps
integrator steps
Figure 7: We use the settings in Figures and and display the estimates (median and interquartile range over 10 training seeds) of DIS, PIS, and the PINN method for log and the expected values of two choices of in (148). We find that the PINN method can provide competitive results.
where c E o0) is a suitably chosen penalty parameter. Using automatic differentiation, one can then compute an approximation of the control as in (121). In Figure 7, we compare this approach to DIS and PIS on a 10-dimensional double well example. While DIS provides better results overall, the PINN approach outperforms PIS and holds promise for further development. A clear advantage of PINN the fact that no time-discretization necessary. On the other hand, one needs compute higher-order derivatives in the PINN objective (122) during training, and one also needs to evaluate the derivative of V to obtain the approximated score during sampling.
oo)
A.11.2 BSDe-based methods and Feynman-Kac formula
To reduce the computational cost for the higher-order derivatives in the PINN objoctive, can also develop loss functionals & for our specific PDEs by considering suitable stochastic representations. This can be done by applying Itô's lemma to the corresponding solutions. For instance, for the HJB equation in (118), we obtain
log ((Xr) = V(E.T) R VV (X) + (X),
(125)
where we abbreviate
Sp(X) = L- Vp) (X,, 9): dB,
(126)
In the above, the stochastic process is given by the SDE
dX, = -xx s) ds + adB., X- ~ E,
(127)
where (E, is again suitable random variable distributed on Rd x [0, Minimizing the squared residual in (125) for an approximation V V, we can thus define a viable loss functional by
LBSDE (V) = E (V(E:.) vV (X) +x(())
+ log log p(Xr))
(128)
see, e.g., Richter (2021); Nüsken & Richter (2021). The name of the loss functional LBSDE stems from the fact that (129) can also be seen as a backward stochastic differential equation (BSDE). A method to combine the PINN loss in A.11.1 with the BSDE-based loss in (128) presented in Nüsken & Richter (2023), leading to the so-called diffussion loss, which allows to consider arbitrary SDE trajectory lengths.
One can argue in similar fashion for the linear PDE in (117), where Itt's lemma establishes that
p(Xr) = P(E.T)+ J. (div(T))) (X., 8) ds S,(X).
(129)
However, for the linear Kolmogorov backward equation, we can also make use of the celebrated Feynman- formula, i.e.,
p(E, = E exp (-f." div())(X.. s) p(XT) (5,T)
(130)
32

---

Published in Transactions on Machine Learning Research (02/2024)
DW (d-20,
marginal histogram of X
1.0
reweighted histogram of
density
0.5
0.0
X1
Figure 8: Illustration of importance sampling in path space. The left panel shows histogram of the first component of the controlled process X4 at terminal time T compared to the corresponding marginal of the larget density Z. While we camnot correct the error caused by pxy R Pyr: we can miligate the error in the control 2 u² by reweighting X7 with the importance weights w", see Section A.12. The right panel shows that this indeed results in an improved approximation of the target density.
to establish the loss
(p) = ( exp (-/," div(F)(X,,)) 5) ds p(Xx)
27
(131)
see, e.g., Berner et al. (2020); Richter Berner (2022); Beck al. (2021). Since the stochastic integral (X) has vanishing expectation conditioned on (E,7), it can also be included in the FK-based loss in (131) to reduce the variance of corresponding estimators and improve the performance, see Richter & Berner (2022).
A.12 Importance sampling in path space
In practice, will usually not have u" available but must rely approximation R u". In consequence, this then yields samples X7 that are only approxinately distributed according to the target density p and thus leads to biased estimates. However, we can correct for this bias by employing importance sampling.
Importance sampling is a classical mothod for variance-reduccion in Monte Carlo approximation that leads to unbiased estimates of corresponding quantities of interest (Liu & Liu, 2001). The idea is sample from a proposal distribution and correct for the corresponding bias by likelihood ratio between the proposal and the target distribution. While importance sampling is commonly used for densities, one also employ it in path space, thereby reweighting continuous trajectories of stochastic process (Hartmann et al., 2017). In our application, u is only an approximation of the optimal u". Note that, in general, it holds for any suitable functional & ([, + R that
can make use of this in order to correct for the bias introduced by the fact that
E p(x"))
|] IE (xx")
dlP X-* (X4)
(132)
dpx
i.c., the expectation can be computed w.r.t. samples from optimally controlled process X even though the actual samples originate from the process X The weights necessary for this can be calculated via
PYr (X.' dP (X") = exp - -R - f-- ((x)) dB, log p(X#)
1
Wh
(133)
where we used the Radon-Nikodym derivative from Proposition A.9. The identity in (132) then yields an unbiased estimator, however, with potentially increased variance. In fact, the variance of importance sampling estimators can scale exponentially in the dimension and the deviation of u from the optimal w", thereby being very sensitive the approximation quality of the control u, see, for instance, Hartmann & Richter (2021).
Note that, in practice, we do not know Z, iee, we need to consider the unnormalized weights
TH t Z.
(134)
33

---

Published in Transactions on Machine Learning Research (02/2024)
We can then make use of the identity
IB [x(x))
IB [
(X*)
(135)
While for normalized weights in (133) the variance of the corresponding importance sampling estimator vanishes at the optimum =u, this is, in general, not the case for the unnormalized weights defined in (134).
In implementations, we introduce bias in the importance sampling estimator by simulating the SDE X using a time-discretization, such as the Marry-Mar scheme, see Section A.13. Moreover, the quantity PYr in (133) typically intractable and we need to use the approximation Pxş R PYr as the loss LDIS, see Corollary 3.1. We have illustrated the effect of importance sampling in Figure 8.
A.13 Details on implementation
In this section, we provide further details on the implementation.
SDE: For the inference SDE Y employ the variance-preservve (VP) SDE from Song et al. (2021) given by
(t) = ->/ I and f(x, t) == -
(136)
with B C( (0, oo)) and n (0, oo). Note that this constitutes Ornstein Ubllebbek process with conditional density
(.))) = N (e (e-a 12 (1-c-20(0)
(137)
where
a(l) = / 3(s)ds.
(138)
In particular, observe that
PY, = E 10rr1x((1130) s N (0. 121)
(139)
for suitable B and sufficiently large T. In practice, choose the schedule
L 1 (1 (1- mmin T T B()==
(140)
with 7 1 and sufficiently large 0 < amin < 0max- This motivates our choice X No T). We can numerically check whether (the unconditional) Yr is indeed close to standard normal distribution. To this end, we consider samples from D and let the process Y run according to the inference SDE in (3). We can now compare the empirical distributions of the different components, which each need to follow a one-dimensional Gaussian. Figure 9 shows that this is indeed the case.
Model: Similar to PIS, we employ the initial score of the inference SDE, ie., the score of the data distribution, log = Voog p, the parametrization of our control. Additionally, our DIS method, we also make use of the (tractable) initial score of the generative SDE Voggxx v (0,,) Specifically, we choose
(x = 10 (z, + (2) (t)(())s(z)
(141)
where
s(x, t) 4 Voggxx (x) (1 (1-2) 1 p(z)
(142)
and (2) and 50 are neural networks with parameters 0. We use the same network architectures as PIS, which in particular uses a Fourier feature embodding (Tancik et al., 2020) for the time variable We initialize (2) ((0) such that pro) = 0 and aco) = 1. The parametrization in (141) establishes that our initial control (approximately) matches the optimal control at the initial and terminal times, iee,
Upro (-, 0) =0) V log p o00) V log u² =²(,0)
(143)
34

---

Published in Transactions on Machine Learning Research (02/2024)
GMM
Punnel
Punnel
histogram of A00, 1)
0.4
0.3
0.2
0.1
0.0
-2
2
-2
0
2
-2
D
Figure 9: Histogram of the first component of YT for the Gaussian mixture model and the first two coordinates of the Funnel distribution (in blue) when starting the corresponding process Y in Yo ~ D and using N 800 steps for the simulation, compared to the density of standard normal (in orange). We do not depict the double well as cannot sample from the ground truth distribution.
and
Ugto (,T) = oT) V logpxy R V logpy, = (.';,))
(144)
In our experiments, we found beneficial to detach the evaluation from the computational graph and to use a clipping schedule for the output of the neural networks.
Training: We minimize the mapping 0 - (u)( using a variant of gradient descent, ie., the Adam
optimizer (Kingma & Ba, 2014), where EDIS is a Monte Carlo estimator of LDIS in (19). Given that we are now optimizing reverse KT. divergence, the expectation LDis(ua) over the controlled process Xa9 in (17) and we need to discretize X and the running costs R. (X")) on a time grid 0 = to < ... < ty = T.
We use the Fuler-Maruyama (EM) scheme given by
8 = 82 + (oto - -)) (8 + +(tx)) (Br... - Brn).
(145)
~N(0.Ar1)
where Al = T/N 1x+1 - in the step size. Given X8 Xy it can be shown that X& convergences X for At - in an appropriate sense (Kloeden & Platen, 1992). This motivates why we increase the number of steps N for methods when the training progresses, see Table 2. A smaller step size typically leads lo better approximation but incurs increased computational costs. We also trained separate models with a fixed number of steps, but did not observe benefits justifying the additional computational effort, see Figure 10.
We compute the derivatives w.r.t 0 by automatic differentiation. In memory-restrieted setting, one could alternatively use the stochastic adjoint sensitimity method (Li et al., 2020; Kidger et al., 2021) to compute the gradients using adaptive SDE solvers. For a given time constraint, we did not observe better performance of stochastic adjoint sensitivity methods, higher-order (adaptive) SDE solvers, or the exponential integrator by Zhang& & Chen (2022b, Section 5) over the Euler-Maruyama scheme. Our training scheme yields better results compared to the models obtained when using the default configuration in the code of Zhang & Chen (2022a), see Figure 11. For our comparisons, we thus trained models for PIS according to our scheme, see also Table 2. Note that in Section A.11 we present an alternative approach of training a control ug based on physics-informed neural networks (PINNs).
Sampling: In order to reduce the variance, we compute an exponential moving average 0 of the parameters ( in Algorithm 1 during training and it for evaluation. We evaluate our model by sampling from the prior X ~ 10) and simulating the generative SDE using the EM scheme, such that X represents an approximate sample from D. Note that we can, in principle, choose the number of steps for the EM scheme independent from the steps used for training. We used an increasing number of steps during training, and our trained model provides good results for range of steps, Figures 4 and 5. Thus, amortized the
35

---

Published in Transactions on Machine Learning Research (02/2024)
GMM
Funnel
DW (d=20,w=3,004
PIS
DIS
0.00
13
ground-truth
(2) (rv)
-0.25
12
-0.50
11
0.11
-0.75
-.2
10
100
200
integrator steps
integrator steps
integralor steps
Figure 10: We use the same setting as in Figure 4, except that for each choice of SDE integrator steps N a separate model is trained, i.e., all gradient steps use the same step size At. This yields qualitatively similar rosults. For our other experiments, we thus amortize the cost and train a model which can be evaluated for a range of steps N.
cost of training different models for different step sizes. Finally, note that similar to Song al. (2021), one could allernatively obtain approximale samples from D by simulating the probabilily flow ODE, originating from the choice À 1 in Theorem A.1, using suitable ODE solvers (Zhang & Chen, 2022b).
Evaluation:
We evaluate the performance of our proposed method, DIS, against PIS on the following two
metrics:
approcimating normalizing constants: We obtain lower bound on the log-normalizing constant. log Z by evaluating the negative control costs, i.e.,
log Z log DKL x (Yy)) CDIs(uu)
(146)
see (20). We note that, similar to the PIS method, adding the stochastic integral (which has zero expectation), i.c., considcring
pxp (X11 R + log p(X7)
+ L-² u(Y.",)) s).,ss.ddl dB,
(T
(147)
yields an estimator with lower variance, see also Corollary A.12. We will compare unbiased estimates using importance sampling in path space, see Section A.12.
approsimaling expeclations and slandard devialions: In applications, one is often interested in computing expected values of the form
IF. (v())),
(148)
where Yo DD follows some distribution and 7 Rd - R specifies an observable of interest. We consider (xx == ||o|² and y(x) = || = Σi= xz|- We approximate (148) by creating samples (K),,))) N k from our model as described in the previous paragraph, where E N specifies the sample sizc. We can then approximate I (148) via Monte Carlo sampling by
P
(
(149)
Given a reference solution of (148), we can now compute the relative error (T/T), which can be viewed as an evaluation of the sample quality on a global level. Analogously, we analyze the error when approximating coordinate-wiee standard deviations of the target distribution Yo using the samples (x (0.00)
-=
We note that an Fuler- Maruyama discretization -LDIS (ug) in (146) can potentially overestimate the log nog-oommalizig constant log Z. We refer Vargas& & Nüsken (2023) for reformulation that guarantee lower bound using backward integration.
36

---

Published in Transactions on Machine Learning Research (02/2024)
GMM
Funnel
DW (d=20,wm5,6-p)
ground-ruut
0.0
0.25
DIS PIS
-0.5
0.00
10
(original)
-0.25
-1.5
0.50
-2.0
-0.75
Figure 11: The boxplot compares our models the results obtained from the repository of Zhang Chen (2022a) using their settings. As they only train with N 100 steps and present approximations of the nog-oomaliziin constants, show corresponding results. The figure shows that our implementation of the PIS method already provides significantly improved results. Compared to their settings, train until convergence with larger batch sizes, use increasingly finer step sizes, a clipping schodule, and an exponential moving average of the parameters.
A.14 Further numerical results
In this section, we present further empirical results and comparisons.
Other divergences on path space: Our principled framework of considering path space measures in Section 2.3 allows us consider divergences other than the (reverse) KT. divergence, which known to suffer from mode collapse (Minka al., 2005; Midgley et al., 2022) or potentially high variance of Monte Carlo estimators (Roeder al., 2017). We conducted preliminary experiments and want to highlight results using the log-veriance divergence that has been suggested in Nüsken & Richter (2021) and often exhibits favorable numerical properties. In fact, we can see in Figure 6 that, for some experiments, this divergence can improve performance significantly sometimes by more than an order of magnitude. We, therefore, suspect that further improvements will be possible by leveraging our path space perspective, see Richter & Berner (2024) for an extensive comparison.
Other parametrizations and initializations: We tried a number of parametrizations and initializations for our neural networks. Specifically, tried to initialize the control such that the generative SDE X
has approximately zero drift (as in PIS),
corresponds to the SDE of the unadjusted Langevin algorithm (ULA),
3. has the correct drift at the initial and terminal times, i.e., the interpolation described in (141).
For Item 1, also experimented with omitting the score Vlog p in the parametrization. We present comparison of the learned densities for the Funnel problem in Figure 12 and GMM problem in Figure 13, where we chose non-symmetric GMM example to emphasize the different behavior. While there are problem-specific performance differences, we observed that all the parametrizations incorporating the score V log yielded qualitatively similar final results. Thus, we opted for the most stable combination, i.c., Item 3, which is detailed (141) above. As further benchmarks and metrics, we also present the results of running ULA, see the next paragraph for details, and the Sinkhorn distance (Cuturi, 2013) to ground truth samples. Finally, provide comparison of the corresponding learned drifts in Figure 14 and Figure 15. We note that the learned drifts are closer to the optimal drift for larger E T]. It would be interesting future direction to try different loss weighting schemes (as have been employed for the denoising score matching objective) to facilitate learning the score for small t.
Note that Langevin-basse based algorithms (which are not model- based), such as ULA or the addiusel Langevin algorithm (MALA), are only guaranteed to converge
Unadjusted Langevin Algorithm (ULA):
15 Our implementation is based httpp:://githbbb htt.a//a..ccneeeallaaaaasssaaaaa
http:////.ttttteeeeeeeea/aaattte sinkhorn with the default parameters.
37

---

Published in Transactions on Machine Learning Research (02/2024)
to the target distribution D in infinite time. In contrast, DIS and PIS are guaranteed to sample from the target after finite time T (given that the learned model converged). This is particularly relevant for distributions 'D for which the convergence speed of the langevin-baser algorithms is slow which often appears if the modes of the corresponding densitics are very disconnected (or, equivalently, if the potential consists of multiple minima, which are separated by rather high "energy barriers"). In fact, by Kramer's law, the time to escape such a minimum via the Langevin SDE satisfies the large deviations asymptotics E|r] x cxp(2A //), where is the exit time from the well, AV is the energy barrier, ie., the difference between the minimum and the corresponding maximum in the potential function, and 17 is temperature scaling (Berglund, 2011). Here, the crucial part is the exponential scaling in the energy barriers, which results in the Langevin diffusion "being stuck" in the potential minima for a very long time. We visualize such behavior in the experiments provided in Figures 12 and 13. In our GMM experiment, only three modes are sufficiently covered, even for large terminal times, such as = 2000, which takes roughly ten times longer than DIS or PIS to sample 6000 points on a GPU. Trying larger stop sizes leads bias and divergence, as can be seen the last plot. We observe similar results for the Funnel distribution, although we need smaller step sizes to even converge. We note that hard to tune the parameters for given problem.
counteract the bias by additionally using MCMC methods (e.g., MALA), but it remains
38

---

Published in Transactions on Machine Learning Research (02/2024)
Table 2: Hyperparameters
DIS SDE inference SDE (corresponding Y)
Variance-Preservinng SDE with linear schedule (Song al., 2021)
min. diffusivity minn
0.1
max. diffusivity max
10
terminal time T
I
initial distribution XØ
10, (truncated 99.99% of mass)
PIS SDE (Zhang Chen, 2022a) uncontrolled process X
scaled Brownian motion
drift f
0 (constant)
diffusivity 5
v0.2 (constant)
terminal time T
5
initial distribution
Dirac delta 6, at the origin
SDE Solver
type
Muler-Maruyamat (Kloeden & Platen, 1992)
steps N (see also figure descriptions)
[100, 200, 400, 800) (each for 1/4 of the total gradient steps K)
Training
optimizer
Adam (Kingma Ba, 2014)
weight decay
10-7
learning rate
0.005
batch size
2048
gradient clipping
1 (e-norm)
clip (
s [-c,c], $62) c
c= 10 (step & 200), 50 (200 < step S 400), 250 (else) 20000 (d S 10), 80000 (else)
gradient steps K
framework
PyTorch (Paszke et al., 2019)
GPU
Tesla V100 (32 GiB)
number of seeds
10
Network lourier-embe (Tancik al., 2020)
dimensions
[0,T" - 12 -> R >> R linear o[p linear] fourier-features GELU (Hendryeks & Gimpel, 2016)
architecture
activation function
Network ( width
(Zhang & Chen, 2022a)
64 (d < 10), 128 (else)
dimensions
Rd x [0,7] + R x RR + RR + Rd
architecture
linear = linear] o [ linear] [gosum] [linear, fourier-embe] GELU (Hendrycks Gimpel, 2016) 0
activation function
bias initialization in last layer
weight initialization in last layer
0
Network ( (Zhang & Chen, 2022a)
dimensions
O. - 1244 -> R&1 I2
architecture
[linear og] o [linear og] fourier-embga GELU (Hendrycks Gimpel, 2016)
activation function bias initialization in last layer
1
weight initialization in last layer
0
Evaluation
exponentially moving average samples M
last 1500 steps (updated every 5-th) with decay 1 14ema_stepp stoppnss_/p/////_ 6000
39

---

Published in Transactions on Machine Learning Research (02/2024)
DIS (Langevin init.)
DIS
PIS (our setting) Sample time=0.3s
(our setting)
(same init. as PIS) Sample time=0.35
Sample time=0.35
Sample time=0.35
Sinkhom dist.=6.51 E[[X[]] error=0.26
Sinkhorn d.5t.=6.52
Sinkharn dist.=6,54
Sinkhorn d
Groundtruth
e[[X[x] error=0.38
=[[X]=] error=0.19
rel. =[[X]]] error=0.55
5
-5
DIS (without Vlog p)
ULA (T=50, At-0.01)
ULA
(without Vlog
(T=5, AC-0.01)
(T=500, 4t=0.1)
Sample uime=0.2s
Sample time=0.2s
Sample time= 4s Sinkhorn dist.=7.39 rel. [[[X]]] error=0.55
Sample time=2.8s
Sample time=2.85
Sinkhor dist.=6.50 rel. [[[X]]] 0.33
Sinkhom dist.=6.74
Sinkharn dist.=7.27
Sinkhom dist.=9.47
rel. [[[X]]] error=0.59
rel. =[[X]]] error=0.41
rel. s[[X]:] error=2.05
5
0
X1
X1
Figure 12: KDE plots visualizing samples from DIS and PIS trained on the Funnel example with different. parametrizations, as described in Section A.13. For comparison, we also present results of ULA with two different settings and the runtimes to sample 6000 points on single GPU.
PIS (our setting) Sample time-0
(our setting) Sample time-0, 4s Sinkhom dist.=0.09 rel. =[[x]]] error-0.02
[Langevin init.)
(same init. as PIS) Sample time-0.44 Sinkhorn dist. =0.t=
Sample time-0 0.4s
Sinkhom dist.=0.07
Sinkhorn dist. 0.84
Groundtruth
rel. [[(]]]] er00
rel. <[[X]a] error=0.01
rel. [[[X]]] 0
5.0
0.0
-.5
-5.0
DIS (without Vlog
PIS (without Vlog Sample time=0.2s
ULA (T-50, At=0.1)
ULA T00000.0 Ar-0.1)
ULA (T-2000, 4t-0 4)
Sample time= 0.25
Sample time=1.00
Sample time=9.00 0.00s
Sample time-9 Os Sinkhorn d E[[[]]] error-0.18
Sinkhom dist. =1.42
Sinkhom d.st.=1.36 rel. E[[X[,] error-0.34
Sinkhom dist.=0.08
Sinkhorn -0.07
rel. [[[X]]] 0.33
E[[X],] error=0.32
E[[X[,] error-0.23
5.0
:
2.5
0.0
-5.0
X1
Figure 13: The same experiment as in Figure 12 for GMM example with means (0, 0), (0, 2), (3. 0),
(-4, 0), (0,-5) and covariance matrix L
40

---

Published in Transactions on Machine Learning Research (02/2024)
5
DIS
setting)
5
DIS
[Langevin Init.)
-10-
(same init. as PIS)
-10
(our setting)
-10-
DIS
5
(without Vlog
-5
-10
PIS
(without vlog
ULA
-10
-10
10
-10
10
-00
10
-10
10
-10
Figure 14: The SDE drifts of the different methods in Figure 12 for I1, 22 E [-10,, 10] and 03 T10
0.
41

---

Published in Transactions on Machine Learning Research (02/2024)
DIS (optimal)
DIS
(our setting)
DIS
(Langevin Init.)
DIS
init. PIS)
PIS
setting)
DIS
(without
P15
(without Vlog.p)
Figure 15: The same visualization as in Figure 14 for the GMM example in Figure 13 and 21, E (,888). We compute the optimal drift for DIS similar to Section A.10.
42


Reviews:
[{"review_id": "ukzoKuBeW6", "reviewer": "Reviewer_ziUJ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n\n- I find interesting the perspective of using different (than the KL) divergences between the path measures of the reference and the constrained diffusion process for crafting alternative algorithms for training. Their experiments suggest that using the log-variance divergence leads to more stable estimates of the log evidence with less valiance.\n\n\n- I find the connections the authors provide with denoising score matching and half-bridge training for solving the Schroedinger bridge problem valuable. However, these parts are probably more fitting for the appendix (especially the former one). In this version they distract the reader from the main message/contribution of the paper.\n\n\n**Weaknesses:**\n\n- While Pavon's work in [7]  already makes the association between diffusion based generative modelling and stochastic control, I acknowledge that the authors provide here a more well-presented and comprehensive version of the association. However they should should acknowledge this existing work of Pavon. They could emphasize their contributions by highlighting their additional contributions.\n\n- On the related work the authors do not make any mention to the extensive optimal stochastic control literature that is a crucial part of their framework, i.e., [1] - [6]. For example [4] uses a forward and time-reversed diffusion process to compute optimal controls for SDEs starting from Dirac delta functions. Their derivation follows similar arguments with minimising the KL divergence between a reference process and the controlled process.\n\n- The association between minimisation of the lower bound for diffusion processes and optimal control has been already proposed for inference of diffusion processes with variational inference in [8] and [9], where in [9] the authors consider a forward and time-reversed process akin to diffusion-based generative modeling to compute the optimal controls that minimise the ELBO.\n\n- I find the writing and presentation style of the numerical experiments a bit confusing. I recommend the authors to rewrite this section to provide initially with a summary and a motivation what they attempt to demonstrate with each experiment and then presentation and walkthrough the results. In the current version, the section 4.1 Examples fits better to an Appendix, while in the results section of the main text I would mention the name of the models employed, why these were selected and then list the results for each.\n\n- The authors do not mention relevant literature on sampling unnormalised densities, i.e. energy based models.\n\n\n[1] Kappen, Hilbert Johan, and Hans Christian Ruiz. \"Adaptive importance sampling for control and inference.\" Journal of Statistical Physics 162 (2016): 1244-1266.\n\n[2] Thijssen, Sep, and H. J. Kappen. \"Path integral control and state-dependent feedback.\" Physical Review E 91.3 (2015): 032104.\n\n[3] Kappen, Hilbert J., Vicenç Gómez, and Manfred Opper. \"Optimal control as a graphical model inference problem.\" Machine learning 87 (2012): 159-182.\n\n[4] Maoutsa, Dimitra, and Manfred Opper. \"Deterministic particle flows for constraining stochastic nonlinear systems.\" Physical Review Research 4.4 (2022): 043035.\n\n[5] Zhang, Wei, et al. \"Applications of the cross-entropy method to importance sampling and optimal control of diffusions.\" SIAM Journal on Scientific Computing 36.6 (2014): A2654-A2672.\n\n[6] Hartmann, Carsten, and Christof Schütte. \"Efficient rare event simulation by optimal nonequilibrium forcing.\" Journal of Statistical Mechanics: Theory and Experiment 2012.11 (2012): P11004.\n\n[7] Pavon, Michele. \"On Local Entropy, Stochastic Control, and Deep Neural Networks.\" IEEE Control Systems Letters 7 (2022): 437-441.\n\n[8] Opper, Manfred. \"Variational inference for stochastic differential equations.\" Annalen der Physik 531.3 (2019): 1800233.\n\n[9] Maoutsa, Dimitra. \"Geometric constraints improve inference of sparsely observed stochastic dynamics.\" International Conference on Learning Representations (ICLR) 2023 - Workshop on Physics for Machine Learning (Physics4ML) (2023).\n\n[10] Zaki, Nikolai, Théo Galy-Fajou, and Manfred Opper. \"Evidence Estimation by Kullback-Leibler Integration for Flow-Based Methods.\" Third Symposium on Advances in Approximate Bayesian Inference. 2020.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors present an optimal control perspective on diffusion based generative modelling. They formulate the minimisation of the evidence lower bound within a stochastic control formalism. By appealing to the Girsanov theorem for the change of measure they compute  the optimal change of measure through the Doob-h transform. This allows to express the evidence lower bound (ELBO) of the generative model in terms of the control costs of the stochastic control problem. \nIn the second part of the paper they proceed with providing an estimator (time-reversed diffusion sampler - DIS) for sampling from unnormalised probability densities. They conduct numerical comparisons between this proposed approach and the recent path integral sampler developed by Zhang & Chen in 2022.\n\nThe paper provides an interesting perspective between diffusion-based generative modeling and stochastic control. While it doesn't unveil anything entirely new, it sheds light on a connection that, as far as I know, hadn't been explicitly articulated before, except for a mention in [7].", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Please include in the relative work some overview on optimal stochastic control and Kulback-Leibler control for stochastic systems, that uses the KL divergence from a reference path measure to compute the controls, similar to the presented approach, and also the other points mentioned in the weaknesses.\n\n- In Figure 4, both PIS and DIS seem to perform particularly bad in approximationg the log evidence for the double well system. Can you provide some intuition behind this. \nAlso how would the method compare with other methods for computing the evidence, like in [10]?", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "wZnfgVYoQU", "reviewer": "Reviewer_A8EE", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n \n- This paper bridges the gap between different fields (generative modelling, approximate Bayesian inference, control theory) to offer a perspective on diffusion models which, to the best of my knowledge, is novel.\n\n- The main contribution is rigorously justified (it is obtained by manipulating PDEs and applying results from control theory) and can be applied to develop practical, theoretically-grounded methodologies, as illustrated with the sampling procedure in Section 3. In that sense, I believe this is relevant for the machine learning community and can foster future interesting work. \n\n**Weaknesses**\n\n- Although the paper is well written overall, it suffers from a lack of clarity in my opinion. The paper is quite dense and high-level/more intuitive explanations on the contributions are missing. For instance, the main results are provided in Sections 2.2 and 2.3, and the authors provide detailed explanations on their mathematical technicality, which I appreciate (e.g., see paragraph below Corollary 2.5) -- but I am missing the big picture, which is: why does it help to view a diffusion model as a controlled SDE? What are the \"valuable new insights\" (Conclusion) on diffusion models provided by control theory? The fact that the authors then parameterize the control $u$ as a neural network, which are black-box algorithms, does not help clarify.\n\n- The experiments are a bit disappointing: given that this paper strongly relies on score-based generative models, I would have expected experiments similar to the ones in (Song et al., 2020b) to be conducted. It would be interesting to understand why the authors have chosen to not explore the consequences of their optimal control perspective on classical diffusion experiments, such as image generation. Besides, regarding the experiments in Section 4, the results are indeed encouraging but no insight is given on why DIS outperforms PIS.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The main contribution of this paper is a connection between diffusion probabilistic models (DPMs) and stochastic optimal control: the authors prove that score-based generative modelling with stochastic differential equations (SDEs) (Song et al., 2020b) amounts to solving a Hamilton-Jacobi-Bellman (HJB) equation (Lemma 2.3). This enables drawing the connection with stochastic optimal control: the HJB equation governing score-based generative models is linked to a controlled SDE (Theorem 2.4). While this result is shown by applying a classical theorem from control theory (the verification theorem), it offers a new perspective on score-based generative models, which are usually associated with the Fokker-Planck equation.  \n\nThe authors then leverage this connection to derive a lower bound on the negative log-likelihood of the considered generative model (Corollary 2.5). This bound is interpreted as an evidence lower bound and justifies another perspective of score-based generative modeling as a variational inference problem on the space of continuous trajectories (Proposition 2.6). Based on this result, the authors finally develop a new procedure called \"time-reversed diffusion sampler\" to sample from distributions with intractable normalizing constant (Section 3) and illustrate its relevance on toy experiments (Section 4).", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Given my points above, I would encourage the authors to clarify the insights offered by their stochastic optimal control approach for diffusion models. In particular,\n\n- There is enough explanations on the mathematical derivations in my opinion: what is missing is intuition. What does this optimal control perspective on diffusion model bring concretely? Why does the proposed sampling procedure outperform PIS? Is the perspective useful for classical experiments for diffusion models, e.g., high-dimensional image generation?   \n\n- The second equality of eq. (5) is crucial for the rest of the paper, but is not proved anywhere. I confirm it is correct (I checked it using identities in Appendix A.2) but I strongly encourage the authors to give the detailed derivations in their paper. \n\n- The authors provide some background on diffusion models in the main text, but not on stochastic optimal control (this is provided in Appendix A.6). Moving some of these elements from the appendix to the main text might help the reader understand why it is useful to use controlled SDEs for diffusion models.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "QSR3Gx8FOi", "reviewer": "Reviewer_JtqT", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: \n- The connections established in the paper may allow for techniques from control theory to be applied to generative models, and vice-versa.\n- The analysis in the paper is quite rigorous and appears correct to the best of my knowledge.\n- The paper offers an interesting approach to sampling unnormalized densities, which is a central task in many areas of applied mathematics. Their method seems to perform well on several multi-modal distributions, even relative to state of the art sampling algorithms.\n\nWeaknesses: \n\n- I am skeptical of the contributions in Section 2. For instance, I don’t quite see the significance in the fact that the ELBO can be derived from control theory (Theorem 2.4, Corollary 2.5). In addition, the connection between the objective of SDE-based generative models and KL divergence between path measures has been well-studied through the lens of Girsanov’s theorem, for instance in [1] and [2]. It is therefore unclear why the insight from Proposition 2.6 is novel or interesting in the generative modeling context (aside from the fact that it is later used to develop the sampling algorithm).\n- The exposition is confusing/difficult to follow at times. It feels like sections 2 and 3 often jump abruptly between several ideas without properly connecting them.\n- The DIS sampling algorithm seems quite similar to the path integral sampler proposed in the work of Zhang and Chen [3]. If I understand both correctly, the main differences are in the initial distributions and running/terminal costs, but the connection between sampling and stochastic control seems to have already been exploited. While I appreciate the authors for thoroughly explaining the relationship between the two methods in Section 3.1, it does take away from the novelty of the algorithm.\n\n[1] Yang Song, Conor Durkan, Iain Murray, Stefano Ermon. Maximum likelihood training of score-based diffusion models. In Advances in Neural Information Processing Systems 34, 2021.\n[2] Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim and Anru R Zhang. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. In International Conference on Learning Representations, 2023.\n[3] Qinsheng Zhang and Yongxin Chen. Path integral sampler: a stochastic control approach for sampling. In International Conference on Learning Representations, 2022a.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors study a connection between stochastic optimal control theory and generative models based on SDEs, such as score-based diffusion models. It is shown that the evidence lower bound (ELBO) for SDE-based models follows from a standard result in control theory, and that the expected variational gap of ELBO can be interpreted in terms of the KL divergence between relevant path measures. The authors employ their results to develop a novel sampling algorithm for applications when the target density itself is known up to normalization, but one does not have access to samples. They demonstrate the utility of their sampling algorithm through several numerical experiments.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please make changes by addressing the comments on the weaknesses above and some further questions/comments below. \n\n- I am interested to know if you can quantify the sampling error of the DIS algorithm in terms of the estimation error of the control term (i.e., if an analogue of Theorem 2 in [1] holds for your sampling algorithm). Is this something you have considered?\n- It is worth to mention that the DIS and PIS algorithms outperform baseline sampling algorithms after having been trained, but that training the models incurs additional run-time. In addition, you might mention any difficulties you encountered during training (and how you dealt with them) in case others want to use your algorithm.\n- There have been several recent works (possibly after your submission) using control theory to study generative models, so it would be good to update the related work section. For instance, in [2], the authors expand on the connection between diffusion models and HJB equations using the framework of mean-field games.\n- I think a table of contents would improve the readability of the appendix.\n\n[1] Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim and Anru R Zhang. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. In International Conference on Learning Representations, 2023.\n[2] Benjamin Zhang and Markos Katsoulakis. A mean-field games laboratory for generative models. Arxiv preprint Arxiv:2034.13534", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "oeeBzLs05J", "reviewer": "Reviewer_QKvk", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\nAlthough the paper is dense and contains a lot of information, I find it well structured and well written. The\nreading is not easy, but I acknowledge that the authors put effort to make it comprehensible. In particular, the details\nprovided in the appendices are useful to understand the paper.\n\nI think that this paper offers an interesting and strongly grounded perspective on diffusion-based generative modeling. In\nparticular, the connection with score-based diffusion methods is well explained, and the links between DIS and the related\nalgorithms is well highlighted.\n\nBesides this, the authors provide lots of details on their numerical implementation, which is key for reproducibility, and\npresent convincing results in practice.\n\nWeaknesses: \n\nAbout the sampling method for un-normalized densities: although the authors claim that DIS may be\napplied to any general forward process (as long as one can sample from the terminal distribution), the authors only consider\nthe OU-VP scheme in their experiments. In this case, DIS reduces DDS (see Appendix A.10.1). Then, it is not clear how\nDIS is a novel contribution in practice.\n\nIn the numerical experiments, a comparison with efficient MCMC methods such as Sequential Monte Carlo is missing.\nAs showed by [3], these methods may give better results than PIS. Moreover, the highest dimension\nconsidered is not very large (d = 50), compared to previous works.\n\n[3] Denoising diffusion samplers. Vargas et al, 2023.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "In this paper, the authors establish a connection between generative models built upon general Stochastic Differential Equations (SDEs) and Stochastic Optimal Control (SOC) theory. Given a target distribution $\\pi\\in \\mathcal{P}(\\mathbb{R}^d)$ and a horizon time $T>0$, the underlying approach consists in using SOC theory to find a drift $\\mu:[0,T]\\times\\mathbb{R}^d \\to \\mathbb{R}^d $, a volatility $\\sigma:[0,T] \\to \\mathbb{R}_+$ and a prior distribution $\\pi^0\\in \\mathcal{P}(\\mathbb{R}^d)$ such that the SDE given by $dX_t=\\mu(T-t, X_t)dt+ \\sigma(T-t)dB_t$ with $X_0\\sim\\pi^0$ verifies $X_T\\sim \\pi$. \n\nTheir method consists in first deriving a Hamilton-Jacobi-Bellman equation for the evolution of the log-densities of the process $(X_{T-t})$. This leads to the formulation of a SOC problem over a specific class of control functions (roughly, additive controls),\nwhich amounts to maximising an evidence lower bound (ELBO) to $\\log \\pi$. Interestingly, this ELBO connects with the framework of score-based diffusion models [1],[4], and its variational gap can be rewritten as a KL divergence between path measures.\n\nThis SOC approach leads to a novel diffusion-based method, referred to as time-reversed diffusion sampler (DIS), to sample from un-normalized densities. In this setting, traditional diffusion-based optimisation tools do not work since one does not have access to samples from $\\pi$. In practice, theDIS procedure relies on a gradient-based optimization scheme along sample trajectories, similarly to Path integral Sampler (PIS) [5] and Denoising Diffusion Sampler (DDS) [3].\n\nThe authors propose numerical experiments where they apply DIS with reference process chosen as the Ornstein Uhlenbeck process (Variance Preserving scheme in [2]). For this setting, their algorithm recovers the concurrent algorithm from [3]. They compare their method with PIS, and provide better numerical results in various settings.\n\n[1] A variational perspective on diffusion-based generative models and score matching. Huang et al., 2021.\n\n[2] Score-based generative modeling through stochastic differential equations. Song et al, 2020.\n\n[3] Denoising diffusion samplers. Vargas et al, 2023.\n\n[4] A connection between score matching and denoising autoencoders. Vincent, 2011.\n\n[5]  Path integral sampler: a stochastic control approach for sampling. Zhang et al.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I recommend the authors to introduce a numerical comparison with (i) the implementation from\n[3], as it may reveal differences in the numerical scheme, and (ii) standard MCMC schemes. Moreover, I\nsuggest the authors to run additional experiments in a higher dimensional synthetic setting: for example, Gaussian mixtures\nwith d = 100, 200.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: All the reviewers agree that the optimal control point of view on diffusion models and related algorithms is valuable for the community. The interaction between diffusion models, sampling and the stochastic optimal control frameworks are very promising. The authors have addressed the main remarks of the reviewers in a convincing manner. I therefore recommend the manuscript for acceptance.  However, I would have liked to see a more detailed comparison between DDS [1] and DIS. In particular, it would be good to highlight where DIS shines and DDS fails.

[1] Denoising diffusion samplers. Vargas et al, 2023.
Decision: Accept (poster)