Diffusion Schrödinger Bridge Matching
Authors: Yuyang Shi
Valentin De Bortoli
Andrew Campbell
Arnaud Doucet
Venue: NeurIPS 2023 poster (2023)
Keywords: diffusion Schrödinger bridge
bridge matching
optimal transport
TL;DR: We introduce a novel theoretical framework for approximating solution of Schrödinger Bridges, and show that the associated numerical scheme scales better and is more accurate than existing algorithms.

Abstract: Solving transport problems, i.e. finding a map transporting one given distribution to another, has numerous applications in machine learning. Novel mass transport methods motivated by generative modeling have recently been proposed, e.g. Denoising Diffusion Models (DDMs) and Flow Matching Models (FMMs) implement such a transport through a Stochastic Differential Equation (SDE) or an Ordinary Differential Equation (ODE). However, while it is desirable in many applications to approximate the deterministic dynamic Optimal Transport (OT) map which admits attractive properties, DDMs and FMMs are not guaranteed to provide transports close to the OT map. In contrast, Schrödinger bridges (SBs) compute stochastic dynamic mappings which recover entropy-regularized versions of OT. Unfortunately, existing numerical methods approximating SBs either scale poorly with dimension or accumulate errors across iterations. In this work, we introduce Iterative Markovian Fitting (IMF), a new methodology for solving SB problems, and Diffusion Schrödinger Bridge Matching (DSBM), a novel numerical algorithm for computing IMF iterates. DSBM significantly improves over previous SB numerics and recovers as special/limiting cases various recent transport methods. We demonstrate the performance of DSBM on a variety of problems.

