L IKELIHOOD T RAINING OF S CHRÖDINGER B RIDGE USING F ORWARD -B ACKWARD SDE S T HEORY
A BSTRACT
Schrödinger Bridge (SB) is an entropy-regularized optimal transport problem that has received increasing attention in deep generative modeling for its mathematical ﬂexibility 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 .
I NTRODUCTION
[Figure]
(SB) transform between two distributions.
SGM requires pre-specifying the data to noise diffusion
SB instead learns the process
Score-based Generative Model (SGM; Song et al.
(2020)) is an emerging generative model class that has achieved remarkable results in synthesizing high-ﬁdelity data (Song & Ermon, 2020; Kong et al., 2020a;b).
Like many deep generative models, SGM seeks to ﬁnd nonlinear functions that transform simple distributions (typically Gaussian) into complex, often intractable, data distributions.
In SGM, this is done by ﬁrst diffusing data to noise through a stochastic differential equation (SDE); then learning to reverse this diffusion process by regressing a network to the score function ( i.e. the gradient of the log probability density) at each time step (Hyvärinen & Dayan, 2005).
This reversed process thereby deﬁnes the generation (see Fig. 1).
Despite its empirical successes, SGM admits few limitations.
First, the diffusion process has to obey a simple form ( e.g. linear or degenerate drift) in order to compute the analytic score function for the regression purpose.
Secondly, the diffusion process needs to run to sufﬁciently 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 (Jolicoeur-Martineau et al., 2021), thereby limiting their practical usages compared to e.g. GANs or ﬂow-based models (Ping et al., 2020; Karras et al., 2020b).
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).
SB – as an entropy-regularized optimal transport problem – seeks two optimal policies that transform back-and-forth between two arbitrary distributions in a ﬁnite horizon.
The similarity between the two problems ( i.e. both involve transforming distributions) is evident, and the additional ﬂexibility from SB is also attractive.
To enable SB-inspired generative training, however, previous works require either ad-hoc multi-stage optimization or retreat to traditional SB algorithms, e.g. Iterative
Proportional Fitting (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-likelihood 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 difﬁculty 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 is characterized by two coupled partial differential equations (PDEs) (Léonard, 2013), the latter integrates over a 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 et al., 2018) – paves an elegant way to connect the two objectives.
The implication of our ﬁndings 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 ﬂow-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.g. a SGM-inspired Langevin corrector (Song & Ermon, 2019).
The resulting method, SB-FBSDE , generates encouraging images on MNIST, CelebA, and CIFAR10 and outperforms prior optimal transport models by a large margin.
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-based objectives (see § 3.2).
Secondly, neither of the prior methods focuses on log-likelihood training, which is the key ﬁnding in SB-FBSDE to bridge connections to SGM and adopt modern training improvements.
Indeed, due to the difference in the underlying SDE classes, 1 their connections to SGM can only be made after time discretization by carefully choosing each step size (De Bortoli et 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-Backward SDEs theory, for computing the log-likelihood 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 ﬂexibility 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 transportinspired baselines on synthesizing high-ﬁdelity images and is comparable to other existing models.
Notation.
We denote p SDE t ( X t ) as the
marginal density
driven by some SDE process X ( t ) ≡ X t until the time step t ∈ [0 , T ] .
The time direction is aligned throughout this article such that p 0 and p T respectively correspond to the data and prior distributions.
The gradient, divergence, and Hessian of a function f ( x ) , where x ∈ R n , will be denoted as ∇ x f ∈ R n , ∇ x · f ∈ R , and ∇ 2 x f ∈ R n × n .
P RELIMINARIES 2. S CORE - BASED G ENERATIVE M ODEL (SGM)
Given a data point X 0 ∈ R n sampled from an unknown data distribution p data , SGM ﬁrst progressively diffuses the data towards random noise with the following forward SDE:
[Equation]
where f ( · , t ) : R n → R n , g ( t ) ∈ R , and W t ∈ R n are the drift, diffusion, and standard Wiener process.
Typically, g ( · ) is some monotonically increasing function such that for sufﬁciently large time steps, we have p (1) T ≈ p prior resemble some prior distribution ( e.g. Gaussian) at the terminal horizon T .
Reversing (1) yields another SDE 2 that traverses backward in time (Anderson, 1982):
[Equation]
where p (1) t corresponds to the marginal density of SDE (1) at time t , and ∇ x log p (1) t 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 t ∈ [0 , T ] ; in other words, p (1) t ≡ p (2) t .
When the drift f is of simple structure, for instance linear
(Ho et al., 2020) or simply degenerate (Song & Ermon, 2019), the conditional score function ∇ x log p (1) t ( X t | X
= x 0 ) ≡ ∇ x log p t | x 0 admits an analytic solution at any time t .
Hence, SGM proposes to train a parameterized score network s ( t, x ; θ ) by regressing its outputs to the ground-truth values, i.e. E [ λ ( t ) (cid:107) s ( t, X t ; θ ) − ∇ x log p t | x 0 (cid:107) 2 ] , where the expectation is taken over the SDE (1).
In practice, λ ( t ) 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 lower-bounded as follows:
[Equation]
where s t ≡ s ( t, x ; θ ) and the expectation is taken over the SDE (1) given a data point X 0 = x 0 .
This objective (3) suggests a principled choice of λ ( t ) :
=
g ( t ) 2 .
After training, SGM simply substitutes the score function with the learned score network s ( t, x ; θ ) to generate data from p prior ,
[Equation]
It is important to notice that p prior needs not equal p (1) T in practice, and the approximation is close only through a careful design of (1).
Notably, designing the diffusion g ( t ) can be particularly problematic as it affects both the approximation p (1) T ≈ p prior and the training via the weighting λ ( t ) ; hence can lead to unstable training (Nichol & Dhariwal, 2021).
In contrast, Schrödinger Bridge considers a more ﬂexible framework for designing the forward diffusion that requires minimal manipulation.
2. S CHRÖDINGER B RIDGE (SB)
Following the dynamic expression of SB (Pavon & Wakolbinger, 1991; Dai Pra, 1991), consider
[Equation]
(5) where Q ∈ P ( p data , p prior ) belongs to a set of path measure with p data and p prior as its marginal densities at t = 0 and T .
On the other hand, P 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 et al. (2021); Pavon & Wakolbinger (1991); Caluya & Halder (2021)) .
Let Ψ( t, x ) and (cid:98) Ψ( t, x ) be the solutions to the following PDEs:
[Equation]
Then, the solution to the optimization (5) can be expressed by the path measure of the following forward (7a), or equivalently backward (7b), SDE:
[Equation]
where ∇ x log Ψ( t, X t ) and ∇ x log (cid:98) Ψ( t, X t ) are the optimal forward and backward drifts for SB.
[Figure]
Similar to the forward/backward processes in SGM, the stochastic processes of SB in (7a) and (7b) are also equivalent in the sense that ∀ t ∈ [0 , T ] , p (7a) t ≡ p (7b) t ≡ p SB t .
In fact, its marginal density obeys a factorization principle: p SB t ( X t ) =
Ψ( t, X t ) (cid:98) Ψ( t, X t ) .
To construct the generative pipeline from (7b), one requires solving the PDEs in (6) to obtain (cid:98) Ψ .
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 ( p data (cid:28) p prior ) with multiple stages, p data (cid:28) p middle (cid:28) p prior , so that each segment admits an analytic solution (Wang et al., 2021), or consider the following half-bridge ( p data ← p prior vs. p data → p prior ) optimization (De Bortoli et al., 2021;
Vargas et al., 2021),
[Equation]
which can be solved with IPF algorithm (Kullback, 1968) starting from Q (0) :
= 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-likelihood of SGM (3).
A PPROACH
We motivate our approach starting from some control-theoretic 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-afﬁne SDEs with additive noise:
[Equation]
It is clear that the control-afﬁne SDE (8) includes all SDEs (1,2,4,7) appearing in § 2 by considering ( A , B , C ) :
= ( f, I , g ) and different interpretations of the control variables u ( t, X t ) .
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 ∇ x log p t | x 0 in SGM or ∇ x log (cid:98) Ψ in SB, that minimizes some objectives, (3) vs. (5), while subjected to some control-afﬁne SDEs, (1,2) vs. (7).
In what follows, we will show that a speciﬁc mathematical methodology in nonlinear SOC literature – called
Forward-Backward SDEs theory (FBSDEs; see Ma et al.
– links the optimality condition of
SB (6) to the log-likelihood objectives of SGM (3).
All proofs are left to Appendix B.
3. F ORWARD -B ACKWARD SDE S (FBSDE S ) R EPRESENTATION FOR SB
The theory of FBSDEs establishes an innate connection between different classes of PDEs and forward-backward SDEs.
Below we introduce the following connection related to our problem.
Lemma 2 (Nonlinear Feynman-Kac; 3 Exarchos & Theodorou (2018)) .
Consider the coupled SDEs
[Equation]
where the functions f , G , h , and ϕ satisfy proper regularity conditions 4 so that there exists a pair of unique strong solutions satisfying (9).
Now, consider the following second-order parabolic PDE and suppose v ( t, x ) ≡ v is once continuously differentiable in t and twice in x , i.e. v ∈ C 1 , 2 ,
[Equation]
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:
[Equation]
and
[Equation]
Lemma 2 states that solutions to a certain class of nonlinear (via the function h in (10)) PDEs can be represented by solutions to a set of forward-backward 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 t 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 t (hence also Z t ) is a time-varying function of X t .
Since the v appearing in the nonlinear Feynman-Kac relation (11) takes the random vector X t as its argument, v ( t, X t ) shall also be understood as a random variable.
Finally, it is known that the original (deterministic) PDE solution v ( t, x ) can be recovered by taking conditional expectation, i.e. v ( t, x ) =
E X t ∼ (9a)
[ Y t | X t = x ] and G ( t, x ) T ∇ x v ( t
[Equation]
Since it is often computationally favorable to solve SDEs rather than PDEs, Lemma 2 has been widely used as a 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 t .
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 are given by two coupled PDEs in (6), remains unclear.
Below we derive a similar FBSDEs representation for SB.
Theorem 3 (FBSDEs to SB optimality (6)) .
Consider the following set of coupled SDEs,
[Equation]
where f and g satisfy the same regularity conditions in Lemma 2 (see Footnote 4), and the boundary conditions are given by X (0)
= x 0 and Y T + (cid:98) Y T = log p prior ( X T ) .
Suppose Ψ , (cid:98) Ψ ∈ C 1 , 2 , then the nonlinear Feynman-Kac relations between the FBSDEs (13) and PDEs (6) are given by
[Equation]
Furthermore, ( Y t , (cid:98) Y t ) obey the following relation:
[Equation]
The FBSDEs for SB (13) share a similar forward-backward structure as in (9), where (13a) and (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 t =
g ∇ x log Ψ .
In other words, these FBSDEs provide a local representation of log Ψ and log (cid:98) Ψ evaluated on the optimal path governed by (7a).
Since Z t and (cid:98) Z t can be understood as the forward/backward policies, in a similar spirit of policy-based methods (Pereira et al., 2020; Schulman et al., 2015), that guide the
SDE processes of SB, they sufﬁciently characterize the SB model.
Hence, our next step is to derive a proper training objective to optimize these policies.
3. L OG - LIKELIHOOD C OMPUTATION OF SB
Theorem 3 has an important implication: It suggests that given a 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-likelihood of the data point x 0 , i.e. E (cid:104) Y
+
[Equation]
Y 0 |
[Equation]
=
[Equation]
=
[Equation]
= log
[Equation]
data ( x 0 ) , where X t is sampled from (13a).
We now state our main result, which makes this observation formal:
Theorem 4 (Log-likelihood of SB model) .
Given the solution satisfying the FBSDE system in (13), the log-likelihood of the SB model ( Z t , (cid:98) Z t ) , at a data point x 0 , can be expressed as
[Equation]
where the expectation is taken over the forward SDE (13a) with the initial condition X 0 = x 0 .
Similar to (3), Theorem 4 suggests a parameterized lower bound to the log-likelihoods, i.e. log p SB 0 ( x 0 ) ≥ L SB ( x 0 ; θ, φ ) where L SB ( x 0 ; θ, φ ) shares the same expression in (16) except that
Z t ≈ Z
( t, x ; θ ) and (cid:98) Z t
≈ (cid:98) Z ( t, x ; φ ) are approximated with some parameterized models ( e.g.
DNNs).
Note that ∇ x log p SB t is intractable in practice for any nontrivial ( Z t , (cid:98) Z t ) .
Hence, we use the divergence-based objective in (16) as our training objective of both policies.
Connection to score-based models.
Recall Fig.
2 and compare the parameterized log-likelihoods of
SB (16) to SGM (3); one can verify that L SB collapses to L SGM when ( Z t , (cid:98) Z t ) :
= ( 0 , g s t ) .
From the SB perspective, this occurs only when p (1) T = p prior .
Since no effort is required in the forward process to reach p prior , the optimal forward control Z t , by deﬁnition, degenerates; thereby making the backward control (cid:98) Z t collapses to the score function.
However, in any case when p (1) T (cid:54) =
p prior , for instance when the diffusion SDEs are improperly designed, the forward policy Z t steers the diffusion process back to p prior , while its backward counterpart (cid:98) Z t 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 log-likelihood training from SGM to SB.
Connection to ﬂow-based models.
Interestingly, the log-likelihood computation in Theorem 4, where we use a path { X t } t ∈ [0 ,T ] sampled from a data point X 0 to parameterize its log-likelihood, resembles modern training of (deterministic) ﬂow-based models (Grathwohl et al., 2018), which have recently been shown to admit a close relation to SGM (Song et al., 2020;
Gong & Li, 2021).
The connection is built on the concept of probability ﬂow – which suggests that the marginal density of an
SDE can be evaluated through an ordinary differential equation (ODE).
Below, we provide a similar ﬂow representation for SB, further strengthening their connection to modern generative models.
Corollary 5 (Probability ﬂow for SB) .
The following ODE characterizes the probability ﬂow of the optimal processes of SB (7) in the sense that ∀ t, p (17) t
[Equation]
One can verify (see Remark 10 in § B) that computing the log-likelihood of this ODE model (17) using ﬂow-based training techniques indeed recovers the training objective of SB derived in (16).
3. P RACTICAL I MPLEMENTATION
In this section, we detail the implementation of our FBSDE-inspired SB model, named SB-FBSDE .
Training process.
We treat the log-likelihood in (16) as our training objective, where the divergence can be can be estimated efﬁciently 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 elserun regression-based objectives.
For low-dimensional datasets, we simply perform joint optimization, max L SB ( x 0 ; θ, φ ) , to train the parameterized policies Z ( · , · ; θ ) and
(cid:98) Z ( · , · ; φ ) .
For higher-dimensional ( 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 a reply buffer and refreshing them in a lower frequency basis (around 1500 iterations).
Although this implies that the gradient path w.r.t.
θ will be discarded, we can leverage the symmetric structure of SB and re-derive the log-likelihood for the sampled noise, i.e. L SB ( x T ) , 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,
[Equation]
Our training process is summarized in Alg.
While the joint training scheme in Alg. 2 resembles recent diffusion ﬂow-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 yet at the cost of introducing memory complexity.
We highlight these ﬂexible training procedures arising from the uniﬁed viewpoint provided in Theorem 4.
Hereafter, we refer to each cycle, i.e. 2 K training steps, in Alg.
3 as a training stage of SB-FBSDE.
Generative process.
While the generative processes for SB can be performed as simply as propagating (7b) given the trained policy (cid:98) Z ( · , · ; φ ) , it has been constantly observed that adopting Langevin sampling to the generative process greatly improves performance (Song et al., 2020).
This procedure, often referred to as the Langevin corrector , requires knowing the score function ∇ x log p t .
For SB, we can estimate its value by recalling (see § 2.2) that Z t + (cid:98) Z t =
g ∇ x log p SB t .
This results in the following predictor-corrector sampling procedure (see Alg. 4 in Appendix D for more details).
[Equation]
where (cid:15) ∼ N ( 0 , I ) and σ t is the pre-speciﬁed noise scales (see (59) in Appendix D).
Limitations & efﬁciency.
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.5 times compared to SGM, we empirically observe that the divergence-based training converges much faster per iteration than standard regression.
As a result, SB-FBSDE admits comparable training time (+6.8% in our CIFAR10 experiment) compared to SGM, yet with a substantially fewer sampling time (-80%) due to adopting nonlinear SDEs.
[Figure]
T=
T=0. T=0.
T=1.
(b) Checkerboard
T=0.
T=0. T=0.
T=1.
Setups.
We testify SB-FBSDE on two toy datasets and three image datasets, i.e. MNIST, CelebA,and CIFAR10.
p prior is set to a zero-mean Gaussian whose variance varies for each task and can be computed according to Song & Ermon (2020).
We parameterize Z ( · , · ; θ ) and (cid:98) Z ( · , · ; φ ) with residual-based networks for toy datasets and consider Unet (Ronneberger et al., 2015) and NCSN++ (Song et al., 2020) respectively for MNIST/CelebA and CIFAR10.
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.
as 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, we set the horizon T = 1.0 and solve the SDEs via the Euler-Maruyama method.
The interval [0 , T ] is discretized into 200 steps for CIFAR10 and 100 steps for all other datasets, which are much fewer than the ones in SGM ( ≥ 1000 steps).
Other details are left in Appendix D.
Toy datasets.
We ﬁrst validate our joint optimization ( i.e. Alg 2) on generating a mixture of Gaussian and checkerboard (adopted from Grathwohl et al.
as the representatives of continuous and discontinuous distributions.
Figure 3 shows how the learned policies, i.e. Z ( · , · ; θ ) and (cid:98) Z ( · , · ; φ ) , construct the vector ﬁelds that progressively transport samples back-and-forth between p prior and p data .
The vector ﬁelds can be highly nonlinear and dissimilar to each other.
This resembles neither SGMs, whose forward vector ﬁeld must obey linear structure, nor ﬂow-based models, whose vector ﬁelds 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 CIFAR10 are presented in Fig.
4, which clearly suggest that our SB-FBSDE is able to synthesize high-ﬁdelity images.
More uncurated images can be founded in Appendix E. Regarding the quantitative evaluation, Table 1 summarizes the negative log-likelihood (NLL; measured in bits/dim) and the Fréchet Inception Distance score (FID; Heusel et al.
on CIFAR10.
For our SB-FBSDE, we compute the NLL on the test set using Corollary 5, in a similar vein to SGMs and ﬂow-based models, and report the FID over 50k
[Table]
[Figure] [Figure]
FID w.r.t the training set.
Notably, our SB-FBSDE achieves 2.98 bits/dim and 3.18 FID score on
CIFAR10, which is 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 a large margin in terms of the sample quality.
More importantly, it enables log-likelihood computations that are 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 ( · , · ; θ ) ≡ Z θ plays an essential role in governing samples towards p prior .
Here, we validate this conjecture by computing the KL divergence between the terminal distribution induced by Z θ , i.e. p (13a) T , and the designated prior p prior .
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 := α t X t ) base drifts, our SB-FBSDE generates terminal distributions that are much closer to p prior .
Note that the values of SGM remain unchanged throughout training since SGM relies on pre-speciﬁed diffusion.
This is in contrast to our SB-FBSDE whose forward policy Z θ gradually shortens the KL gap to p prior , thereby providing a 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 CIFAR10 throughout training.
Since the SDEs are often solved via the Euler-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 SB t .
We emphasize this improvement as the beneﬁt gained from applying modern generative training techniques based on the solid connection between SB and SGM.
C ONCLUSION
In this work, we present a novel computational framework, grounded on Forward-Backward 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 ﬁndings 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 CIFAR10, showing encouraging results in synthesizing high-ﬁdelity samples while retaining the rigorous mathematical framework.
A CKNOWLEDGMENTS
The authors would like to thank Ioannis Exarchos and Oswin So for their generous involvement and helpful supports during the rebuttal.
The authors would also like to thank Marcus A Pereira and
Ethan N Evans for their participation and kind discussion in the early stage of project exploration.
A UTHOR C ONTRIBUTIONS
The original idea of solving the PDE optimality of 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 B 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 to code development.
R EPRODUCIBILITY S TATEMENT
Our training algorithms are detailed in Alg. 1, 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 sufﬁcient information for readers of interests to reproduce our results.
As we strongly believe in the merit of open sourcing, we 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.
A I NTRODUCTION OF S CHRÖDINGER B RIDGE
In this subsection, we provide a brief review for Schrödinger Bridge (SB) and some reasonings for
Theorem 1.
The SB problem, at its classical form, considers the following optimization (Dai Pra, 1991; Pavon & Wakolbinger, 1991),
[Equation] [Equation]
where U := { u : [0 , T ] × R n (cid:55)→ R n |(cid:104) u , u (cid:105) < ∞} .
The optimization (22) characterizes a standard stochastic optimal control (SOC) programming with energy ( i.e. 12 (cid:107) u (cid:107) 2 ) minimization except with an additional terminal boundary condition.
The optimality conditions to (22) are given by
[Equation]
where ψ ( t, x ) ∈ C 1 , 2 is known as the value function in SOC literature and p ∗ ( t, x ) ∈ C 1 , 2 is the associated optimal marginal density.
∆ denotes the Laplace operator.
Equations (23a) and (23b) are respectively the Kolmogorov’s backward and forward PDEs, also known as Hamilton-Jacobi-Bellman and Fokker-Planck equations.
The SB system can be obtained by applying the Hopf-Cole (Hopf, 1950; Cole, 1951) transformation ( ψ, p ∗ ) (cid:55)→ (Ψ , (cid:98)
[Equation]
In this work
[Equation]
we consider a recent generalization of (22) and (24) to an SDE class with nonlinear drift, afﬁne control
[Equation]
and time-varying diffusion.
We synthesize their results below.
Theorem 6 (SB optimality; Caluya & Halder (2021)) .
Consider the following optimization
[Equation]
where g ( t ) is uniformly lower-bounded and f ( t, X t ) satisﬁes Lipschitz conditions with at most linear growth in x .
Then, the Hopf-Cole transformation to (25) becomes
[Equation]
with the same boundary conditions Ψ(0 , · ) (cid:98) Ψ(0 , · ) = p 0 , Ψ( T, · ) (cid:98) Ψ( T, · ) =
p T .
The optimal control to (25) is thereby given by
[Equation]
t, X t ) .
Proof.
See Section III and Theorem 2 inf Caluya & Halder (2021).
Theorem 6 is particularly attractive to us since its SDE corresponds exactly to the one appearing in score-based generative models.
One can recover Theorem 1 by
(i) Following Pavon & Wakolbinger (1991), we know that the objective in (25) is equivalent to D KL ( Q || P ) by an application of Girsanov’s Theorem.
(ii) Equation (26) is exactly (6) with (cid:15)
= 12 .
Furthermore, substituting the optimal control (27) to the stochastic process in (25) yields the optimal forward SDE in (7a).
(iii)
Finally, reversing the SDE (7a) from forward to backward following Anderson (1982),
[Equation]
and recalling the factorization principle, log p SB t ( · ) = log Ψ
( t, · )
+ log (cid:98) Ψ( t, · ) , from Equation (4.15) in Chen et al.
(2021) yield the optimal backward SDE in (7b)
B P ROOFS AND R EMARKS IN S ECTION
In this section, we provide proofs for all of our theorems.
We following the same notation by denoting p SDE t ( X t ) as the marginal density driven by some SDE process X ( t ) ≡ X t until the time step t ∈ [0 , T ] .
Gradient and Hessian of a function f ( x ) , where x ∈ R n , will respectively be denoted as ∇ x f ∈ R n and ∇ 2 x f ∈ R n × n .
Divergence and Laplace operators will respectively be denoted as ∇· and ∆ .
Note that ∆ = ∇ · ∇ .
For notational brevity, we will only keep the subscript x for multivariate functions.
Finally, Tr( A ) denotes the trace of a square matrix A .
We ﬁrst 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 d W t 2 and d t are of the same scale and keeps the expansion up to O (d t ) .
Lemma 7 (Itô formula; Itô (1951)) .
Let v ∈ C 1 , 2 and let X t be the stochastic process satisfying
[Equation]
Then, the stochastic process v ( t, X t ) is also an Itô process satisfying
[Equation]
Next, the following lemma will be useful in proving Theorem 3.
Lemma 8.
The following equality holds at any point x ∈ R n such that p ( x ) (cid:54) = 0 .
[Equation]
Proof.
[Equation]
Assumptions Before stating our proofs, we provide the assumptions used throughout the paper.
These assumptions are adopted from stochastic analysis for SGM (Song et al., 2021; Yong & Zhou, 1999; Anderson, 1982), SB (Caluya & Halder, 2021), and FBSDE (Exarchos & Theodorou, 2018;
Gorodetsky et al., 2015).
(i) p prior , p data ∈ C 2 with ﬁnite second-order moment.
(ii) f
and g are continuous functions, and | g ( t ) | 2 > 0 is uniformly lower-bounded w.r.t.
t .
(iii) ∀ t ∈ [0 , T ] , we have f ( t, x ) , ∇ x log p t ( x ) , ∇ x log Ψ( t, x ) , ∇ x log (cid:98) Ψ( t, x ) , Z ( t, x ; θ ) , and (cid:98) Z ( t, x ; φ ) Lipschitz and at most linear growth w.r.t.
x .
(iv) Ψ , (cid:98) Ψ ∈ C 1 , 2 .
h , and ϕ are continuous functions.
h satisﬁes quadratic growth w.r.t.
x uniformly in t .
(v) ∃ k > 0 : p SB t ( x ) =
O (exp −(cid:107) x (cid:107) 2 k ) as x → ∞ .
Assumptions (i) (ii) (iii) are standard conditions in stochastic analysis to ensure the existenceuniqueness of the SDEs; hence also appear in SGM analysis (Song et al., 2021).
Assumption (iv) allows applications of Itô formula and properly deﬁnes the backward SDE in FBSDE theory.
Finally, assumption (v) assures the exponential limiting behavior when performing integration by parts.
Now, let us begin the proofs of Theorem 3, 4, and Corollary 5.
Theorem 3 (FBSDEs to SB optimality (6)) .
Consider the following set of coupled SDEs,
[Equation]
where f and g satisfy the same regularity conditions in Lemma 2 (see Footnote 4), and the boundary conditions are given by X (0)
= x 0 and Y T + (cid:98) Y T = log p prior ( X T ) .
Suppose Ψ , (cid:98) Ψ ∈ C 1 , 2 , then nonlinear Feynman-Kac relations between the FBSDEs (13) and PDEs (6) are given by
[Equation]
Furthermore, ( Y t , (cid:98) Y t ) obey the following relation:
[Equation]
Proof.
Similar to how the original nonlinear Feynman-Kac ( i
[Equation]
e. Lemma 2) can 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 t ) w.r
[Equation]
t. the optimal forward SDE (7a).
[Equation]
From the PDE dynamics (6)
[Equation]
we know that
[Equation]
The ﬁrst term in the RHS can be readily canceled out with the related f -term in (32).
The second term can also be canceled out using the fact that ∇ 2 x log Ψ
=
1Ψ ∇ 2 x Ψ − 1Ψ 2 ∇ x Ψ ∇ x Ψ T .
Hence, we are left with
[Equation]
Likewise, applying Itô lemma to log (cid:98) Ψ( t, X t ) , where X t follows the SDE in (7a),
[Equation]
[Equation]
but now noticing that the dynamics of ∂ log (cid:98)
[Equation]
become
[Equation]
Only the ﬁrst term in the RHS will be canceled out in (34).
Hence, we are left with
[Equation] [Equation]
Notice that the trace terms above can be simpliﬁed to
[Equation]
where the last equality follows by Lemma 8.
Substituting this result back to (35)
[Equation]
we get d log (cid:98) Ψ =
[Equation]
Finally
[Equation]
by rewriting (33) and (36) with the nonlinear Feynman-Kac in (31) yields
[Equation]
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 not exist (Pardoux & Peng, 1992).
For the completeness, one shall understand them in the sense of v ( t, x ) =
lim (cid:15) →∞ v (cid:15) ( t, x ) uniformly in ( t, x ) over a compact set.
Here v (cid:15) ( t, x ) is the classical solution to (10) with ( f (cid:15) , G (cid:15) , h (cid:15) , ϕ (cid:15) ) converge uniformly toward ( f, G, h, ϕ ) over the compact set.
We refer readers of interests to Exarchos & Theodorou (2018);
Negyesi et al.
(2021), and their references therein.
Theorem 4 (Log-likelihood for SB models) .
Given the solution satisfying the FBSDE system in (13), the log-likelihood of the SB model ( Z t , (cid:98) Z t ) , at a data point x 0 , can be expressed as
[Equation]
where the expectation is taken over the forward SDE (13a) with the initial condition X 0 = x 0 .
Proof.
[Equation]
which recovers (39).
Finally, notice that with integration by part, we have
[Equation]
where we adopt common practice and assume the limiting behavior of p SB t ; in other words
[Equation]
∃ k > 0 : p SB t ( x ) =
O (exp −(cid:107) x (cid:107) 2 k ) as x → ∞ .
With (41)
[Equation]
we can rewrite the related parts in (40) as
[Equation]
Hence
[Equation]
we also recover (38)
[Equation]
Corollary 5 (Probability ﬂow for SB) .
The following ODE characterizes the probability ﬂow of the optimal processes of SB (7) in the sense that ∀ t, p (17) t
[Equation]
Proof.
The probability ODE ﬂow (Song et al., 2020; Maoutsa et al., 2020) suggests that the equivalent
ODE model for the SDE (1) is given by
[Equation]
We can adopt this result to the SDEs of SB (7a) by considering f ← f + g Z t and p (1) t ← p SB t .
This yields
[Equation]
Applying the the factorization principle (Chen et al., 2021) with g log p SB t =
Z t + (cid:98) Z t concludes the proof.
Remark 10 (Connection between SB-FBSDE and ﬂow-based models) .
To demonstrate how applying ﬂow-based training techniques to the probability ODE ﬂow of SB (43) recovers the same loglikelihood objective in (39), recall that given an ODE d X t = F ( t, X t )d t with X 0 = x 0 ∼ p data , ﬂow-based models compute the change in log-density using the instantaneous change of variables formula (Chen et al., 2018):
[Equation]
which implies that the log-likelihood of x 0 can be computed as
[Equation]
Now, consider the probability ODE ﬂow of SB in (44),
[Equation]
Substituting this vector ﬁeld F SB to (45) yields
[Equation]
where (*) is due to integration by parts (recall (41)) and (**) again uses the factorization principle
Z t + (cid:98) Z t = g ∇ x log p SB t .
One can verify that (46) indeed recovers (39).
Theorem 11 (FBSDE computation for L SB ( x T ) in SB models) .
With the same regularity conditions in Theorem 3, the following FBSDEs also satisfy the nonlinear Feynman-Kac relations in (31).
[Equation]
Given a backward trajectory sampled from (47a), where X T = x T and x T ∼ p prior , the log-likelihood of x T is given by log p prior ( x T ) =
E (cid:104) Y T + (cid:98) Y T | X T = x T (cid:105) :
= L SB ( x T ) .
In particular,
L SB ( x T ) =
E [log p T ( X 0 )] − (cid:90) T
[Equation]
Proof.
Due to the symmetric structure of SB
[Equation]
we can consider a new time coordinate
[Equation]
Under this transformation
[Equation]
the base reference P appearing in (5) is equivalent to
[Equation]
The corresponding PDE optimality becomes
[Equation]
and the optimal forward/backward policies are given by
[Equation]
By comparing (50) with (7), one can notice that the new SB system (Φ , (cid:98) Φ) s corresponds to the original system (Ψ , (cid:98) Ψ) t via
[Equation]
and
[Equation]
Equation (51) shall be understood as the forward policy in t -coordinate system corresponds to the backward policy in s -coordinate system, and vise versa.
Following similar derivations in the proof of Theorem 3, we can apply Itô lemma to expend the stochastic processes d log Φ and d log (cid:98) Φ w.r.t.
(50a).
This yields the following FBSDE system.
[Equation]
Similar to
[Equation] [Equation]
Y
[Equation]
Z
[Equation]
s ) relate to the original FBSDE system (30) by
[Equation]
Changing the coordinate from s to t and applying (53) readily yield (47).
Finally, the expression in (48) can be carried out similar to (40):
[Equation]
We conclude the proof.
C C OMPARISON 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 is in contrast to our SB-FBSDE, which instead utilizes a divergence-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.
[Equation]
and X k ∼ N ( B k +1 ( X k +1 ) , 2 γ k +1 I ) , where F k ( x ) :
= x + γ k +1 f k
x ) and B k +1 ( x ) :
= x
+ γ k
+1 b
k +1 ( x ) are solved alternately via
[Equation]
In what follows, we focus mainly on the connection between (55a) and Theorem 4, yet similar analysis can be applied to (55b).
Now, expanding (55a) with the deﬁnition of ( B k +1 , F k )
[Equation]
which resembles the term (cid:107) Z t + (cid:98) Z t − g ∇ x log p SB t (cid:107) 2 in (15).
While 2 (cid:13) indeed corresponds to our backward policy (cid:98) Z ( k +1 , X k +1 ) after time discretization, 1 (cid:13) slightly differs from Z ( k +1 ,
X k +1 ) in how time is integrated, γ k +1 f
( k, X k +1 )
vs. Z ( k +1 , X k +1 ) .
On the other hand, 3 (cid:13) may be seen as an approximation of g ∇ x log p SB t , which, crucially, is not utilized in SB-FBSDE training.
Since ∇ x log p SB t is often intractable (nor does SB-FBSDE try to approximate it), SB-FBSDE instead uses the divergence-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,
[Equation]
which, crucially, differs from the SDE concerned by both our SB-FBSDE and SGM,
[Equation]
in that the diffusion g ( t ) is a time-varying function.
This implies that the connection between classical
SB models and SGM can only be made in discrete-time after choosing proper step sizes.
For instance,
De Bortoli et al.
(2021) considers the Euler-Maruyama discretization (see their § C.3),
[Equation]
58) where (cid:15) ∼ N ( 0 , I ) .
In order for (58) to match the discretization of (57), where g ( t ) is often a monotonically increasing function, the step sizes { γ k } Nk =1 must also increase monotonically.
However, since ∇ x log p SB t 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-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 ﬂawlessly); hence unaffected by the choice of numerical discretization or step sizes.
Model parametrization.
While Vargas et al.
(2021) utilizes non-parametric models, e.g. 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 SBFBSDE can also be trained with IPF (see Alg. 3), we stress that it is also possible to train both policies jointly whenever the computational budget permits.
Interestingly, this joint optimization – which is not presented in concurrent methods – resembles the training scheme of the recently-proposed diffusion ﬂow-based model (Zhang & Chen, 2021).
We highlight these ﬂexible training procedures arising from the uniﬁed 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 a promising future work.
Corrector Sampling.
While both De Bortoli et al.
(2021) and SB-FBSDE implement corrector sampling, they corresponds to different quantities.
Speciﬁcally, our SB-FBSDE relies on the same predictor-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 “ ∇ log p t ”.
In SB, this term corresponds exactly to adding the outputs of our networks " Z + (cid:98) Z ".
This computation differs from the corrector sampling appearing in De Bortoli et al.
(2021), which relies on single network ( i.e. either
Z or (cid:98) 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).
D E XPERIMENT D ETAILS
Table 1 with other models.
[Table]
Dataset
[Table]
p
[Table]
Toy 2e1.
Mnist 2e-
[Figure] CelebA 2e- [Figure]
CIFAR 1e2500.
Z t ( · , · ; θ ) and # of parameters (cid:98) Z t ( · , · ; φ ) and # of parameters
[Table]
NCSN++ (62.69M [Table] Unet (
[Table]
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 al.
(2021), the alternate training scheme may substantially accelerate the convergence under proper initialization.
Speciﬁcally, when Z t is initialized with degenerate outputs ( e.g. by zeroing out its last layer), training (cid:98) Z t at the ﬁrst K steps can be made in a similar SGM fashion since p SB t now admits analytical expression.
As for the proceeding stages, we resume to use (18, 19) since ( Z t , (cid:98) Z t ) no longer have trivial outputs.
Data pre-processing.
MNIST is padded from 28 × 28 to 32 × 32 to prevent degenerate feature maps through Unet.
CelebA is resized to 3 × 32 × 32 to accelerate training.
Both CelebA and CIFAR10 are augmented with random horizontal ﬂips to enhance the diversity.
Input: p prior , policies Z ( · , · ; θ ) and (cid:98) Z ( · , · ; φ ) Sample X
T ∼ p prior .
for t = T to ∆ t do
Sample (cid:15) ∼ N ( 0 , I ) .
Predict X t, 1 ← X t + g (cid:98)
Z t ∆ t + √ g ∆ t (cid:15) .
for i = 1 to N do
Sample (cid:15) i ∼ N ( 0 , I ) .
Compute ∇ x log p SB t,i ≈
[ Z ( t, X t,i )+ (cid:98) Z ( t, X t,i )]
/g .
Compute σ t,i with (59).
Correct X t,i +1 ← X t,i + σ t,i ∇ x log p SB t,i + √ 2 σ t,i (cid:15) i .
end for
Propagate X t − ∆ t ← X t,N . end for return X
Sampling.
The sampling procedure is summarized in Alg. 4.
Given some pre-deﬁned signal-to-noise ratio r (we set r = 0.05 for all experiments), the
Langevin noise scale σ t,i at each time step t and each corrector step i is computed by
[Equation]
Network architectures.
Table 8 summarizes the network architecture used for each dataset.
For toy datasets, we parameterize Z ( · , · ; θ ) and (cid:98) Z ( · , · ; φ ) with the architectures shown in Fig. 9.
Speciﬁcally, FCBlock represents a 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 al.
(2020) and Song et al.
Remarks on Table 1.
We note that the values of our SB-FBSDE reported in Table 1 are computed without the Langevin corrector due to 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 CIFAR10, despite already being encouraging, may be further improved with the
Langevin corrector.
Remarks on Fig.
To estimating KL( p T , p prior ) , we ﬁrst compute the pixel-wise ﬁrst and second moments given the generated samples X T at the end of the forward diffusion.
After ﬁtting a diagonal
Gaussian to { X T } , we can apply the analytic formula for computing the KL divergence between two multivariate Gaussians.
Remarks on Fig.
To accelerate the sampling process with the Langevin corrector, for this experiment we consider a reduced Unet (see Table 8) for CelebA.
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 A DDITIONAL E XPERIMENTS
Comparison to De Bortoli et al.
(2021) under same setup.
To demonstrate the superior performance of our model, we conduct experiments with the exact same setup implemented in De Bortoli et al.
Speciﬁcally, we adopt the same network architecture (reduced U-net), image preprocessing (center-cropping 140 pixel and resizing to 32 × 32), step sizes ( N =50), and horizon (0.
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 our model have higher diversity ( e.g. color skin, facing angle, background, etc) and better visual quality.
We conjecture that our performance difference may come from (i) the (in)sensitivity to numerical discretization between our divergence objectives and their mean-matching regression, and (ii) the foundational differences in how diffusion coefﬁcients are designed.
[Figure]
[Figure]
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 ﬁrst train the backward policy with SGM’s regression loss for a sufﬁcient long iterations (50k) until the FID roughly converges.
Then, we switch to our alternate training (Alg. 3) using the divergence-based objectives.
Crucially,
[Table]
with only 5k iterations of our divergence-based training, we drop the FID dramatically down to 13.35 from 33.68.
With a full stage of training (
), the FID decreases even lower to 11.85.
The qualitative results are provided in Fig.
Comparing Fig.
11a (corresponds to “SGM/50k” in
Table 3) and Fig.
11b (corresponds to “SGM/50k + SB/b/5k” in Table 3), it can be seen that the visible ﬂaw and noise have been substantially improved.
Additional Figures