Likelihood Training of Schrödinger Bridge using Forward-Backward SDEs Theory
Authors: Tianrong Chen
Guan-Horng Liu
Evangelos Theodorou
Venue: ICLR 2022 Poster (2022)
Keywords: Schrödinger Bridge
score-based generative model
optimal transport
forward-backward stochastic differential equations
stochastic optimal control

Abstract: Schrödinger Bridge (SB) is an entropy-regularized optimal transport problem that has received increasing attention in deep generative modeling for its mathematical flexibility compared to the Scored-based Generative Model (SGM). However, it remains unclear whether the optimization principle of SB relates to the modern training of deep generative models, which often rely on constructing log-likelihood objectives.This raises questions on the suitability of SB models as a principled alternative for generative applications. In this work, we present a novel computational framework for likelihood training of SB models grounded on Forward-Backward Stochastic Differential Equations Theory – a mathematical methodology appeared in stochastic optimal control that transforms the optimality condition of SB into a set of SDEs. Crucially, these SDEs can be used to construct the likelihood objectives for SB that, surprisingly, generalizes the ones for SGM as special cases. This leads to a new optimization principle that inherits the same SB optimality yet without losing applications of modern generative training techniques, and we show that the resulting training algorithm achieves comparable results on generating realistic images on MNIST, CelebA, and CIFAR10. Our code is available at https://github.com/ghliu/SB-FBSDE.

Published as a conference paper at ICLR 2022
LIKELIHOOD TRAINING OF SCHRÖDINGER BRIDGE USING FORWARD-BACKWARD SDES THEORY
Tianrong Chen", Guan-Horng Liu", Evangelos A. Theodorou
Georgia Institute of Technology, USA (t (tianrong. chen, ghliu, evangelos svangelosstheeddroou @gatechh edu
ABSTRACT
Schrödinger Bridge (SB) is an entropy-regularin optimal transport problem that has received increasing attention deep generative modeling for its mathematical flexibility compared to the Scored-base Generative Model (SGM). However, it remains unclear whether the optimization principle of SB relates to the modern training of deep generative models, which often rely constructing log-likelihood objectives. This raises questions on the suitability SB models as a principled alter- native for generative applications. In this work, present novel computational framework for likelihood training of SB models grounded on Berward-Backweer Stochastic Differential Equations Theory a mathematical methodology appeared in stochastic optimal control that transforms the optimality condition of SB into of SDEs. Crucially, these SDEs can be used construct the likelihood objectives for SB that, surprisingly, generalizes the ones for SGM as special cases. This leads to a new optimization principle that inherits the same SB optimality yet without losing applications of modern generative training techniques, and we show that the resulting training algorithm achieves comparable results on generat- ing realistic images on MNIST, CelebA, and CIFARIO. Our code is available at https: htt///gtttuaa/ com/ ghliu/ /SB -B- FBSDE.
INTRODUCTION
Score-based Generative Model (SGM; Song et al. (2020)) is emerg- ing generative model class that has achieved remarkable results in synthesizing high-fidelity data (Song & Ermon, 2020; Kong et al., 2020a;b). Like many deep generative models, SGM seeks to find nonlinear functions that transform simple distributions (typically Gaus- sian) into complex, often intractable, data distributions. In SGM, this done by first diffusing data noise through stochastic differential equation (SDE); then learning to reverse this diffusion process by regressing network the score function (i.e. the gradient of the log probability density) at each time step (Hyvärinen & Dayan, 2005). This reversed process thereby defines the generation (see Fig. 1).
Date-to-noise (diffusion) SDE
Data
Noise
Noise-to-data (generation) SDE
SGM
Given
Learned
Leamed
Leamed
Figure 1: Both Score-based Generative Model (SGM) and Schrödinger Bridge (SB) transform between two distributions. While SGM requires pre-specifying the data-to-noise diffusion, SB instead learns the process.
Despite its empirical successes, SGM admits few limitations. First, the diffusion process has to obey a simple form (e.g. linear degenerate drift) in order to compute the analytic score function for the regression purpose. Secondly, the diffusion process needs to run to sufficiently large time steps so that the end distribution is approximate Gaussian (Kong& & Ping, 2021). ) For these reasons, SGM often takes a notoriously long time in generating data [Jolicocur-Martii compared to e.g. GANs or flow-based models (Ping et al., 2020; Karras et al., 2020b).
et al., 2021), thereby limiting their practical usages
In the attempt to lift these restrictions, a line of recent works inspired by Schrödinger Bridge (SB; Schrödinger (1932)) has been proposed (De Bortoli et al., 2021; Wang et al., 2021; Vargas et al., 2021). optimal transport problem seeks two optimal policies that transform back-and-forth between two arbitrary distributions in finite horizon. The similarity between the
SB as entropy-regularicn
"Equal contribution. Order determined by coin flip. See Author Contributions section.

---

Published as a conference paper at ICLR 2022
two problems (i.e. both involve transforming distributions) is evident, and the additional flexibility from SB is also attractive. To enable SB- 3B-inspired generative training, however, previous works require either ad-hoo tho multi-stage optimization retreat to traditional SB algorithms, e.g. Iterative Proportional Fitring (IPF; Kullback (1968)). The underlying relation between the optimization principle of SB and modern generative training, in particular SGM, remains relatively unexplored, despite their intimately related problem formulations. More importantly, with the recent connection between SGM and log- lolleeelioood computation (Song et al., 2021), it is crucial to explore whether there exists an alternative way of training SB that better respects, or perhaps generalizes, modern training of SGM, so as to solidify the suitability of SB as a principled generative model.
In this work. we present a fundamental connection between solving SB and training SGM. The difficulty arises immediately as one notices that the optimality condition of SB and the likelihood objective of SGM are represented by merely two distinct mathematical objects. While the former characterized by two coupled partial differential equations (PDEs) (Léonard, 2013), the latter integrates over notably complex SDE that resembles neither its diffusion nor reversed process (Song et al., 2021). Nevertheless, inspired by the recent advance on understanding deep learning through the optimal control perspective (Li & Hao, 2018; Liu et al., 2021a;b), we show that Forward Backward SDEs a mathematical methodology appearing in stochastic optimal control for solving nonlinear PDEs (Han al., 2018) paves an elegant way connect the two objectives. The implication of our findings is nontrivial: It yields an exact log-likelihood expression of SB that precisely generalizes the one of SGM (Song et al., 2021) to fully nonlinear diffusion, thereby providing novel theoretical connections between the two model classes. Algorithmically. our framework suggests rich training procedures that resemble the joint optimization for diffusion flow-based models (Zhang & Chen, 2021) or more traditional IPF approaches (Kullback, 1968; De Bortoli et al., 2021). This allows one to marry the best of both worlds by improving the SB training with e. a SGM- (GMM-ispired Langevin corrector (Song& & Ermon, 2019). The resulting method, SB-FBSDE, generates encouraging images on MNIST, CelebA, and CIFARIO and outperforms prior optimal transport models by a large margin.
Forward-Backward
Our method differs from the concurrent SB methods (De Bortoli et al., 2021; Vargas et al., 2021) in various aspects. First, while both prior methods rely on solving SB with mean-matching regression, our SB-FBSDE instead utilizes a divergence-bastt -based objectives (see $3.2). Secondly, neither of the prior methods focuses on liglliieldooo training, which is the key finding in SB-FBSDE to bridge connections to SGM and adopt modern training improvements. Indeed, due the difference in the underlying SDE classes, their connections to SGM can only be made after time discretization by carefully choosing each step size (De Bortoli al., 2021). In contrast, our theoretical connection is derived readily in continuous-time; hence unaffected by the choice of numerical discretization.
In summary, we present the following contributions.
We present a novel computational framework, grounded on Forward-Backwark SDEs theory, for computing the log-lik lileeikiihood objectives of Schrödinger Bridge (SB) and solidifying their theoretical connections to Score-based Generative Model (SGM).
Our framework suggests a new training principle that retains the mathematical flexibility from SB while enjoying advanced techniques from the modern generative training of SGM.
We show that the resulting method named SB-FBSDE outperforms previous optimal transport- inspired baselines on synthesizing high-fidelity images and is comparable to other existing models.
Notation. We denote pE x))) as the marginal density driven by some SDE process X(t) X, until the time step E [0, 7]. The time direction is aligned throughout this article such that po and pr respectively correspond to the data and prior distributions. The gradient, divergence, and Hessian of
a function of fxx), where x e R"," will be denoted as Vaf E IR"", Va f E R, and 22 E Raxa
2
PRELIMINARIES
2.1
SCORE-BASED GENERATIVE MODEL (SGM)
Given data point X, E R" sampled from an unknown data distribution Pdata SGM first progressively diffuses the data towards random noise with the following forward SDE:
dX, ft, X,)dt + g(f)dW,, 8, Xo Pdata:
We adopt the recent advance in SB theory (Caluya & Halder, 2021) that extends classical SB models (used in prior works) the exact SDE class appearing in SGM. See Appendices and C. for more details.
2

---