Diffusion Schrödinger Bridge Matching
Yuyang Shi' University of Oxford
Valentin De Bortoli'
Andrew Campbell University of Oxford
Arnaud Doucet University of Oxford
ENS ULM
Abstract
Solving transport problems, i.e. finding map transporting one given distribution to another, has numerous applications in machine learning. Novel mass trans- port methods motivated by generative modeling have recently been proposed, e-g- Denoising Diffusion Models (DDMss and Flow Matching Models (FMMs) im- plement such a transport through a Stochastic Differential Equation (SDE) or an Ordinary Differential Equation (ODE). However, while it is desirable in many applications to approximate the deterministic dynamic Optimal Transport (OT) map which admits attractive properties, DDMs and FMMs are not guaranteed provide transports close to the OT map. In contrast, Schrödinger bridges (SBs) compute stochastic dynamic mappings which recover entropy regularized versions OT. Unfortumately, existing numerical methods approximating SBs either scale poorly with dimension accumulate introduce Iterative Markovian Fitting (IMF), a new methodology for solving SB problems, and Diffusion Schrüdinger Bridge Matching (DSBM), novel numerical algorithm for computing IMF iterates. DSBM significantly improves over previ- ous SB numerics and recovers as special/limiting cases various recent transport methods. We demonstrate the performance of DSBM on a variety of problems.
across iterations. In this work, we
Introduction
Mass transport problems are ubiquitous in machine learning (Peyré and Cuturi, 2019). For discrete measures, the Optimal Transport (OT) map can be computed exactly but is computationally inten- sive. In landmark paper, Cuturi (2013) showed that an entropy-regulate computed more efficiently using the Sinkhorn algorithm (Sinkhorn, 1967). This has enabled the use of OT techniques variety of applications ranging from biology (Bunne et al., 2022) shape correspondence (Feydy et al., 2017). However, applications involving nigh-dimensional continuous distributions and/or large datasets remain challenging for these techniques.
version of OT can be
One of such data-rich applications is generative modeling, central transport problem in machine learning which requires designing a deterministic or stochastic mapping transporting reference "noise" distribution to the data distribution. For example, Generative Adversarial Networks (Good- fellow et al., 2014) define static, deterministic transport map, while Denoising Diffusion Models (DDMs) (Song et al., 2021b; Ho et al., 2020) build a dynamic, stochastic transport map by simulating Stochastic Differential Equation (SDE), whose drift is learned using score matching (Hyvärinen, 2005; Vincent. 1). The excellent performances of DDMs have motivated recent developments of Bridge Matching and Flow Matching models, which are dynamic transport maps using SDEs (Song et al., 2021 la; Peluchetti, 2021; Liu, 2022; Albergo et al., 2023) or ODEs (Albergo and Vanden- Eijnden, 2023; Heitz al., 2023; Lipman al., 2023: Liu al., 2023b). Compared DDMs, Bridge and Flow Matching methods do not rely forward "noising" diffusion converging to the reference distribution in infinite time, and are also more generally applicable as they can approximate transport maps between two general distributions based on their samples. Nonetheless, these transport maps
Equal contribution.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Sets for alternating projections
Preserved properties
DSBM
IPF
Po
To: Pr = TT M: R(O)
M.. R(Q) Po No. Pr = RT
Bridge Matching
IMF
Demoising Diffusion
Flow Matching
Table 1: Comparison between Iterative Markovian Fitting (IMF) and Iterative Proportional Fitting (IPF). The Schrödinger Bridge is the unique IP s.t. Po *0, Pr *r, P M, P R(Q) simul- taneously by Proposition 5. M is the space of (regular) Markoy measures and R(Q) the space of reciprocal measures of Q
Relationship between DSBM and existing methods.
Figure I:
are not necessarily close to the OT map minimizing the Wasserstein-2 metric, which is appealing for its many attractive properties (Peyré and Cuturi, 2019; Villani, 2009).
In contrast, the Schrödinger Bridge (SB) problem is a dynamic version of entropy-reeullrizeed OT (EOT) (Föllmer, 1988; Léonard, 2014b). The SB is the finite-time diffuuson which admits as initial and terminal distributions the two distributions of interest and is the closest in Kullback-Leibler divergence reference diffusion. Numerous methods to approximate SBs numerically have been proposed, see e.g. (Bernton et al., 2019; Chen et al., 2016; Finlay et al., 2020; Caluya and Halder, 2021; Pavon et al., 2021), but these techniques tend to be restricted to low-dimessional settings. Recently, novel techniques using diffusion-basse
ideas have been proposed in (De Bortoli et al., 2021; Vargas et al., 2021: Chen et al., 2022) based on Iterative Proportional Fitting (IPF) (Fortet, 1940; Kullback, 1968; Rüschendorf and Thomsen, 1993), continuous state-space extension of the Sinkhorn algorithm (Essid and Pavon, 2019). These approaches have been shown to scale better empirically, but numerical errors tend to accumulate over iterations (Fernandes et al., 2021).
In this paper, our contributions are three-fold. First, we introduce Iterative Markovian Fitting (IMF), new procedure to compute SBs which alternates between projecting on the space of Markov processes and on the reciprocal class, i.e. the measures which have the same bridge as the reference measure of SB (Léonard et al., 2014). We establish various theoretical results for IMF. Contrary to IPF, the IMF iterates always preserve the initial and terminal distributions. The differences between IPF and IMF are presented in Table 1. Second, we propose Diffusion Schrödinger Bridge Matching (DSBM), novel algorithm approximating numerically the SB solution derived from IMF. DSBM requires at each iteration solving a simple regression problem in the spirit of Bridge and Flow Matching, and does not suffer from the dine-discrttization and "forgetting" issues of previous DSB techniques (De Bortoli et al., 2021; Vargas et al., 2021; Chen et al., 2022). Finally, we demonstrate the performance of DSBM on a variety of transport tasks.?
Notations. We denote by P(C), the space of path measures, i.e. P(C) P(C TJ. Rd) where T> 0. The subset of Markov path measures associated with an SDE of the form dX, vn (X, o, with o, locally Lipschitz, is denoted M. For any Q E M. the reciprocal class of Q is denoted R(Q), see Definition 3. We also denote Qu its marginal distribution at time t, Q the joint distribution at times and t, Qstt the conditional distribution at time given state at time and Qjo. P(C) its diffusion bridge. Unless specified otherwise, all gradient operators V are w.r.t. the variable with time index t. Let (X. and (Y, )2) be probability spaces. Given Markov kernel K X x y [0, 1] and probability measure fu defined on X. we write μK the probability measure on y' such that for any A E y we have μk(AA = Sx K(1r, A)dp(z). In particular, for any joint distribution I, over Rd x Rd. we denote the mixture of bridges measure as II Ho E P(C), which is short for II()) = Jgax Qj, (Zo, xr Io, (dæy, dzr).
2 Dynamic Mass Transport Techniques
2.1 Denoising Diffusion and Bridge Matching Models
Denoising Diffusion Models (Song et al., 2021b; Ho et al., 2020) are a popular class of generative models. They define forward noising process Q E M using the SDE dX, - + dB, on the time-interval 00, T), where X, I is drawn from the data distribution no and (B.)cejo,rj is a d-
Code be found at https: sttps://githube com comm/uyang--aae whsddssmpptmtooscc......m pytorch.

---

dimensional Brownian motion. This diffusion' converges towards the standard Gaussian distribution N(,, Id) T + oc. A generative model is given by its time-reversal (o. (xT-t)tejo.77 where Yn ~ Qr and dY, = Y + V log Q (Y,) + dB, (Anderson, 1982; Haussmann and Pardoux, 1986). In practice, (Y, I is initialized with Yo- मे' N(00, Id), and the Stein score V log (z+) = a V log Q o (X, X, = xı] is approximated using neural network de st xz) minimizing the denoising score matching loss Eg.., [Il Qoo(XX,(Xo) sg X=)
An alternative to considering the time-reversal between the two distributions and learn mimicking diffusion process. This approach generalizes DDMs and allows for more flexible choices of sampling processes. We call this framework Bridge Matching and adopt presentation similar to Peluchetti (2021); Liu et al. (2022b), where æy is the data distribution." We denote Q E .M the path measure associated with the following process
forward noising process is to "build bridges"
dX. fxxx)) o, X Qo.
Consider now the distribution of this process pinned down at an initial and terminal point zo. TT- dr).
denoted Q Q. (.,r(1/2,,,0 (r)) Under mild assumptions, the pinned process 2-o.r(1/am (o, is diffusion bridge and is given by
dxo. = ( + of V log (T (x )dd + a,dB
X No.
which satisfies X = ør using Doob h-transform theory (Rogers and Williams. 2000). Next, we define an independent coupling Ho, = To @ NT, and let II 1o.r This path measure is mixture of bridges. We aim to find Markov diffusion dY, = (,f((,,)) + (Y. + a,dB, on [0, T'] which admits the same marginals as II; i.c. for any E (). Y, ~ II, so YT ~ TT. For such vr, a generative model for sampling data distribution कт is obtained by simulating (Y)) can be verified that indeed Y, ~ II, for (2+) = Er (V log (X, present the theory behind this idea more formally using Markovian projections in Section 3.1. In practice, we do not have access to oy and it is learned using neural networks with regression loss
X = zz]. We
En r [ log (X, (1, x.)l
For ft = and o3 = o, Q 0,7 is a Brownian Bridge and we have
XO. far + (1 ) + oe (B, (BT), ax. ((IT x1 - /(T- +o,dB, (4)
with (B, - Br) N(O, t(1 Id).
Id). The regression loss (3) associated with (4) is given by En, ( [X,, X.) - t) vo (1, x.)1
(5)
Letting + 0, we recover Flow Matching models (see Appendix A.1 for further details).
2.2 Schrödinger Bridges and Optimal Transport
The Schrödinger Bridge (SB) problem (Schrödinger, 1932) consists in finding a path measure PSB P(C) such that
SSU (KLLi))((KLPP(Q)) Po no. Pr TT).
(6)
where Q E P(C) is reference path measure. In what follows, we consider Q defined by the diffusion process (1) which is Markov, and without loss of generality, we assume Qo To. Hence pSS is the path measure closest to Q in terms of Kullback-Lebbler divergence which satisfies the initial and terminal constraints PaSB
no and PSB nr.
Another crucial property of JPS is that it can also be defined as a mixture of bridges PSB_ = IT Cjo.T. where ITSB_ argmin (KT.(Tl,,)
Ho To, Ty тT} is the solution of the static SB problem (Léonard, 2014b). In particular, for associated with (oB))re(o.7
have
-o argmin (LEn X, 12 2T H(ll,r) : 1I No. II, ,r),
This known as the Omstein Uhlenbeck (OU) process VPSDE (Song et al., 2021b). "To keep notations consistent with existing works, no is the data distribution in the context of DDM and SB, whereas Ar is the data distribution in Bridge Matching. However, both SB and Bridge Matching methods allow transfer between arbitrary distributions mo, Tr, so this distinction is not important.

---

where H(μ) denotes the entropy, i.c. Hor is the solution of the autropy-regulane OT problem. In this case, the SB can also be obtained theoretically by solving the following problem (Dai Pra, 1991)
USB = argmin, 1.So Ee, [lle X.)l dx; ett, )dt + adBr, Pu To: Pr = .
Then pss is given by the SDE with drift usn initialized with Xo To- For a 0, we recover the classical OT problem and the Benamou-Brenier formula (Benamou and Brenier, 2000).
common approach to solve (6) is the Iterative Proportional Fitting (IPF) method (Fortet, 1940; Kullback, 1968; Rüschendorf, 1995) defining a sequence of path measures (PP") where
p2=+1 argming { (LL(PIP2)) PT тT}. p2m+2 argming (KL(P)(2(t) Po mo (7)
with initialization 150 = Q. This procedure alternates between projections on the set of path measures with given initial distribution To and terminal distribution πT. It can be shown (De Bortoli et al., 2021)
that (P") are associated with diffusions and that for any E N, p2an+1 is the time-reversal of 12t with initialization Ty, and p2a++2 is the time-reversal of [p2++1 with initialization To. Leveraging this property, De Bortoli et al. (2021) proposed Diffusion Schrödinger Bridge (DSB), an algorithm which learns the time-reversals iteratively. In particular, DDMs can be seen as the first iteration of DSB.
3 Iterative Markovian Fitting
3.1 Markovian Projection and Reciprocal Projection
Markovian Projection. Projecting on Markov measures is key ingredient in our methodology and in the Bridge Matching framework. This concept was introduced multiple times in the literature (Gyöngy, 1986; Peluchetti, 2021; Liu et al., 2022b). In particular, we focus on Markovian projection of path measures given by a mixture of bridges II Qo,, E P(C).
Definition 1. Assume that Q is given by (1) and that for any (Io, E R.. Qja. (20, ØT) is associated with (x!-") [e given by dx. = ) (T)
)d) Jdt ood
with [0, T] + (0, +o0). Then, when it is well-defined, we introduce the Markovian projection of 11, MI* = proj (1I) E M, which is associated with the SDE
dx; = (f(()) + of (Xt + o,dB,
of (xx) = of En V log (X, X, = 24..
Note that in our definition o; > 0 so V log (TT(T))))) is well-defined, but Flow Matching can be recovered as the deterministic case in the limit o, + O. In the following proposition, we show that the Markovian projection is indeed projection for the reverse Kullback-Leibler divergence, and that it preserves marginals of II,
Proposition 2. Assume that oz > O. Let M = proj (II. Then, under mild assumptions, we have
M- argming MMI : M E M}.
KL(MI(*) = 1/2 SC" Er,,, [llo,E (V log Qr Xo, x.) (x.))l1ooo
In addition, we have that for any t E [0, M. = II. In particular. M7. = 1l7.
Reciprocal Projection. While the Markovian projection ensures that the obtained measure is Markov, the associated bridge measure is not preserved in general, i.e. Mr * II = Qjo... Measures with same bridge as Q are said to be in its reciprocal class (Léonard et al., 2014).
Definition 3. IT E P(C) is in the reciprocal class R(Q) of E M if = Ho Qo, We define the reciprocal projection of E P(C) as II+ projæ(Q) (P) = Qu.r
Similarly to Proposition 2, we have the following result, which justifies the term reciprocal projection. Proposition 4. Let P E P(C), IT* proja(o) Then, IT argmin (KI.(P]IT) II E R()).
The reciprocal projection 11f of a Markov path measure M does not preserve the Markov property in general. fact, the Schrödinger Bridge is the unique path measure which satisfies the initial and terminal conditions, is Markov and is in the reciprocal class of Q. see (Léonard, 2014b).
Proposition 5. Let TP be a Markov measure the reciprocal class of Q such that Po = To. PP nT. Then, under assumptions on Q, wo and nT, P is unique and is equal to the Schrödinger Bridge pSB

---

3.2 Iterative Markovian Fitting
Based on Proposition 5, we propose novel methodology called Iterative Markovian Firting (IMF) to solve Schrödinger Bridges. We consider a sequence (P" NEN such that
P g projm (P²"),
p22n+2 projR(q) (1p2++1),
(8)
with ppo such that PB No. Pt тт and pd R(Q). These updates correspond to alternatively performing Markovian projections and reciprocal projections.
Combining Proposition 2 and Definition 3, we get that for any E N. P8 = To and Pq. = RT. This property is contrast to the IPF algorithm (7) for which the marginals at the initial and final times are not preserved. We highlight this duality between IPF (7) and IMF (8) in Table 1.
We conclude this section with a theoretical analysis of IMF. First, we start by showing a PattwwingaPthhtgggoo Pythagorean theorem for both the Markovian projection and the reciprocal projection.
Lemma 6. Under mild assumptions, if M M, II R(Q) and KL(IIMM) < to0, we have
ML = (roooj (II) (L))
If IT <
+00, we have
KL(MII) (M (M)) + ()(proiR(a)(M)]
Using Lemma 6, we have the following proposition.
Proposition 7. Under mild assumptions, we have [L[[P++1 pSs) < KT.(PPNIPBB) < oc, and
li (TP(IP [p=-1)= = O.
Hence, for the IMF sequence (R"), the Markov path measures (I22++1) nen are getting closer to the reciprocal class, while the reciprocal path measures (I2n+2) nEN are getting closer the set of Markov measures. Proposition should be compared with (Rüschendorf, 1995, Proposition 2.1, Equa-
tion (6)) which shows that, for the IPF sequence (IPm) nee, we have lim (fn+11 This result is similar to Proposition but for the forward Kullback-Leibhe divergence.
( Ii(m) 0.
Using Proposition 7, we finally prove the convergence of the IMF sequence (P") to the Schrödinger Bridge. This result was first shown in the concurrent work (Peluchetti, 2023, The- orem 2). We present a simpler proof in Appendix C.6.
Theorem 8. Under mild assumptions, the IMF sequence (P"")E admits a unique fixed point 1' PSB and limm-++00 KL(P"IP** O.
4 Diffusion Schrödinger Bridge Matching
In this section, we present Diffusion Schrödinger Bridge Matching (DSBM), practical algorithm for solving the SB problem obtained by combining the IMF procedure with Bridge Matching.
Iterative Markovian Fitting practice. IMF alternatively projects on the Markov class M and the reciprocal class RQ). We denote M = p22+1 E . and II = p2n R(Q). Assuming we know how to sample from the bridge o given the initial and terminal conditions, sampling from the reciprocal projection projR(q) (MI) is simple: First, sample (Xo, Xx) from the joint distribution Mo,r Then, sample from the bridge Qjo. (X,,Xx) The bottleneck IMF is in the computation of Markovian projections. By Definition 1, M proj (II) is associated with the process
dX, = (f((X.) + IT V log Qryy X Xx) X.]]dt + a,dB,,
Xo ~ To.
By Proposition 2, we can learn MI using M given by
d.X, (fx( + uæ- (t, ]tt dB Xo ~ To:
(9)
0+ argming 5 En. r ||lo log (X vo x,)1
θ),
(10)
"I practice, we sample the SDE associated with MI and save a batch of joint samples (Xo, X²). This is similar to the trajectory caching procedure in De Bortoli et al. (2021), but only retain initial and final samples.
5

---

where {uo A E } is a parametric family of functions, usually given by a neural network. The optimal lig. (t, m4) = F V log Qrje (rnn(X)(XX)) X, = = for any t E [0, T and Ty E IR4.
With the above two procedures for computing ( (M1) and proj (II), we can now describe a numerical method implementing IMF (8). Let II Qior where D1 = no. ILg = иT. Learn M' R pM (II)) given by (9) with Vg- given by (10). Next, sample from II' proj (M')
M by sampling from My and reconstructing the bridge Qj,. We iterate the process to obtain a sequence (TI", "P,,++)))ee In practice, this algorithm performs poorly (see Figure 3), since the approximate minimization (10) for computing M may not admit Mr++1 = Wy exactly as in Proposition 2. Instead, we incur a bias between MG-+1 and тт which accumulates for each n E N.
To mitigate this problem, we alternate between forward Markovian projection and a backward Markovian projection. This procedure is justified by the following proposition.
Proposition Assume that II Qoo. with Q associated with dX; = fx(X))dt + o, Under mild conditions, the Markovian projection M projm is associated with both
dX;- = (X, + E V log Q1 x dt d, X ~ Ho.
(11)
dY, (-fr- Y + o Bu V log QT - YY)YT)) Y.J)dt + or dB. Yo IT2.
In Proposition 9, (11) is the definition of the Markovian projection, see Definition 1. However, (12) an equivalent representation as time-reversal. In practice, (Yt)eEO,TT) approximated with
dY, (-))- + ug (T - Y.) + dd
Y ~ TT.
(13)
6' = argmin 0.10 Eno, [Ilo? log Q:1o X Xo) - (t X, x,)
De $).
(14)
The optimal Ug- (f) of B V log Q:jo (X, (Xo) X, = ma] for any t [0, T and E Rd.
Note that X, To in the forward projection, while Yo ~ Tт in the backward projection. Therefore, us- ing the backward projection removes the bias on कт accumulated from the forward projection. Leveraging the time-symmetry of the Markovian pro- jection and alternating between (13) and (9) yields the DSBM methodol- ogy summarized in Algorithm 1.
Algorithm Diffusion Schrödinger Bridge Matching
Input:
Joint distribution 11,,,, tractable bridge Qio,
number of outer iterations N' E N.
2: Let ITo = -T0 1
3: for E {0. NO. 1} do
4: Learn vg- using (14) with 11 112".
5: Let M be given by (13).
6:
Let 2 = M 0 0.7-
7:
Learn vg- using (10) with II 112
8: Let 2 be given by (9).
is also possible to learn both the for- ward and backward processes each step, and enforce that the backward and forward processes match. We ex- plore this in Appendix G.
9: Let 1T2m+2 = M 2 p.
10: end for
11: Output: Ugr, Ug-
Initialization coupling. We relate Algorithm to the classical IPF and practical algorithms such as DSB (De Bortoli et al., 2021). Instead of initializing DSBM with Io_ given by a coupling between To. TT, if we initialize it by To Qo where Qo no and Qrio is given by the reference process defined in (1), then DSBM also recovers the IPF iterates used in DSB.
Proposition 10. Suppose the families of functions {ua 0 Q} and {us : é 4) are rich enough so that they can model the optimal vector fields. Let (II", ,MM++) be the optimal DSBM sequence in Algorithm 1 initialized with no Qu.T, and let (IP") NEN be the optimal DSB sequence given by the IPF iterates in (7). Then for any N, 2 we have pm.
We will thus call DSBM-IPF, the DSBM algorithm initialized with the joint distribution given by the forward reference process Io Q and DSBM-IMF, the DSBM algorithm initialized with an independent coupling no = Ro mr. However, the training procedure of DSBM-IPF is very different from the one of (De Bortoli et al., 2021: Chen et al., 2022). In existing works, ji++1 is

---

obtained as the time-reversal of Pr which requires full trajectories from P", see e.g. (De Bortoli et al., 2021, Proposition 6). In contrast, in Algorithm we only use the coupling M to create the bridge measure 11" M Qjo.r. .. By doing so, (i) the losses (10) and (14) can be easily evaluated at any time E [0. : (ii) the trajectory caching procedure in DSBM is more computationally and memory efficient; (iii) while every IPF iteration Tpm is also supposed to be in R(Q). practice one can observe forgetting of the bridge Q o, (Fernandes et al., 2021). In DSBM, this effect is countered by explicit projections on the reciprocal class. See Appendix F for more details.
At equilibrium of DSBM, we have that (Yt) given by (13) is the time reversal of (Xe)xgjo.21 given by (9) and are both associated with the optimal Schrödinger Bridge path measure Pt. As a result, we have that Es- (t, x) = -Ug- (t, x)) + ofV log (x). Hence, a probability flow (Zi [e(o,T) such that Law(Z)
Probability flow ODE.
P7 for any E [0, T] is given by
dz;- (f) + 110- (t, Z7) V (t. )d)
Z no.
See also De Bortoli et al. (2021); Chen et al. (2022) for derivation of this result. Note however that the path measure induced by (Z) [0 does not correspond to P"; in particular, (Zy, Z7-) is not an entropic OT plan. However, since for any E [). Z has marginal distribution IP;, we can compute the log-likelihood of the model (Song et al., 2021b; Huang et al., 2021).
5 Related Work
Markovian projection and Bridge Matching. The concept of Markovian projection has been rediscovered multiple times (Krylov, 1984; Gyöngy, 1986; Dupire, 1994). In the machine learning context, this was first proposed by Peluchetti (2021) define Bridge Matching models. More recently, Liu et al. (2022b) derived theoretical properties of the Markovian projection in Proposition 2, first part of Lemma 6, and applied Bridge Matching for learning data on discrete and constrained domains.
Bridge and Flow Matching. Flow Matching corresponds to deterministic bridges with determinis- tic samplers (ODEs) and has been under active study (Liu et al., 2023b; Liu, 2022: Lipman et al., 2023; Albergo and Vanden-Eijnden, 2023; Heitz et al., 2023; Pooladian et al., 2023; Tong et al., 2023). Denoising Diffusion Implicit Models (DDIM) (Song et al., 2021 la) can also be formulated as discrete-time version of Flow Matching, see Liu et al. (2023b). These models have been extended to the Riemannian setting by Chen and Lipman (2023). Recently, Albergo et al. (2023) studied the influence of stochasticity in the bridge, through the concept of stochastic interpolants. Liu et al. (2023a); Delbracio and Milanfar (2023) used Bridge Matching to perform image restoration tasks and noted benefits of stochasticity empirically. Closely related to our work is the Rectified Flow algorithm of Liu al. (2023b), which corresponds an iterative Flow Matching procedure in order to improve the straightness of the flow and thus eases its simulation. An iterative rectifying procedure using stochastic interpolants is also proposed in (Albergo al., 2023, Section 3.5). Our proposed DSBM-IMF algorithm is closest to Rectified Flow, which can be seen as the deterministic limiting case of MSBM as a O. However, there are few important theoretical and practical differ- ences. Most notably, we adopt the SDE approach which is crucial for the validity of Proposition 5 as well as for the empirical performance of DSBM. We discuss further distinctions between DSBM and Rectified Flow in Appendix A.3.
Diffusion Schrödinger Bridge. Schrödinger Bridges (Schrödinger, 1932) are ubiquitous in probability theory (Léonard, 2014b) and stochastic control (Dai Pra, 1991; Chen et al., 2021). More recently, they have been used for generative modeling: De Bortoli et al. (2021) introduced the DSB algorithm and Vargas et al. (2021); Chen et al. (2022) introduced similar algorithms. The case Dirac delta terminal distribution to solve conditional simulation and more general control problems (Shi et al., 2022; Thornton et al., 2022; Liu et al., 2022a; Chen al., 2023; Tamir et al., 2023). In Somnath et al. (2023), SBs learned using one Bridge Matching iteration, assuming access to the true Schrödinger static coupling. Our proposed method DSBM-IPF is closest to DSB, but with improved continous-time training and projections on the reciprocal class which mitigate two limitations of DSB. Concurrently with our work, Peluchetti (2023) independently introduced the DSBM-IMF approach (named IDBM therein).
investigated by Wang al. (2021). These methods were later extended
7

---

Experiments
2D Experiments. We first show our proposed methods can generate correct samples and learn lower kinetic energy transport maps some 2D examples. We compare our method DSBM with flow-based methods including Flow Matching (FM) (Lipman et al., 2023), Conditional Flow Matching (CFM), OT-CFM (Tong et al., 2023), and Rectified Flow (RF) (Liu et al., 2023b); and other SB methods including DSB (De Bortoli et al., 2021) and SB-CFM (Tong al., 2023). OT-CFM and SB-CFM utilizes sample-based mini- mani-batch OT or EOT solvers (Fatras et al., 2021; Flamary et al., 2021) to define an approximate OT or SB static coupling 10 or fisH see also Pooladian et al. (2023); Stromme (2023). We can also utilize this idea the DSBM-IMP framework, which corresponds to using the initialization coupling ng. = iTSR, in Algorithm 1. This approximate SB coupling n also satisfies T To and Mīş #r but can provide better initialization than the independent coupling 01. T To @ RT. We name this approach DSBM-IMF+. The rest of the methods do not use OT solvers. DSB and DSBM directly learn the EOT map as the solution of the diffusion process.
In Table 2, we show the 2- Wasserstein distance between the true and generated samples, as well as the integrated path energy defined as EEI [| Z,)l [dd] where is the learned drift along the ODE trajectory Z, For direct comparability, we report for DSBM using its probability flow ODE. Lower path energies represent shorter (and potentially easier integrate) trajectories. We find that in this low dimensional setting. OT-CFM performs the best by utilizing OT solvers, but DSBM outperforms FM and CFM when OT solvers are not used. Further, DSBM outperforms DSB on all datasets, suggesting DSBM solves the SB problem with higher accuracy. The results also show that among SB methods, DSBM-IMF+ can achieve lower sampling error than DSBM-IPE and DSBM-IMF. It also performs better than SB- SB-CFM on 3 of the datasets and achieve lower path energy on all datasets. Finally, we find Rectified Flow achieves lower sampling error than DSBM except for the moons- -aaas--gassssans task, for which DSBM is significantly more accurate. Since RF can be informally seen as DSBM in the case 0 - O, this suggests the optimal a varies for each task and between generative and general transfer tasks. Figure 2 visualizes how affects the straightness and sample quality of learned transport maps between two mixture distributions.
High- Hiigh-Dimessionl Gaussian Experiment. We next perform the Gaussian transport experiment in De Bortoli et al. (2021) with dimension d = 50 to verify the scalability of our proposed approach. The true SB can be computed analytically in this case (Bunne et al., 2023). In Figure 3, we plot the convergence of the learned mean [[] variance Var (Xo), and covariance Cov (Xo, Xr) between times 0, T' We also consider RF and related baseline IMF-b, which performs IMF numerically but only in the backward direction. All methods converge approximately to the correct mean. However,
2-Wussersttei (Enler steps)
energy
Doasse!
scurve
Ngaussians
Ngaussians
moons
Scurve
Rgaussians
moons- Sgaussäams
DSDM-IPF
0.14020.000
0.140+0.024
0.31520.079
0.81220.092
1.59810.034
2.11010.059
14.9110,310
1..16
DSBM-IME
0.144=0.024
0.14520.037
0.338#0.091
0.83840.098
1.50H000036
2099200053
148
41.00
DSBM-IMF+
0.123=0.014 0.130:0.025
0.276:0:000
0,802+0.172
1.59/99.00
2.116H00188
14.88+0.252
$1.09=7.206
DSB
0.190=0.049
0.27260.065
0.411=0.084
0.98740.324
SB-CFM
0.129=0.24
0.136+0.000
0,238:0.044
0.84340.099
L61910.035
2,141000
15.08.0.209
45.6920.661
PM
0.212=0.025
0.161+0.033
0.351#0.066
2,22710.056
2.95010.044
18.1210.116
CFM
0.215-0.028
0,171+0,023
0.370H0.49
1285+0.114
2.39110.043
3.07140:026
18.0010.090
116.522.633
RF
0.129=0022 0.136e0019
0.26740047
1,522H
1.78520.052
1633400.044
14.8420.441
37.61=3.906
OT-CFM
0.111=0.005
0.102:0.013
0.253:0.040
0,716:0.187
1.178:0.020
1.57759.036
15. 1050.215
30.50el0.626
Table 2: Sampling quality as measured by 2- Wasserstein distance and path energy for the 2D experiments. +1 SD over 5 seeds. Best values are in bold and second best are italicized.
0=000
0=011
8=03
o=
3.0
Figure 2: Learned SB probability flow between two mixtures of Gaussians (green + yellow).
8

---

Mean
Variance
Covariance
x10"10"
d=b
0.00
DSB 3.2611660
13 13.0033449
32.81122 28
DGRMUUPF
1.0
0.00
-0.0 0.05
TIcRS- IMF
SB-CHM 14560.73
12.3e1.47
49.4+3.91
RF
DSBM-IPF 1.23:0.33
4.4210.76 8.7540.87
0.10
0.50
DSBM-IME 1.340051 5.0510955 9.7511.67
Iteration
Inoration
lioration
Table 3: Average (L(PP,IPP at 21 uniformly spaced t.
Figure 3: Convergence of Gaussian experiment in d = 50.
C 12//2322
2
4
$
LISR EIGRA IFF FIGRM IMF
KM
C 9 2 2 0
0 0oo0
0
6
2
2
4
D
9 6
3
2 C 0
7 2
2
3000
. 1 6 0
a 2 2303
2 4 2
9
7
1
a
2040 444066500
$ 1 0 50052200 22222
b 187888852
7 A
6
20
6J
200 00855
8 8
10
(a) OT-CFM
(b) DSB
(c) DSBM-IPP
30
Training step (H1000))
100
1.0
200
230
Figure 4: Samples of MNIST digits transferred from letters.
Figure 5: FID vs iteration.
the variance estimates become inaccurate for RF and IMF-b. Among SB methods, DSB and IMF-b also gave inaccurate SB covariance estimates as the number of iteration increases. On the other hand, DSBM does not suffer from this issue. In Table 3, we further quantify the accuracy and compare with SB-CFM (Tong et al., 2023) by computing the KL divergence between the marginal distributions of the learned process Pr, and the true SB P. Our proposed methods achieve similar KL, divergence as SB-CFM in dimension d = 5, but are much more accurate in higher dimensions.
MNIST, EMNIST transfer. We test our method for domain transfer between MNIST digits and EMNIST letters as in De Bortoli al. (2021). We compare DSBM a direct substitute of DSB, and also with Bridge Matching (BM) (Peluchetti, 2021; Liu et al., 2022b), CFM, OT-CFM and RF. We plot some output samples from different algorithms in Figure and the convergence of FID score in Figure 5. We find that OT-CFM becomes less applicable in higher dimensions and produces samples of worse quality (Figure 4a). On the other hand, image quality deteriorates during training of DSB and RF DSBM achieves higher quality samples visually, and does not suffer from deterioration. It is also about 30% more efficient than DSB in terms of runtime.
CelebA transfer. Next, we evaluate and perform some ablations of our method on transfer task on the CelebA 64 x 64 dataset. We consider the images given by the tokens male/old and female/young. In Figures and 7, we show that as increases, the quality of the images (as measured by the FID score) increases until a is too high, but the alignment (as measured by LPIPS) between the generated image and the original sample decreases. Additionally, we investigate the dependency between a and image dimension in Figure 8. In particular, for the same a 1, the outputs of DSBM for CelebA 128 x 128 are better aligned with the original data than for Celeb 64 x 64. This is in agreement with the observations of Chen (2023); Hoogeboom et al. (2023) that the noise schedule in diffusion models should scale with the resolution.
AFHQ transfer. We demonstrate the scalability of our method on an additional transfer experiment on the AFHQ 512 x 512 dataset between the classes cat and wild. The results are shown in Figure 9, On this higher-dimensional problem, we observe that DSBM can also generate realistic samples which are similar to the input.
Unpaired Fluid Flows Downscaling. Finally, we apply DSBM to perform downscaling of geo- physical fluid dynamics, i.e. super-resolution of low resolution spatial data. We use the dataset in (Bischoff and Deck, 2023), which consists of unpaired low (64 x 64) and high (512 x 512) resolution fields. As shown in Figure 10, DSBM is able learn high resolution reconstructions by only slightly noising the low resolution input. In contrast, Bischoff and Deck (2023) use two diffusion models in forward and backward directions (Diffusion-fb) based on Meng et al. (2022), which improves over the Random baseline. Figure 11 shows that DSBM-IPF and DSBM-IMF achieve much lower l

---

70
60
0.15
40
0.10
30
0.05
10
10²
Figure 6: Left to right: initial and generated samples (64 64) ob- tained after 20 DSBM-IME itera- tions for a² {0.01, 0.1, 10).
Figure 7: FID (blue) and LPIPS (red) scores (lower is better for both) vary o².
Figure 8: Top to bottom: DSBM (a 1) 64; 128 x 128; original images.
(a) Transfer between classes cat and wild.
(b) Transfer between classes wild and cat.
Figure 9: DSBM domain transfer results on the AFHQ 512 x 512 dataset.
Superottantion
Vonticty
1.0
(a)
(b)
Data waxenamber
Figure 10: (a) Left to right: source low resolution sample, intermediate state and final reconstruction of DSBM-IPF; (b) unpaired high resolution sample.
Figure 11: & distance between low olution source and high resolution recon- structed fields.
distances for all frequency classes in the dataset than Diffusion- ffb (and thus Random), indicating DSBM able reconstruct high resolution fields consistent with the low resolution source.
Discussion
In this work, we introduce IMF, a new methodology for learning Schrödinger Bridges. IMF is an alternative to the classical IPF and can be interpreted as its dual. Building on this new framework, we present two practical algorithms, DSBM-IPF and DSBM-IME for lcarning SBs. These algorithms mitigate the sime-discretization and bias accumulation issues cxisting methods. However, DSBM still has some limitations. First, results suggest DSBM is most effective for solving general transport problems. For generative modeling, we only find minor improvements compared to Bridge and Flow Matching CIFAR-10 (see Appendix I.6). Second, while DSBM is more efficient than DSB. it still requires sampling from the learned process during the caching step. Finally, the EOT problem becomes more difficult solve numerically for small values of
In future work, we would like to further investigate the differences between DSBM-IMF and DSBM- IPF. IMF also appears useful for developing a better understanding of the Rectified Flow algorithm (Liu et al., 2023b), as IMF minimizes a clear objective (6) and Rectified Flow can be seen as a limiting case of it. Finally, Rectified Flow has also been extended to solve OT problems with general convex costs by Liu (2022), and would interesting to derive a SB version of this extension.
10

---

Acknowledgements
YS acknowledges support from the Huawei UK Fellowship. AC acknowledges support from the EP- SRC CDT in Modern Statistics and Statistical Machine Learning (EP/S023151/1). AD acknowledges support of the UK Dstl and EPSRC grant EP/R013616/ This part of the collaboration between US DOD, UK MOD and UK EPSRC under the Multidisciplinary University Research Initiative. He also acknowledges support from the EPSRC grants CoSines (EP/R034710/)) and Bayes4Health (EP/R01 8561/1).
References
Albergo, M. S., Boffi, N. M., and Vanden-Eijnden, E. (2023), Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.085: 0879.7
Albergo, M. S. and Vanden Bijnden, E. (2023). Building normalizing flows with stochastic inter- polants. In International Conference on Learning Representations.
Anderson, B. (1982). Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313-326
Banerjee, A., Guo, X., and Wang, H. (2005). On the optimality of conditional expectation Bregman predictor. IEEE Transactions on Information Theory, 51(7):2664-260 2669..
Barczy, M. and Kern, P. (2013). Representations of multidimensional linear process bridges. Random Operators and Stochastic Equations, 119291959-18.
Benamou, J.-D. and Brenier, Y. (2000). A computational fluid mechanics solution to the Monge- Kantorovich mass transfer problem. Numerische Mathematik, 84(3):375-393.
Bernton, E., Heng, J., Doucet, A., and Jacob, P. E. (2019). Schrödinger bridge samplers. arXiv
preprint
Bischoff, T. and Deck, K. (2023). Unpaired downscaling of fluid flows with diffusion bridges. arXiv preprint arXiv:2305.01822
Bogachev, V. I., Krasovitskii, T. I., and Shaposhnikov, S. V. (2021). On uniqueness of probability solutions of the Pokker-Plance Talannkk-ooooooooo equation. Shornik: Mathematics, 212(6):745.
Bunne, C., Hsieh, Y.-P., Cuturi, M., and Krause, (2023). The Schrödinger bridge between Gaussian measures has closed form. In International Conference Artificial Intelligence and Statistics.
Bunne, C., Papaxanthos, L., Krause, A., and Cuturi, M. (2022). Proximal optimal transport modeling of population dynamics. In International Conference on Artificial Intelligence and Statistics, pages 6511-6528. PMLR.
Caluya, K. and Halder, A. (2021). Wasserstein proximal algorithms for the Schrödinger bridge prob- lem: Density control with nonlinear drift. IEEE Transactions on Automatic Control, 67(3): 1663- 1178.
Chen, R. T. and Lipman, Y. (2023). Riemannian flow matching on general geometries. arXiv preprint Ntti::2302.036660
Chen, T. (2023). On the importance of noise scheduling for diffusion models. arXiv preprint arXiv:2301.10972.
Chen, T.. Liu, G.-H., Tao, M., and Theodorou, E. A. (2023). Deep momentum multi-marginal Schrödinger bridge. arXiv preprint arXiv:2303.01751.
Chen, T., Liu, G.-H., and Theodorou, E. A. (2022). Likelihood training of Schrödinger bridge using forward-backwand SDEs theory. International Conference on Learning Representations.
Chen, Y., Georgiou, T., and Pavon, M. (2016). Entropic and displacement interpolation: a computa-
tional approach using the Hilbert metric. SIAM Journal Applied Mathematics, 76(6):2375-25
11

---

Chen, Y., Georgiou, T. T., and Pavon, M. (2021). Optimal transport in systems and control. Anmual Review of Control, Robotics, and Autonomous Systems,
Choi, Y., Uh, Y., Yoo, J., and Ha, J.-W. (2020). Stargan v2: Diverse image synthesis for multiple domains. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition.
Chung, K. L. and Walsh, J. B. (2006). Markov processes, Brownian motion, and Time Symmetry, volume 249. Springer Science & Business Media.
Csiszár, (1975). I-divergence geometry of probability distributions and minimization problems. The Annaly of Probability, 3(1):146-1558
Cuturi, M. (2013). Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems.
Dai Pra, P. (1991). A stochastic control approach to reciprocal diffusion processes. Applied Mathematics and Optimization, 23(1):313-329
De Bortoli, V., Thornton, J., Heng, J., and Doucet, (2021). Diffusion Schrödinger bridge with applications to score-based generative modeling. In Advances in Neural Information Processing Systems.
Delbracio, M. and Milanfar, P. (2023). Inversion by direct iteration: alternative to denoising diffusion for image restoration. arXiv preprint ar345
Dupire, B. (1994). Pricing with smile. Risk, 7(1):18-20)
Essid, M. and Pavon, M. (2019). Traversing the Schrödinger bridge strait: Robert Fortet's marvelous proof redux. Journal of Optimization Theory and Applications, 181(1):23-60.
Fatras, K., Zine, Y., Majewski, S., Flamary, R., Gribonval, R., and Courty, N. (2021). Minibatch optimal transport distances; analysis and applications. arXiv preprint arXiv:2/01.01792.
Fernandes, D. L., Vargas, F., Ek, C. H., and Campbell, N. D. (2021). Shooting Schrödinger's cat. In Fourth Symposium on Advances in Approximate Bayesian Inference.
Feydy, J., Charlier, B., Vialard, F.-X., and Peyré. G. (2017). Optimal transport for diffeomorphic registration. In Medical Image Computing and Computer Assisted Intervention- MICCAI, pages 291-299.
Finlay, C., Gerolin, A., Oberman, A. M., and Pooladian, A.-A. (2020). Learning normalizing flows from entropy- Kantorovich potentials. arXiv preprint arXiv:2006.0603 06033.
Flamary, R., Courty, N., Gramfort, A., Alaya, M. Z., Boisbunon, A., Chambon, S., Chapel, L.,
Corenflos, A., Fatras, K., Fournier, N., Gautheron, L., Gayraud, N. T., Janati, H., Rakotomamonjy. A., Redko, I., Rolet, A., Schutz, A., Seguy, V., Sutherland, J., Tavenard, R., Tong, A., and Vayer, T. (2021). Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1-8.
Föllmer, H. (1988). Random fields and diffusion processes. École d' d'Ééé de Probabilités de Saint-Flour XV-XVII, 1985-87, pages 101-203. Springer.
Fortet, R. (1940). Résolution d'un système d' d'uuuuttios de M. Schrödinger. Journal de Mathématiques Pures et Appliqués, 1:83-105.
Goodfellow, I. J., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A.. and Bengio, Y. (2014). Generative adversarial networks. arXiv preprint arXiv: (olo:4406.2111
Gyöngy, I. (1986). Mimicking the one-dimensional marginal distributions of processes having Itô differential. Probability Theory and Related Fields, 71:501-516.
Haussmann, U. G. and Pardoux, E. (1986). Time reversal of diffusions. The Annals of Probability. 14(4): 188-1205.
Heitz, E.,, Belcour, I.., and Chambon, T. (2023). Iterative aleddinbedddnng a minimalist deterministic diffusion model. In SIGGGRAPH.
12

---

Heng, J., Bortoli, D., Doucet, A., and Thornton, J. (2021). Simulating diffusion bridges with score matching. arXiv preprint arXiv:2111.07243.
Ho, J., Jain, A., and Abbeel, P. (2020). Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems.
Hoogeboom, E., Heek. J., and Salimans, T. (2023). simple diffusion: End-to-end diffusion for high resolution images. arXiv preprint arXiv:2301. 110..
Huang, C.-W., Lim, H., and Courville, A. C. (2021). A variational perspective on diffusioo-basee generative models and score matching. In Advances in Neural Information Processing Systems.
Hyvärinen, A. (2005). Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4).
Krylov, N. (1984). Once more about the connection between elliptic operators and Itô's stochastic equations. In Statistics and Control of Stochastic Processes, Steklov Seminar, pages 214-229.
Kullback, S. (1968). Probability densities with given marginals. The Annals of Mathematical Statistics, 39(4):1236-1243
Kullback, S. (1997). Information Theory and Statistics. Dover Publications, Inc., Mincola, NY. Reprint of the second (1968) edition.
Léonard, C. (2012). Girsanov theory under a finite entropy condition. In Séminaire de Probabilités XLIV, pages 429-465. Springer.
Léonard, C. (2014a). Some properties of path measures. In Séminaire de Probabilités XLVI, pages 207-230. Springer.
Léonard, C. (2014b). A survey of the Schrödinger problem and some its connections with optimal transport. Discrete & Continuous Dynamical Systems-A, 34(4):1533-1574
Léonard, C., Rælly, S., Zambrini, J.-C., al. (2014). Reciprocal processes. a measure-theoretiaal point of view. Probability Surveys, 11:237-269.
Lipman, Y., Chen, R. T., Ben-Hamu, H., Nickel, M., and Le, M. (2023). Flow matching for generative modeling. In International Conference Learning Representations.
Liu, G.-H., Chen, T., So, O., and Theodorou, E. A. (2022a). Deep generalized Schrödinger bridge. In Advances in Neural Information Processing Systems.
Liu, G.-H., Vahdat, A., Huang. D.-A., Theodorou, E. A., Nic, W., and Anandkumar, A. (2023a). PSB: image-to-image Schrödinger bridge. arXiv preprint arXiv:2302.08877
Liu, Q. (2022). Rectified flow: A marginal preserving approach to optimal transport. arXiv preprint arXiv:2209.145:
Liu, X., Gong, C., and Liu, Q. (2023b). Flow straight and fast: Learning to generate and transfer data with rectified flow. In International Conference Learning Representations.
Liu, X., Wu, L., Ye, M., and Liu, Q. (2022b). Let us build bridges: Understanding and extending diffusion generative models. arXiv preprint arXiv:2208.1 2e0e:22988..499
Liu, Z., Luo, P., Wang, X., and Tang, X. (2015). Deep learning face attributes in the wild. In International Conference Computer Vision.
Meng, C., He, Y., Song, Y., Song, J., Wu, J., J.-Y, and Ermon, S. (2022). SDEdit: Guided image synthesis and editing with stochastic differential equations. In International Conference on Learning Representations.
Palmowski, Z. and Rolski, T. (2002). A technique for exponential change of measure for Markov processes. Bernoulli, pages 767-785.
Pavon, M., Trigila, G., and Tabak, F. (2021). The data-driven Schrödinger bridge. Communications on Pure and Applied Mathematics, 14:5533-1
13

---

Peluchetti, S. (2021). Non-denoising forward-time diffusions. https://openreviev wttp/iiwwwwe//ormm id=oVfIKuhqfC.
Peluchetti, S. (2023), Diffusion bridge mixture transports, Schrödinger bridge problems and genera- tive modeling. arXiv preprint arXiv:2304.00917
Peyré, G. and Cuturi, M. (2019). Computational optimal transport. Foundations and Trends@ in Machine Learning, 1(5-6))355-607.
Pooladian, A.-A., Bam-H H.,, Domingo- C., Amos, B., Lipman, Y., and Chen, R. (2023).
Multisample flow matching: Straightening flows with minibatch couplings. In International Conference on Machine Learning.
Rogers, L. G. and Williams, D. (2000). Diffusions, Markov processes, and Martingales. Vol. 2. Cambridge Mathematical Library. Cambridge University Press, Cambridge. Itô calculus, Reprint of the second (1994) edition.
Rüschendorf, L. (1995). Convergence of the iterative proportional fitting procedure. The Annals of Statistics, 1304)41160
Rüschendorf, and Thomsen, W. (1993). Note on the Schrödinger equation and i-projections. Statistics & Probability letters, 1765):369-375
Schrödinger, (1932). Sur théorie relativiste de l'électron et e l'interprétation de la mécanique quantique. Amnale de l'Institut Henri Poincaré, 2(4):266-3) 310.
Shi, Y., De Bortoli, V., Deligiannidis, G., and Doucet, A. (2022). Conditional simulation using diffusion Schrödinger bridges. In Uncertainty in Artificial Intelligence.
Sinkhorn, R. (1967). Diagonal equivalence matrices with prescribed and column sums. The American Mathematical Monthly, 74(4):402-405.
Somnath, V. R., Pariset, M., Hsieh, Y.PP., Martinez, M. R., Krause, A., and Bunne, C. (2023). Aligned diffusion Schrödinger bridges. arXiv preprint arXiv:2302.1144:
Song, J., Meng, C., and Ermon. S. (2021a). Denoising diffusion implicit models. In International Conference on Learning Representations.
Song, K.-U. (2022). Applying regularized Schröddnggr--bridge-aaedd stochastic process in generative modeling. arXiv preprint wrX1v:2088.07331.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. (2021b). Score- based generative modeling through stochastic differential equations. International Conference on Learning Representations.
Stromme, A. (2023). Sampling from Schrödinger bridge. In International Conference on Artificial Intelligence and Statistics.
Tamir, E., Trapp, M., and Solin, (2023). Transport with support: Data-conditional diffusion bridges. arXiv preprint arXiv:2301.
Thornton, J., Hutchinson, M., Mathieu, E., De Bortoli, V., Teh, Y. W., and Doucet, A. (2022). Riemannian diffusion Schrödinger bridge. arXiv preprint arXiv:2207.03024
Tong, A., Malkin, N., Huguet, G., Zhang, Y., Rector-Brooks, J., Fatras, K., Wolf, G., and Bengio, Y. (2023). Conditional flow matching: Simulation-frea dynamic optimal transport. arXiv preprins 2.0000230..00482..
van Erven, and Harremoes, (2014). Rényi divergence and kullback-leibler divergence. IEEE Transactions Information Theory, 60(7):3797-3820.
Vargas, F, Thodoroff, P., Lamacraft, A., and Lawrence, N. (2021). Solving Schrödinger bridges via maximum likelihood. Entropy, 23(9):1134.
14

---

Villani, C. (2009). Optimal transport, volume 338 of Grundlehren der Mathematischen Wis- senschaften [Fundamental Principles of Mathematical Sciences]. Springer- Verlag, Berlin, Old and new.
Vincent, P. (2011). A connection between score matching and denoising autoencoders. Neural Computation,
Wang. G., Jiao, Y., Xu, Q.. Wang. Y., and Yang, C. (2021). Deep generative learning via Schrödinger bridge. In International Conference Machine Learning.
15

---

Outline of the Appendix
In Appendix A, we first clarify the relationship between different methods in the existing literature and our proposed DSBM framework. In Appendix B, we focus on the family of linear SDEs, and draw link between the parameterization of bridges in this paper and the stochastic interpolant in Albergo et al. (2023). In Appendix C, we give proofs for results in the main text. In Appendix D, we present additional theoretical results for IMF in the Gaussian case. In Appendix E, we derive the discrete-time version of Markovian projection. In Appendix F, we explain the benefits of DSBM compared to DSB in more detail. In Appendix G, we describe a method for learning the forward and backward processes jointly and propose consistency loss between the forward and backward processes. In Appendix H, we present additional methodological details for a practical scaling of the loss function reduce variance, similar to standard Denoising Diffusion Models. In Appendix I, we give further details for all experiments and additional experimental results. Finally, we discuss broader impacts of our work in Appendix J.
A Discussion of Existing Works
A.1 Bridge Matching and Flow Matching Models
In this section, we clarify the relationship between variants of Flow Matching and show that they are equivalent under some conditions. We follow the nomenclature of Tong et al. (2023). We refer to the algorithm originally proposed Lipman et al. (2023) using linear probability paths and described in (Tong et al., 2023, Section 4.1) as Flow Matching (FM), and the algorithm proposed in (Tong et al., 2023, Section 4.2) as Conditional Flow Matching (CFM). There is a small constant parameter miin in both algorithms, which controls the smoothing of the modeled distribution. We consider the case min = 0. Then CFM recovers exactly the 1st iteration of Rectified Flow (Liu et al., 2023b). Furthermore, FM, CFM and the 1st iteration of Rectified Flow are all equivalent when performing generative modeling with a standard Gaussian uy- We refer to them collectively as Flow Matching models (FMMs) as they only differ in the smoothing method. We also present them all under the Bridge Matching framework. These models can also be interpreted in the context of stochastic interpolants (Albergo and Vanden- Eijinenn, 2023; Albergo et al., 2023). Finally, present recent applications of Bridge Matching and show that some of the objectives in Somnath et al. (2023); Liu al. (2023a); Delbracio and Milanfar (2023) are identical.
Flow Matching and Conditional Flow Matching. al., 2023, Equation (21)) is
In Flow Matching (FM), the objective (Lipman
En -X) -) vo (t. x,11
where 17, is given by Fr (X, N(X,; XT,, (1-
4)²)-
In Conditional Flow Matching (CFM), X..T = - + (1 -1-3) 1 Xo. with Xo ~ N(O, Id) and the objective (Tong al., 2023, Equation (16)) is given by
E [I (Xr Xo)/T vo (1, ))2)
(15)
This is the same as (Liu et al., 2023b, Equation (1)). Furthermore, (X, - x1- (Xr
XT) - t) (XT- ) / t), so the CFM objective equivalent to
En, r [1((Tr x 0.7 )/( -) -) vo (t, X. )12).
(16)
The optimal ve xx) = (Bur X X, = xzc ITT - t). In the case of generative modeling, To is a standard Gaussian distribution and Hor is given by N(Xu; 0, Id) Thus, Ir is also given by NT (Xr(X) ( T Xx, (1 ))²). Therefore, the FM (Lipman et al., 2023) and CFM (Tong et al., 2023) objectives are exactly the same. However, CFM is also applicable when To is not Gaussian distributed, so CFM is a generalized version of FM'.
'In the case min > O. FM and CFM are indeed different in how smoothing is performed, and we refer to Tong et al. (2023) for more detailed analysis.
16

---

Stochastic Interpolant. In (Albergo and Vanden- Eijnde, 2023; Albergo et al., 2023), the concept of stochastic interpolant introduced. In Albergo and Vanden-Eijned (2023), the interpolation is deterministic (not necessarily linear), of the form (T 2T) a(() + (() while in Albergo et al. (2023), the interpolation is stochastic given by [2 TT) a(() + (() + Z for Z ~ ((,, Id). In Albergo and Vanden-Eijnden (2023), an ODE is learned and the associated velocity field uy is obtained by minimizing the following objective (Albergo and Vanden-Eijnden, 2023, Equation (9))
BI,, (1l,,11((((x,..... XT) vo X )[())
x
Hence, if (100, xT) = 7000 + (1 7 )eT, we recover (15).
Link with Bridge Matching. When Q is associated with the Brownian motion (oB,))e(o.r) and + 0 in Bridge Matching, we recover the same objective (5) as the Flow Matching objective (16), since log (Xr(X,) (X (T)
t))- Bridge Matching can also be applied to general distributions To, Tr; i.e. does not have to be restricted to a Gaussian. Therefore, Bridge Matching a generalized version of Flow Matching, see also (Liu et al., 2022b, Equation (10)).
Inverse problems and interpolation. Somnath et al. (2023) and Liu et al. (2023a) present Bridge Matching algorithms between aligned data (Xo: (x):X)) T,, The objectives (Somnath et al., 2023, Equation (8)) and (Liu et al., 2023a, Equation (12)) are equivalent to the Bridge Matching objective (3). The main difference between Liu et al. (2023a) and Somnath al. (2023) resides in the choice of Ho.r. In the case of Somnath et al. (2023), this choice is motivated by the access to aligned data with applications in biology assuming they are distributed as the true Schrödinger static coupling, iee. Hor = IS In the case of Liu et al. (2023a), Ho, corresponds to pairing between clean and corrupted images, e.g. with Hh To the distribution of clean images and ΠT = Tr the distribution of corrupted images obtained from the clean images using the degradation kernel NToo--
Finally, in (Delbracio and Milanfar, 2023, Equation (5)) the authors consider a reconstruction process of the form
dX2- = (Biooj Xo X,] - ttdt,
XT ~ 1lys
(17)
where here we have replaced F( t) by En,.. Xo X, m²). This is justified if the ||-1Ip norm in (Delbracio and Milanfar, 2023, Equation (4)) is replaced by 11 (or any Bregman Loss Function, see Banerjee et al. (2005)). In Delbracio and Milanfar (2023). Пo, corresponds to the joint distribution of clean and corrupted images as in Liu et al. (2023a). Exchanging the role of Hl and Ty, (17) can be rewritten equivalently
dx,
(T) t)dt, (FIry X x,
X Too
We thus obtain the optimal Flow Matching vector field vp(t, = (Biri Xr X; = xx)- t) in (16). Note that Delbracio and Milanfar (2023) also incorporates a stochastic version of their objective (Delbracio and Milanfar, 2023, Equation (7)). It remains open question whether this objective can be understood as special instance of the Bridge Matching framework.
A.2 On DSBM and Existing Works
In this section, we show that the DSBM framework recovers the above existing algorithms for different choices of bridges Q and couplings I Algorithm 1. For the independent coupling
n = To 8 Ty and Brownian bridge Qjo (4) with diffusion parameter o a, the loss function (10) recovers the Brownian Bridge Matching loss (5). Letting -> O, we recover Flow Matching (Lipman et al., 2023). In this case, further iterations repeating lines 7-9 in Algorithm (with only forward projections) recover Rectified Flow (Liu al., 2023b). If the coupling 11. is given by an estimation of the OT map between To and тT, then the first iteration recovers OT-CFM (Tong al., 2023; Pooladian et al., 2023). Finally, for general bridges Q if we are given the optimal Schrödinger Bridge static coupling To m then the DSBM procedure converges in one iteration and we recover Somnath et al. (2023).
A.3 DSBM and Rectified Flow
We discuss the differences in more detail between our proposed DSBM method and Rectified Flow. Both Rectified Flow and DSBM are general frameworks for building transport maps between
17

---

two general distributions To. Rr. However, there are a few important theoretical and practical differences. Firstly, we adopt the SDE approach as opposed to the ODE approach in Rectified Flow. This distinction is crucial in theory, as Proposition 5, which guarantees the uniqueness of the characterization of SB, valid only when ot > O. Consequently, Rectified Flow is not guaranteed to converge to the dynamic optimal transport solution (see e.g. counterexample in Liu (2022)). In following work, Liu (2022) established formal connections between Rectified Flow and OT when restricting the class of vector fields to gradient fields. In DSBM, the connection to OT is obtained by considering its entropy-regularize observe significant improvements of sample quality in our experiments when performing transport between two general distributions. This is in line with the theoretical analysis in Albergo et al. (2023). On the other hand, while Bridge Matching also achieves high sample quality using the SDE approach, the transported samples are much more dissimilar to the input data (see e.g. Figures 14, 15, 19). Lastly, Rectified Flow also performs Markovian projections iteratively, but only in the forward direction. Consequently, the bias in the learned marginals Py is accumulated and cannot be corrected in later iterations, i.e. the first iteration of RF will achieve the most accurate marginal P.-- Subsequent iterations can improve the straightness of the flow, but at the cost of sampling accuracy of 17. We observe in practice that this becomes particularly problematic if the first iteration of Rectified Flow (which is equivalent to CFM) fails to provide a good transport and learn an accurate JP7, e.g. in the case of moons-8 -ggaussinnss (Table 2), Gaussian transport (Figure 3), and MNIST, EMNIST transfer (Figure 5 and Figure 13). As Rectified Flow cannot recover from this issue, we observe the accuracy of l7 only deteriorates in further iterations as n increases. In our methodology, we leverage Proposition 9 to perform forward and backward Bridge Matching, and we observe that the marginal accuracy is able to improve with iteration.
version. Furthermore, by adopting the SDE approach, we
B
The Design Space of Brownian Bridges
B.1
Relationship to Stochastic Interpolants
From stochastic interpolants to Brownian bridges. In this section, we draw a link between our parameterization of bridges and the one used in Albergo et al. (2023). In Albergo al. (2023), stochastic interpolant is defined as
X, = م,50 + 3,OT + 7e
(18)
where Z N(0, Id). Since their methodology and analysis mainly relies on the probability flow, they work with (18), which is easier to analyse. In our setting, as we deal mostly with diffusions, it is natural to parameterize Brownian bridges as follows
dX, {- Bar Jdl + d.
(19)
The goal of this section is to derive explicit formulas between the parameters ā,, 3r and 5e of (18) and the parameters at 3r and 7. of (19). Consider (Xi)ee00,1) given by (19). We have that for any I [0,T]
X, expl-Ai)ag + So B. exp[A, AijdsT + So explA. A,JdB,
where A, So a,s.. Therefore, we have that
āe =xeeppp- S0 a,ds).
B = f0 B. expl- 1 andujds,
7/ = S0 13 expl-2 2 andudds
(20)
of 04
Be = 3 + Borr
x2 = (5: + 27/0t ar = 27t + 25,7
(21)
Using this relationship, we get that the Markovian projection, see Definition 1, is given by
dx = ff + ddB
fx = Enr -ax + 3,Xr X, zz).
We have that
(x.) = Engu az X X, = =x]]
= [Eumm (a, + X, + 7) + 3, X, = xz).
18

---

Using (21), we get that
(xt) = Li + 3 + fi Z X, = Fz..
In Albergo al. (2023), it is shown that V log M2 (1r) = = -o. [ X; = /t///t. where M is
the Markovian projection. The probability flow associated with (Xi) given by
dz: (T((Z))
z V log (2 2
(E o6/00
+
3 x + (-a + 22 Z X = Z,]]dt
2%
(Eulory o
x r + 7 Z
X dt
Hence, recover (Albergo al., 2023, Theorem 2.6).
Non-Markov path measures. A natural question is whether (19) arises as the bridge measure of some Markov measure. For instance, Q is associated with (zg Bt)xe (o,77, then pinning the process at 2T at time T, we get that the associated bridge measure Qi0, is given by
dx. = (øT x,) - t)dt + ddB,
1 and T/ = 1. Using (20), we get that 0; T Therefore, we recover (19) with a; = Br = T-t
3, f and 5.- (T t) We recover (4), upon noting that B, Br is Gaussian with zero mean and variance (T - t)t/T.
More generally, we consider Markov measure Q associated with (X,) (X,) such that
dX, = a, X,dl + adB
X = =00
We now derive the associated bridge measure Q 0.7
X, = expl + A ff cs explA,, Ar]dB.,
with A, = =0² a,d.. We have that
a Vn,, n, log Qr1((ar/a))) = (c, exp(A Ar// 1ff ct exp/2(A, - Ar)]ds)or
- (c, 2(A - Ar) Jtt ce exp|2(A, - Ad)sss
Therefore, combining this result and (2), we get that Q associated with
a, = a: + ef expl-2 ST a,ds)/ J," et expp1 S a,dulds, A. = a expl- Jt a,d8// S exxp1- anddudss
7t = Ct.
In that case (a,, G) entirely parameterize (at, 3r (E).T Hence, in the Omntein-UUllnneeck setting, if QQ, the bridge of Markov measure, it is fully parameterized by two functions while in the non-Markov setting it is parameterized by three functions.
In this paper, present framework in the Markovian setting the Schrödinger Bridge problem is usually defined with respect to Markov reference measures. However, our methodology could be extended in a straightforward fashion to the non-Markovian setting. This would allow for a further exploration of the design space of DSBM.
B.2 Linear SDE and Bridge Matching
In this section, study further the diffusion bridge of linear SDEs. Arbitrary Markov measures can be chosen to build bridges; however, we want to be able to compute some representations of the bridge in an explicit way. More precisely, denoting (x) 0.7) )1e the diffusion bridge with To, Sr the initial and final condition, we want to have access to the following:
integral sampler: we want to have formula to sample X.TT for any / E [0, T] without having to run stochastic process forward backward.
19

---

forward sampler: we want to have a forward SDE for X with explicit coefficients, which might depend on Ir,, running in a forward fashion terminating at 27.
backward sampler: we want have backward SDE for Y X 0.7 with explicit coefficients, which might depend on To. running in a backward fashion terminating at zo-
We focus on linear SDEs of the form dX, -0,,, + 0B aB, which are particularly amenable, where (E) is schedule with /3 E C( 7), (0, +00)).
B.2.1 Brownian motion
First, consider the Brownian motion setting and let (X))tejo.ry be associated with Q with dX, -B We consider (X. )eeo, conditioned at both ends X zo and X7 IT First, using (Barczy and Ker, 2013, Theorem 3.3), we have that for any /. E [0, T)
X 0,7 B(LT)
RO,T) To HOOT (IT XT) X,,
R(0,0)
with R(s,t) () S22 du = (B - -B),, where for any t E [0, B, = S0 3,ds. Therefore, we get that
x = (1
B(T) mm +
B(t)
B(t) B(T) (2T X X
(22)
(22) defines the integral sampler. In addition, using (Barczy and Kern, 2013, Theorem 3.2), we have
x 20.
dx. 0.T Be/ntt to² + t a / aB.
where
(s, (B((t() B(s)). Therefore, get that
Xo_ To
8. dx 0.T {- B(T) X + B(T) BO)
pdt ]et + 0 dB.
2r
(23)
(23) defines the forward sampler. Finally, derive the backward sampler by considering the time-reversal of the forward unconditional process (initialized at 2o). Following Haussmann and Pardoux (1986),
Yo IT
dY!.T o² Br Vog fo O.T Ir + 0372 o
(24)
In addition, we have that
X, = xo + B(t
~~ N(O,
Hence, we get that for any Combining this result and (24), get
[0, and Rdd
log @1)()))0)) = -(x (
YOT UT,
dy =
Br YO. B(T
8r .T zo + o
(25)
Combining (22), (23) and (25), get
XOT A1To + -
- -XT) + X..
OOT 10:
dx..7 = { x.x v xr )el + o/222BBBB
YOT FT:
dvo,t = -K YDT + , roldd 3// dB,
with
X = 1
B[T
B(t)
6:
B(T)
B(0)
B
v
8
K
B[T)
B
B(T
bb
B.
T
81
B(O
20

---

B.2.2
Second, we consider the Orastein-Uhlenbeck setting and let (Xx) with dX, adt 00//2 dB, with a & 0. We consider (X 0. )ae(0,7]. the stochastic process (Xx) conditioned both ends XD_ In and x TT. First, using (Barezy and Kern, 2013, Theorem 3.3), have that for any E [0, T]
0.T RI.
ROTT To HOTT (TT XT) + X.,
R1000)
with R(s,t) (xp)a(B(t)) ((s))lr(sstt) t), with 7(s,t) S: 02 exp[-2a(B(f) ((u))]du. In particular, we have
7(8 () = = (1 1eppl-2a(B(() - B())), R(s,() = =2 =intaa= sinh(a(B(()) B())))-
(26)
Therefore, we get that
x (B()((((()))()))
Bit)(a sinb(aB(T]) To siab(a((Ti) (IT XT) + X.,
(27)
(27) defines the integral sampler. In addition, using (Barczy and Ker, 2013, Theorem 3.2) and (26), have x 2a and
dxo. (-a,,X,-
2 expl 2B(T B(())
X + a expl -a((B((-B(())))
FT dt + 0 /
{(,T)
7(.2)
1-a 2a (xp)2a(((T)-B(0)
2cr9,
2a..
ctt + o
expl -B)
{-a exp (o()) - B(])
0.T ((TT)
x
+ 27) )t
2n
/ dB,
expl- a B(t))
exp
a-a, coth(a B(t)) X
+ aB3 (scc((a(B(BB)) (()))
dB.
In the formula, coth is the hyperbolic cotangent function defined as coth(x) Laub(x) sinb(a) and esch the hyperbolic cosecant function defined esch(x) (26), we get that
cosh
"Sn((z] Combining this result and
0.T vo, dx = (-a (LL 0.T +aß, (cl(
dB,
(28)
(28) defines the forward sampler. Finally, derive the backward sampler by considering the time-reversal of the forward unconditional process (initialized at zo). Following Haussmann and Pardoux (1986),
Y 0.T ET,
dY (aBr-oY! + Vogg--ooo@rr-oooo (=)))
[oo)] o
(29)
In addition, we have that
X = B(t] + a expl exp)= I(0,
/2a
Hence, we get that for any E [0, T) and Rdd
V log ( (x|20) = 22(( 2xpplaa8(t)aa/lddd1oxpp-------01
Combining this result and (29), we get Y = xT and
dy = {aßr- 0.T
2asr
0,T
20 exp
T-exp! 2a ) dt + o
1
T-expl
(T-)
(-0.3T- 1 1texp|-2aB(12
-20 y U.T
2007
)ll 033//,,n
- 21/2 0
+
Therefore
VO.T 0 = ST, dYO.T_ (-ao0--cah(aBBT-11)) hhtp:ooooaaaooo""""oooooooccoeed (c)))))((())) dB,
(30)
21

---

Combining (27), (28) and (30), we get
T Ar + Fi(TT - XT) + X..
x 20:
dx. = {N/ x.. + V dd + 0 d
Y JT,
dYp." ( Y T + +0 20 dt + o , aB
with
A ((T))))(T)-B(())
9t sinh(a
sin((aB(T))
m/ a3, coth(((((())) B()))) w/ aBt -
Ky aBt coth(a
vi = aB, csch(a
Using that tanh(x) ~ z and sinh(z) ~ x for x + O, we recover the Brownian motion setting by letting + 0. Note that Albergo al. (2023) show that given integral sampler, which does not necessarily comes from a Markovian process, a forward sampler with the same marginals can be defined, although it does not necessarily satisfies the fact that the paths have the same distribution.
C Proofs
C.1 Proof of Proposition 2
We refer the reader to Chung and Walsh (2006); Rogers and Williams (2000) for an introduction to Doob h-trranffrmm Our theoretical treatment of the Doob h-transform closely follows Palmowski and Rolski (2002).
First, we introduce the infinitesimal generator given for any f C ((0, x Rdd R), [0, T)
and x E Rd by
xf(t, = (f((z), +
+ at(,
(tf
(31)
The following assumption ensures that the diffusion associated with Q as well as its Markovian projections well-defined.
A1. f. a and xx) + ExTryel log Qr (X)(X,) X, = xa] are locally Lipschitz and there exist C O, E C(10, T), IR) such that for any E [0, T] and no, Ie E Rd, have
(x:))l - C(1 ||x
C N az Z I/C,
[8 [ V log
X, z,]ll S Cc(()(1+ ||xe|]).
We consider the following assumption, which will ensure that we can apply Doob l-ttransform techniques.
A2. For any Do E R4,, Iroo is absolutely continuous w.r.t QTio- For any ay E Rd, let Prio be given for any E Rd by Prro(orr/z)) (zo) = (oorr////drmppp To E Rd. Fr (pro((rlao)) is bounded. For any zo Rd, let lelo given for any In Rd and
and assume that for any
t E [0, by
(o)) Sye (r (Trn Finally, we assume that for any Do E Rdd, (t) + (=+ and (1, + Apqoo(z((aa) are bounded.
(32)
This means that for any Zo E ; (t, 2x) + (e (mn) is a good function and Rolski, 2002, Proposition 3.2). Note here that these assumptions could be relaxed case-by- case basis. We leave this study for future work.
function in the sense of (Palmowski
The following lemma is direct consequence of A2 and (32). It ensures that the h-function PPoj satisfies the backward Kolmogorov equation.
Lemma 11. Assume A2. Then, & E C²,2 T) Rd R) and Aipoo 0.
Using (31), we have that for any zo E Rd and f E C (00 T Rd R). E [0,7] and x, Rdd
(A(fv10) fAppo)(t / Af o ft, v log (1)))1))))))))
22

---

Finally, we consider the following assumption, which will ensure that the Doob h-transform is well-defined.
3. For any zo E Rd, there exists C 2 such that for any E [0, and It E Rd |/ log (=jo( C(1 + || [
We now ready state and prove Proposition? 2. Note that the Markovian projection is defined in Definition 1. Finally, we define M the space of path measures such that P E M if IP is associated with dx, (X,, 2 (X,) + a,dB,, with o, v locally Lipschitz. This restriction of Markov measures allows us to apply the entropic version of the Girsanov theorem (Léonard, 2012). It has no impact on our methodology.
Proposition. Assume A1, A2, A3. Let M' = ( Then,
M = ( (T]M) M E EM}
KL((TT(M') = Lo' Eno, ( (V log Qry X.) Xo, X,] og /
In addition, we have that for any t [0, Mc II.. In particular, Mi- 1l7.
Proof. First, we recall that II is given by II = Q with Po,r = d In particular, we have aQo
all 1h QoProo where pryjo aQrip Therefore, using Lemma 11, (Palmowski and Rolski, 2002,
Lemma 3.1, Lemma 4.1), the remark following (Palmowski and Rolski, 2002, Lemma 4.1), A1, A2
and A3, get that g is Markov and associated with the distribution of (X))x{(o, given for any
t e [0, Tj by
X, x- So
ofV Po (x,, (x,)d))dds + So adB
(33)
where for any E [0, Fu, v2 E Rdd we recall that
Prjo (22 (mo) fR+ 4T (2))
(34)
First, we have that for any E [0, 17, Iz, 2o Rd
(=)0 (2)) = Sg. Qrjo. (2, Ir, fo) o( (2))
Therefore, we get that for any E [0, T) and F2, Fo E Rdd
(r = (=)
(35)
aQ u(<rlao)
In addition, we have the following identity for any E [0, T) vo, F4, or E Rd Qrjo
rT) = Q (2x (rur((())))) Using (34), this result and (35), we get that for any / E [o, T) and To, E IRd log Pejo (xo)
S oo (mn (xrr()))))))
log Quy (IT (xT
(o (mm (m,
or (wo Ceoo) (mx (so,eT)
V log (oggl
@njo (m)) (a))ccioo (o (mo)
= SM 11 (mo)
Il. (ara) V log (2T(((TT
= Sg V log Qrj
2o).
Hence, combining this result and (33), we get
X, [ + of of V log Q77 Xx, X,) X, Xo) ds + fo a,dB.. Let MI be Markov defined by dX,, = (X, ) + v { + d, such that KL(IIMM) < to0 with o, v locally Lipschitz. Using (Léonard, 2012, Theorem 2.3), we get that
[L( = 2 Je Eno, [ V log @T X X,, Xo] ue (X, (1
In addition, we have that for any E [0,
En,, !|lo, n.a VV log QT (X, X X, X 1/ (
> En... lllo 1|10 / log Qr (Xy X,) X,, X.] - 0j (X,)
where of (1x) = En,, log (X,, X, 2r} which concludes the first part of the proof. For the second part of the proof, we show that for any / E [0, T). we have M; II, First, we have that MI; and II, satisfy the same Fokker- Planck equation, see (Peluchetti, 2021, Theorem 2) for instance. We conclude by uniqueness of the solutions of the Fokker-Pland equation under A and A3, see Bogachev al. (2021) for instance.
D
23

---

C.2 Proof of Proposition 4
Proof. By the additive property of KL divergence (Léonard, 2014a), [L(PI]I) = KL Ep [KL(P (.. Restricting Ho = Q directly gives that the KL minimizer is II with oor FO,T, and thus II* Qo which we recall is short for IT)
Jaxxa (-)or xr (rrrrr)oo (orr)rrrrrd drr).
C.3 Proof of Proposition 5
This result is direct consequence of (Léonard, 2014b, Theorem 2.12), which we recall here for completeness.
Proposition. Assume that Q E M, that Qo Qr Qu that for any Io. 2T R ( IT) 2 expl=A(zo)- --(x))r with A 0 measurable, fRexR= expl-B(xx) (xo)= B(rr))dQ (zo, IT) < +oo with B 2 0 measurable. Assume that there exists to E (0, and X
measurable such that Qto (X) > 0 and for all x E X. Qo.T & Qu.T (1X X x). In ad-
dition, assume that (L ( +o,, KL(TT(Q) t00, SR + ()
to0,
[ga ()(rr)drr (xT) +00.
Then there exists a unique Schrödinger Bridge PSB. addition let I be Markov measure in the reciprocal class of Q such that Po = and Pr Tr. Assume that (L(PP QQ) < +o0. Then IP is the unique Schrödinger Bridge PSB
Proof. The first part of the proof is a consequence of (Léonard, 2014b, Theorem 2.12(a)). The second part a consequence of (Léonard, 2014b, Theorem 2.12(b)) and (Léonard et al., 2014, Theorem 2.14).
C.4 Proof of Lemma
Lemma. Let M E M and IT e R(Q) and assume A1, A2, A3. If KT((IIMM) too and (T,(pojj, 400 we have
M = (L)) + KL(proi,
(36)
For any P E P(C). III +00, have
(L(P|IT) KT. (P)) + (T(
(37)
Proof. We start with the proof (36). Similarly to Proposition 2, where we have M. E M to ensure that we can apply (Léonard, 2012, Theorem 2.3), we get
(T((IIMM) of Ent, (X.)) Bn V log (XT Xo. X.1131///
In addition, we have
(T.(PrjjM 2 1 En, (X,), af (V log (TT: | /
Finally, using Proposition 2, we have that
(I))
2 S En.. I1 Eo V log Q(Xr(X;) Xo. X.]
L V Q(xr Xx
2 fo (E Em Qlg Q(x
Xn, X.) [] En, Em (V log Qo X.) X.
24

---

Using this result, we have
[T (T)) 2KT.(pro
=1o (Bul, [ 11r
V log Q( X Xo. X. En, Bu | log Q(X+((.)) x, x
+ So En, (X:)) Eur, V log Q (X, X.) dott
1o (BBoo,.llETT, V log ( X X ]|]
[ Q(X7((( o)
+ So (En, [x)])//))x Em, B oog ( offl
2 1od En, E log x x,l)lo,dt
Lo' Em... II log
X | + En, (nn,llo((x)///f1l/oood
fo En, ((, (X.) ,,, log X X.)
=Jo' En. [I ETyoo. log
x X oo
1od En
2 5o' BT IN,, (X,) // E En,o. log Qr X
X X. 2KL(II)M),
which concludes the first part of the proof.
For (37), define 11" pro (P") Po,r Q/o.r. Using (Csiszár, 1975, Equation 2.6), we have
[T KT.(PITT*) + Se log( aml (d)P()
KL(P + JR4 xR_ log( dil
a (Do, x1 21))dPPPuup x1)
(L(P)IT') + SRaxRd log( llg(, (zo, x1 2 (2), KL(PIII") + (LIIIIII)
dill.r
which concludes the proof.
C.5 Proof of Proposition 7
Proposition. Assume that the conditions of Proposition 5 and Lemma apply for P for every N and for the Schrödinger Bridge TSS we have KT. (PpBt < KT. PSB) < co, and lom [Tp(PP 1pn-1) =0 O.
Proof. We follow the technique of Rüschendorf (1995) but for the reverse HellbackkLeeiblle diver- gence. Applying Lemma 6, we get for any N E N
(T(B))BBB) = (II(PO)P)) + KLLPPIPSS = E KL(PP(Pttt) + KI(PNNIIISB)
which concludes the proof.
C.6 Proof of Theorem 8
Theorem. Assume that the conditions of Proposition 5 and Lemma 6 apply for 1 for every n E N for the Schrödinger Bridge pSn the IMF sequence (IP") admits a unique fixed point 1=* = psn
and lim (L [P**)
Proof. By Proposition 7, (L(P" < (I < oo for all E N. Now, using the coercivity of KL(-IPSR) (this is where our analysis differs from the one of the IPF), we have that the IMF sequence (IPm) eee and its subsequences (M.*+)) )eEN and (IT") are subsets of (TP E P(C)= KT.(PP pSB) < (PTIBB))))) which (relatively) compact. Thus, (MM++)) ) contains convergent subsequence M + M. j - and (TI%) (TT)) contains further convergent subsequence I - IT weakly as k + sc. As the Markov and the reciprocal classes closed under weak convergence, M E MM and 11* R(Q). Now, by the lower semi- continuity of KL divergence in the weak topology (van Erven and Harremoes, 2014, Theorem 19). 0 KL(MIIII' < lim inf KL(M [1 as 1P E M n RQQ). Also, P satisfies P = my and P7 = Ry as is satisfied by all IP". By Proposition 5, the unique Schrödinger Bridge PsB Finally, limm too (T.(P")PP)) using limp--oc (L [F) 0 and the monotonicity of (L(PIIpSB by Proposition 7. D
O. Hence, MI II which we denote
follows
25

---

C.7 Proof of Proposition 9
Proof.
The proof similar the one of Proposition 2.
In particular, the time-reversal of (20 ur) is associated with
dY = (- - v log
Y..." |mo)
( Jdt + or-
YOT
TT- (38)
0
One can view both (11) (12) as SDEs with drift defined as the conditional expectation of the drift of (2) (38) under Ho,r in the forward and backward directions respectively.
C.8 Proof of Proposition 10
Proof. We proceed by induction. Firstly, for N0.7 Qu,r. Il = pu = Q at initialization. We can also define M Q, such that MI 100 and IT projR(e) (M") for all n E N. By (De Bortoli al., 2021, Section 3.5), the optimal DSB sequence [p" is Markov and Ipn Q where Pa 0.7 is the IPF sequence of the static SB problem. In other words, Fi E MMR((Q.
Suppose M j2n+1 By definition, M I22+22 To. i.e. both forward processes are initialized at mn- In DSB, by De Bortoli et al. (2021), p2n+2 is defined as the time-reversal of f22+11 such that 1 1P00+1 Hence, j22+22 Tolp 1b2++1
In DSBM, we first perform reciprocal projection 112 projR(Q)) (M2n+1) M Qjo. Since M p2n+1 E R(Q), however, have that 112 M Furthermore, since M p2n+1 E M, (
(MP2 M Thus, M2 given by M We conclude that M TnM2n-1 = p2n+2 Similar arguments holds for the the reverse projection (13). Therefore,
(9) is such that M To and M projM(II?*t)
M fmm for all n N.
C.9 The set of Markov measures is not convex
The result of Lemma should be compared with the information geometry result of (Csiszár, 1975, Theorem 2.2), which states that if C is convex set and IP E C. then under mild conditions, KL(PIQ) = (P|ojjrrjc (Q)) (Q))), where proje (Q) argmin, (KL(I : P C} is the projection of Q on C. Note that, contrary to Lemma 6, (Csiszár, 1975, Theorem 2.2) is given for the forward Kullback-Leibler divergence whereas Lemma given for the reverse KI. divergence. In addition, (Csiszár, 1975, Theorem 2.2) requires the projection set C to be convex which not satisfied for the space of Markov measures M. We give a simple counter-example proving that the set of Markov measures is not convex.
Let pi (zo, 31, x2) (1()))p1 (21l (12) and 22(20,01,22) p=(co)ps((r/()))ps(aa)p)) on (0, 1)3 such that
Pi (1 = 1) = ao:
pi(21 = 1/20) = aa,
P1 = 1|x1) = 02
Additionally, we set
P2(20 1) Bo,
p2(21 1220) 31,
P2(#2 1|m)) = B2.
Finally, we set (1/2)p: + (1/2)pg- Consider g(0² = 1/2 = 1, 1) = q(02 = 1, =
1, Zi = / = 1, xg = 1) and (xx2 = 1|2| = 1) = ( 1, = 1)/g(=1 = 1). Let
A 4lg(02 1, 1, 1)((21 1) ((1- = 1,x, 1) 1)]
(0pa)a0 +
+ 31) +
+ -031)
ana13,a2 - 80,1
(0, - on) + 018100(000 - 3o)
(3, (B) 02).
q is Markov if and only if A = 0. Therefore g is not Markov soon on * 8o and ag # B2.
26

---

C.10 Impact of the resolution on the entropic regularization
In this section, we show how the resolution of the input affects the entropic regularization. First, we consider the Schrödinger bridge problem between and (T = 1) with Q associated with (o)oo)e)...1- In that case the static Schrödinger Bridge is given by Pū, such that
F0 = argmin (S==x ||z0 - x1 (xo; x1) - <KT([P]m To @ mI) Po = =oo Pi Pi=m}
where ε = 20 We now describe the solution of a higher dimensional problem given by the upsampling of the marginals. Let E N with 2 1 and down W + Rd such that for any k E {l..... .d.. and æ E IRd,, down(z) = I with ã E R' and In = Ixf. We also denote up Rd + fdd such that for any k E {l...., d). and f E Rd, up(Ø) x with I E R and F& = 1) for (l.....dd .d. Note that down up Id.
We denote To u no and मं up We extend up and down to Fd Ra and IR x 18da by letting for any I, y E F and z, E Rd up( ū 5) (up(#), up()),
down(æ, v) (down(r)) down(y)).
First, note that since uponown- ld on (p( we get that for any supported on up((R)) x up())). using (Kullback, 1997, Theorem 4.1)
(L(P|=0' @ πi HP) = KL P Plmo @ TT).
Second let P such that Pp To and P1 "l Then, IP is supported on up(Rd) x up((R").
Therefore, for any IP such that Po up and P1 = x
(F a KL(down
TT).
Finally, for any P such that Po TD up and P1 = we have
uh
SEraxRaa I| I1 /|2d x1) f S 1:o - I1 [Pd(ddon
Therefore, for any such that Pg = uP and P: . we have
JRitxRra 11:00 -
= 1: ddwwn P) (1) (e/ m)))
Therefore, we have the following result.
Proposition 12. Let O. F is the solution of the static Schrödinger bridge with marginals To मh and regularization E if and only if dlownyP Ix is the solution of the static Schrödinger bridge with marginals To, 31 and regularization e/f²
This means in particular that the Schrödinger Bridge not invariant via upsampling. In Appendix I.4, we confirm these results visually upon noting that for the same running DSBM at resolution 64 x 64 and 128 x 128 gives different results, even after downsampling of the 128 x 128 results.
D Convergence of IMF in the Gaussian setting
In this section, we study the IMF in an one-dimensional Gaussian case. We consider T 1, Th III N(0, (1/23²)) and associated with (oB:)ue(o.1) where 0. In what follows, let
(1 x2
(t 1).
0 = (1/282)
(39)
where e E [0, 1]. We also denote Ø² = 202 We start with the following result which gives an explicit expression of some marginals of the reciprocal projection.
Lemma 13. Ler I1,1 = N(0, Σ²) with Σd given by (39). Ler I1o = H8,1 Qjo. For any E [0, 1], we
have that IT 1. = N(0. Σ) with
2²
a
Σ = (1/288)
2
1
a
1-1
0
k
where we have
of - t to?
- +07)
27

---

Proof. Let E [0, 1] and x
H1 =/o,1 = 1/ Using (Barezy and Kern 2013, Theorem 3.3), we
get that
X
0.1 (1 ()xx + ( + +(((1 ())
with Z (d,,d)) independent from (X., X1). Hence, we get that E [X and
Cov(Xu; x E| [ = (1 (00-t)/222) ()/ + to /233) - (
Similarly, we get that Cov (X,, x 0.1 = / (2 Finally, we get that
Var(X? 0.1) TE[((1 1) +1, (1)-)
=(1 ) + - t)e (2 + +22 (23() + ()((1))
(1 21 212 + +2((1 t)e + (2(((1)))((233
1) = (1+ t(1 - 1)) (2)))
which concludes the proof.
Leveraging Lemma 13, we can give an explicit expression of the drift term in the Markovian projection.
Lemma 14. Let N0,1 = N(O, Σd) with E given by (39). Let ITo Mg Q 0.1- For any t E [0,
and Ft E Rd, we have that
o² o'Eng V log Qij X, xr] I+1(1 (]2(c 112
21 -1)-6="0--0
Hence, the Markovian projection of II,, denoted M. is associated with (X.) with
dx, (1-20(c2-1)--) 11- 1=(1-0[(c-- 1)+02 X, + adB.
Xg No,
(40)
Proof. Using (Barczy and Kern, 2013, Theorem 3.2), we get that Qjo is associated with
dx = V log Qup( X + +od
x no,
where for any / [0, 1). we have and E RR, o² log (2j((21 = (œ1 xx./(1 Therefore, get that MI is associated with (X))cejo.1) such that
dx;- Eng, 2 log Q
(X, (X, X,) adB,,
X₂ Zg.
Therefore, we get that
dX, = (Bng X X - xz) - t)dt + odBr
Xo = Z0.
Using Lemma 13, we have that for any E [0, 1] and x, E Rd, ETg addition, we have for any E [0,1]
X X, = xi] = a //,1/b/,... In
1-t //- - 1 (t + (1 - t)e² - 1 t(1
- 1)) -- ()) 1) + +0)
= -(1- - (e - 1) -1111
1(1 1))(()))(2( - (1)++))/0++(111102200000 - -1) σ2))
(1 ()((2 - 1) - 1000000-))+1111-(-----------1111
(1 () - 21) -1) - ((1
which concludes the proof.
Note that since o > 0 and e² [0, we get that for any E -
0
and therefore the drift is well-defined, smooth and sublinear. In particular, (40) admits a unique strong solution. In what follows, we denote G [0. 1) (0, 1] RR given for any E [0, 1] by
G(t, 1
(1 (1-29)62-1) 11 82x
We have the following useful lemma.
28

---

Lemma 15. L [0, 8 > 0 and p = 2c² + 0.. We distinguish three cases:
(a) Ifp then
((1,6)==07 -14-pp)) 1
1/2 tan ((4- p2)2/2/ /p).
(b) P 2, then
G(L, = -02
(c) P > 2, then
((1,))==000 1// tanh ((p² -2)12/1 4)1/2/p).
This lemma is useful combined with the following proposition, which gives the analytical update formula of Gaussian IMF covariances as a function of c(),c²).
Proposition 16. Ler ng. = N(0, with I given by (39). Then I1,1 = N(O, with
(1 e
)
'' (1/239)
= f(c))
with /: [0, [0, 1] given for any E [0, 1] by
f(e2) = exp|G(1,
Proof. We have that X, exp[G((, Xo + M1. where M, is Gaussian random variable with zero mcan independent from Xg. Therefore, we get that Cov(X,, X1) = exp) In addition, have that E [ E [X]] 1/2,32, which concludes the proof.
Iterating the procedure in Proposition 16, we obtain a sequence of IMF covariances (c ) satisfying 1+1 = f(.)). Finally, we show that this iterative procedure recovers the true SB coupling 11ş5)
N(,, E as a fixed point. The formula of S is given e.g. in (Bunne al., 2023, Equation (2))
which we use below.
Proposition 17, Let IT = N(O, ESB be the true static SB solution, with
SB = (1/28²)
1
-V4+0++-0
CSB
2
Then I is a fixed point of the iterative procedure in Proposition 16, i.e. f((3)) = C.B
Proof. By straightforward calculations, PsB = 20 - /4+ I 0, psB = and thus
G(1, (1.))))) = /2 0. Hence f((5) exp (1,, S5
2 I 0, we have psn >> 2
SB
and G(I, (a)) tanh (02 a4). Hence, (( exp( )3s)) } (V4 -4-
S.. We thus correctly recover the static SB solution 1 fixed point of IMF
We visualize the convergence of this fixed point procedure for a variety of parameter settings in Figure 12. The convergence appears be very fast in only two or three iterations.
Chevepoor
Conergonec
Convegpact
Chrmeppoe
4.1
014
31-
10
65
03
04
45
41
60
62
46
(a) 0.1, c8
(b) 7 1,
(c) 7 1,
(d) 7
5, of
Figure 12: Convergence of IMF in the analytic case given by Proposition 16.
29

---

E Discrete-Time Markovian Projection
We derive in this section discrete-time version of the Markovian projection and show that, some limiting case, we recover the continuous-time time projection. In the discrete case, we let
( x IN) I=== 9++110.,.. TR.
= (ro, .N) I 9k+ (2 (T,
We consider Markovian measure p given by p(ro.N) p(zo) In Pk+1 ( Now let us compute KL(T/p). We have
2 J KL 2xw
+ (ooopl)) +
In what follows, we denote
Lo (Loo)) LN =
L = S KL(g++1JA,N (I EN jdz.NNN
We have the following proposition. Proposition 18. The minimizer Ps+1|k of Lx+1 1 is given by
Pk+ (T&+1 [xx) = JRa 9k+ (I Tki EN (TN (x)ddx
(41)
!f po No, then for any & E {0. (0.....N 11, Ps = Tu- In addition, assume that PI+1 (x (1xx) =
expl - |x T2- (2x) (27)) and 14+1jk,N (a 1 = expl-Il
T& (x))
d/2, Finally, assume that || axll & 222 Then, have that
f(xx) = Sgs J (1V, NN) w (x + 1(72
(42)
Proof. The proof' of (41) and (42) are straightforward and left to the reader. We now prove that Po mg. then for any k (1..... N). pk Th. First, have that for any l: ,O,...,NN 1}.
< Tk41, T (g) llx ZN).
Assume now that Pa = Tk. then we have
(244(1 = SRe Plo (Tx JP ll (2 lzx)d
SRexxe Ph (T 19 up, (FR+ Ima 2ND Nij wa dry fox drRdNN SSs # (Ir Th+1, aN dayday = (x41(14+1).
Su Jue xud Th (a 11 (1844 (mn aN NT )
which concludes the proof.
In particular, in the previous proposition, if f(z,, V log i.e. we have a discretiza- tion of the bridge then f(xx) = Jg- (x) dar, which recovers the Markovian projection in continuous-time.
F Comparing DSBM-IPF and DSB
We analyze further the differences between DSBM-IPF proposed here and DSB proposed in De Bortoli al. (2021) and related algorithms in Vargas et al. (2021); Chen et al. (2022). All algorithms solve the SB problem using the IPF iterates. However, DSB-type algorithms solve for the IPF iterates using time-reversals, whereas DSBM solves for the iterates using Markovian and reciprocal projections. A comparison between these two methodologies made Section 4.
We investigate here further benefit (iii) of DSBM in Section 4, i.e. the benefit of explicitly projecting onto the reciprocal class Q Intuitively speaking, we directly incorporate the reference measure ? in the training procedure inductive bias. More formally, suppose we have at the current
30

---

IPF iteration 2 (Markov diffusion in the forward direction) and want to learn M (Markov diffusion in the backward direction). Due to training error and the forgetting issue (Fernandes et al., 2021), however, M no longer has the correct bridge Q Q Now suppose we first perform IMF for M and learn M2 in the forward direction. That is say, repeat alternative reciprocal and Markovian projections and obtain a sequence (IVI"..")) mEN in the forward direction converging to M Then M now has the correct bridge 1 by Proposition 7, since M is the SB between Mi and M Theoretically, M = Mil for = 0, T' but due to training error accumulating it may be that M o M However, M8 = M6 since M in the forward direction and starting from samples from wy- As a result, we can obtain a Markov forward diffusion M²",,, which has Mg = To and the correct bridge Qjo.-- These are the same set of properties that the reference measure Q has. As a result, replacing Q with M in (6) results in the same SB solution.
Consequently, now continuing the IPF iterations from Mi m it is as if we restart IPF afresh using modified SB problem
s argmin (g : Po = No: - PT.t.
If M,2.. is closer than Q MI in the sense of KL divergence, then we obtain better initialization of the IPF procedure. As proposed in Algorithm 1, DSBM performs the Markovian and reciprocal projection only once before switching between the forward and backward directions. However, it is still beneficial compared to DSB with less bias accumulation in the bridge.
Algorithmically, one main difference between DSB- type algorithms and DSBM due the above distinction occurs in the trajectory caching step. In DSB, fixed discretization of SDE needs to be chosen, and all intermediate samples from the discretized Euler- Maruyama simulation of the SDE need to be saved. Furthermore, a second set of drift evaluation needs to be performed for all datapoints in the trajectory (De Bortoli al., 2021, Equations (12), (13)). The IPML algorithm in Vargas et al. (2021) is also similar to DSB, but Gaussian processes are used to fit the drifts of forward and backward SDEs instead of neural networks. In Chen al. (2022), the implicit score matching loss is used instead, but all intermediate points in the SDE trajectory also need to be saved. On the contrary, DSBM does not require intermediate samples during trajectory caching and only retains the joint samples at times 0, T. Then, the intermediate trajectories are reconstructed using the reference bridge Qjo.r.
G Joint Training of Forward and Backward Processes
We recall below the DSBM algorithm given in Algorithm 1.
Algorithm 2 Diffusion Schrödinger Bridge Matching
1: Input: Joint distribution n7,r, tractable bridge Qo,T number of outer iterations N N.
2: Let IIo = Q
3: for n E (O...., N - 1} do
4: Learn ug- using (14) with 11 112".
5: Let M be given by (13).
6:
Let 1 = M o
7:
Learn va. using (10) with II 112
8: Let M be given by (9).
9: Let 1I = M2 or
10: end for
11: Output: Ug-, Ug.
Our main observation comes from Proposition 9. In particular, under mild assumptions, we have that the Markovian projection MI proj (II) is associated with
dx,- (X. + o 0 V log Qr (XT x.]]dt ad X. no:
dY, = (Y,) + oz Eugr. V log 10 (Y, Yr) Y,]]dt + dr Yo ~ WT-
31

---

Considering the following losses,
8" argmin, S.6' Eu,, r |1 log ( X
()
(43)
6- = a En,, Ill | log 2 X Xo) volt, (11/ood de d).
(44)
If the families of functions (va 0 E } and {us : 0 { are rich enough, we have for any I. E [0, 7) and xy E RR, Ug- (1, = nf V log ( X, = ma and ve (x) = B. V log (X, X, = zz). In practice, this means that the Markovian projection can be computed in a forward or backward fashion equivalently.
Therefore, given coupling II = 1122,, we can update both vg and ue. This means that we train the forward and backward model jointly. We then consider M associated with (13) and M2
associated with (9). Note that if the families of functions {wo 0 0 [ and (vo 0 E Φ) are rich enough then M 2 M!?+1
Mixture from forward and backward. Once we have obtained both the forward update and the backward update, our next task is define the new mixture of bridge II In Algorithm 1, since train only the backward model M Mr define 1 M ij,, In the case of joint training, we have access to M and M2 One way to define a new mixture of bridge to compute II (M Qjo.T + Qio. This choice ensures that in the case where
M = Mg we have
11 -M M Qjo.r M Qjo.T-
It also ensures that all the steps in the joint DSBM training algorithms symmetric. We leave the study of an optimal combination of 2 and M for future work.
Consistency loss. In addition to the losses (43) and (44), we also consider an additional consistency loss. A similar idea was explored in Song (2022). DSB (De Bortoli al., 2021; Chen al., 2022) and DSBM, see Algorithm 1, the processes parameterized by (forward) and vs backward are identical only at equilibrium. Thus imposing the forward and the backward processes match each step of DSB or DSBM would lead to some bias. However, this is not the case in the joint training setting. Indeed, in that case, we have M M if the families are rich enough. Therefore, we get that
dY, = (-fr- + (T - Y.)) + or-,dB, the time reversal of
Yo TT,
(45)
dx, = (ff( + volt, dt) od
Xo
(46)
Computing the time-reversal of (45), have
dX,- = {f((X)) - u (t, X,) of v log (I,""x))))dl dt o. Identifying (47) and (46), we get that for any E [0, 7" and 2; Rdd
Xy ~ no-
(47)
20 - ro (4, mr) of log 11," (x)).
(48)
We highlight that letting az - 0 for any E [0, Tj. we get that Ug = - Uns which confirms that the time-reversal of an ODE is simply given by flipping the sign of the velocity. Therefore, propose the following loss which links the parameters and Φ
Lenns (0, (b)
1o" Engn [l|vo( (1, X.) + (1, X.) - Vf (x.)
Leveraging tools from implicit score matching (Hyvärinen, 2005) and the divergence theorem, we get that
Lcons (0, Eng (t, X,) + (t, X,)l /o? + 2div X.) no(1, X.)) + C,
where C 0 a constant which does depend on 0 and d. Alternatively, (48) shows that
v log (x+) = Engn [V log Qric X = xx] + Enga (V log XX
X, = xz].
32

---

We thus also have the following denoising score matching consistency loss Lcous (8, Φ) So Eng. I|le tust X,) - log ( (x) log (Xx //o/dd
The advantage of this DSM loss is that it does not rely any divergence computation.
Below, we recall the two losses used estimate the Markovian projection (43) and (44)
()) = S Rn, r V log ( X,) up (t, x, /o, ( fo" o log Q:jo [X, Xx) ut X, /offd..
The complete loss we consider in the joint training of the algorithm is of the form
(x の) (()) (()) + ACoons Φ),
(49)
where > 0 is an additional regularization parameter. We now state a version of DSBM which performs joint training in Algorithm
Algorithm 3 Diffusion Schrödinger Bridge Matching (Joint Training)
1: Input: Coupling H6 tractable bridge Q N
2: Let IT = -m0 1
3: for E {0, 1} do
4: Learn vg-, lig- using (49) with II".
5: Let Mr be given by (9).
Let Mg++1 be given by (13).
6:
7: Let M (M + (Mg+)
8: Let I1+ = My QI.
9: end for
10: Output: Ug-, Ug-
H Loss Scaling
Similar to the loss weighting in standard diffusion models (Song et al., 2021b; Ho et al., 2020), we derive similar weighting reduce the variance our objective. We focus on the forward direction of Markovian projection in this case, and the backward case can be derived similarly. Our forward loss the DSBM framework is given by (10), where the inner expectation is given by
En, r Illo? log QTT (X, Xx)) - (t, X,)
Letting Q be Brownian bridge with diffusion parameter and assuming T 1, this becomes
-xx. X - Xxo - aaotoo//11110 t)Z vo(t, x! 0.7 )1?1
with X :X1 +
t)x, (1- / Z)Z. When 1, we see that the regression target dominated by the noise term o /t/(1 ttz which needs to be predicted based on information contained within x. The loss will have an approximate scale of o²t/(1
when t Rs 1 which will very large. To avoid these large values affecting gradient descent, we can downweight the loss by + 22t t) (we can add to effectively cause no loss scaling when is close to 0)
(1+ o 22/(1 t)) -1 " Exx..x. )~ (o.ia) [||
Xo - //ttt//11111111////11 t)Z vo X, 0,7 >)
Similar arguments can be applied to the backward loss (14)
Eu.olloo log Q X,) (t, (t, 11)
= Exx X, (o.La Xn - XI -a - (1-11 x
which we then downweight by + 0² t)/t
(1+03(1 -))ttt/) "E,, |Xo XX - - - (t, )1
33

---

I Experiments
In this section, we present further details of the experiment setups as well further experiment results. In all experiments, we use Brownian motion for the reference measure Q with corresponding Brownian bridge (4) and T We use the Adam optimizer with learning rate 10 and SiLU activations unless specified otherwise. The experiments are run on computing clusters with a mixture of both CPU and GPU resources.
I.1 2D Experiments
For the 2D experiments, we closely follow Tong et al. (2023) and the released code?, and use the same synthetic datasets and the 2-Wasserstein distance between the test set and samples simulated using probability flow ODE as the evaluation metric. However, we use 10000 samples in the test set since we find the 2-Wasserstein distance can vary greatly with only 1000 samples (which can be as high as 0.3 even between two set of samples both drawn from the ground truth distribution). We use simple MLP with hidden layers and 256 hidden units to parameterize the forward and backward drift networks. We use batch size 128 and 20 diffusion steps with uniform schedule at sampling time. Each outer iteration is trained for 10000 steps and we train for 20 outer iterations. As the initial coupling is more optimal in ASBM-IMF44 we reduce the number of outer iterations to 4 and train for 50000 steps in each outer iteration. For flow methods, we train for 200000 steps total. For Table 2 we use o; = 1 in all cases, except the moons-8gaussian dataset where we use o, 5. Note that FM cannot used for the moons-8gaussian task since it requires a Gaussian source, but CFM is applicable. The experiments are using 1 CPU and take approximately 200 minutes (for both training and lesting).
L.1.1 Variance of the reference measure Q
We comment further on the effect of oe in the reference path measure Q. We assume time- homogeneous oe = a for simplicity. In Figure 2, we vary a and visualize the learned transport for a 3gaussians problem of transporting between two Gaussian mixtures. In Table 4 we show the 2- 2-Wassrsttein distance between the test set and generated samples for this 3gaussians problem as well as the moons-Sgaussiins problem. We find that large values of result in increasingly curved transport paths, and correspondingly reduced performance when a is excessively large. Conversely, also find reduced performance when is excessively small. We conjecture this is due increased optimization difficulty and bias accumulation. Firstly, the EOT problem becomes more difficult to solve as is taken to O, which would require higher number of outer iterations. Further, the introduction of noise also decreases optimization difficulty by smoothing the intermediate marginals between the two terminal distributions. The benefit of setting was also observed in Albergo et al. (2023): Delbracio and Milanfar (2023). Finally, we conjecture setting a >> could also increase the diversity of sampled couplings and may alleviate some bias accumulation issues in the outer iterations. When O, these issues result in the artifacts observed in the transferred samples (marked using yellow points) in Figure 2. The appropriate value for o depends on the spatial scaling of the problem as shown in Table 4, where the optimum or is larger for the larger scale moons-88 moons- 8gaussiaas problem.
and using a stochastic sampler
3ganssians
moons- Soons-8gmussians
2- 2-Wasserztein
2- Wusssersteie
0.64620.028
14499000008
07
0.1
0.724x0.039
a
1.28540.346
0 = 0.3
0.546a0.169 169
Φ 2.0
0.91629.292
0 1.0
0.439:0.072
8 4.0
0.81840.249
= 3.0
0,34320.078
8.0
0.98920.179
Table 4: 2- 2-Wasserstein distance for varying value of a used in the DSBM-IMF method.
(code released under MIT license)
34

---

1.2 Gaussian Experiment
Similar to the 2D experiments, we use a simple MLP with 2 hidden layers and 256 hidden units to parameterize the forward and backward drift networks. This is a smaller network compared to the "large" network in De Bortoli et al. (2021). We use batch size 128 and 20 diffusion sampling steps with uniform time schedule at inference time. Each outer iteration is trained for 10000 steps and we train for 20 outer iterations. The experiments are also run using CPU, and for the case d = 50 finish in approximately 200 minutes. We note that DSB, IMF-h, and DSBM methods all solve for the Schrödinger Bridge transport, whereas Rectified Flow does not and so is not plotted on the covariance convergence plot in Figure 3 since it is not comparable.
For Table 3, we assume the marginals of the learned process Pr are also independently Gaussian distributed in each dimension. We thus estimate the KL divergence from the sample mean and variance of each dimension of P& using the analytic KL formula between Gaussian distributions.
1.3 MNIST Transfer Experiment
We follow De Bortoli et al. (2021) closely for the setup of this experiment. We use the set of first 5 letters of EMNIST (A-E, a-e) such that both domains have 10 classes in total. We use the same U-Net architecture, batch size 128 and 30 diffusion sampling steps. The network size is approximately 6.6 million parameters. Each outer iteration is trained for 5000 steps. We refresh the cache dataloader every 1000 steps with 10000 new samples. Contrary to De Bortoli et al. (2021), in our experiments we find that we obtain better sampling quality for both DSB and DSBM using a uniform noising schedule. We simply choose o, = 1 for all t and T = 1 in our experiments.
For DSBM-IPF, we train for most 50 outer iterations (i.e. 250000 total number of steps). For DSBM-IMF and Rectified Flow, since their first iteration corresponds to Bridge Matching and Flow Matching respectively, we first pretrain the forward and backward networks for 100000 steps using the Bridge Matching or Flow Matching losses. DSBM-IMF then switches to iterative training with 5000 steps per outer iteration, the same as DSBM-IPE. For Rectified Flow, we train for 50000 steps per outer iteration. The experiments are performed using 2 GPUs and take approximately one day.
We provide further experiment results in Figures 13, 14 and 15. In Figure 13, we show samples generated using different algorithms and at different points of convergence. Samples generated using CFM and Rectified Flow with 2 rectified steps in (a) and (b) appear be less clear and identifiable. OT-CFM improves upon CFM slightly in (c), but many samples still appear to be unclear. For DSB, the algorithm has not converged after 10 iterations, and many samples in (d) still appear be letters. After 20 iterations, there are still letter-like like samples in Figure 4b, and the digit classes also appear to be unbalanced with many instances of digit '0', After 30 iterations, however, the sample quality of DSB becomes very poor in (e). On the other hand, as shown in (f-(i), we observe that DSBM-IPF converges faster than DSB, with more accurate samples even in iterations 10 and 20, and the sample quality continues to improve until the end of training after 50 iterations.
We present some additional trajectory samples at the end of training in Figures 14 and 15 in both the forward and backward directions. We observe DSBM is able to transfer samples across the two domains faithfully, and the output samples preserve interesting similarities compared to the input, whereas Bridge Matching preserves much less similarity. The Mean Squared Distance (MSD) between the initial and final samples also confirm that DSBM methods transfer more closely to the original inputs, and DSBM-IMF further achieves the best FID score out of the three methods.
1.4 CelebA Transfer Experiment
For CclebA (Liu et al., 2015), we test DSBM on resized 64 x 64 and 128 x 128 resolution images. We evaluate the dependency of the results with respect to a > 0 on CelebA 64 x 64. We also showcase the scalibility of our method by training DSBM on CelebA 128 128. In both cases, the dataset is split between male e/old and remale/yougg We gather 20000 samples of each class on which we perform classical data augmentation such as horizontal flipping.
For our ablation study w.r.t. the value of a >> O, we run DSBM for 20 iterations (note that the loss and the quality was still improving after 20 DSBM iterations but by stopping the runs early we were allowed to draw comparisons with more values of a). We use U-Net architecture with 4 resolution levels and 2 residual blocks per resolution level. The batch size is fixed to 64 and the EMA rate is
35

---

5
D
0
3
5 5 8 22222222
O
4
6
J 6 47 0
2
9
1
3 Probasdd 22322302
&
1493
110451
5 159466142 020=CL/25/
N
5 :
/
4
h PRSSSSPAD4
(a) CFM
(b) Rectified Flow @ 3
OT:CMM
(d) DSB @ 10
(C)
DSB
@ 30
o 8 7
6
8 0 8 < 2
8 4
3 3 4 0 1.0444000
0
5 8 6
0 5
44 3
1 0
2
6 6
0 4
2
2 4
2 2
0
2
0
1 2
2 6 3 a 3 4 7
3 30003 00082610 2 27
5
9
2 2 2293338
9
2 3
8 3
58
10 23 6
6 1998770
7
4 0 56775
1710196288
9 5 2
5 900388662 9633
à 6 2 5 0
r04 5093869514
7 5 5 260
7/5580093
4
7 8 A 82023
0 0760585000
017700000 076606623
2 n 8 & C000
3 85 1551888
2 79324 22487230
4
6
053013177
(l) DSBM-IPF @ 10 (g) DSBM-IPF @ 20 (h) DSBM-IPF @ 30 DSBM-IPF @ 40
DSBM-IPF @ 50
Figure 13: Samples of MNIST digits transferred from the EMNIST letters using different methods.
@ indicates the progressed number of outer iterations n +
fixed to 0.999, We refresh the cache dataloader every 500 steps with 10000 new samples. The SDE sampler is chosen to be the modified Euler-Maruyama sampler, see Heng et al. (2021) for instance, with constant schedule for the stepsizes. We use 100 sampling steps at inference time and to refresh the cache. For each outer DSBM iteration we train the model for 20000 iterations.
We provide additional transfer results in resolution 128 128 in Figure 16. We do not change the training setting for this experiment.
1.5 AFHQ Transfer Experiment
For AFHQ (Choi et al., 2020), we test DSBM between classes cat and wild with 512 x 512 resolution images. Each class contains approximately 5000 samples. We first pretrain the networks using Bridge Matching for 100000 steps, then run DSBM for 20 iterations with 25000 steps per outer iteration. We follow Liu et al. (2023b) and use the same U-Net architecture The batch size is 4 and the EMA rate is 0.999, We choose o² = 5 and again we use 100 sampling steps with constant stepsizes.
I.6 CIFAR-10 Generative Modeling Experiment
We also test our method in the standard generative modeling framework on the CIFAR-10 dataset. We again use a U-Net architecture with 4 resolution levels and 2 residual blocks per resolution level. The network size is approximately 39.6 million parameters. The batch size is fixed to 1 128 and the EMA rate is fixed to 0.9999. The Adam optimizer is used for this task. For DSBM-IMF and Rectified Flow, again we first pretrain the networks using the Bridge Matching or Flow Matching losses, then switch to DSBM-IMF or RF training with 100000 steps per outer iteration. We find that 1 or 2 additional outer iterations appear sufficiently effective this task, and additional outer iterations can cause sample quality to drop. For our main experiment, the pretraining stage ran for approximately days, and DSBM-IMF ran for approximately 4 additional days using 4 V100 GPUs.
The best results during training for each method are reported in Table 5 and Figure 17, where we compute the FID score between 50000 samples in the CIFAR- 10 training set and 50000 generated samples following standard practice for this task. We observe that DSBM-IMF can clearly improve upon Bridge Matching the same value of a, which suggests that further outer iterations in DSBM is beneficial for improving sample quality. This is contrary to Rectified Flow, which causes the FID score to worsen compared to Flow Matching after only I rectified iteration. However, increases,
'nttps://github. wtt.g/baaa//tttiiieedeeeeeeeeeee hwshh/oooololloo httpb//bbb///mm////////// Immggeeenrettiionnn /naggg/Geerrrattoooonnig//////// wwwiifiidd_oow/fffff hatthhh__attpptttteee wtt_paatocc___tt_____ee auuuuuutorcc.r__aauussaaaa py (code released under Apache-2.0 2.0 license)
36

---

10874
/
FID
17.14
8
9 99 17357745
MSD
0.579
0 63302231
D
5 2
0
9 5505
(a) Bridge Matching
8
R
0
e
96 A
Y
9 7
6 5 3
FID
15.27
D
8
O
2
8
T
77554
MSD
0.354
5 0
2
6
2
8
6
O
7
2
0
(b) DSBM-IPF
4
8
5
6
8
5
3
222222
8
C
2
0
FID
10.59
2
22
5
0
MSD
0.375
A 8 8
0
M
0
f
4
A
Aa
8
20022000
9
3
(c) DSBM IMF
Figure 14: Left: EMNIST to MNIST sample trajectory with 30 diffusion steps at / = 0, 1/3, 2/3, 1. Right: FID score final samples, and Mean Squared Distance between initial and final samples.
7
0 l
9
0
9
8 2
8 9
eee
2
3
02
02
A
0
Leas
FID
9.402
4
6
5
8
X
C88 8
MSD
0.561
8
7
0
&
8
(a) Bridge Matching
2 4
4 A
5
A 4 7
6 8
A
6 5
0
#
1
4/
2 41 2
A
A 1
13.59
2
2
A
60
6
&
6
C
MSD
0.359
8
3
saec
D
2
1
3
3
F
9
/
(b) DSBM-IPF
0 lo
5 9
8
O
5
C
cabb
0
0
4.
D
e.
2
3
4
C
C C44
FID
8.990
7
13
C
N
U e
6
L:
H
C
R:
C de
MSD
0.375
8
h
34
C
6
8
4
6
C
7
(c) DSBM IMF
Figure 15: Left: MNIST to EMNIST sample trajectory with 30 diffusion steps at - 0,3, 2/3, Right: FID score of final samples, and Mean Squared Distance between initial and final samples.
37

---

X
IX
Transfer results between images given by the tokens f female/young and male/old. Top row: original images (left) and generated images (right). Bottom row: original images (left) and generated images (right).
Figure 16:
we observe the FID score worsens for both Bridge Matching and DSBM as more stochasticity is introduced in the sampler. The best result of DSBM-IMF is obtained using o² = 0.2 and is slightly better than FM (i.e. with x² = 0) using 100 Fuler steps. the other hand, using the dopri5 ODE solver, FM achieves FID of 4.055 with on average 148 integration steps. In Figure 17, we observe that both RF and DSBM-IMF are very effective in improving sampling quality at low number of diffusion steps, i.c low number of function evaluations (NFEs), compared to Bridge and Flow Matching as well as OT-CFM which improves upon CFM slightly. DSBM-IME also achieves lower FID score than RF as the NFEs are taken higher. Additional strategies such as distillation and fast SDE solvers can also be useful for improving few-step sampling quality further.
1.7 Fluid Flows Experiment
We use the fluid flows dataset" from Bischoff and Deck (2023). The dataset consists of unpaired low (64 x 64) and high (512 x 512) resolution fields, as well as context field with local information for the high resolution field. The data fields consist of two channels representing supersaturation and vorticity. The context field is dependent the wavenumber ki = ky E (1, 2, 4, 8, 16), which
hhttts://gitee httpg/iht....//..iaaaa. wtt.cooco///c.ee/e. /Cta///taaaaaaattss jl (code released under MIT license)
38

---

FM
RF
BA CPM
RM CPM
4,931
6.010
OT-EPM
BM
DSBM-IMF
RD
DIRM
0.2
5.427
4.511
0.5
8.274
6.896
12.749
NEE
NEE
Table 5: FID results on the CIFAR-10 train set using 100 Euler(-Maruyama) steps.
Figure 17: FID vs number of diffusion steps (NFE) with NFE between and 1000.
Figure 18: Zoomed-in version of Figure 17 with NFE be- tween 100 and 1000.
(a) Low-res
(b) Diffusion-ft
(c) Bridge Matching
(d) DSBM-IPF
(e) DSBM-
Figure 19: (a) Source low resolution sample; (b)(c) intermediate state and final reconstruction of each algorithm, for wavenumber lig
ky 2.
specifies the frequency of the saturation specific humidity modulation. We follow Bischoff and Deck (2023) for data processing and network architecture, which is given by U-Net with additional spatial mean bypass network given by an MLP. The network size is approximately 1 11.3 million parameters in total. We train using batch size 4, learning rate 10 for 5000 steps per outer iteration and N = 12 outer iterations. We refresh the cache dataloader every 2500 steps with 1250 new samples. The training takes approximately 20 hours
single RTX GPU. We used o² 0.3 and sample using 30 diffusion steps without finetuning these parameters. For the Diffusion-fb method in Bischoff and Deck (2023), we use the released code without modifying any parameters.
We visualize intermediate and final reconstruction samples for different algorithms in Figure 19. We that DSBM-IPF and DSBM-IMF provide consistent samples with the low resolution source, whereas Diffusion-fb and Bridge Matching produce dissimilar samples. We also follow Bischoff and Deck (2023) for comparable performance as Diffusion-fh in terms of these statistical profiles, and can be comparatively more accurate e.g. in the tails of the distributions in Figure 20, and for the case kg ky 4 in Figure 21 for which the power spectrum of supersaturation is correctly captured by DSBM-IMF but by other methods. Comparing this analysis with Figure 11, DSBM-IMF is also significantly more accurate in terms of conditional consistency than Diffusion-fh fb. On the other hand, DSBM- appears less accurate in terms of these unconditional statistics than Diffusion-ft and DSBM-IMF, but achieves lower & distances from the input sources in Figure 11. This suggests that DSBM-IPF and DSBM-IMF exhibit different empirical biases before convergence, and DSBM-IME more preferable when the accuracy of the samples are important. This is in line with IMF theory as the marginals no. Ty are preserved IMF but not in IPF.
refined statistical analysis in Figures 20, 21, 22. DSBM-IME achieves
J Broader Impact
Our work focuses on theoretical and methodological research and is intended to bring closer the fields of generative modeling and optimal transport. It can useful for learning transport maps between general distributions with high accuracy and high scalability, which can have useful applications
10 "nttps://githuu wttpttttuw..oo//////aaaaaaa wwtffu/iiiii.fioonnoooooo bridge- deadgg--donssalling ing (code released under Apache-2.0 2.0 li-
cense)
39

---

167
167
The Dappoa
10'4
10:1
DIDAGET
Dansts
Sopersaturation
2
ko m1
Ry
Au
Tour
Isarus Bap
10--
39
10
10777
10
10"11
10'11
-30
10
33
-20
-30
10
-37
-10
-20
-30
Vortictty
Figure 20: KDE estimates of values in supersaturation and vorticity fields.
167
16"
Denatty
10-12
10
101
10
107
107
'2'
107
2, Supersatoration
R.
6. = =
16
10'1
20
10:1
Daniy
The jon-rs
10
Vorikity
Figure 21: Spectral density estimates of supersaturation and vorticity fields.
1,=2,=d
1.0
1.0
0.5
Dillaae
0.5
02
0.1
Supersaturation
Figure 22: KDE estimates of spatial means of the supersaturation field. The shaded areas denote 99% confidence interval obtained using 10000 bootstrap samples.
40

---

in machine learning, but also natural science areas such as physics, biology and geosciences in which optimal transport maps with theoretical guarantees are appealing. Our fluid flows experiment demonstrates such potentials. However, as is the case for generative models as whole, intentional malicious use could cause detrimental societal impacts.
41


Reviews:
[{"review_id": "Ly5QaBT1Oa", "reviewer": "Reviewer_Axbu", "summary": "This paper proposes IMF, an iterative method for solving Schrodinger Bridge problems where the solution at each iteration preserves the correct marginal distributions at times 0 and T. This is in contrast to the existing method IPF which only satisfies this in the limit (and is very difficult to reach in practice). ", "questions": "  - \"Most notably, we adopt the SDE approach which is crucial for the validity of Proposition 5\" - Can the authors expand on what breaks when there is no stochasticity? From what I understand, RF should also preserve marginal distributions and converge to the optimal transport solution, which is discussed more in [1]. The solutions are also unique similar to any Flow Matching-based approach.\n\n\n  - All of the experiments in the main paper are currently only for transferring between two data distributions, but being able to find optimal transport maps is useful for generative modeling as well. I think it could really make the paper stronger by having experiments in generative modeling (Gaussian -> data) and showing some empirical metrics from e.g. [2] such as sampling cost and consistency, since these are important problems in generative modeling at the moment. [Edit: I found the experiments in I.5. I see that sampling efficiency is on par with RF but Table 6 seems pretty good! Although it does look like if we extrapolate to NFE > 100, the other methods are descending quite a bit faster in FID...]\n\n\n  - Can the authors expand on the differences between the two proposed methods, DSBM-IPF and DSBM-IMF? If I understood correctly, the main difference between DSBM-IPF and DSBM-IMF is only in the first iteration, where the pairs (x0, xT) are sampled from (either an OU process Q or the independent joint distribution). And the real difference is that Q has a slight bias at time T because it may not have reached its stationary distribution. Practically, this difference seems very negligible, yet there seems to be differences in empirical results (e.g. Table 3 and Figure 5). Theoretically, since Q does not match the correct marginal at T, this is also problematic for the proof of convergence to SB right?\n\n\n[1] \"Rectified Flow: A Marginal Preserving Approach to Optimal Transport\" https://arxiv.org/abs/2209.14577\n\n\n[2] \"Multisample Flow Matching: Straightening Flows with Minibatch Couplings\" https://arxiv.org/abs/2304.14772", "limitations": "The authors do mention limitations such as marginal improvement, computational issues, and difficulties when sigma --> 0. I think these are sufficient.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "  - An interesting marginal-preserving algorithm for solving Schrodinger bridge problems. The popular method IPF can have convergence problems in practice, while this seems to fix some of that.", "weaknesses": "  - Presentation. While I eventually understood the resulting IMF algorithm, I felt the presentation can be much more simplified. The algorithm is actually quite simple and closely related to priors works on SB and OT. While I understand the need to prove convergence and showcase novelty, the heavily front-loaded mathematical notation made reading on the first pass very difficult. I imagine the current paper may be off-putting for many readers who are not extremely familiar with related works. \n\n\n  - Novelty. The proposed algorithm seems very similar to Rectified Flow but just with an additional noise parameter.\n\n\n  - Limitation. The proposed algorithm only works for the standard formulation of Schrodinger bridge with quadratic costs. However, in this setting, it isn't clear to me why the extra stochasticity is useful, as opposed to solving deterministic maps (optimal transport problems).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thorough evaluation of our work. We appreciate their interest and their thoughtful questions. We would like to address the raised questions here which hopefully would clarify the role of stochasticity further.  \n\n**“Algorithm presentation can be simplified.”**:\nWe have tried to make the terminology clearer with the\nintroduction of a Notation section. The introduction of some mathematical\nconcepts like the Markovian and reciprocal projections seems\nimportant to us, as it allows us to describe within a same framework many\ndifferent existing algorithms ([1,2,3,4,5] can be seen as special case of the\nMarkovian projection procedure for instance). However, we understand the\nreviewer's concerns and will try to improve the clarity of the work in the revised version of the paper. \nIf there are specific presentation issues the reviewer would like to point out, \nwe would be happy to change them as well.\n\n**“Role of stochasticity”**: We thank the reviewer for this question. Indeed, it can be unclear at first why stochasticity is important. We would like to provide three motivations regarding the use of stochasticity:\n\n* Firstly, from a theoretical point of view, stochasticity is important for\nestablishing the convergence of the IMF algorithm to the SB solution.\nThe IMF methodology provably converges to the unique EOT/SB solution, see [8, Theorem 2] (concurrent work) for a proof.\nWe have also found a short proof of this which we will include in the revised version of the paper.\nIn contrast, such results appear to be more difficult to establish for\nRectified Flow (RF).\nThe RF algorithm can only provably solve the OT problem when the learned vector fields $v_t(x_t)$ are restricted to gradient\nfields, and counterexamples exist if this condition does not hold [7].\nEven with the above restriction on the vector fields, the\nstrongest convergence results for RF as of now appear to be Theorem 5.6 and\nCorollary 5.7 in [7], which rely on a surrogate loss and does not\ndirectly prove the convergence of the learned coupling to the OT map.\n\n* Empirically, we also encountered difficulties with the original RF formulation\nfor transfer tasks, and we observe adding stochasticity can improve\nperformance significantly. Qualitatively speaking, in the case of image\ntransfer, when $\\sigma$ is close to zero (i.e. our approach is closer to RF)\nthe learned process only changes minimally the input samples.\nAs a result, the output samples might look unrealistic, especially if the two domains are very different (for example interpolating\nbetween the classes \"horse'' and \"plane'' in CIFAR-10, see the additional\nresults in the one-page PDF).\n\n* Finally, we would like to make a comparison between non-regularized OT and\nthe SB problem. In itself, the Schrödinger Bridge problem is an important\ntheoretical problem, with many recent works tackling it. Unlike\nnon-regularized OT, the SB formulation can be useful when the solution does\nnot necessarily have a one-to-one mapping structure. For example, in inverse\nproblems (such as super-resolution), there are multiple solutions\ncorresponding to a given condition, and so the solution mapping is one-to-many\nand the SB formulation is a more natural framework.\n\nWe also would like to emphasize that another crucial difference with RF is\nthat DSBM-IMF applies forward-backward training, whereas RF only trains the flow\nmatching process iteratively in one direction. The latter will cause errors in\nthe marginal distributions to accumulate, and marginal accuracy to deteriorate\nin each iteration, whereas DSBM-IMF can still improve marginal accuracy (as\nshown in e.g. Figure 5).  Therefore, we believe our proposed method is a novel\nand relevant contribution, which also helps to\nunify the flow matching and SB research areas.\n\nWe do hope that our clarification of the main differences between our algorithm\nand RF [6], as well as our justification of the use of stochasticity, have resolved the concerns raised in the\nreview.\n\n**“Differences between DSBM-IPF and DSBM-IMF”**: Your understanding is indeed\ncorrect. The DSBM-IPF and DSBM-IMF algorithms only differ in the first\niteration, and when the OU process is used their differences should be\ncomparatively small. We use Brownian process in our work motivated by optimal\ntransport, and their difference is more significant. However, both\nalgorithms converge theoretically to the true SB solution. In the case of DSBM-IPF,\nalthough the process $\\mathbb{Q}$ does not match with $\\pi_T$ at time\n$T$, by Proposition 9 DSBM-IPF can recover the true IPF sequence\n$(\\tilde{\\mathbb{P}}^n)$ (see lines 118-125), for which a different, classical\nproof of convergence of IPF exists. To summarize, both algorithms converge to\nthe true SB (assuming that the score is perfectly learned),\nhowever the theoretical frameworks are very\ndifferent. We will highlight this in the revised version of the paper.\n\nOverall, we would like to thank the reviewer for their insightful questions\nwhich have helped us to clarify several points of the paper, in particular the\nimportance of stochasticity.\n\n[1] Delbracio, Milanfar (2023) -- Inversion by Direct Iteration: An Alternative to Denoising Diffusion for Image Restoration\n\n[2] Liu, Vahdat, Huang, Theodorou, Nie, Anandkumar (2023) -- I2SB: Image-to-Image Schrödinger Bridge\n\n[3] Lipman, Chen, Ben-Hamu, Nickel, Le (2022) -- Flow Matching for Generative Modeling\n\n[4] Heitz, Belcour, Chambon (2023) -- Iterative $\\alpha$-(de)Blending: Learning a Deterministic Mapping Between Arbitrary Densities \n\n[5] Albergo, Boffi, Vanden-Eijnden (2023) -- Stochastic Interpolants: A Unifying Framework for Flows and Diffusions\n\n[6] Liu, Gong, Liu (2022) -- Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow\n\n[7] Liu (2022) -- Rectified Flow: A Marginal Preserving Approach to Optimal Transport\n\n[8] Peluchetti (2023) -- Diffusion Bridge Mixture Transports, Schrödinger Bridge Problems and Generative Modeling"}, {"review_id": "EHPyGKPSdc", "reviewer": "Reviewer_J5AJ", "summary": "The paper introduces Iterative Markovian Fitting (IMF) as a new method to compute Schrödinger Bridges (SBs), which are dynamic versions of entropy-regularized optimal transport. IMF alternates between projecting on the space of Markov processes and the reciprocal class, and it preserves the initial and terminal distributions. The paper also proposes Diffusion Schrödinger Bridge Matching (DSBM) as an algorithm to approximate SB solutions derived from IMF. DSBM overcomes issues of previous techniques and solves a simple regression problem at each iteration. The performance of DSBM is demonstrated in various transport tasks. The paper provides theoretical results, notations, and definitions related to the topics discussed.", "questions": "1, How does the Iterative Markovian Fitting (IMF) method differ from the traditional Iterative Proportional Fitting (IPF) approach in computing Schrödinger Bridges (SBs)?\n2, Can you provide more insights into the theoretical results established for IMF in the paper?\n3, How does Diffusion Schrödinger Bridge Matching (DSBM) overcome the time-discretization and \"forgetting\" issues of previous diffusion-based techniques?", "limitations": "As stated in \"Weakness\"", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1, The paper introduces Iterative Markovian Fitting (IMF) as a new procedure for computing Schrödinger Bridges (SBs). This approach provides a fresh perspective and contributes to the field of optimal transport and generative modeling.\n2, The paper establishes various theoretical results for IMF, demonstrating the validity and effectiveness of the proposed method. These results enhance the understanding of SBs and their applications in machine learning.", "weaknesses": "1， While the paper mentions that IMF and DSBM address scalability issues, it lacks a detailed analysis of their scalability in terms of computational resources and dataset sizes. Understanding the scalability limitations of the proposed methods is crucial for their applicability to real-world, large-scale problems.\n2, The paper assumes locally Lipschitz continuity of drift and limited settings for the proposed methods. These assumptions may limit the applicability of IMF and DSBM to certain scenarios. A discussion of the limitations imposed by these assumptions would provide a more comprehensive view of the proposed methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for taking the time to review our submission and for their constructive feedback. We address here the main points from the review.\n\n**“Analysis of scalability”**: Our proposed DSBM method leverages tools from the recent flow/bridge matching literature, which are highly scalable generative models.\nIn particular, the first stage of DSBM is the same as bridge matching, and so DSBM can be viewed as refinements of bridge matching in later stages.\nWe have demonstrated our proposed method in a number of high-dimensional tasks, such as CelebA ($128 \\times 128$) and fluid flows downscaling ($512\\times512)$. We have additionally performed an experiment on the AFHQ dataset at $512 \\times 512$ resolution using DSBM-IMF, and have included it in the one-page response PDF.\nOur experimental setup is fully described in the appendix of our paper (see Appendix I). \nRegarding the computational resources, our MNIST experiments were run on two 2080Ti GPUs, \nwhile the CelebA and fluid flows downscaling experiments were run on a single A100 or RTX GPU. \nIn the case of image experiments, we use the standard dataloaders provided by the torchvision package. \nFor the fluid flow experiment, further details on the climate dataset and preprocessing can be found in [1]. \nWe hope this clarifies the reviewer's concerns and are happy to provide further clarifications if needed.\n\n**“Limitations of locally Lipschitz continuity assumptions”**: We thank the reviewer for pointing out this\nlimitation. First, we would like to highlight that in most applications the\nquantities of interest are locally Lipschitz (which is one of the weakest\nrequirement regarding the regularity of the drift and diffusion\nmatrix). Dropping this requirement might lead to technical difficulties (for\ninstance it is known that the solutions of SDEs with locally Lipschitz\ncoefficients have a unique strong solution, which is not the case if the drift\nis only assumed to be continuous [2]). However, it is likely that the local Lipschitzness assumption could be dropped and\nreplaced with finite entropy conditions. Currently, we\nmake the Lipschitzness assumption to use known results regarding the\nwell-posedness of the Doob $h$-transform [3]. However, there is another line of work that studies Doob $h$-transform\nunder finite entropy conditions [4]. These conditions might be more amenable\nto the study of Schrödinger Bridges. We plan to study further the theoretical properties of the IMF under this finite entropy setting in a future work.\n\n**“How does IMF differ from IPF”**: IPF and IMF schemes project on different classes of path measures (see Table 1). We refer to the “Relationship between IPF and IMF” section in our response to all reviewers for a more detailed comparison.\n\n**“More insights into the theoretical results established\nfor IMF in the paper”**: In what follows, we give a high level explanation of\nthe main results of the paper. Our main contribution is to show that the two\nalternate projections considered in the IMF are 1) indeed projections\nunder the Kullback--Leibler divergence (Proposition 2 and 4);\n2) that they satisfy a Pythagorean theorem (Lemma\n6). The fact that these two projections satisfy a Pythagorean theorem is\nsurprising and one of the important results of our work. Our results\ndiffer from classical information theoretic results [5] in that our\nPythagorean theorems are stated for the backward KL divergence\nand not the forward one. This is a key difference which greatly simplifies the\nanalysis. As a result, the concurrent work [6] was able to prove the\nconvergence of IMF to the Schrödinger Bridge. We would like to emphasize\nthat since the submission of this paper we have found a proof of the\nconvergence of IMF to the SB which is shorter while similar in spirit to the one of [6]. \nWe summarize this proof below and\nand will include it in the revised version of the paper.\n\nUsing the Pythagorean theorems, we can show that the KL\ndivergence $\\mathrm{KL}(\\mathbb{P}^n|\\mathbb{P}^\\star)$ between the IMF iterates and the SB is\nfinite. Using the coercivity of $\\mathrm{KL}(\\mathbb{P}^n|\\mathbb{P}^\\star)$\n(this is where our analysis differs from the one of the IPF) we have that\nthe IMF sequence is relatively compact. Then, we use the space of Markov\npath measure and the reciprocal class are closed under weak convergence to\nshow that the only limiting point of IMF must be Markov and in\nthe reciprocal class. Since they met the marginal constraint at endpoints, the unique limit point is the Schrödinger Bridge (Proposition 5). \n\n**“How does DSBM overcome the time-discretization and 'forgetting' issues”**: DSBM straightforwardly overcomes the time-discretization issue, as both the diffusion bridge $\\mathbb{Q}_ {|0,T}$ (e.g. equation (4)) and the training objectives (e.g. equations (9)(13)) are derived in continuous time. As for the \"forgetting\" issue, previous diffusion SB methods only use the reference process $\\mathbb{Q}$ in the first iteration of the algorithm which leads to the \"forgetting\" issue. On the contrary, DSBM directly uses the diffusion bridge $\\mathbb{Q}_{|0,T}$ in each iteration. This can be intuitively understood as incorporating $\\mathbb{Q}$ as inductive bias in the DSBM algorithm. See also Appendix F for a more detailed discussion on how DSBM overcomes the \"forgetting\" issue. \n\nWe hope that our answers have resolved the reviewer's concerns. \n\n[1] Bischoff, Deck (2023) -- Unpaired Downscaling of Fluid Flows with Diffusion Bridges  \n\n[2] Ikeda, Watanabe (1996) -- Itô’s Stochastic Calculus and Probability Theory\n\n[3] Palmowski, Rolski (2002) -- A technique for exponential change of measure for Markov processes\n\n[4] Leonard (2011) -- Stochastic derivatives and generalized h-transforms of Markov processes  \n\n[5] Csiszar (1975) -- I-Divergence Geometry of Probability Distributions and Minimization Problems\n\n[6] Peluchetti (2023) -- Diffusion Bridge Mixture Transports, Schrödinger Bridge Problems and Generative Modeling"}, {"review_id": "ISIvFepOWl", "reviewer": "Reviewer_24jx", "summary": "The submission suggests a numerically effective approach for solving the Schroedinger Bridge (SB) problem and illustrates its potential for generative modeling. The approach generalizes in some sense some recent flow matching methods. Furthermore, it provides a more efficient (no full trajectory caching) algorithm with less error accumulations that better remembers the prior reference measure for solving SB problems compared to previous works. Empirical results indicate for instance that it can yield better generative performance for MNIST-EMNIST transfer compared to previous work, or better reconstructions for downscaling geophysical fluid dynamics.", "questions": "In Figure 16, it appears that FID can be improved by using more than 100 steps. Does DBSM obtain better FID values compared to Bridge matching for more than 100 steps?\n\nIt was not clear to me what to take away from some parts of the appendix, such as as E Discrete-Time Markovian Projection or the Brownian/OU bridge representations in B.2. Are the results in B.2 helpful for saying anything about how to specify the SDE in the forward process $X$ for the suggested approach?\n\n", "limitations": "Yes", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Originality: The proposed methodology and algorithm is novel as far as I am aware. The work provides additional contributions compared to cited concurrent work [Peluchetti (2023)] that shares some similar ideas. Related literature and approaches are discussed in detail.\n\nQuality: Claims are adequately supported by proofs, and numerical experiments seem to support this. \n\nClarity: The paper is well written. The presentation is quite rigorous. \n\nSignificance: In my opinion, this work is of interest to the community as the method (i) generalizes recent flow matching approaches with non-degenerate noise and (ii) improves experimentally previous Diffusion Schrödinger Bridge approaches and (iii) provides some theoretical guarantees for solving the SP problem (under idealized assumptions).", "weaknesses": "It was not clear to me how the computational cost of the suggested method compares to related works per iteration or diffusion step (e.g. in Figures 5 and 16).\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their thorough evaluation of our work. We appreciate your acknowledgment of the paper’s merits. Here we would like to expand on the other comments raised in the review. \n\n**“Computational cost of the suggested method”**: At training time, our proposed method has comparable computational cost with previous SB methods such as DSB [1], but is more costly than flow matching [2] and bridge matching [3,4,5,6] (however, the first iteration of DSBM corresponds to bridge matching and so has comparable cost). \nDSBM can be seen as a refinement of bridge matching, and so pre-trained bridge matching models can be used to initialize DSBM. \nThe additional training cost comes from the trajectory caching procedure in subsequent iterations (see footnote 4), which requires sampling from the SDE model trained in the previous iteration. \nFor instance, at iteration $n+1$ we need to sample from $\\mathbb{M}^n$ and save a batch of joint samples $(\\mathbf{X}_0, \\mathbf{X}_T)$. \nThis batch of samples can then be cached and looped over for a number of epochs to sample different $\\mathbf{X}_t$ at different values of $t$. \nUnder reasonable assumptions on the SDE simulation cost and the number of epochs, the training cost averaged per iteration is about 1.5-2 times the training cost of flow matching and bridge matching (which is the case in e.g. Figure 5). \nAlso, as both the forward and backward models are trained, the total training cost is doubled. \nOn the other hand, DSBM is about 30\\% more efficient than DSB, as the trajectory caching procedure in DSBM requires less NFEs than DSB (see lines 944-954).  \n\nAt sampling time, the computational costs are all equal among all methods per diffusion step. In Figure 16, all methods have the same sampling cost for each vertical slice. \nTherefore, DSBM can attain similar sampling performance as flow and bridge matching with significantly less sampling cost. \nWe also note that FID score is not the only metric for well-learned transport. Indeed, while the FID score quantifies information about the accuracy of the marginal distribution, we also need to check for the similarity between the samples $(\\mathbf{X}_0, \\mathbf{X}_T)$, see the Unpaired Fluid Flows Downscaling experiment for instance.\n\n**“Not clear what to take away from some parts of the appendix ...”**: In\nAppendix B, we showcase that the DSBM methodology can cover a large class of\nlinear drifts. Our goal is to derive the OU bridge and push the methodology to\nhave an equivalent of VPSDE used in standard diffusion models. We think such a\nformulation could be beneficial for generative modeling (when\n$\\pi_T=\\mathcal{N}(0,I)$) but we haven't thoroughly tested it. In our\npreliminary experiments, we did not observe significant benefits for using the\nOU bridge.\n\nRegarding Appendix E, we have\nderived a discrete-time version of all of our results for readers who would\nprefer to introduce the methods without having the knowledge of SDE. Another\ngoal of Appendix E is to show that there is nothing in our method that is\nintrinsic to the continuous-time  framework and that our methodology could be recovered as\nthe limiting case of a fully discrete numerical scheme.\n\n[1] De Bortoli, Thornton, Heng, Doucet (2021) -- Diffusion Schrödinger Bridge with Applications to Score-Based Generative Modeling \n\n[2] Lipman, Chen, Ben-Hamu, Nickel, Le (2022) -- Flow Matching for Generative Modeling\n\n[3] Delbracio, Milanfar (2023) -- Inversion by Direct Iteration: An Alternative to Denoising Diffusion for Image Restoration\n\n[4] Liu, Vahdat, Huang, Theodorou, Nie, Anandkumar (2023) -- I2SB: Image-to-Image Schrödinger Bridge\n\n[5] Heitz, Belcour, Chambon (2023) -- Iterative $\\alpha$-(de)Blending: Learning a Deterministic Mapping Between Arbitrary Densities \n\n[6] Albergo, Boffi, Vanden-Eijnden (2023) -- Stochastic Interpolants: A Unifying Framework for Flows and Diffusions"}, {"review_id": "gKD0BMDgcJ", "reviewer": "Reviewer_7Dnn", "summary": "Flow matching and alpha blending have achieved tremendous attention in the matching problems. Although the straight trajectories based on $X_t = (1-\\alpha)X_0+\\alpha X_1$ yield **fast inference**, it doesn't necessarily mean they are efficient in score estimations in training. To make the training theoretically more efficient, rectified flow leverages the optimal transport ideas to minimize the transport cost. However, the stochastic alternative is still lacking (stochastic interpolant may not be optimal in terms of optimal transport). To tackle this issue, the authors proposed the stochastic version of flow matching/ alpha blending with entropic optimized transport. A novel scheme IMF that resembles the IPF iteration is proposed.\n", "questions": "\n1. DDMs can be seen as the first iteration of DSM, which implies that we can use DDMs model weights to initialize the model weights of DSM. Does it mean that empirically we can also use weights from stochastic interpolant / the stochastic version of flow matching to train the first step of DSBM?\n\n2. In line 993 in the appendix, the authors mentioned that \"The advantage of this DSM loss is that it does **not rely on any divergence computation**\". I believe this is quite an important part of scalability, but why it is only emphasized in the joint training instead of stating it in the main paper?\n\n\n\n\n ", "limitations": "NA", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "1. Despite the advances of stochastic interpolant, the transport from marginals may not be optimal in general. The **optimal-transport**-based stochastic version of flow matching and alpha blending is **still missing** in the community.\n\n2. In my understanding, the proposal of the IMF algorithm is the key/ most important contribution of this paper.  \n\n3. solid experiments and comparisons are studied, which verify the promising potential of DSBM.\n\n4. Connections to flow matching, rectified flow, and stochastic interpolant are studied, which makes the authors easier to understand the benefit of this paper.", "weaknesses": "1. Clearer clarification between IPF and IMF would be required to facilitate the understanding. IPF is based on forward K; IMF is based on backward KL, this is good in math, but more detailed information is lacking. for example, if we solely check proposition 2, we can derive almost the same loss based on forward KL (which leads to the IPF algorithm) except that the underlying measure is not based on $\\Pi_{0, t}$. So **what is the real benefit/ advantages of simulating $\\Pi_{0, t}$ to train the loss in proposition 2 would clarify our concerns**\n\n2. Algorithm 1 is not clear enough. The authors may consider providing us more details in the learning of $\\upsilon_{\\phi^{\\star}}$ in Alg 1. For example, in the training of alpha-blending (line 60-64 in https://github.com/tchambon/IADB/blob/main/iadb.py), given some samples of $x_1$ and $x_0$, very simple code would be used to train the models. However, when it comes to DSBM, how is it implemented? pseudocode would work if it clarifies our concerns.\n\n3. My concern is that this algorithm may be quite expensive.\n\n\n\n\nMinor \n\n1. missing one related work:\n\n    Discussions on the accumulations of IPF errors based on general optimal-transport cost functions studied in [1] would be appreciated.\n\n    [1] Provably Convergent Schrödinger Bridge with Applications to Probabilistic Time Series Imputation. ICML'23.\n\n\n2. writing of Equation between line 637 and line 638 is weird, because given X_t=x_t and x_T, what is the randomness there and why do we need expectation w.r.t. $\\Pi_{T|t}$?  \n\n3. derivation of eq after line 113 is not clear to me. why $E_{\\Pi_{0, T}}[||X_0-X_T||]^2$? is there any reference for that?\n\n4. first part of section 4 in the main paper, the need the alternating schemes is still not well motivated\n\n5. what is IMF-b?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for your thorough review. We will incorporate your suggestions regarding the clarity of Algorithm 1 and the related work. We now address your comments in more detail.\n\n**“Clearer clarification between IPF and IMF”**: Like you mentioned, we have related IPF with the forward KL, and IMF with the backward KL; however, they are not simply interchangable by switching the direction of KL in Proposition 2. We refer to the “Relationship between IPF and IMF” section in our response to all reviewers for details.\n\n**“Algorithm 1 is not clear enough...”**: We were constrained by space in the main paper to describe the algorithm fully. However, we plan to revise the paper to include a more detailed and down to earth description of the algorithm. In what follows, we include a more complete description of the algorithm. We start by describing a function DatasetUpdate which updates the dataset.\n\n>Input: direction $d$ (forward or backward), drift $v$\n\n>If $d$ is forward do\n\n>>Sample $\\mathbf{X}_0\\sim\\pi_0$\n\n>>Simulate $\\mathrm{d}\\mathbf{X}_t=v(t,\\mathbf{X}_t)\\mathrm{d}t+\\sigma\\mathrm{d}\\mathbf{B}_t$\n\n>>Output Dataloader($(\\mathbf{X}_0,\\mathbf{X}_T)$)\n\n>If $d$ is backward do\n\n>>Sample $\\mathbf{Y}_0\\sim\\pi_T$\n\n>>Simulate $\\mathrm{d}\\mathbf{Y}_t=v(t,\\mathbf{Y}_t)\\mathrm{d}t+\\sigma\\mathrm{d}\\mathbf{B}_t$\n\n>>Output Dataloader($(\\mathbf{Y}_T,\\mathbf{Y}_0)$)\n\nEquipped with this function we are now ready to describe the full pseudocode of DSBM-IMF.\n\n>Input: forward drift network $v_\\theta$, backward drift network $v_\\phi$\n\n>Initialize $d$ = backward and PairedDataset with independent samples $\\mathbf{X}_0\\sim\\pi_0$ and $\\mathbf{X}_T\\sim\\pi_T$\n\n>For $n\\in\\{0,\\dots,N-1\\}$ do\n\n>>While not converged do\n\n>>>Sample $t\\sim\\mathrm{Unif}([0,T])$\n\n>>>Sample $(\\mathbf{X}_0,\\mathbf{X}_T)\\sim$ PairedDataset and $\\mathbf{Z}\\sim\\mathrm{N}(0,\\mathrm{Id})$\n\n>>>Compute $\\mathbf{X}_t=(1-t/T)\\mathbf{X}_0+(t/T)\\mathbf{X}_T+\\sigma\\sqrt{t(T-t)/T}\\mathbf{Z}$\n\n>>>If $d$ is forward $v=v_\\theta$; update $\\theta$ using ADAM on the loss function $\\|v_\\theta(t,\\mathbf{X}_t)-(\\mathbf{X}_T-\\mathbf{X}_t)/(T-t)\\|^2$\n\n>>>If $d$ is backward $v=v_\\phi$; update $\\phi$ using ADAM  on the loss function $\\|v_\\phi(T-t,\\mathbf{X}_t)-(\\mathbf{X}_0-\\mathbf{X}_t)/t\\|^2$\n\n>>Update PairedDataset using DatasetUpdate$(d,v)$\n\n>>If $d$ is forward change it to backward; if $d$ is backward change it to forward\n\nIn practice, we can use caching to reuse samples from PairedDataset. We hope that this pseudocode resolves the\nconcerns of the reviewer and we are happy to provide more details if more\nclarifications are needed. \n\n**“My concern is that this algorithm may be quite expensive.”**: For the first stage of DSBM, there is no computational difference with bridge matching. Later stages are trained similarly to Rectified Flow. The main computational complexity is that the coupling $(\\mathbf{X}_0,\\mathbf{X}_T)$ is not independent but given by the previous iteration. This requires running a sampling algorithm and is shared between DSBM, DSB and Rectified Flow. In practice, the computational load of this step can be greatly reduced using caching strategies.\n\n**“Missing one related work...”**: We thank the reviewer for pointing us to [1]. In [1], the authors study the bias accumulation for IPF when the potentials are approximated. The results can be applied to provide error bounds on the marginals. However, in the IMF procedure, there does not exist potentials associated with the path measure in the sense of [1]. It would be extremely interesting to derive similar results for the IMF sequence. We leave this study for future work.\n\n**“Writing of Equation between line 637 and line 638 ...”**: Thank you for this question. The variables $x_0$ and $x_T$ in lines 637-642 should all be changed to $\\mathbf{X}_0$ or $\\mathbf{X}_T$. We will make this modification in the paper.\n\n**“Why $\\mathbb{E}\\_{\\Pi_{0,T}}[||X_0-X_T||]^2$?”**: One reference for this result is [2] around equation (1.2). This is a standard result when $\\mathbb{Q}$ is a Brownian process $\\sigma \\mathbf{B}_t$.\n\n**“The need of alternating schemes...”“what is IMF-b?”**: As detailed in lines 186-192, a naive version of IMF can be implemented by iteratively performing Markovian projections in the forward-time direction using equation (9). Due to time-symmetry (Proposition 8), IMF can also be derived in the reverse-time direction using equation (13). They are essentially the same algorithm, and only differs in the direction of transport we learn. We name the latter \"IMF-b\" which stands for backward-only IMF. However, as shown in Figure 3, IMF-b suffers from increasing error as the number of iterations $n$ increases. This is in line with the analysis in lines 190-192. On the other hand, Figure 3 shows DSBM-IMF does not suffer from the same issue. This indicates that the forward-backward scheme is highly useful for avoiding bias accumulation and improving accuracy in the learned marginals.\n\n**“Can use pretrained weights in first step of DSBM?”**: That is correct. The first iteration of DSBM coincides with bridge matching, and so DSBM can also be interpreted as further refinements of bridge matching.\n\n**“The advantage of this DSM loss is that it does not rely on any divergence computation...”**: We did not mention this advantage in the main paper because some previous SB methods (such as DSB) also don't rely on divergence computation. However, we agree with the reviewer that any loss based on divergence (like the ISM loss) may suffer from scalability issues.\n\nWe thank the reviewer for your suggestions to clarify our contribution and the pseudocode. We hope that our rebuttal resolves the main concerns of the reviewer.\n\n[1] Chen, Deng, Fang, Li, Yang, Zhang, Rasul, Schneider, Nevmyvaka (2023) -- Provably Convergent Schrödinger Bridge with Applications to Probabilistic Time Series Imputation\n\n[2] Léonard (2013) -- A survey of the Schrödinger problem and some of its connections with optimal transport"}, {"review_id": "fC389gXWrn", "reviewer": "Reviewer_Lk9k", "summary": "This paper formulates a new framework of a probabilistic model called Diffusion Schrödinger Bridge Matching (DSBM) and an iterative algorithm which is called iterative Markovian fitting (IMF). Firstly, the paper has contribution of pointing out a relationship between famous generative models (score matching and flow matching) in the form of DSBM. IMF is also novel and might be more suitable to be trained in neural networks than well-established algorithms such as IPF. To achieve these, the authors utilize a tractable bridge structure (Brownian bridge), which enables preserving key properties, i.e., densities of marginal distribution at the start and end of the process. The proposed algorithm often shows better generation results in certain tasks, such as image transfer. \n ", "questions": "1. Does $f_t$ is predefined in all tasks? \n2. Is $\\sigma_t$ constant and **not** scheduled? The authors of [1] mentioned time-symmetric scheduling the diffusion term, as suggested by prior SB models.\n3. Could you explain how the IMF scheme is fundamentally better than IPF algorithms? \n\n\n\n[1] Guan-Horng Liu et al. (2023), I2SB: Image-to-Image Schrödinger Bridge. ", "limitations": "Although I think overall paper is good, but I also think the proposed experimental results are not quite sufficient; the paper need more empirical evidence to show the excellence of the model. I hope to see response from the authors for this point.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-written despite dealing with a complex subject. As far as I understand, most of the arguments seem to be theoretically sound.\n2. SB algorithms have suffered from numerical errors and computational complexity. Since the IMF approach is akin to standard matching algorithms for neural networks, the provided empirical results show that it scales well. \n3. The paper provides a novel perspective on the diffusion bridge problem with a novel algorithm. As a result, DSBM-IMF shows promising results in certain tasks.", "weaknesses": "1. Significance. In Tables 2, 3 and Figure 10, It is difficult for me to observe that DSBM-IMF excels DSMB-IPF. Therefore, I would say significance of the proposed algorithm is not clearly shown in the experiments. The paper needs more empirical evidence of the actual benefits of choosing this algorithm, such as performance boost and scalability, in more challenging, or large-scale settings.\n2. Limited applications. I believe the actual implementation of the algorithm heavily relies on the well-known Brownian bridge. Although I think this is actually a valid contraint for certain applications such as image generation, at the same time it means that the algorithm is not always applicable to general optimal transport problems. Therefore, I recommend putting more effort on analyzing more classes of $\\mathbb{Q}$ in the experiment (if there are any), or putting clear justification on certain classes of $\\mathbb{Q}$ used in the experiment.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback and their positive evaluation. We appreciate their interest in our work. \nWe address here the reviewer's concerns raised in the review.\n\n**“Difficult to observe DSBM-IMF excels DSMB-IPF ...”**: We would like to first make a clarification that DSBM-IPF is also a proposed novel algorithm in our work. While DSBM-IPF is based on the classical IPF theory, it is different from previous diffusion-based algorithms based on IPF such as DSB which we compare to in the paper. Comparing DSBM-IPF with DSB, our proposed methodology achieves significantly better results, see Figure 5. On the other hand, comparing DSBM-IPF with DSBM-IMF, they exhibit a comparatively minor difference like you mentioned. Algorithmically, the two DSBM methods only differ in the initialization of $\\Pi^0$ in the first iteration. To summarize, DSBM-IMF and DSBM-IPF are both original contributions of the paper (differing mostly on their theoretical foundations while being very close in practice).\n\n**“Actual implementation of the algorithm heavily relies on the well-known Brownian bridge ...”**: We thank the reviewer for this comment. We refer to the \"Choice of diffusion bridge\" section in our response to all reviewers for a detailed response.\n\n**“More empirical evidence”**: We have demonstrated in our paper the usefulness of our proposed method in a number of high-dimensional tasks, such as CelebA ($128 \\times 128$) and fluid flows downscaling ($512 \\times 512)$. We have additionally performed an experiment on the AFHQ dataset at $512 \\times 512$ resolution using the unmodified DSBM-IMF algorithm, and have included it in the one-page response PDF. Additionally, we note that the scalability of our method is comparable to the one of bridge matching, as our method can be viewed as a refinement of it.\n\n**“Is $f_t$ predefined ...?”**: Yes, the reference process $\\mathbb{Q}$ has a similar role as the “forward noising process” in standard diffusion models. Two standard choices of $\\mathbb{Q}$ include the Brownian motion for which $f_t(x_t)=0$, and the Ornstein-Uhlenbeck (OU) process for which $f_t(x_t)=-\\frac{1}{2}x_t$.\n\n**“Is $\\sigma_t$ constant and not scheduled? ...”**: We have used a constant schedule for the diffusion coefficient $\\sigma_t$ in our work. We have experimented with a symmetric schedule of integration step sizes similar to the related works you mention, but did not observe any improvement. We agree that the schedules for both the noising term and the integration steps may be important directions to study for future work.\n\n**“Could you explain how the IMF scheme is fundamentally better than IPF algorithms?”**: In short, the IMF scheme always preserves $\\pi_0, \\pi_T$ as the marginals at the initial and final times, whereas the IPF scheme only reaches the SB solution $\\mathbb{P}^\\star$ which satisfies both marginal constraints in the limit. For extra details, we refer to the “Relationship between IPF and IMF” section in our response to all reviewers."}]

Author rebuttal:
We would like to thank all reviewers for their time and their very helpful feedback. 
We will take care to address all suggestions on improving clarity and minor typos in the next version of the paper. 
We would like to take the opportunity here to address a few common questions raised by several reviewers. 

**“Relationship between IPF and IMF”**: Several reviewers are interested in a more detailed discussion on the relationship between the classical IPF and the novel IMF schemes. In short, IPF and IMF schemes use *different projections* on *different classes* of path measures (see Table 1 for a summary). In more detail, IPF uses alternative forward KL projections onto the path measure classes $\\{\mathbb{P}: \mathbb{P}_0=\pi_0\\}$ and $\\{\mathbb{P}:\mathbb{P}_T=\pi_T\\}$. IMF uses alternative backward KL projections onto the classes $\mathcal{M}$ (Markovian class) and $\mathcal{R}(\mathbb{Q})$ (reciprocal class). 
The true Schrödinger Bridge solution $\mathbb{P}^\star$ is the unique path measure in all 4 classes. 
A fundamental difference is that the IMF iterates always admit $\pi_0, \pi_T$ as the marginals at the initial and final times. In contrast, each IPF iterate only satisfies one of the marginal constraints (odd iterate satisfies $\tilde{\mathbb{P}}^{2n+1}_T=\pi_T$, even iterate satisfies $\tilde{\mathbb{P}}^{2n+2}_0=\pi_0$, see line 119), and they only reach the SB solution $\mathbb{P}^\star$ which satisfies both marginal constraints in the limit as $n\to\infty$. Therefore, the IMF scheme is more preferable when the marginal accuracy of the samples at the initial and final times are important, and we would like to obtain more accurate samples before the algorithms have converged. 

From a practical point of view, IMF also gives rise to more stable training procedures. The reciprocal projection can be carried out empirically by simulating the modeled SDEs, while the Markovian projection leverages tools from the flow/bridge matching literature [1,2,3,4,5]. In contrast, in DSB there is an additional approximation that the transition density is Gaussian for small stepsizes. This assumption is only true for small stepsizes and can make the loss of DSB quite unstable. In addition, DSB suffers from bias accumulation throughout the IPF iterations. The IMF scheme circumvents this issue since we interpolate between the endpoints using the Brownian bridge.

**“Choice of diffusion bridge”**: We use a Brownian motion as our reference process $\mathbb{Q}$ in our experiments, as it corresponds to solving the standard Wasserstein-2 EOT problem (see line 113). It is also the choice used in previous bridge matching papers [1,2,3,4]. We have analyzed a larger class of linear diffusion bridges in Appendix B. Extending bridge matching procedure to non-linear processes would be non-trivial and is still an open question, 
as this requires solving and sampling from an intractable diffusion bridge.
However, even if the diffusion bridge is not available in close form, one might be able to first learn the diffusion bridge using methodologies such as [6], and then apply DSBM using the learned diffusion bridge. 
This extension should also be motivated by specific applications, which we leave for future work. In our case, the choice of Brownian motion is motivated by the optimal transport point of view.

**“FID with more than 100 steps”**: We have performed further evaluations for the CIFAR-10 generative modeling task in Appendix I.5 with NFEs higher than 100 as suggested by several reviewers. We have included the result in our submitted one-page response PDF. We observe that at higher NFEs, DSBM-IMF still achieves better FID scores than bridge matching (BM) and rectified flow (RF), but achieves slightly worse FID than conditional flow matching (CFM) and OT-CFM for NFE ≥ 300. Overall, DSBM-IMF is effective in improving sample quality when the NFE is very low, but unlike RF, DSBM-IMF still results in comparably low FID when the NFE is very high. 

**“Scalability of DSBM”**: The scalability of our method should be comparable to bridge matching, as our method can be viewed as a refinement of it. For the first stage of DSBM, there is no computational difference between our proposed scheme and [5]. 
Later stages correspond to refinements of the process using improved couplings $(\mathbf{X}_0, \mathbf{X}_T)$, which involves a sampling process similar to DSB and Rectified Flow. The computational load of this step can be
greatly reduced using caching strategies. 
We have demonstrated in our paper the usefulness of our proposed method in a number of high-dimensional tasks, such as CelebA ($128 \times 128$) and fluid flows downscaling ($512 \times 512)$. We have additionally performed an experiment on the AFHQ dataset at $512 \times 512$ resolution using DSBM-IMF, and have included it in the one-page response PDF.

[1] Delbracio, Milanfar (2023) -- Inversion by Direct Iteration: An Alternative to Denoising Diffusion for Image Restoration

[2] Liu, Vahdat, Huang, Theodorou, Nie, Anandkumar (2023) -- I2SB: Image-to-Image Schrödinger Bridge

[3] Heitz, Belcour, Chambon (2023) -- Iterative $\alpha$-(de)Blending: Learning a Deterministic Mapping Between Arbitrary Densities 

[4] Albergo, Boffi, Vanden-Eijnden (2023) -- Stochastic Interpolants: A Unifying Framework for Flows and Diffusions

[5] Lipman, Chen, Ben-Hamu, Nickel, Le (2022) -- Flow Matching for Generative Modeling

[6] Heng, De Bortoli, Doucet, Thornton (2021) -- Simulating Diffusion Bridges with Score Matching

Meta-review: The paper proposes a novel numerical approach for solving the Schroedinger Bridge (SB) problem and 
show its potential impacts for generative modeling. Experimental results illustrate the the benefit
it bring for generative modelling.

Overall, all reviewers are happy about the paper  and the rebuttal and  propose to accept the paper.
Decision: Accept (poster)