Published as a conference paper at ICLR 2022
where f(.tt) : R" R", g(t) R, and We E R" are the drift, diffusion, and standard Wiener process. Typically, g()) is some monotonically increasing function such that for sufficienlly large time steps, we have ()) x Pprior resemble some prior distribution (e.g, Gaussian) at the terminal horizon T. Reversing (1) yields another SDE² that traverses backward in time (Anderson, 1982):
dx,- [f g2 V logp: " )dt + g dWe, XT -
where Pt corresponds to the marginal density of SDE (1) at time t. and V logp," is known as the score function. These two stochastic processes are equivalent in the sense that their marginal densities are equal to each other throughout / [0, 7): in other words, P. =
When the drift f is of simple structure, for instance linear (Ho al., 2020) or simply degenerate (Song Ermon, 2019), the conditional score function v log P: (1) (X, X - zn) = v log Pt/a, admits an analytic solution at any time Hence, SGM proposes train a parameterized score network s(t, æ;; 0) by regressing its outputs to the ground-truth values, i.e. [[ (tt, X,; 0) - V log Ptja.. 112), where the expectation is taken over the SDE (1). In practice, x(() is some hand-designed weighting function that largely affects the performance. Recent works (Song et al., 2021; Huang et al., 2021) have shown
that the log-likelihood of SGM, despite being complex, can be
as follows:
SGM (300) 2 2CsOM (200 0) = E [log pT (XT)] L.
+ V (g's; - -5) -f))aa dl, (3)
log Po
E
2 g
= E [log pT (XT))
E
|| V log Pt|mull
1|
- 1 log moll 12 Vo f dt,
where S, (t,, x;; 8) and the expectation is taken over the SDE (1) given a data point X. 2p. This objective (3) suggests a principled choice of A(t) := g(t)?. After training, SGM simply substitutes the score function with the learned score network s(1, æx; 0) to generate data from Ppriors
g2 stt, x:; dt + dW., XT dx, = if
Pprior.
(4)
(1) It is important to notice that Porior needs not equal pr in practice, and the approximation is close only through a careful design of (1). Notably, designing the diffusion gt) can particularly problematic as it affects both the approximation PT s Pprior and the training via the weighting A(t); hence can lead to unstable training (Nichol & Dhariwal, 2021). In contrast, Schrödinger Bridge considers more flexible framework for designing the forward diffusion that requires minimal manipulation.
(1)
2.2 SCHRÖDINGER BRIDGE (SB)
Following the dynamic expression of SB (Pavon & Wakolbinger, 1991; Dai Pra, 1991), consider
min QEP(Paa Pasio))
(KL(Q P),
(5)
where Q E P(Patta Ppricr) belongs to set of path measure with Paata and Pprior as its marginal densitics at t = 0 and T. On the other hand, IP denotes a reference measure, which we will set to the path measure of (1) for later convenience. The optimality condition to (5) is characterized by two PDEs that are coupled through their boundary conditions. We summarize the related result below.
Theorem 1 (SB optimality; Chen al. (2021); Pavon & Wakolbinger (1991); Caluya & Halder (2021)). Let (1, x) and (t, æ²) be the solutions to the following PDEs:
0 V / f VV
s.t. VOO,)) =Paata V (T.))((TT.)) = Pprior
(6)
00
= Va TT
Then, the solution the optimization (5) can be expressed by the path measure of the following forward (7a), or equivalently backward (7b), SDE:
dX, = [f - g" Va log V (t, X,)]dt + dW,, X, ~ Poatus
(7a)
dx, [f g" V log V (t. ) + g dW., Xy Ppriors
(7b)
where Vae log (t, X,) and V log (t, Xx) are the forimal and backward drifts for SB.
Heereafter, we will sometimes drop f ft, X.) and g = g({) for brevity.
3

---

Published as a conference paper at ICLR 2022
Score-based Generative Model (SGM)
Schrödinger Bridge (SB)
dx,
/
dt g gdW,
Stochastic Optimal Control Perspective
dx,-
[ V/''"V 4 of gdW,
Data
Noise
Data
X.
Noise
Control-affine SDEs
dx.-
[s V
V.logg. gdW,
dx.-
f -ig Vzlog V dt + gdW,
Log-likelihood
FBSDEs Theory (Theorem & 4)
Objective (3)
PDEs Optimality (6)
Schematic diagram of the our stochastic optimal control interpretation, and how it connects the objective of SGM (3) and optimality of SB (6) through
Figure 2:
SDEs theory.
Similar to the forward/backward processes in SGM, the stochastic processes SB in (7a) and (7b) are also equivalent in the sense that Vt [0, T] Pt (T)) = ((bb) = S fact, its marginal density obeys factorization principle: pr (X.) = (1, X (1, X,).
To construct the generative pipeline from (7b), one requires solving the PDEs in (6) to obtain T. Unfortunately, these PDEs are hard to solve even for low-dimensional systems (Renardy & Rogers, 2006); let alone for generative applications. Indeed, previous works either have to replace the original Schrödinger Bridge (Peata = Pprior) with multiple stages, Pdata = Pmiddle = Pprior. so that each segment admits analytic solution (Wang et al., 2021), or consider the following half-bridge (Pdaa f Peiio VS. Pdata + Pprice) optimization (De Bortoli et al., 2021; Vargas et al., 2021).
Q(1)
arg min DKL(Q Q Q(),,
n(0)
arg min DKL ( || (1)))
=
(EP((-Prur)
CEP(Pan::)
which can be solved with IPF algorithm (Kullback, 1968) starting from Q()) == P. In the following section, we will present a scalable computational framework for solving the optimality PDEs in (6) and show that it paves an elegant way connecting the optimality principle of SB (6) to the parameterized log-likeliiee likeliihoood of SGM (3).
3 APPROACH
We motivate our approach starting from some theoreec observation (see Fig. 2). Notice that both SGM and SB consist of forward and backward SDEs with similar structures. From the stochastic control perspective, these SDEs belong to the class of control-affine SDEs with additive noise:
dx, A(t, X, )d (I, X,) X,)dt + C(t) dW,-
It is clear that the control-affine SDE (8) includes all SDEs (1,2.4.7) appearing $2 by considering (A, B, C') == (f, I, g) and different interpretations of the control variables u(t, X.). This implies that the optimization processes of both SGM and SB can be aligned through the lens of stochastic optimal control (SOC). Indeed, both problems can be interpreted as seeking some time-varying control policy, either the score function V Ptjmg in SGM or V, log T in SB, that minimizes some objectives, (3) vs. (5), while subjected to some control-affine SDEs, (1,2) vg. (7). In what follows, will show that specific mathematical methodology in nonlinear SOC literature called Backaardd SDEs theory (FBSDEs; see Ma et al. (1999)) links the optimality condition of SB (6) to the ligeeiill objectives of SGM (3). All proofs are left Appendix B.
3.1 FORWARD-BACKWAN SDES (FBSDES) REPRESENTATION FOR SB
The theory of FBSDEs establishes innate connection between different classes of PDEs and -awkard-bckwward SDEs. Below we introduce the following connection related to our problem.
Lemma (Nonlinear Feynman-Kac; Exarchos Theodorou (2018)). Consider the coupled SDEs
dxx (1, X, )al G(I, X,),,,
Xo =F0
(9a)
dY, = -ht, X,, Y, Zt)dt + Z(1, T.W, Y, = (({)
(9b)
'I Lemmm 2 can be viewed as the nonlinear extension of the celchrated Feynman- Kac formula (Karatzas & Shreve, 2012), which characterizes the connection between linear PDEs and forward SDEs.
4

---

Published as a conference paper at ICLR 2022
where the functions f. G. h, and ip satisfy proper regularity conditions that there exists a pair of unique strong solutions satisfying (9). Now, consider the following second-ore order parabolic PDE and suppose v(t, x) is once continuously differentiable in t and twice in i.c. C....
de Öt + Tr G + litt, x. U, G ('v))=== =0, v(T,z)) = ((x),
(10)
then the solution to (9) coincides with the solution to (10) along paths generated by the forward SDE (9a) almost surely, i.e., the following stochastic representation (known as the nonlinear Feynman- Kac relation) is valid:
(t, X.) YY;
and G(t, (T..X)))Vaell (t, Ze.
(11)
Lemma 2 states that solutions a certain class of nonlinear (via the function h in (10)) PDEs can be represented by solutions to a set of horward-bacwwand SDES (9) through the transformation (11), and this relation can be extended to the viscosity case (Pardoux & Peng (1992); see also Appendix B). Note that Y, is the solution to the backward SDE (9b) whose randomness is driven by the forward SDE (9a). Indeed, it is clear from (11) that Y, (hence also Z,) is a time-varying function of X, Since the appearing in the nonlinear Feynman-Kan relation (11) takes the random vector X, as its argument, ett X,) shall also be understood as a random variable. Finally, it is known that the original (deterministic) PDE solution vtt,, æ) can be recovered by taking conditional expectation, i.e.
(12) Since it is often computationally favorable to solve SDEs rather than PDEs, Lemma 2 has been widely used as scalable method for solving high-dimensional PDEs (Han et al., 2018; Pereira et al., 2019). Take SOC applications for instance. their PDE optimality condition can be characterized by (11) under proper conditions, with the optimal control given in the form of Z, Hence, one can adopt Lemma 2 to solve the underlying FBSDEs, rather than the original PDE optimality, for the optimal control. Despite seemingly attractive, whether these principles can be extended to SB, whose optimality conditions given by two coupled PDEs in (6), remains unclear. Below we derive similar FBSDEs representation for SB.
v(4,, a:) = (xa)) Y, X, = xx] and G(I, T VVa
Va 0(, x)) = Ex [2, X, = =x]-
Theorem (FBSDEs SB optimality (6)). Consider the following set of coupled SDEs,
dx; = (f + gZze) dt + gdW,
(13a)
dY, 1
2 Z +Z.dW,
(13b)
1 -, V. -19Z)-1) Z, dt + Z, d.W.
dÝ,
(13c)
where f and g satisfy the same regularity conditions in Lemma 2 (see Foomote 4), and the boundary conditions are given by X(0) 3g and YT Y = log Pprior (XT). Suppose V, v E C then the nonlinear Feynman-Kae Kac relations between the FBSDEs ( (13) and PDEs (6) are given by
Y, = ( x,) = log (t, X.),
Z: Ztt, X,) vv
- log (t, 2, 2(t X.) v. log (t, X.).
(14)
Ý. (t,,x,) = log Xt),,,
Furthermore, (Y,, Y.) obey the following relation:
Y, +,+Y, = log SB (X.).
The FBSDEs for SB (13) share a similar (13b,13c) respectively represent the forward and backward SDEs. One can verify that the forward SDE (13a) coincides with the optimal forward SDE (7a) with the substitution Z, vVa log W. In other words, these FBSDEs provide a local representation of log v and log v evaluated on the optimal path governed by (7a). Since Z, and Z can be understood the forward/backwart policies, in similar spirit of policy-based methods (Pereira et al., 2020; Schulman et al., 2015), that guide the SDE processes of SB, they sufficiently characterize the SB model. Hence, our next step is derive proper training objective to optimize these policies.
structure as in (9), where (13a) and
Yong & Zhou (1 (1999); Kobylanski (2000) require f. G. h, and yp be continuous, and G to be uniformly L Lipschitz in æ, and h to satisfy quadratic growth condition in
5

---

Published as a conference paper at ICLR 2022
3.2 LOG-LIKELIHOOU COMPUTATION OF SB
Theorem 3 has important implication: It suggests that given path sampled from the forward SDE (13a), the solutions to the backward SDEs (13b,13c) at t = 0 provide an unbiased estimation of the log-liielihood of the data point Fo. i.e. E Yo + Y Xo- =0 log pd (xo) log Pdata (20). where X, is sampled from (13a). We now state our main result, which makes this observation formal: Theorem ligeliiil of SB model). Given the solution satisfying the FBSDE system in (13), the log-tikelihood of the SB model (Ze, Zi), at data point Vo, can be expressed as
{
log log (xoo) = E [log pr (XT)]
E
Zel
Z, - v log SB + 1
2
2
[1 log SB -2.12- -V f dt (15)
[ pT (XT)]
E
Z 12
2.11 + V (gZ, - f) + 2,,,, Z, dt, (16)
2
where the expectation is taken over the forward SDE (13a) with the initial condition Xo= 20-
Similar to (3), Theorem suggests a parameterized lower bound to the log-likelihoods. i.e. SR
log Po (zo) 2 SSB (Iu; 0 6) where LsB (Zu; 6, Φ) shares the same expression in (16) except that Z, (t, x; 0) and Ze Z(t, xx; Φ) are approximated with some parameterized models (e.g. DNNs). Note that v log is intractable in practice for any nontrivial (Z,, Z.). Hence, we use the divergence- hiveggenee-based objective in (16) as training objective of both policies.
Connection to basee models. Recall Fig. 2 and compare the parameterized log-likelihoode of SB (16) to SGM (3); one can verify that LsB collapses to LsOM when (Z., 2.) == (0, 8:). From the SB perspective, this occurs only when PT (1) Pprior- Since no effort is required in the forward process to reach Pprior. the optimal forward control Z, by definition, degenerates; thereby making the backward control Z, collapses to the score function. However, in any case when pr + Pprior- for instance when the diffusion SDEs are improperly designed, the forward policy Z, steers the diffusion process back to Pprior. while its backward counterpart Z, compensates the reversed process accordingly. From this view, SB alleviates the problematic design in SGM by enlarging the class of diffusion processes to accept nonlinear drifts and providing an optimization principle on learning these processes, Moreover, Theorem 4 generalizes the ligeliiildood training from SGM to SB.
Connection to flow-based models. Interestingly, the log-likelihood computation in Theorem 4, where we use a path (Xrrrejo sampled from a data point Xo to parameterize its log-likelihood, resembles modern training of (deterministic) flow-based models (Grathwohl et al., 201 8), which have recently been shown admit a close relation to SGM (Song et al., 2020; Gong & Li, 2021). The connection is built on the concept of probability flow which suggests that the marginal density of an SDE can be evaluated through an ordinary differential equation (ODE). Below, we provide a similar flow representation for SB, further strengthening their connection to modern generative models.
Corollary 5 (Probability flow for SB). The following ODE characterizes the probability flow of the optimal processes of SB (7) in the sense that Vt, (17) = (7) = pr.
dx,-
f (t, 29 (Z Xx) + Z(t, X()) dt
(17)
One can verify (see Remark 10 in (B) that computing the log-likelihood of this ODE model (17) using flow-based training techniques indeed recovers the training objective of SB derived in (16).
3.3 PRACTICAL IMPLEMENTATION
In this section, we detail the implementation of our FBSDE-inspired SB model, named SB-FBSDE.
Training process. We treat the likelihoooo in (16) training objective, where the divergence can be can be estimated efficiently following Hutchinson (1989). This immediately distinguishes SB-FBSDE from prior SB models (De Bortoli et al., 2021; Vargas et al., 2021), which instead rely on

---

Published as a conference paper at ICLR 2022
Algorithm
Likelihood training of SB-FBSDE
Algorithm 3 Alternate (IPF-based) training
Input:
boundary distributions Paura and Ppriors
Input: Caching frequency M
parameterized policies [-:,;;) )) and ((:1))
for 1 to K do
repeat
k%M then
memory resource affordable then run Algorithm 2.
Sample Xieju.ry from (13a) where zo rv Plasa (computational graph discarded).
else
end if
Compute Ess (æo: Φ) with (18).
run Algorithm 3.
end if until converges
Update with gradient V (wo; o).
end for
for l = 1 to K do
Algorithm 2 Joint (diffusion flow-based) training
if A%M == 0 then
Sample X from (7b) where æx Pacior (computational graph discarded).
for = 1 K do
Sample X from (13a) where Bo rv Plaa (computational graph retained).
end if
Compute (xr; 0) with (19).
Compute Lsg (Zo; 8, o) with (16).
Update with gradient V VeLsa (2r; 0).
Update (8, Φ) with Ss [oo 0 o).
end for
end for
regression- based objectives." For low-dimensional datasets, we simply perform joint optimization. max SsB (ao;: 8, Φ). to train the parameterized policies -(:::;) 6) and ( Φ). For higher-dimensionas (e.g. image) datasets, however, it can be prohibitively expensive to keep the entire computational graph. In these cases, we follow De Bortoli et al. (2021) by caching the sampled trajectories in reply buffer and refreshing them in lower frequency basis (around 1500 iterations). Although this implies that the gradient path w.r.t. 0 will be discarded, we can leverage the symmetric structure of SB and re-derive the log-like loglliielhoodd for the sampled noise, iee. (2n(2r), ) based on the backward trajectories sampled from (7b). We leave the derivation to Theorem 11 in $B due to space constraint. This results in an alternate training between the following two objectives after dropping all unrelated terms,
EsB (Do: ②) = L/ Ex,~(7a)
Z(t, X x) + vV (t, Xx: z, X1; Φ) dt, (18)
2
/'" (7b)
T
[ Z(t X,; 6) v Z(t 0, + 2, Zt, X,; 0) dt. (19)
Esn (xr;
2
Our training process is summarized in Alg. 1. While the joint training scheme in Alg. 2 resembles recent diffusion flow-based models (Zhang& Chen, 2021), the alternate training in Alg. 3 relates to the classical IPF (De Bortoli et al., 2021), despite differing in the underlying objectives. Empirically, the joint training scheme can converge faster yct at the cost of introducing memory complexity. We highlight these flexible training procedures arising from the unified viewpoint provided in Theorem 4. Hereafter, we refer to each cycle, i.e. 2K training steps, in Alg. as a training stage of SB-FBSDE.
Generative process. While the generative processes for SB can be performed simply propagat- ing (7b) given the trained policy Z(, :; :)),, it has been constantly observed that adopting Langevin sampling to the generative process greatly improves performance (Song et al., 2020). This procedure, the Langevin corrector, requires knowing the score function V log Pt. For SB, we can estimate its value by recalling (see (2.2) that Z, + Z, = gV a log . This results in the following predictor-comrector sampling procedure (see Alg. 4 in Appendix D for more details).
often referred
Predict step: X: X, + Z, VgDte
(20)
Correct step: X,
X, + E [Z 2,) + /20,
(21)
where ~~ I), and or is the pre-specified noise scales (see (59) in Appendix D).
Limitations & efficiency. The main computational burden of our method comes from the computation of the divergence and maintaining two distinct networks. Despite it typically increases the memory by 2~2.55 times compared to SGM, we empirically observe that the divergence-based training converges much faster per iteration than standard regression. As a result, SB- F--FBDD admits comparable training time (+6.8% in our CIFAR 10 experiment) compared to SGM, yet with a substantially fewer sampling time (-80% due to adopting nonlinear SDEs.
'In fact, their regression targets may he recovered from (15) under proper transformation; see Appendix C.

---

Published as a conference paper at ICLR 2022
4 EXPERIMENTS
(a) GMM
(b) Checkerboard
T=0
T=0.33
T=0.66
T=1.0
T=0.00
T-0.33
T=0.66
T=1.00
Figure 3: Validation of our SB-FBSDE model on two synthetic toy datasets that represent continuous and discontinuous distributions. Upper: Generation (Peata + Pptior) process with the backward vector
field ())+))). Bottom: Diffusion (Pdata + Pprior) process with the forward vector field ()::00).
0
6
2
2
7
1
9
1
7
4
7
/
9
3
3
4
3
1
6
3
8
4
8
9
b
2
5
8
1
6
4
1 /
2
2
2
6
3
Y
7
7
2
1
7
8
5
9
1 4 S 3
3
2
/
7
4
1
Figure 4: Uncurated samples from our SB-FBSDE models trained on MNIST (left), resized CelebA (middle) and CIFAR10 10 (right). More images can be found in Appendix E.
Setups. We testify SB-FBSDE on two toy datasets and three image datasets, iee. MNIST, CelebA, 6 and CIFARIO. Pprior set zero-mean Gaussian whose variance varies for each task and can be computed according to Song & Ermon (2020). We parameterize Z(- (0) and Z(,, : () with residual-based networks for toy datasets and consider Unet (Ronneberger al., 2015) and NCSN+ (Song et al., 2020) respectively for MNIST/CelebA and CIFARIO. All networks adopt position encoding and are trained with AdamW (Loshchilov Hutter, 2017) on a TITAN RTX. We adopt VE-SDE (i.e. f == 0: see Song et al. (2020)) our SDE backbone, which implies that in order to achieve reasonable performance, SB must learn a proper data-to-noise diffusion process. On all datasets, set the horizon T = =1.0 and solve the SDEs via the Euler- Maruyama method. The Euler-Maruyam
interval [0. 7] is discretized into 200 steps for CIFAR 10 and 100 steps for all other datasets, which much fewer than the ones in SGM (> 1000 steps). Other details are left in Appendix D.
Toy datasets. We first validate our joint optimization (i.e. Alg 2) on generating mixture of Gaussian and checkerboard (adopted from Grathwohl al. (2018)) the representatives of continuous and discontinuous distributions. Figure shows how the learned policies, i.e. Z).-:0) and Z(--; o). construct the vector fields that progressively transport samples back-and-forth between Pprior and Pdata- The vector fields can be highly nonlinear and dissimilar to each other. This resembles neither SGMs, whose forward vector field must obey linear structure, nor flow-based models, whose vector fields are simply with opposite directions. We highlight this as a distinct feature arising from SB models.
Image datasets. Next, we validate our alternate training (i.e. Alg 3) on high-dimensional image generation. The generated images for MNIST, CelebA, and CIFAR 10 are presented in Fig. 4, which clearly suggest that our SB-FBSDE is able to synthesize high-fidelity images. More uncurated images can be founded Appendix E. Regarding the quantitative evaluation, Table 1 summarizes the negative ligeliieliood (NLL: measured in bits/dim) and the Fréchet Inception Distance score (FID; Heusel et al. (2017)) on CIFAR10. For our SB-FBSDE, we compute the NLL on the test set using Corollary 5, in similar vein SGMs and flow- models, and report the FID over 50k
We follow similar setup of prior SB models (De Bortoli al., 2021) and resize the image size 32.
8

---

Published as a conference paper at ICLR 2022
Table 1: CIFARIO evaluation using negative log-likelihood (NLL; bits/dim) on the test set and sample quality (FID score) w.r.t. the training set. Our SB-FBSDE outperforms other optimal transport baselines by large margin and is comparable to existing generative models.
Model Class
Method
NLL +
FID +
SB-FBSDE (ours)
2.96
3.01
DOT (Tanaka, 2019)
15.78
Optimal Transport
Multi-stage SB (Wang et al., 2021)
12.32
DGflow (Ansari et al., 2020)
9.63
SDE (deep, sub- Song et al. (2020))
2.99
2.92
ScoreFlow (Song et al., 2021)
2.74
5.7
SGMs
VDM (Kingma et al., 2021)
2.49
4.00
LSGM(Vahdat et al., 2021)
3.43
2.10
MNIST
Sample Quality (FI)
(resized)
CIFAR10
Linear drift If- 9.51)
Degenerate diift 1/-
KLOD KLOD..Pol Pyo)
SGH
wile Langevi correcter Langevin correcter
FBSDE [ours]
10
50- 58-FBSDE Train Stage
FBSDE Train
Figure 5: Validation of our SB-FBSDE on
Figure Ablation analysis where we show that adding Langevin corrector to SB-FBSDD uniformly improves the FID scores on both CelebA and CIFARIO training.
learning forward diffusions that are closer (in KL sense) to Pprior compared to SGM.
samples w.r.t the training set. Notably, our SB-FBSDE achieves 2.98 bits/dim and 3.18 FID score on CIFARIO, which comparable to the top existing methods from other model classes (e.g. SGMs) and outperforms prior Optimal Transport (OT) methods (Wang et al., 2021; Tanaka, 2019) by large margin terms of the sample quality. More importantly, it enables log-likelihood computations that otherwise infeasible in prior OT methods. We note that the quantitative comparisons on MNIST and CelebA are omitted as the scores on these two datasets are not widely reported and different pre-processing (e.g. resizing of CelebA) can lead to values that are not directly comparable.
Validity of SB forward diffusion. Our theoretical analysis in $3.2 suggests that the forward policy Z(::; 0) Ze plays an essential role in governing samples towards Ppeier. Here, validate this conjecture by computing the KL divergence between the terminal distribution induced by Zo. i.e. PT
(13a) and the designated prior Pprior- We refer readers to Appendix D for the actual computation.
Figure 5 reports these values over MNIST training. For both degenerate (f := 0) and linear (f (f:=;; a,) base drifts, our SB-FBSDE generates terminal distributions that are much closer to Ppeioe- Note that the values of SGM remain unchanged throughout training since SGM relies on pre-specified diffusion. This is in contrast to our SB-FBSDE whose forward policy Ze gradually shortens the KL gap to Ppriors thereby providing better forward diffusion for training the backward policy.
Effect of Langevin corrector. In practice, we observe that the Langevin corrector greatly affects the generative performance. As shown in Fig. 6, including these corrector steps uniformly improves the sample quality (FID) on both CelebA and CIFAR 10 throughout training. Since the SDEs are often solved via the Murer-Maruyama method, their propagation can be subjected to discretization errors accumulated over time. These Langevin steps thereby help re-distributing the samples at each time step t towards the desired density p.hh SB We emphasize this improvement as the benefit gained from applying modern generative training techniques based on the solid connection between SB and SGM.
5 CONCLUSION
In this work, we present a novel computational framework, grounded on Forward- -ackkarrd SDEs theory, for computing the log-likelihood of Schrödinger Bridge (SB) a recently emerging model that adopts entropy-regularized optimal transport for generative modeling. Our findings provide new theoretical insights by generalizing previous theoretical results for Score-based Generative Model, and facilitate applications of modern generative training for SB. We validate our method on various image generative tasks, e.g. MNIST, CelebA, and CIFAR1O, showing encouraging results in synthesizing figh-iideiity samples while retaining the rigorous mathematical framework.

---

Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Published as a conference paper at ICLR 2022
ACKNOWLEDGMENTS
The authors would like to thank loannis Exarchos and Oswin So for their generous involvement and helpful supports during the rebuttal. The authors would also like to thank Mareus A Pereira and Ethan N Evans for their participation and kind discussion in the early stage of project exploration.
AUTHOR CONTRIBUTIONS
The original idea of solving the PDE optimality SB with FBSDEs theory was initiated by Tianrong. Later, Guan derived the main theories (i.e. Theorem 3, 4, 11, and Corollary 5) presented in Section 3.1, 3.2 and Appendix with few helps from Tianrong- Tianrong designed the practical algorithms (e.g. stage-wise optimization and Langevin-corrector) in Section 3.3 and conducted most experiments with few helps from Guan. Guan wrote the main paper except for Section 4, which were written by both Tianrong and Guan. Both Guan and Tianrong contributed code development.
REPRODUCIBILITY STATEMENT
Our training algorithms are detailed in Alg. 2, and 3, with the training objectives given in the same section (see (16, 18, 19)). Other implementation details (e.g. data pre-processing) are left in Appendix D. This shall provide sufficient information for readers of interests to reproduce our results. As strongly believe in the merit of open sourcing, intend to release our implementation upon publication. On the theoretical side, all proofs are left to Appendix B due to space constraint. We provide the assumptions in the same section.
REFERENCES
Applications, 12(3)3313-32, 1982.
Abdul Fatir Ansari, Ming Liang Ang, and Harold Soh. Refining deep generative models via discrimi- nator gradient flow. arXiv preprint arXiv:2012.00780 2020.
Kenneth Caluya and Abhishek Halder. Wasserstein proximal algorithms for the schrödinger bridge problem: Density control with nonlinear drift. IEEE Transactions on Automatic Control, 2021.
Jianfei Chen, Cheng Lu, Biqi Chenli, Jun Zhu, and Tian Tian. Vflow: More expressive generative flows with variational data augmentation. In International Conference on Machine Learning, pp. 1660-1669. PMLR, 2020.
Tian Qi Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary dif- ferential equations. In Advances Neural Information Processing Systems, pp. 3583, 2018.
Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. Stochastic control liaisons: Richard sinkhorn meets gaspard monge on schrödinger bridge. SIAM Review, 63(2):249-313, 2021.
Rewon Child. Very deep vaes generalize autoregressive models and can outperform them on images. arXiv preprini arXiv:2011.10650 2020.
Julian D Cole. On a quasi-linear parabolic equation occurring in aerodynamics. Quarterly of applied mathematics, 933):225-236 1951.
Paolo Dai Pra. A stochastic control approach to reciprocal diffusion processes. Applied mathematics and Optimization, 23(1):31--299, 1991.
Valentin De Bortoli, James Thornton, Jeremy Heng, and Amaud Doucet. Diffusion schrödinger bridge with applications to ccore-based generative modeling. arXiv preprint 11X.::2106.01355 2021.
10

---

Published as a conference paper at ICLR 2022
Ioannis Exarchos and Evangelos A Theodorou. Stochastic optimal control via forward and backward stochastic differential equations and importance sampling. Automatica, 87:159- 165, 2018.
Wenbo Gong and Yingzhen Li. Interpreting diffusion score matching using normalizing flow. arXiv preprint arXiv:2107.1001: 2021.
Xinyu Gong, Shiyu Chang, Yifan Jiang, and Zhangyang Wang. Autogan: Neural architecture search for generative adversarial networks. In Proceedings of the IEEE/CVF International Conference Computer Vision, pp. 2224-3234, 2019.
Alex A Gorodetsky, Sertac Karaman, and Youssef M Marzouk. Efficient high-dimensional stochastic optimal motion control using tensor-train decomposition. In Robotics: Science and Systems, 2015.
Will Grathwohl, Ricky TQ Chen, Jesse Betterncourt, Ilya Sutskever, and David Duvenaud. Ffjord: Free-form continuous dynamics for scalable reversible generative models. arXiv preprint 1rXiv:1810.01167, 2018.
Jiequn Han, Arnulf Jentzen, and E Weinan. Solving high- high-dimensional partial differential equations using deep learning. Proceedings of the National Academy of Sciences, 115(34)8505-85110 2018.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. arXiv preprint arXiv:2006. 11239, 2020.
Eberhard Hopf. The partial differential equation ut+ uux= μax. Communications on Pure and Applied mathematics, 3(3):201-25 -230, 1950.
Chin-Wei Huang, Laurent Dinh, and Aaron Courville, Augmented normalizing flows: Bridging the gap between generative flows and latent variable models. arXiv preprint arXiv:2002.07101, 2020.
Chin-Wei Huang, Jae Hyun Lim, and Aaron Courville. A variational perspective on diffusion-based generative models and score matching. arXiv preprint arXiv:2106.02866, 2021.
Michael F Hutchinson. A stochastic estimator of the trace the influence matrix for laplacian and Computation, 185)):059-1076.
smoothing splines. Communications 1989.
Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
Kiyosi Itô. On stochastic differential equations, volume 4. American Mathematical Soc., 1951.
Alexia Jolicoeur-Martinenna Ke Li, Rémi Piché- Taillefre,, Tal Kachman, and Ioannis Mitliagkas. Gotta go fast when generating data with score-based models. arXiv preprint urXiv:2105. 14000 2021.
Ioannis Karatzas and Steven Shreve. Brownian motion and stochastic calculus, volume 113, Springer Science & Business Media, 2012.
Tero Karras, Miia Aittala, Janne Hellsten, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Training generative adversarial networks with limited data. arXiv preprint arXiv:2006. 06676, 2020a.
Tero Karras, Samuli I Laine, Miika Aittala, Janne Hellsten, Jaakko Lehtinen, and Timo Aila. Analyzing and improving the image quality of stylegan. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 81 8110-81 19, 2020b.
Diederik P Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. arXiv preprint 10t.::2277.0006.. 2021.
Magdalena Kobylanski. Backward stochastic differential equations and partial differential equations with quadratic growth. Annals of probability, pp. 558 658, 2000.
11

---

Published as a conference paper at ICLR 2022
Jungil Kong, Jachyeon Kim, and Jackyoung Bae. Hifi-gam. Generative adversarial networks for efficient and high fidelity speech synthesis. arXiv preprint arXiv:0010.06646, 2020a.
Zhifeng Kong and Wei Ping. On fast sampling of diffusion probabilistic models. arXiv preprint arXiv:2106.00132. 2021.
Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: A versatile diffusion model for audio synthesis. arXiv preprini arXiv:2009.0976/ 2020b.
Solomon Kullback. Probability densities with given marginals. The Annals of Mathematical Statistics, 39(4):1236 1243, 1968.
Christian Léonard. A survey of the schrödinger problem and some of its connections with optimal transport. arXiv preprint arXiv:1308.022: 0215, 2013.
Qianxiao Li and Shuji Hao. An optimal control approach to deep learning and applications to discrete-weight neural networks. arXiv preprint 1rX.iv/003.0122 2018.
Guan-Horn Liu, Tianrong Chen, and Evangelos A Theodorou. Ddpnopt: Differential dynamic
programming neural optimizer. In International Conference on Learning Representations, 2021a.
Guan-Horng Liu, Tianrong Chen, and Evangelos Theodorou. Dynamic game theoretic neural optimizer. In International Conference Machine Learning, 2021b.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 17Xiv:1711.051 101, 2017.
Jin Ma. J-M Morel, and Jiongmin Yong. Forward-backwer stochastic differential equations and their applications. Number 1702. Springer Science & Business Media, 1999.
Lars Maaløe, Marco Fraccaro, Valentin Liévin, and Ole Winther. Biva: A very deep hierarchy of latent variables for generative modeling. arXiv preprint arXiv:1902.02102, 2019.
Dimitra Maoutsa, Sebastian Reich, and Manfred Opper. Interacting particle solutions of fokker- planck equations through gradient-log-density estimation. Entropy, 2288):802, 2020.
Balint Negyesi, Kristoffer Andersson, and Cornelis W Oosterlee. The step malliavin scheme: new discretization of bsdes implemented with deep learning regressions. arXiv preprint arXiv:21 10.0542, 2021.
Alex Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. arXiv preprint arXhv:2102.09672, 2021.
Etienne Pardoux and Shige Peng. Backward stochastic differential equations and quasilinear parabolic partial differential equations. In Stochastic partial differential equations and their applications, pp. 200-217. Springer, 1992.
Jeeseung Park and Younggeun Kim. Styleformer: Transformer based generative adversarial networks with style vector. arXiv preprint arXiv:2 000037 2021.
Michele Pavon and Anton Wakolbinger. On free energy. stochastic control, and schrödinger processes. In Modeling. Estimation and Control of Systems with Uncertainty, pp. 334-348. Springer, 1991.
Marcus Pereira, Ziyi Wang, Ioannis Exarchos, and Evangelos A Theodorou. Neural network architectures for stochastic control using the nonlinear feynman-kac lemma. arXiv preprint arXiv:1902.03986, 2019.
Marcus Pereira, Ziyi Wang, Tianrong Chen, Emily Reed, and Evangelos Theodorou. Feynman-kac neural network architectures for stochastic control using second-order fbsde theory. In Learning for Dynamics and Control, pp. 728-738. PMLR. 2020.
Wei Ping, Kainan Peng, Kexin Zhao, and Zhao Song. Waveflow: A compact flow-based model for raw audio. In International Conference Machine Learning, pp. 7706-7716. PMLR, 2020.
12

---

Published as a conference paper at ICLR 2022
Prajit Ramachandran, Barret Zoph, and Quoc V Le. Searching for activation functions. arXiv preprint arXiv:1710.05941, 2017.
Michael Renardy and Robert C Rogers. An introduction to partial differential equations, volume 13. Springer Science & Business Media, 2006.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In International Conference on Medical image computing and computer- assisted intervention, pp. 234 241. Springer, 2015.
Erwin Schrödinger. Sur la théorie relativiste d'éeectron l'interprétation de mécanique quantique. In Annales de l'institut Henri Poincaré, volume 2, pp. 269-310, 1932.
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International conference on machine learning, pp. 1889-1897, 2015.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. arXiv preprint 1907v::907.00550 0000000 2019.
Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. arXiv preprint wwXiv:2006.09011 2020.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13346 2020.
Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. arXiv e-prints, pp. 2r0
Akinori Tanaka. Discriminator optimal transport. arXiv preprint arXiv:1910.06832, 2019.
Arash Vahdat and Jan Kautz. Nvae: A deep hierarchical variational autoencoder. arXiv preprint arXiv:2007.03888, 2020.
Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. arXiv preprint arXiv:2106.0500 2021.
Francisco Vargas, Pierre Thodoroff, Neil D Lawrence, and Austen Lamacraft, Solving schrödinger bridges via maximum likelihood. arXiv preprint arXiv:2106.00 020000 2021.
Gefei Wang. Yuling Jiao, Qian Xu, Yang Wang, and Can Yang. Deep generative learning via schrödinger bridge. arXiv preprint arXiv:2106 10410, 2021.
Jiongmin Yong and Xun Yu Zhou. Stochastic controls: Hamiltonian systems and IJJB equations, volume 43. Springer Science & Business Media, 1999.
Qinsheng Zhang and Yongxin Chen. Diffusion normalizing flow. arXiv preprint arXiv:211000757 2021.
13

---

Published as a conference paper at ICLR 2022
A INTRODUCTION OF SCHRÓDINGER BRIDGE
In this subsection, we provide a brief review for Schrödinger Bridge (SB) and some reasonings for Theorem The SB problem, its classical form, considers the following optimization (Dai Pra, 1991; Pavon & Wakolbinger, 1991).
1/'0 2 [ [1)((1,,,)11
min uedl
(22)
dX,
u(t, X,)dt + V2c dW,
s.t.
Xo (X): Xr ~ pT (X)
where U == {u: [0, T] x R + (u, u) < oo). The optimization (22) characterizes a standard
stochastic optimal control (SOC) programming with energy (i.e. |||u](?) minimization except with
an additional terminal boundary condition. The optimality conditions to (22) are given by
Oup / - AA..
(23a)
Öt
Op" Ōt
V (p v +e
(23b)
where (t, v) E C is known as the value function in SOC literature and (t, x) E C is the associated optimal marginal density. A denotes the Laplace operator. Equations (23a) and (23b) respectively the Kolmogorov's backward and forward PDEs, also known as Hamilton-Jacobi-Blllle and Tokker-Planchs equations. The SB system can be obtained by applying the Hopf-Cole (Hopf, 1950; Cole, 1951) transformation (vh,p") + V V)
04
AD
8t
s.t. (0, (0, )) Po, (TT,)) Pr
(24)
aqj с T0
0t
In this work, we consider recent generalization of (22) and (24) to an SDE class with nonlinear drift, affine control, and time-varying diffusion. We synthesize their results below.
Theorem 6 (SB optimality; Caluya & Halder (2021)). Consider the following optimization
min E
2 Nu(t,x,))
(25)
dx, = [S(1, X.) + g(1) m(t, X,) + 22 g(() dW,
s.t.
Xo ~ PX(),, Xx (X)
and (t1, x,) satisfies Lipschitz conditions with at most linear growth æ. Then, the Hopf-Cole transformation to (25) becomes
where gt) is uniformly
- V -- Vv.
dt
(26)
00
with the same boundary conditions (O,,) po. V V(T,-)V(T, (T, pr. The optimal control to (25) is thereby given by
u (t, X,) = -rg V log (t, X,).
(27)
Proof. See Section and Theorem inf Caluya & Halder (2021).
Theorem 6 is particularly attractive to since its SDE corresponds exactly to the one appearing in score-based generative models. One can recover Theorem by
(i) Following Pavon Wakolbinger (1991), we know that the objective in (25) is equivalent to DKL. (Q PP) by application of Girsanov's Theorem.
14

---

Published as a conference paper at ICLR 2022
(ii) Equation (26) is exactly (6) with = 1 Furthermore, substituting the optimal control (27) to the stochastic process in (25) yields the optimal forward SDE in (7a).
(ii)) Finally, reversing the SDE (7a) from forward to backward following Anderson (1982),
dx. [f + g2 Væ log (t. X g2 V log SB dt + dW..
(28)
and recalling the factorization principle, log p (-) log (t. log (t,-)), from Equation (4.15) in Chen et al. (2021) yield the optimal backward SDE in (7b).
B PROOFS AND REMARKS IN SECTION
In this section, we provide proofs for all of our theorems. We following the same notation by denoting Pt pSDE (X,)
(X₂) as the marginal density driven by some SDE process X(t) X, until the time step t E (O.) Gradient and Hessian of function f(x), where R", will respectively be denoted as v., Ra and / f E R.xx Divergence and Laplace operators will respectively denoted as V and A. Note that A = V V. For notational brevity, we will only keep the subscript for multivariate functions. Finally, (r( denotes the trace of a square matrix A.
We first restate the celebrated Itō lemma, which is known as the extension of the chain rule of ordinary calculus to the stochastic setting. It relies on the fact that dW, 2 and dl are of the same scale and keeps the expansion up to (dt))
Lemma (Itō formula; Itô (1951)). Let E C22 and let X, be the stochastic process satisfying dX;- Xt)dt G(t, d.W.
Then, the stochastic process vt,, X₂) is also an Itō process satisfying
d.) = Ou(t,
Ōt dt + Xt,
+ Tr [G'V (,) dt
1 2
+ Vou(t, X.) ((t,tt,)))G(t(,)/))) Gtt, dW..
(29)
Next, the following lemma will be useful in proving Theorem 3. Lemma 8. The following equality holds at any point a E R such that p(x) #
1 p(
Tr (V)p(z))
|| logp(>>)11 + Alogp(æ)
Proof.
Tr (v'p(z)) Ap(x) = V Vp(x)
V (p(x)Vlogg((n)
-o(g)\l VVlgp(<>) + logp(()) p(x) (lV logg((())
Assumptions Before stating our proofs, we provide the assumptions used throughout the paper. These assumptions are adopted from stochastic analysis for SGM (Song al., 2021; Yong & Zhou, 1999; Anderson, 1982). SB (Caluya & Halder, 2021), and FBSDE (Exarchos & Theodorou, 2018; Gorodetsky al., 2015).
(i) Pprior: Panta E C with finite second-order moment.
(ii) f and g are continuous functions, and (2 > is uniformly lower-bounded w.r.t.
(iii) Vt E [0, T], we have f (t, x), Va (x:), V log (tt, æ). V log (t, æx) Z(t, æ; 0), and Ž ax; Φ) Lipschitz and at most linear growth w.r.t. 2.
15

---

Published as a conference paper at ICLR 2022
(iv) V. T E C h, and yp are continuous functions. h satisfies quadratic growth w.r.. uniformly in t.
(v) 3k (
Pi
(lxxl as x oc.
Assumptions (i (ii) (iii) standard conditions in stochastic analysis to ensure the existence- uniqueness of the SDEs; hence also appear in SGM analysis (Song et al., 2021). Assumption (iv) allows applications of Itō formula and properly defines the backward SDE in FBSDE theory. Finally, assumption (v) assures the exponential limiting behavior when performing integration by parts.
Now, let begin the proofs of Theorem and Corollary
Theorem (FBSDEs SB optimality (6)). Consider the following of coupled SDEs,
dx; (f + +Zze) df gdW,
(30a)
dY. 1 - Z, d.W
(30b)
dY, - 2 (Z Vx (9 f) Z, dt + Z dW.
(30c)
where f and g satisfy the same regularity conditions in Lemma 2 (see Footnote 4), and the boundary conditions are given by X(0) nonlinear Feynman- Kac relations between the FBSDEs (13) and PDEs (6) are given by
and YT - Yx- log Pprior Xr). Suppose V E C then
Y, (t,, log (t, X,),
Z: Z(t, X.) g V log (t, X.),
(31)
Y, Yt, = log (t, Xt),
Ż, (t, x,) vV log (t. X.).
Furthermore, (Y,, Y.) obey the following relation:
Y, + Y, = log (X.).
Proof. Similar to how the original nonlinear Feynman-Kac (i.e. Lemma 2) be carried out by an application of Itô lemma (Ma et al., 1999). We can apply Itô lemma 7 to the stochastic process log (t, X,) w.r.t the optimal forward SDE (7a).
d log = log v
1 dt + V log VT + gf log v) 99 Tr V2 log w dt + [gV log vT dW,
Øt
(32)
From the PDE dynamics (6), we know that
log V
-V
V V
Ot
V
2
- log T
7 Tr v2
The first term in the RHS can readily canceled out with the related f-term in (32). The second term also be canceled out using the fact that v log v are left with
v2 22 V VV. VT. Hence,
1
dlog V
V a log 1// 9 Tr
V V.
dt vv log v law,
122
マ
lo log 2/ + v log T dw.
(33)
Likewise, applying Itô lemma to log (t, X.), where X, follows the SDE in (7a).
d log = 0t dt + V log v (f / log v - 2
Bog v
1
Tr
V log T
dt + VV log TT] dw.
(34)
16

---

Published as a conference paper at ICLR 2022
now noticing that the dynamics of log v dt
become
i log T
- V (Tf) + (
Öt
log v. f 2 Tr -
Only the first term in the RHS will be canceled out in (34). Hence, we are left with
d log v0 = V f f+ 2
V
dt
(35)
+ V log
Va log + 2
Tr
V2 log
v
di + gVa log dww.
Notice that the trace terms above can be simplified to
T
V TV
V
log
V
Tr
V V
v log 11/²
2
V log L I + A log VT
2
where the last equality follows by Lemma 8. Substituting this result back (35), get
1 =l1l log 1 + 92 A log TF + g V log v7 V log v
log v
Va f
dt vV log vT dw,
= Va: (g²V log ύー 5) + Illl log 111 + V log v V log v
dl vv log V dw,
(36)
Finally, by rewriting (33) and (36) with the nonlinear Feynman-Ko Kac in (31) yields
dX, = (ʃ + (zz di + gdW,
dY, 1 z Zi Zidt + Z,,W,
dY. = 2, + V (922-
- +)+ 2. de + 2,dW:
This concludes the proof.
Remark 9 (Viscosity solutions). These FBSDE results can be extended to viscosity solutions in the case when the classical solution does exist (Pardoux & Peng. 1992). For the completeness, one shall understand them in the sense of v(t, x)) lim (t,t, x) uniformly in (t, æ) over a compact set. Here of (t, æ) the classical solution to (10) with (fe, (fe hes iP.) converge uniformly toward (J, G, h, sp) over the compact set. We refer readers of interests to Exarchos & Theodorou (2018);
Negyesi et al. (2021), and their references therein.
Theorem 4 (Log-likelihoo for SB models). Given the solution satisfying the FBSDE system in (13), the log-likelioood of the SB model (Ze, Z.), at data point By, can be expressed as
/'
Ss (mx) E. (logpr((X)))
Z 2 Z,
Vv log + Z
0
1
V
logr Zu 2-12--2 Vz
(38)
= EE ([oo))((XT)]]
E
2+
2, V (gZ, f) 2, Z. dt, (39)
where the expectation taken over the forward SDE (13a) with the initial condition X, Do-
17

---

Published as a conference paper at ICLR 2022
Proof.
LsB
(20)
E
Yo
Yo =
Yr !'¹
T
(- + (9 - f) 2, di
=E
2 [/12211?
dt Y
do Xo 20
- | E T
TT
=E
Yr + Y Xo = wo
+ V (gZ: - f) + Z Xo 20 dl
T
1
=E(log Xx))
E
Z
112
12
1
(gZ, f) Z Zu dt,
(40)
2
which recovers (39). Finally, notice that with integration by part, we have
E Ex,
gV. 2.|
/ (gV Z,) ps dX,
sn dX, / g2, VaPi
(41)
/ (12 vo log SB ps dX,
Ex
[-92, V log Pt
SB
where we adopt common practice and assume the limiting behavior of pr. in other words, El: > 0: (ar) O(exp O(exp-I1al1)
||
as a
xi With (41), we can rewrite the related parts in (40) as
E
Zz + Z. 2.
E
12. 2
27 ( log (B) + Z, Z.
2
E
2
1
VV
log SB Z. 2 All log SB
1
Z.1
(42)
Hence, we also recover (38).
Corollary (Probability flow for SB). The following ODE characterizes the probability of the optimal processes of SB (7) in the sense that Vt, (17) (
= =p
SB
dx, f 1 gZ X.) 29 (Z(t, x,) + Z(t, Xi),
dt
(43)
Proof. The probability ODE flow (Song et al., 2020; Maoutsa et al., 2020) suggests that the equivalent ODE model for the SDE (1) given by
dx, = [ - 5 Vn log !1) dt.
We can adopt this result to the SDEs of SB (7a) by considering f + f + gZ, and p, <- ps. This yields
1 dx, = f+ sz gZz / VV Va logr," SB
(44)
Applying the the factorization principle (Chen et al., 2021) with g logp, Z, 2, concludes the proof.
n
18

---

Published as a conference paper at ICLR 2022
Remark 10 (Connection between SB-FBSDE and flow-based models). To demonstrate how applying flow-based training techniques to the probability ODE flow of SB (43) recovers the same log- likelihood objective in (39), recall that given ODE dX, F(t, dt with X, 2n Pdaa- flow-based models compute the change in log-density using the instantaneous change variables formula (Chen al., 2018):
b Ologp(X, p(Xi) V F,
Öt
which implies that the log-likellng oogllikelhhod of an can be computed as
log p(2) logp(wo) /-" V dt.
(45)
Now, consider the probability ODE flow of SB in (44),
Fsn = f + gZ. 25(2 (Z, + 2n) = f + -0 Zz - 2.).
Substituting this vector field FSB to (45) yields
logpr pr (x) (XT)) = pooo)( - S0 V
f
2 g (Z. -2,)
T
f+ 17 70(Zz - -2,))
E [log po (zo) = E pI Pr (XT))
+
E
V x
(gZ, - f) -ー 59V= (Ze + En) dt
E [rogpr(XT))
E
V
T
1
E [log PT (x)]
E
Vx
(gZz
f) +
2 7g + 2.)² (V log
dt
T
E log pr T])
E
(gZ, 0) + 2 Z. + 2.)² dl,
(46)
where (*) is due to integration by parts (recall (41)) and (**) again uses the factorization principle SB
Z. + Z, V lo One can verify that (46) indeed recovers (39).
Theorem 11 (FBSDE computation for ( in SB models). With the same regularity conditions
in Theorem 3, the following FBSDEs also satisfy the nonlinear Feynman-Kan relations in (31).
dx, = (5 gZ.) dt + gdW,
(47a)
dY, = Z Z. -v (9 + ff++ +)++ Z, dt + Z,dW,
(47b)
dY, = 2,
2,dt + 2. a.W.
(47c)
2
Given a backward trajectory sampled from (47a), where Xx = FT and F. ~ Putions the log- likelihood of vo is given by log Pprice (xr)) (x = TF
Y + V XT = 2T = L.g (FrT). In particular,
/ E
(xT) = [l(ooggT(Xo)) -
Z. 12 Zell V. (ghh +
+ f)+
dt, (48)
Proof. Due the symmetric structure of SB, consider a new time coordinate
Under this transformation, the base reference appearing in (5) is equivalent to = fs, xd + dW..
dx, =
The corresponding PDE optimality becomes
00 V Tr
a
= V (0/)+ Tr
s.l. (00,-))(() -) Pprior, (T, .)
Pdatas
(49)
84
19

---

Published as a conference paper at ICLR 2022
and the optimal forward/backward policies given
dX; = | + g2 V log (s, x,) + 9 dW
X ~ Ppriors
(50a)
dX,- I g2 V a log ( x ds + 9 dW.,
XT Pdatu-
(50b)
By comparing (50) with (7), one can notice that the new SB system ( D). corresponds to the
original system (V, V) via
(s x -TT-t,X--1
and (s. = -1,, Xy-
(51)
Equation (51) shall be understood as the forward policy in l-coordinate system corresponds to the backward policy in s-coordinate system, and vise versa. Following similar derivations the proof of Theorem 3, we can apply Itó lemma to expend the stochastic processes d dlog & and doo w.r.t. (50a). This yields the following FBSDE system.
dx. (gZ, -f) ds + gdW.
(52a)
dY' 1 ; 11,,,1 2ds Z law,
(52b)
dY = (
2
1
|2 + V (g + f) 2, Tz ds + T dd,
(52c)
Similar to (51), (Y, Y, Z Z',) relate the original FBSDE system (30) by
(Y Zg 22)) (T T -1: -1)
(53)
Changing the coordinate from to and applying (53) readily yield (47). Finally, the expression in (48) can be carried out similar to (40):
Ss (IT))
=E
YT
YT IT
Yo (-¹ (1122111)
T
/''
( |z,1| A (gZ.
=E
dt d Y
+f f)+ 2T2, ()+
dt( IT
=E
Yo + Yn Xx = xr
2, 2 1 2. 1 + V (g2, + 1) 2,2,
Xx 2T dl
00
2
=Ellog (o (Xo)
E
Z,
11
2
V (gZ, + f) 2 dl.
(54)
We conclude the proof.
C COMPARISON WITH PRIOR SB WORKS
Our method is closely related to two concurrent SB models (De Bortoli et al., 2021; Vargas et al., 2021), yet differs in various aspects. Below we enumerate some of the differences.
Training loss. Both concurrent methods rely on solving SB mean-matching regression between the current drift and (estimated) optimal drift. This in contrast to our SB-FBSDE, which instead utilizes a divergence-ber based objective (16). However, the regression objectives are in fact captured by Theorem 4. To see that, recall the forward and backward transition models considered in De Bortoli et al. (2021),
X ~ (XF((X))) 216)1))))) and X ~ ((B++(XX++) 27)-1
where (x( = 7k+ (x) and B (x) x Y dy (x) are solved alternately via
B arg min E B B241 (x+1)) (X++1 Fa (X+) F& (x((+))))
(55a)
B
Fb t arg min E [[I - (X+ + (x+1 Bx+1(X)))11
(55b)
Ft
20

---

Published as a conference paper at ICLR 2022
In what follows, focus mainly the connection between (55a) and Theorem 4, yet similar analysis can be applied to (55b). Now, expanding (55a) with the definition of (Bx+1, Fi))
|B++ (X++1)x+)))))+)) - (X++1 + Fa( - F& 11l
=I1(x++1 + 1 (X)+ (X+++ + X + (A - Xx+1 -k+ Th+1 (X =I =Ilm+1Se(X + 1 (X++1) (XA + (xx:) - Xx+1)l??
(56)
which resembles the term [Z + Ze 9 Vo log log "h 12 in (15). While 2 indeed corresponds to our backward policy 1(k+1, 1, Xx+1) after time discretization, 1 slightly differs from 2( in how time is integrated, [k, X vs. 2(k+1,X, Xx+1). On the other hand, 3 may be seen as an approximation of g v log log D² which, crucially, is nor utilized in SB-FBSDE training. Since V x log pr is often intractable (nor does SB-FBSDE try to approximate it), SB-FBSDD instead uses the divergence-buce based objective (16), which does not appear in their practical training.
SDE model class. It is important to recognize that both concurrent methods are rooted in the classical SB formulation with the following SDE model,
dx. f(t, dt) + 22 dW 2
which, crucially, differs from the SDE concerne by both our SB-FBSDE and SGM,
dx, = ft, )t + g(t)dW,,
(57)
in that the diffusion (t) is vire-varying function. This implies that the connection between classical SB models and SGM can only be made in discrete-ting after choosing proper step sizes. For instance, De Bortoli et al. (2021) considers the Euler Mallya-aauyymaa discretization (see their $C.3),
X - Xx + 7 X&) + 2 6,
(58)
where (O, I). In order for (58) to match the discretization of (57), where g(() is often a monotonically increasing function, the step sizes { must also increase monotonically. However, since V. logp, is approximated in De Bortoli et al. (2021) using the states from two consecutive steps (see (55)), this may also affect the accuracy of the regression targets.
In contrast, our SB- SB-FBSDE is grounded on the recent SB theory (Caluya & Halder, 2021), which considers the same SDE model as in (57). As such, connection between SB-FBSDE and SGM is made directly in continuous-time (and can be extended to discrete-time flawlessly); hence unaffected by the choice of numerical discretization or step sizes.
Model parametrization. While Vargas et al. (2021) utilizes non-parametric models, egg. Gaussian processes (hence are not directly comparable), both De Bortoli et al. (2021) and SB-FBSDE use DNNs to approximate the SB policies.
Training algorithm and convergence. Both concurrent methods rely on solving SB with IPF algorithm, which performs alternate training between the forward/backward policies. While SB- FBSDE can also be trained with (see Alg. 3), we stress that is also possible to train both policies jointly whenever the computational budget permits. Interestingly. this joint optimization which not presented in concurrent methods resembles the training scheme of the recently-proposed diffusion flow-based model (Zhang & Chen, 2021). We highlight these flexible training procedures arising from the unified viewpoint provided in Theorem 4. Finally, with the close relation between Alg. 3 and IPF (despite with different objectives and SDE model classes), convergence analysis from classical IPF can be applied with few efforts. We leave it as promising future work.
Corrector Sampling. While both De Bortoli et al. (2021) and SB-FBSDA implement corrector sampling, they corresponds to different quantities. Specifically, our SB-FBSDE relies on the same predicorr-corrector scheme proposed in SGM (see Sec 4.2 in Song et al. (2020)), where the "corrector" part is made with a Langevin sampling using the desired optimal density "V log In SB, this term corresponds exactly to adding the outputs of our networks "Z+ 2" This computation differs from the corrector sampling appearing in De Bortoli et al. (2021), which relies on single network (í.e. either Z or Z). Crucially, this implies that the two methods is approaching different target distributions; hence leading to different training results. Notably, it has been reported in De Bortoli et al. (2021) that corrector sampling only gives negligible improvement (see $J.2 in De Bortoli et al. (2021)), yet in our case we observe major quantitative improvement (up to 4 FID).
21

---

Published as a conference paper at ICLR 2022
D EXPERIMENT DETAILS
Table with other models.
Table 2: CIFARIO evaluation.
Model Class
Method
NLL +
FID
SB-FBSDE (ours)
2.96
3.01
DOT (Tanaka, 2019)
15.78
Optimal Transport
Multi-stage SB (Wang et al., 2021)
12.32
DGflow (Ansari et al., 2020)
9.63
SDE (deep, sub- Song et al. (2020))
2.99
2.92
ScoreFlow (Song et al., 2021)
2.74
5.7
SGMs
VDM (Kingma al., 2021)
2.49
4.00
LSGM(VVhdat et al., 2021)
3.43
2.10
VDVAE (Child, 2020)
2.87
NVAE (Vahdat Kautz, 2020)
2.91
23.49
VABs
BIVA (Maaløe et al., 2019)
3.08
FFJORD (Grathwohl et al., 2018)
3,40
Flows
VFlow (Chen et al., 2020)
2.98
ANF (Huang et al., 2020)
3.05
AutoGAN (Gong al., 2019)
12.42
GANs
StyleGAN2-ADA (Karras al., 2020a)
2.92
LeCAM (Park & Kim, 2021)
2.47
Figure 7: Training Hyper-parameters
Dataset
learning rate
time steps
batch size
variance of Pprior
Toy
2c-4
100
400
1.0
Mnist
2e
100
200
1.0
CelebA
2e-4
100
200
900.0
CIFAR10
1e-5
200
64
2500.0
Figure 8: Network Architectures
Dataset
Z. 0) and # of parameters
Zz(-::)) and # of parameters
Toy Mnist
Fe-ResNet (0.76M)
FC-ResNet (0.76M)
reduced Unet (1.95M)
reduced Unet (1.95M)
CelebA
Unet (39.63M)
Unet (39.63M)
CIFAR10
NCSN+ (62.69M)
Unet (39.63M)
Training. We use Exponential Moving Average (EMA) with the decay rate of 0.99. Table 7 details the hyper-parameters used for each dataset. As mentioned in De Bortoli et (2021), the alternate training scheme may substantially accelerate the convergence under proper initialization. Specifically, when Z, is initialized with degenerate outputs (e.g. by zeroing out its last layer), training Z, the first K steps can be made a similar SGM fashion since pf now admits analytical expression. As for the proceeding stages, we resume to use (18, 19) since (Z,, Z.) no longer have trivial outputs.
Data pre-procss,, MNIST padded from 28 28 x22 to prevent degenerate feature maps through Unet. CelebA is resized to 3x x33 to accelerate training. Both CelebA and CIFAR10 are augmented with random horizontal flips to enhance the diversity.
22

---

Published as a conference paper at ICLR 2022
ResBBlock
ReaBllckk
ReaBlock
FCBlock
output
Posttion encoding
FCBlock
Figure 9: Network architecture for toy datasets.
Sampling. The sampling procedure is summarized Alg. 4. Given some pre-defined signal-to-noise ratio (we set r =0.05 for all experiments), the Langevin noise scale t,, at each time step and each corrector step is com- puted by
Algorithm 4 Generative Process of SB-FBBDE
Input: Pprce policies :(::;0) and [
Sample Xr ~ Pprior- for = T to At do
Sample ~ (0, I). Predict X 1 X. + Z.AL + VgÂTe.
for i = 1 to N do
2r 2r
Sample c, (O,
Tr
[[( + (t, X
Compute V log Pt. [Z(t, xx,)) +Z(t, x,,)\/9-
Compute de,: with (59).
(59)
Correct X,,+11 Xx, + o V log log + /2
end for Propagate X.-AI f X,,. end for return X.
Network architectures. Table 8 summa- rizes the network architecture used for each dataset. For toy datasets, we param- eterize :; () and :(:;:d)) with the architectures shown in Fig. 9. Specifically, FCBlock represents fully connected layer followed by a swish nonlinear activation (Ramachandran et al., 2017). As for MNIST, we consider a smaller version of Unet (Ho et al., 2020) by reducing the numbers of residual block, attention heads, and channels respectively to 1, 2. and 32. Unet and NCSN+ respectively correspond to the architectures appeared in Ho et (2020) and Song al. (2020).
Remarks on Table 1. We note that the values of our SB-FBSDE reported in Table are computed without the Langevin corrector due the computational constraint. For all other experiments, we adopt the Langevin corrector as it generally improves the performance (see Fig. 6). This implies that our results on CIFARIO, despite already being encouraging, may be further improved with the Langevin corrector.
Remarks on Fig. To estimating KL(pr> Pprior), we first compute the pixel-wise first and second moments given the generated samples X at the end of the forward diffusion. After fitting diagonal Gaussian to { we can apply the analytic formula for computing the KL divergence between two multivariate Gaussians.
Remarks on Fig. 6, To accelerate the sampling process with the Langevin corrector, for this experiment consider a reduced Unet (see Table 8) for CclebA. The FID scores on both datasets are computed with 10k samples. We stress, however, that the performance improvement using the Langevin corrector remains consistent across other (larger) architectures and if one increases the FID samples.
E ADDITIONAL EXPERIMENTS
Comparison to De Bortoli al. (2021) under same setup. To demonstrate the superior perfor- mance of our model, we conduct experiments with the exact same setup implemented in De Bortoli (2021). Specifically, we adopt the same network architecture (reduced U-net), image pre- processing (center-cropping 140 pixel and resizing to 32 x 32), step sizes (N=50), and horizon (0.5 second) for fair comparison. Comparing our Fig. 10b to De Bortoli et al. (2021) (see their Fig. 6), it is clear that images generated by model have higher diversity (e.g. color skin, facing angle, background, ett) and better visual quality. We conjecture that our performance difference may come from (i) the (in)sensitivity numerical discretization between our divergence objectives and their mean-man matching regression, and (ii) the foundational differences in how diffusion coefficients are designed.
23

---

Published as a conference paper at ICLR 2022
(a) Ground Truth
(b) SB-FBSDE Generated Image
Figure 10: Comparison between images generated by ground truth and SB-FBSDE reduced CelebA. Our SB-FBSDE trained under the same data pre-processing. network architecture and stepsizes implemented in De Bortoli al. (2021).
(a) SGM/50k
(b) SGM/SOk + SB/b/5k
(c) SGM/50K SB/F/5k SB/b/5k
Figure 11: Qualitative results at the different stages of training. (a) Results after 50k training iterations using SGM's regression loss. (b) Refine the results of Fig. 1la by training the backward policy using (18) with 5k iterations. (c) Refine the results of Fig. lla with full SB-FBSDE stage using (18.19).
SGM regression training + SB divergence- based training. Table 3 reports the FID (using 10k samples, without corrector steps) at different stages of CIFAR10 training. We first train the backward policy with SGM's 's's regression loss for a sufficient long iterations (50k) until the FID roughly converges. Then, switch to our alternate training (Alg. 3) using the divergence-based objectives. Crucially,
24

---

Published as a conference paper at ICLR 2022
Table 3: SGM regression training + SB divergence-bast training. We denote "SGM/50K" as "training 50k steps using SGM loss", and "SB/(f.b//5"" "training forward/backward policy with 5k steps using our divergence loss", and etc.
SGM/50k SB/b/5k
SGM//0K SB/f/5k + SB/b/5k
initialization
SGM/1OK
SGM/20K
SGM/50K
FID
448
41.37
35.47
33.68
13.35
11.85
with only 5k iterations of our divergence-based training, we drop the FID dramatically down to 13.35 from 33.68. With full stage of training (last column), the FID decreases even lower to 11.85. The qualitative results are provided in Fig. 11. Comparing Fig. la (corresponds to "SGM//OK" in Table 3) and Fig. 1lb (corresponds to "SGM/50k + SB/b/5k" in Table 3), it can be seen that the visible flaw and noise have been substantially improved.
Additional Figures
0
2
2
7
6
}
9
3
3
3
8
4
7
5
+
/
3
7
5
3
4
7
/
1
9
1
2
5
3
0
3
7
3
9
9
9
3
3
1
4
3
)
/
4
Y
2
7
~
9
2
/
S
9 7 4
3
8
4
8
9
2
9
5
/
3
9
2
1
1
4
5
8
/
4
2
6
9
$
8
4
6
9
8
2
2
6
3
7
/
2
8
7
7
6
2
7
8
2
4
(
7
7
2
3
7
7
8
3
8
o
7
6
7
2
9
6
1
5
9
/
4
S
3
3
8
5
8
5
5
9
9
2
8
3
7
4
9
1
5
2
3
5
3
/
6
9
7
3
2
3
7
/
5
9
7
5
6
7
7
7
1
9
4
3
7
1
1
2
7
9
0
/
5
4
5
5
5
8
9
7
5
8
7
7
/
8
7
0
8
4
3
2
7
3
7
8
9
7
3
f
A
7
3
2
3
3
3
2
5
4
9
7
3
3
7
3
7
2
9
,
3
5
5
S
7
=
7
7
1
/
6
>
3
9
5
1
/
9
d
4
9
3
4
4
Q
o
5
8
3
6
o
>
9
8
7
3
7
7
5
2
3
3
6
a
9
8
D
6
/
4
8
1
3
1
8
3
7
1
5
/
7
7
}
5
f
7
3
3
6
9
4
1
5
8
3
6
7
7
3
7
'
3
5
3
7
/
1
7
3
5
5
5
7
8
7
7
/
1
4
4
Figure 12: Uncurated samples generated by our SB-FBSDE on MNIST.
25

---

Published as a conference paper at ICLR 2022
Figure 13: Uncurated samples generated by our SB-FBSDE on resized CelebA
26

---

Published as a conference paper at ICLR 2022
Figure 14: Uncurated samples generated by our SB-FBSDE on CIFARIO.
27


Reviews:
[{"review_id": "kFJq094YaLA", "reviewer": "Reviewer_dQMT", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "STRENGTHS:\n\n-Overall the presentation of the paper is good with a clear presentation of the\nscore-based generative models and the Schrodinger bridge problem. \n\n-The idea of using the FBSDE framework is new and original. I think that the\nidea of using Feynman-Kac based representations to solve (non)linear PDEs is\ndefinitely a good idea. \n\n-The toy experiments and the generative modeling experiments are satisfactory.\n\nWEAKNESSES:\n\n-My first concern is with the formulation of the loss function that I find\nopaque to say the least. Looking at (14) it turns out that the loss function is\nin fact given by\n$\\int_0^T \\mathbb{E}[\\| \\hat{z}_t - g \\nabla \\log p_t^{SB} + z_t \\|^2]$ where\nthe expectation is taken over the path measure induced by the forward SDE. It is\nnot clear how $\\nabla \\log p_t^{SB}$ is computed. The classical ideas of [1,2]\ncannot be used since the $\\nabla \\log p_t^{SB}$ is not given by a\nOrnstein-Ulhenbeck process. Therefore the authors might rely on the idea used in\nthat approximate $\\nabla \\log p_t$ for general processes using their Langevin\ndiscretization. The authors should clearly state what loss function they use in\npractice (discretization of the time + approximation of $\\nabla \\log p_t$). Maybe they use the first part of (14) to compute the loss function but this is not clear either. The\npseudocode provided by the authors is not helpful and there is no anonymous\nrepository to check out the code.\n\n-My second concern is a consequence of the first one. The authors claim that\ntheir work is different from [3,4] but in practice Algorithm 1 is the same\ntraining algorithm as the one derived in [3]. Indeed because the authors propose\nan alternate minimization of the loss function they derive the loss function\nthey minimize at each step corresponds exactly to compute the IPF step\n(Algorithm 1 in the current paper is almost the same as [3, Algorithm 1], see also [4, Algorithm 1], with the only difference being in how the loss is computed which is not clear from the current paper, see previous comment) . In this sense\nthe algorithm provided by the authors corresponds to a rewriting of [3,4]. Hence I\nthink the claims of the authors concerning the originality of the method is overstated and misleading. Similarly it is misleading to state that no\nconnection with Langevin based sampling as been established for Schrodinger\nbridge models, see [3,4]. To summarize there is a clear overlap between the\nmethod proposed by the authors and the work of [3,4] which is not acknowledged in\nthe paper. \n\n-Another concern I have is related to the derivation of the algorithm. Given how\nthe authors derive Algorithm 1 there is no guarantee that 1) the procedure is\ngoing to converge and 2) if it converges that the limiting bridge is indeed the\nSchrodinger bridge. Both these facts can be obtained using the IPF approach, see\n[3, Proposition 6]. As of now Theorem 4 is only valid at equilibrium, i.e. for\nthe optimal set of controls $(z_t, \\hat{z}_t)$. For arbitrary controls \nTheorem 4 might still be valid (or at least in a lower-bound form) but this is not\nwhat is currently stated in the paper. The way the paper is written it looks like Algorithm 1 is motivated and justified by Theorem 4 which is not the case. \n\n\n-I have a minor concern with the contribution of the authors to the FBSDE\ntheory. It seems that Lemma 2 and the current FBSDE is not enough to provide a\nnon-linear Feynman-Kac representation. I find that this contribution is not\nclearly presented by the authors. Also the proof seems to be very heuristic, the\nauthors do not check the necessary regularity conditions to apply Ito\nformulas. They refer to the ``same regularity conditions in Lemma 2'' but Lemma\n2 is not informative because it is only stated that the functions ``G, f, h and\n$\\varphi$ satisfy mild regularity conditions''. The concept of viscosity\nsolution is not reintroduced in the paper (even in the supplementary material).\n\n-Finally I have also concerns with the experiments. In particular there is no\ncomparison with [3,4] which are currently the concurrent for score-based\ngenerative modeling using Schrodinger bridge. In particular, since Algorithm 1\nand [3, Algorithm 1] appear almost identical it is unfortunate that the authors do not\nprovide any comparison with this approach in similar setting (same number of\nsteps, stepsize and without correcting step). \n\nCOMMENTS:\n\n-It would be interesting to observe what is the generative model obtained after\none iteration of the algorithm, i.e. when the backward network is updated for\nthe first time. This should recover the generative model obtained using SGM and\nfurther steps are a refinement of the method. \n\n-It would be interesting to precisely quantify the influence of the corrector\n(which seems to constitute the only difference between Algorithm 2 and the sampling\nprocedure proposed in [3]).\n\n-``Poof'' and ``Remakrs'' in the title of the sections in the appendix.\n\n-In Algorithm 1 do the authors really use a gradient descent algorithm? Or is a\nmore efficient like ADAM is used to train the neural network?\n", "summary_of_the_paper": "In this paper the authors introduce a score-based generative model which relies\non a Schrodinger bridge formulation. More precisely the authors derive a loss\nfunction for the control of a forward/backward SDE. By alternating the\nminimization of the loss function between the forward control and the backward\ncontrol they obtain an approximation of the Schrodinger bridge. This loss\nfunction is obtained using the forward backward SDE (FBSDE) framework to obtain\na non-linear Feynman-Kac representation of the evolution of the\nlog-potentials. This theoretical/methodological contribution is accompanied by a\nexperiments in generative modeling on the MNIST/CelebA/CIFAR-10 datasets.", "summary_of_the_review": "To conclude I think that the idea of using FBSDE to derive the Schrodinger bridge\nis original. However the method has a lot of overlap with the\nworks of [3,4] which is not acknowledged by the authors. The authors do not\nprecise which loss function they use in practice and because they do not give\naccess to the code (or at least to a detailed pseudo code) it is hard to\ndetermine what is the original methodological contribution of the paper.\nI think that the paper is not mature enough and that a true justification of Algorithm 2 with FBSDE is still missing. Experimentally speaking the authors did not compare their work with existing SB methods.\nBased on these comments I recommend the rejection of the paper.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "ZldHoWLv0NO", "reviewer": "Reviewer_8tG7", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Strengths\n1. The proposed method is motivated from optimal stochastic control in a principled way. It bridges the gap between $p_T$ and the prior distribution in existing score-based techniques.\n2. Compared to other works based on the Schrodinger bridge, the approach described in this paper seems to be more scalable, without the need to hand design part of the transformation as in Wang et al., 2021, or iterative proportional filtering as in De Bortoli et al. 2021.\n\n## Weaknesses\n1. The efficiency of the proposed method is unclear in current writing. In Theorem 4, the loss function is described as integrals of expectations. However, it is unclear with respect to which random variable is the expectation taken. How do you estimate the expectation in the loss function? If the expectation is over $x_t$, don't you need to simulate the forward SDE in equation (12a) for each datapoint? Since equation (12a) depends on $z_t$, which is parameterized by a deep neural network, there is no easy way to sample $x_t$ without solving the SDE numerically. By contrast, in score-based generative modeling, $x_t$ is sampled as a noise-perturbed Gaussian with a closed form.\n\n2. There are many errors in writing that affect reading and understanding. For example, in page 3, $p_t^{(2)}$ is never defined. In Lemma 2, I am not sure how to understand the expression $v(t, x) = y(t, x)$. As the solution to an SDE, shouldn't $y$ be a stochastic process? How can a stochastic process equal to $v(t, x)$, which is a deterministic function? What's the definition of $y(t, x)$? Similarly, in Theorem 3, how is $y_t \\equiv y(t, x_t)$, what's the definition of $y(t, x_t)$, and how is $y_t$ a function of $x_t$? In Corollary 5, why does an ODE contain the stochastic term $g d w_t$? In Algorithm 1, why are there references to (23a) and (24)? Do you mean other equations?\n\n3. Authors reported better performance compared to prior optimal transport methods like Wang et al., 2021. However, is this because the network architecture used in this paper (NCSN++) is better than those used in previous methods? Will it be more fair to compare with the same network architecture?\n\n4. Isn't the framework proposed in this work exactly the same as score-based generative modeling? If you let the forward SDE be $d x_t = (f + gz_t )dt + gdw_t$, and parameterize the score network as $z_t + \\hat{z}_t$, then the reverse SDE is the same as (7b). In this case, you can derive Theorem 4 and Corollary 5 automatically from the theory of score-based generative models.", "summary_of_the_paper": "Inspired by recent work on score-based generative modeling, this paper proposes to solve Schrodinger bridges for generative modeling. Different from other Schrodinger bridge works, this paper connects the training to maximum likelihood, and provides a way to compute the log-likelihood of the model. The resulting method can shrink the gap between $p_T$ and the prior distribution, and produce high quality image samples and likelihood comparable to score-based generative models.", "summary_of_the_review": "The paper proposes a useful Schroding bridge framework for generative modeling that is simpler than previous counterparts. However, there are concerns on efficiency issues, writing errors, and fairness in experimental comparison. It is also unclear how much difference is this method compared to the original formulation of score SDEs.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "EFNuiY5Yfc4", "reviewer": "Reviewer_4oR9", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper introduced a new way to look at the SB problem via a non-linear Feynman-Kac representation of a PDE that defines the SB optimal path measure. This overall framework seems to provide significant new insight on how SB can be solved practically, which is a key merit of the work, and I applaud the authors for the finding. \nThe strengths of the paper are:\n* Proposing an interesting framework to approach the SB problem, via solving the SB optimization problem using a PDE to define a pair of equivalent forward and backward SDEs whose path measures correspond to the SB solution, and then solving this SDE using a Feynman-Kac type of representation. \n* Making a connection with prior work, i.e. generalizing the likelihood-based training of score-based generative models. \n\nNovelty aside, I do find the paper requires some major improvement in terms of clarity and a clearer discussion and analysis of limitations and the similarity/difference with prior work. Below are the weaknesses of the paper. \n* Clarity needs work: the narrative at times is not very clear, which makes it harder to grasp what is being solved. (See more detailed questions below)\n* Differences with some prior work (such as iterative proportional fitting) are not not well explained, even though the proposed likelihood training of the forward SDE and the backward SDE seems very similar. \n* Limitations are not properly discussed: training the inference process is done at the cost of an increased cost of computation, which is arguably one of the most important features of SGM. Convergence property of the proposed method not discussed, i.e. whether or not the proposed likelihood training stage will lead to convergence towards the SB solution. A thorough discussion of these limitations and perhaps an experiment on compute cost analysis would be needed. \n\nSome more detailed comments and questions:\n* Throughout the introduction, it is not very clear what “computing the log-likelihood objective of Schrodinger Bridge” means. The narrative needs more work. More precisely, it wasn’t very clear to me what the “model” is until I finished reading section 3, and it took me multiple passes. For example theorem 4 refers to the log-likelihood of SB, it’s the log-likelihood of which model exactly? Fixing the inference SDE (i.e. $z_t$) and looking at the marginal likelihood induced by $\\hat{z}_t$? Is it a likelihood or a lower bound?  The (bold) $z$ used in the previous theorem seems to not depend on any parametric form (as it solves the SB problem), but here it suddenly becomes parametric. \n* End of 2.1, it is not clear why having a more flexible framework will help mitigate the instability problem just mentioned.\n* Is the order of the arguments of $h$ in (11) incorrect?\n* The presentation of Lemma 2 is a bit confusing, is it correct to say that despite the randomness induced by the Brownian motion used in the ito integral, the solution of y (an SDE) will still be a deterministic and smooth function, as it is after all a solution of the PDE? What’s confusing is that v is a solution of a PDE and y is a solution of an SDE, which is random by nature. \n* Presentation of theorem 3 also needs work: is the bold $z$ related to the regular $z$ used in Lemma 2? The notation hasn’t been introduced.\n* Last paragraph of page 5: what does the new interpretation of optimal control bring to us? Is there any practical benefit of it? \n* Last paragraph of 4.1, I am not sure how meaningful the comparison is with some other OT methods, especially since the SB problem is connected to “entropy-regularized” OT, and is not exactly OT.  \n\nMinor points / typos:\n* End of page 6, is it a maximization (of likelihood) instead of min?\n* The last sentence of page 7: can be founded -> found\n", "summary_of_the_paper": "Inspired by stochastic control and forward-backward SDE theory, this paper proposed an iterative algorithm to approach the Schrodinger bridge (SB) problem, generalizing variational likelihood-based training of score-based generative models (SGM). The proposed method is then used to train both the generative process and the inference process of SGMs, recasting the generative modeling problem as the problem of SB. The author then performed experiments on standard image datasets. \n", "summary_of_the_review": "The paper introduces a new framework for likelihood-based training of the forward and backward SDEs that are inspired by the Schrodinger bridge problem, which is quite novel. But the paper is not very well written and the practical limitations of the proposed method are not sufficiently addressed. Therefore I do not vote for acceptance given its current form. \n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "43Mmf5TGku5", "reviewer": "Reviewer_Sm2V", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "[Novelty and Significance]\n\nThe paper presents an elegant framework for likelihood training of Schrödinger bridge (SB)-based generative models. SB generative modeling is a developing area in the field of diffusion generative models and this paper makes a significant contribution to the area. It draws interesting connections to (and generalizes) the framework of previous diffusion-based generative models (e.g., SBGM). \n\n[Writing and Clarity]\n\nThe paper is written fairly well, albeit for the expert reader. There exist some clarity issues:\n\n- For Theorem 1, please provide a proof-sketch and/or cite the exact theorem number. In particular, it is unclear to me how to arrive at 7(a) and 7(b) from the Kolmogorov equations in (6). Is this after time-reversal?\n- For the proofs, please include a brief description of each step. In the current state, some steps of the proofs are unclear.\n    - In the proof of Theorem 3, why do you begin with (1) as the reference measure $\\mathbb{P}$ and not 7(a)?\n    - In several steps, $dt^2$ and $dtd\\mathbf{w}_t$ terms have been dropped. Please mention this in the proof text for clarity.\n    - In Eq. (20), I do not understand how you go from $\\nabla \\cdot(g \\hat{z}_t-f)$ to $-\\hat{z}_t^{\\top}(g \\nabla \\log p_t^{\\mathrm{SB}}) - \\nabla \\cdot f$.\n\nCorrections:\n\n- Eq. (15) is not an ODE and looks like a typographical mistake — the $g\\mathrm{d}\\mathbf{w}_t$ term should not be there.\n- \"... which can be _probability_ expensive on high-dimensional datasets\" — *prohibitively*.\n- Section A.1: ito —> Ito.\n- Section A.1 before Eq. (17). What is $b$? I believe it should be $f$.\n\n[Empirical Evaluation]\n\nThe empirical evaluation, although not extensive, is reasonable.\n\nOne thing that is unclear from Table 2 are the primary baselines; in my opinion, they should be Multi-stage SB and SGMs. It's good that the authors have reported results from several previous works; however, most of the results are not directly comparable, so it is important to specify what the main baselines are. A particular example is DOT which is not a generative model in itself but a sample improvement technique that operates on a pretrained GAN. \n\nThe authors are missing the following work in Table 2 in the optimal transport model class:\n\nAnsari, Abdul Fatir, Ming Liang Ang, and Harold Soh. \"Refining deep generative models via discriminator gradient flow.\" *arXiv preprint arXiv:2012.00780* (ICLR 2021).\n\nOn reproducibility: From the paper, it appears that the practical implementation has several moving parts and is not straightforward. The paper does not contain a reproducibility statement and also does not provide enough details for an expert reader to reproduce the results. I encourage the authors to release their code as supplementary material and include further details of the practical implementation in the Appendix for the benefit of the research community.\n\n[Questions]\n\n- Can you elaborate what exactly is meant by \"... SGM by enlarging the class of diffusion processes to accept *nonlinear* drifts ...\"?\n- How is $\\nabla\\log p^\\mathrm{SB}_t$ computed?\n- What is the form of the function $f$ in your practical framework/implementation?\n\n[Suggestions]\n\n- It may be better to move the paragraph on \"Connection to flow-based models\" before \"In practice, we parameterize the forward...\". The current structure breaks the flow of the reader.\n- To make the paper more accessible to readers unfamiliar with Schrödinger bridges, it would be helpful if you provide a brief review in the Appendix.", "summary_of_the_paper": "This paper presents a framework for likelihood-based training of Schrödinger bridge-based generative models using the theory of Forward-Backward SDEs. In doing so, the paper draws relations to score-based generative models (SBGMs) and shows that the proposed framework is a generalization of the SBGM framework. The proposed framework also provides additional control to the forward SDE to reach the prior distribution unlike the case of SBGM. The authors propose a practical training algorithm which alternates between likelihood training of the forward and backward controls. Experiments on multiple image generation benchmarks demonstrate that the proposed SB-FBSDE performs well as a generative model of the data.", "summary_of_the_review": "The paper presents a novel framework for likelihood training of Schrödinger bridge (SB)-based generative models. The technical contribution of this paper is significant: it generalizes the SBGM framework and provides a practical algorithm for likelihood training of SB models. Although the paper has some clarity issues, I believe these can be fixed in the revision. Overall, I think this is a good paper and I recommend acceptance.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper presents a new computational framework, grounded on Forward-Backward SDEs theory, for the log-likelihood training of Schrödinger Bridge and provides theoretical connections to score-based generative models. The presentation of the results is not satisfactory (the algorithm should be clarified in several places and the notation is not accurate which raises doubts about the soundness of the method). The paper is thus very hard to read for the non-experts on the subject. Furthermore, some reviewers raise concerns about the similarity of this method to other algorithms that were never cited in the paper. Finally, the empirical analysis, as of now, is limited.

In the rebuttal the authors carefully addressed lots of the comments. However paper's presentation still needs to be substantially improved (de-densification of the paper would be extremely important since now the main narrative is very convoluted). The authors made several changes in the manuscript, but detailed discussion regarding training time complexity still seems to be missing (main body and the Appendix) in the new version of the manuscript, even though this was one of the main raised concerns. Overall, the manuscript requires major rewriting. Since the comments regarding the content were successfully addressed (the reviewers are satisfied with detailed answers given by the authors), the paper satisfies the conference bar and can be accepted.
Decision: Accept (Poster)