Schrodinger Bridge Flow for Unpaired Data Translation
Authors: Valentin De Bortoli
Iryna Korshunova
Andriy Mnih
Arnaud Doucet
Venue: NeurIPS 2024 spotlight (2024)
Keywords: unpaired data translation
diffusion models
Schrodinger bridge
bridge matching
stochastic interpolant
flow matching
Primary area: diffusion_based_models
TL;DR: We identify a flow of measures whose discretisations define online versions of the DSBM algorithm which are easier to train.

Abstract: Mass transport problems arise in many areas of machine learning whereby one wants to compute a map transporting one distribution to another. Generative modeling techniques like Generative Adversarial Networks (GANs) and Denoising Diffusion Models (DMMs) have been successfully adapted to solve such transport problems, resulting in CycleGAN and Bridge Matching respectively. However, these methods do not approximate Optimal Transport (OT) maps, which are known to have desirable properties. Existing techniques approximating OT maps for high-dimensional data-rich problems, including DDMs-based Rectified Flow and Schrodinger bridge procedures, require fully training a DDM-type model at each iteration, or use mini-batch techniques which can introduce significant errors. We propose a novel algorithm to compute the Schrodinger bridge, a dynamic entropy-regularized version of OT, that eliminates the need to train multiple DDMs-like models. This algorithm corresponds to a discretization of a flow of path measures, referred to as the Schrodinger Bridge Flow, whose only stationary point is the Schrodinger bridge. We demonstrate the performance of our algorithm on a variety of unpaired data translation tasks.

Schrödinger Bridge Flow for Unpaired Data Translation
Valentin De Bortoli Google DeepMind
Iryna Korshunova Google DeepMind
Andriy Mnih Google DeepMind
Arnaud Doucet Google DeepMind
Abstract
Mass transport problems arise in many areas of machine learning whereby one wants to compute map transporting one distribution to another. Generative modeling techniques like Generative Adversarial Networks (GANs) and Denoising Diffusion Models (DDMs) have been successfully adapted to solve such transport problems, resulting in CycleGAN and Bridge Matching respectively. However, these methods do not approximate Optimal Transport (OT) maps, which known to have desirable properties. Existing techniques approximating OT maps for high-dimensional data-rich problems, such DDM-based Rectified Flow and Schrödinger Bridge procedures, require fully training DDM-type model at each iteration, or use mini-batch techniques which can introduce significant errors. We propose novel algorithm to compute the Schrödinger Bridge, dynamic entropy- regularised version of OT. that eliminates the need to train multiple DDM-lik, models. This algorithm corresponds to a discretisation of a flow of path measures, which we call the Schrödinger Bridge Flow, whose only stationary point is the Schrödinger Bridge. We demonstrate the performance of our algorithm on a variety of unpaired data translation tasks.
Introduction
The problem of finding map to transport one probability distribution another has numerous applications in machine learning. In particular, it is at the core of generative modeling where the idea transform a noise distribution into the data distribution, and is also central to transfer learning tasks such as image-to-image translation. For discrete probability distributions, it is possible to compute the Optimal Transport (OT) map but this computationally expensive (Peyré et al., 2019). By showing that an entropy-rege regullaree version of OT, the Entropic OT (EOT), could be computed much more efficiently using the Sinkhorn algorithm, Cuturi (2013) has enabled transport ideas to be used in numerous applications (Ge et al.. 2021: Zhou et al., 2022). However, the computational complexity of Sinkhorn algorithm quadratic in the sample size, which makes application to very large datasets impractical. Mini-batch versions have been proposed, see e.g. (Genevay et al., 2018), but tend to introduce significant errors in high dimensions (Sommerfeld al., 2019).
In the context of generative modeling, Denoising Diffusion Models (DDMs) (Song et al., 2021a; Ho variety of domains. DDMs define forward
al., 2020) have shown impressive performance process progressively noising the data, and sample generation is achieved by approximating the time- reversal of this diffusion. In order to leverage the iterative refinement properties of DDMs the OT setting, methods exploiting the equivalence between the static versions of (E)OT and their dynamic counterparts (Benamou and Brenier, 2000; Léonard, 2014) have been developed. A procedure to approximate the dynamic OT is considered by Liu et al. (2023b), while techniques to approximate the dynamic equivalent to EOT, the Schrödinger Bridge (SB), have been proposed in (De Bortoli et al., 2021; Vargas et al., 2021; Chen et al., 2022; Peluchetti, 2023; Shi et al., 2023). These techniques are
"Equal contribution.
Submitted to 38th Conference on Neural Information Processing Systems (NeurIPS 2024).

---

expensive however, as they require training multiple DDM-typp type models. Mini-batee versions of OT and Sinkhorn (Pooladian et al.. 2023; Tong et al., 2024b) combined with bridge or flow matching have also been proposed to approximate the OT path and SB, but they optimise minibatch OT objective that can introduce significant errors in high dimensions: the error Wasserstein-1 distance of order D(B-1/(24) where d is the dimension of the problem and B the minibatch size, see (Sommerfeld et al., 2019, Corollary 1).
In this paper, we propose a novel approach to computing the SB. Similarly to Iterative Markovian Fitting (IMF) and its practical implementation, Diffusion Schrödinger Bridge Matching (DSBM) (Shi et al., 2023; Peluchetti, 2023), it leverages the fact that the SB is the only Markov process with preseribed marginals at the endpoints which is in the reciprocal class of the Brownian motion, i.e. it has the same bridge the Brownian motion (Léonard, 2014); Markov processes and the reciprocal class. Compared to DSBM, our approach is casier to implement as it does not require caching samples, alternating between optimising two different losses, and, optionally, uses one neural network instead of two. In Section 3, we start by introducing flow of path measures whose -incressretiitionn yields a family of algorithms called a-IMF and presented in Section 4. Notably, we show that a-IMF converges to the Schrödinger Bridge for any E (0. 1]. Additionally, for special value of the discretisation stepsize = 1, we recover the IMF procedure (Peluchetti, 2023; Shi et al., 2023), while < corresponds to online versions of IMF. We implement a parametric version of the a-IMF as an online DSBM procedure, called a-DSBM. We illustrate the efficiency of our approach in umpaired image-to-image translation settings in Section 6.
Section 2 for more details on
Notation. We denote the space of path measures by P(C). i.e. P(C) - P(C((O, 1], RR)), where C([0, Rd) is the space of continuous functions from [0, 1] to Fd The subset of Markov path mea- sures associated with a diffusion of the form dX,- = (X,) df + J,dB,, with Ø, u locally Lipschitz, is denoted M. For Q induced by (VEB,)ceoo.)) with & > 0 and (Bt) a d-dimensional Brownian motion, the reciprocal class of Q is denoted R(Q), see Definition 2. For any IP E P(C). we denote by P., its marginal distribution at time I, Pa., the joint distribution at times 8, I, Pst the conditional distribution at time given the state at time 1, and P/o.1 E P(C) the distribution of the path on time interval (0, 1) given its endpoints; e.g. Q10.1 is a scaled Brownian bridge. Unless specified otherwise, all gradient operators V are w.r.t. the variable H, with time index t. Given probability spaces (X, X) and y), Markov kernel K xxy - [0, 1). and probability measure μ defined on X, write uK for the probability measure on y' such that for any A E y we have (u((A) = fx K(rr, A)) In particular, for any joint distribution Il over IRd x Rdd we denote the mixture of bridges mea- sure as II Io, Finally, we define the Kullback Leibler (KL) divergence between two probability measures To, E PP() as (L(To = Jx (dTo/ / (z)) (r) if To is absolutely continuous w.r.t. T and KL (o/ I = +oo otherwise.
1. po.1 PCC). which is short for II..) Jg"x2= Po., P. 21 )dHl,, (Zo, 21).
2 Optimal Transport and Schrödinger Bridge
Unpaired Transfer and Optimal Transport. Given unpaired data samples from To and T1, where To. Ti are two distributions on Rd, we are interested in designing a transport map from To to TT. This corresponds to an unpaired data transfer task. We can formulate this problem as finding distribution on Rd x Rd with marginals Ho = To and II, = m so that if X, ~ Fo then X, X ~ ( satisfies X1 wj. Among an infinite number of such so-called coupling distributions II, we are here interested in finding the Entropic Optimal Transport (EOT) coupling IT* defined as
W
1 1 secevllee aal v) - sH(M): To = mp IT -h=mp m } (1) xe
IT = argminnep(R*xR)) X
where H(II) is the differential entropy of II and E > 0 is a regularisation hyperparameter (Peyré et al., 2019). For & = 0, we recover the standard OT.
In order to leverage the recent advances in generative modeling, and in particular the concept of iterative refinement central to DDMs, we turn to a dynamic formulation of EOT known as the Schrödinger Bridge problem (Léonard, 2014). It is defined as follows: find IP* e P(C) such that
F argminpep(c) (KL(P(Q) ; P. To. P. PI=T1),
(2)

---

with Q E P(C) induced by scaled d-dimensional Brownian motion (eeBo) The term dynamic here refers to the fact that (2) defined on path measures, i.e. on (stochastic) processes, in contrast to the static problem (1) which is defined on measures on the space IRd x IR. In Section 3, we show that solving (2) is equivalent to optimising the vector field of a stochastic process using objectives similar to the ones of bridge matching (Peluchetti, 2021; Albergo and Vanden-Fijnden, 2023; Lipman et al., 2023; Liu et al., 2023a). Under mild assumptions, it can be shown that 16 = II*, see e.g. (Léonard, 2014; Pavon et al., 2021). Hence solving (1) reduces to solving (2). Once we have found P* associated with (E1) we can sample from P by first sampling Xj ~ no and then sampling the trajectory (X)))((.1)) which yields (X;: Xi) ~ IT"
Reciprocal and Markoy projections. To introduce our methodology, it is necessary to recall the notions of reciprocal and Markov projections. We refer to Shi et al. (2023) for more details. For practitioners, more intuitive explanation of these projections is given in Appendix E.
Definition 2.1 (Reciprocal projection): IP E PC()) is in the reciprocal class R(Q) Qf
Po. Q We define the reciprocal projection of E P(C) as IP* = pproi (P) - Po, Q/0,1- We will write projæ instead of projæ(Q) to simplify notation.
In other words, P is in the reciprocal class of Q if the conditional distribution of a path given its endpoints is identical under and Q see (Rællyy, 2013). Sampling from the reciprocal projection of P can be achieved by sampling path (X,)) from P,, keeping only the values of the endpoints, say Xo, X1, and then sampling new value for the bridge (X)))E(0,1) from Q).
Definition 2.2 (Markov projection): Assume that Q is induced by (VeB,) 1) for E > 0. Then, when it is well-defined for any P E R(Q), the Markovian projection M = proj (P) E M is the path measure induced by the diffusion (X, )eE(o.1) with for any t E [0, 1]
dx; = (X;) + VEdB
of (x1) = (E [ X, = z1} - 24) / t),
X ~ Po.
In practice, implementing Markovian projection requires solving regression problem to approx- imate EE,,, [X] X = 2r). similar to the one appearing in bridge matching and flow matching.
One key property of the Markovian projection is that Mr = P, for all t E [0, 1], ie. the Markovian projection preserves the marginals; see (Peluchetti, 2021) for instance.
Iterative Markovian Fitting. Leveraging the reciprocal and Markovian projections, Peluchetti (2023) and Shi et al. (2023) concurrently introduced IMF. Starting from 1 = (To @ ) o, measure where endpoints are sampled independently from To and иj and then interpolated using a (scaled) Brownian bridge, they define a sequence of path measures (IP", fom) )eT) where pm (roj)((P)) and pa+1 = projæ (P"). This ensures that Ph = Tu, PT = Ti for all n, and can be shown that the sequence (IP"), nn converges to the SB, (Peluchetti, 2023, Theorem 2). The practical implementation of this algorithm proposed by Shi et al. (2023) is called DSBM. Implementing DSBM poses challenges, as each Markovian projection requires training neural network to approximate the relevant conditional expectations by minimising a bridge matching loss. Furthermore, in practice, generated model samples are stored cache order to train the next iterations of DSBM. This introduces additional hyperparameters that require tuning. In Section 3 we propose a-IMF, an algorithm which can be interpreted as the discretisation of flow of path measures. This leads to a- D-DSBM, an algorithm that is computationally much more efficient than DSBM as it does not rely on a Markovian projection at each step.
3 Schrödinger Bridge flow
We will now introduce a flow of path measures )22), and show that the dineeesscrttistionnn of this flow with an appropriate stepsize a E (0, 1] yields a family of procedures called a-IMF, which all converge to the Schrödinger Bridge. While a = 1 yields the classical IMF, E (0, yields an increenntal version of IMF. In Section we show that a-IMF can be implemented as an online version of DSBM
3

---

3.1 A flow of path measures
Let (IP*, IPr)s20 be alow of path measures defined for any 2 0 by
p = (To @ n1 Q
0,6² projæ(proi (P=)) - pe. = 3=projm (IP"),
(3)
which we assume is well-defined. Note that for any 2 0, Tpe Markov while pas is in the reciprocal class of Q. Crucially, the only fixed point of (3) is the SB. Indeed, let P be a fixed point of (P)) in (3). Then, we have that IP projæ (projи((P)
M
Hence, we get P projæ (proju (.. (projæ (proj, M Hence, under mild assumptions, IP is limit point of IMF and therefore P is the SB IP* given by (2), see (Peluchetti, 2023, Theorem 2).
20
Next, for any a E (0, 1]. we define the following discreti- sation of (3) called a-IMF:
R(Q)
Figure Illustration of the SB Flow and comparison with IMF, IP* is the SB, (P"))ett the IMF sequence and (P*).20 the flow we consider. See Appendix for the analysis of this example.
pa+1 = (1 - a)) + aprojæ (ppr)jm(P"))) (Pm)),
(4)
and IPn pr Note that for any N, 1m R(Q). This recovers the IMF procedure (Shi et al., 2023;
Peluchetti, 2023) when a 1. Using the definition of the sequence (P") nnEN. is possible to analyse the sequence (PP))ne using the properties of the KL divergence well as the Pythagorean identities derived in (Shi et al., 2023; Peluchetti, 2023). We first introduce some assumptions on the Schrödinger Bridge problem. We recall that the differential entropy of a probability measure æ is given by
H(x) - / [og((dd/dLeb))
π admits a density with respect to the Lebesgue measure and too otherwise. Recall that Q is associated with (VeBe)xejo.1) and assume that Qu Leb. Let To, Ti P(RR) such that
[ [<])) < +00,
H(m)) +00,
f E {0, 1}. Under these assumptions, we can use the characterisation of the SB as the only path measure that preserves To, æi, and is both Markov and in the reciprocal class of Q (see e.g. (Léonard, 2014, Theorem 2.))).. We get the following result.
Theorem 3.1 (Convergence of o-IMF): Let (0, 1] and (IPm fTm) defined by (4). Under mild assumptions, we have that lim, ppm = 1i* where IP* is the solution of the Schrödinger Bridge problem (2).
3.2 Discretisation and non-parametric loss
We show here that a- a-IMF associated with an incremental version DSBM for & E (0, 1).
Iterative Markovian Fitting. For any v:: [0, 1] IRd + Rd, we introduce the loss function
L(e,P)= = - 1 L11, [/
(Raya 1 [[0(
- 2z "I dPo (2o, 2n ]dQro, (aa x2 )df
(5)
where we recall that Q is induced by (VeBe)etoo.) for some > 0. This loss was already considered in (Peluchetti, 2021; Lipman al., 2023; Liu et al., 2023a; Liu, 2022; Shi et al., 2023). We also define the path measure P, E PCC)) associated with
dX; (x(x))tt + VEdB,,
Xo mo.
(6)

---

Consider first the sequence (v") defined by
r+ argmin,,(e,.F...
Using Definition 2.2, we have that pe = proj (projR )), which corresponds to pa+11 the IMF sequence. Therefore we have that tot Pun IP* under mild assumptions (Peluchetti, 2023, Theorem 2).
Functional gradient descent. We now introduce a relaxation of (7), where, instead of considering the argmin, we update the vector field with one gradient step. To define this relaxation, we recall that for functional F F + R. where F is appropriate function space, its functional derivative (Courant and Hilbert, 2008) with reference measure y is denoted V,FF and given for any F, when it exists, by
lim,- + 70)
X; = z| -=1-)/(( 2)/(1- t), where Po (mo @ Tı Q1o,1 now introduce sequence of vector fields (v") nn") This corresponds to training bridge matching model (see e.g. Liu et al. (2023a); Albergo al. (2023)), giving Foo proiu(((") Then for n E N, let
Initialised with (x) (Bjg [
i/e
it (x) = (x) ón v Cx (z),
with 5, > O and μ² E P(C). The parameters (6n, μz" nEN will be made explicit in Proposition 3.2. We emphasize that, in contrast to the IMF procedure, in the online update (8) we not need to solve a Markovian projection problem at every step: instead we simply take gradient step the loss (5).
Connection with o-IME The following proposition shows that (PoN))EN defined by (8) is associated with a-IMF defined in (4).
Proposition (Non-parametric updates are a-IMF): Let a 10, (Pm, f. (Pm,,""") )nex as in (4), on = a and nº = (1 a)r n + aprojz (P")). Then, under mild assumptions, we have P = T for all N.
Combining Theorem 3.1 to Proposition 3.2, we get that lime Pur P i.e. the non-parametric procedure converges to the SB.
o-Diffusion Schrödinger Bridge Matching
From DSBM to a-DSBM. In Section 3, we introduced a-IMF, scheme which defines sequence of path measures converging to the SB for all a E (0, 1]. For a = 1, this corresponds to the IMF, whose practical DSBM implementation (Shi al., 2023) requires repeatedly solving an expensive minimisation problem (7). In contrast, for < we are only required to take one (non-parametric) gradient step update the vector field, (8). This suggests the following practical implementation of a- a-IMF, called a- a-DSBM: First, pretrain a bridge matching model so that for t E [0, 1] and æ e Rd
(x) (Egg X X, x] x1 t), where 100 = (no क Q Then, perform the
parametric version of the update (8):
0 0 (-o-aVoL((,Pa)) L(O,P) = - (3:rja
31 1-1-11 o. (.))))
1-t d (1o, (x x1 21)d
(9)
where Py is a stop-gradient version of P.... In Appendix D.2, we give a theoretical justification for this parametric equivalent of (5) and (8) by showing that, as ++ 0, the update on the velocity fields o given by (9) corresponds to direction of descent for the non-parametric loss (8) on average. Once again, we emphasize that if replace the gradient step in (9) with the minimisation f- LrgmingLO0, Pg), we recover DSBM.
Bidirectional online procedure. As with DSBM, directly implementing (9) leads to error quickly accumulating, see Appendix for details. One way to circumvent this error accumulation is- sue is to consider a bidirectional procedure, in which we train both forward and backward
5

---

model. This is possible because the Markovian projection coincides for forward and back- ward path measures, see (Shi et al., 2023, Proposition 9). This suggests considering the loss C(U,, uf, F,,,,) = =Jo Li et, Pt, jtt, which an extension of (5), where
/ (x(e)
I1 Zir || di 0.1 (No, x1 )d(roo, (12 21)
Ce(17, of, P., )P))
(10)
(Bad3a
1 1-t
Tu xz dl p. (2o. xy )dQ.o (2 (1).
(Rays
Similarly to (6), define Pe., Por, associated with (E))) and (Y 1E respectively, which are defined by forward and backward SDES
(fwd): dX; 0X( VEdB,, X no. (bwd): dy, (t) Yo ~ Tι. (1I)
Similarly (8), we define non-parametric updates for any n E N. [0, 1] and Rd
(z). 0t + (+++((=)))))))))))))))))(()))( (x)) (v.' (x)) - V,, Li (z)), ii (z), P, (z).
We have the following proposition which ensures our bidirectional procedure is still valid and that the results of Proposition 3.2 still hold.
Proposition 4.1 (Bidirectional updates): Let a E (0, 1l. For any n E N, define (P", iPn nne by
(4). Then, under mild assumption and assuming that on = a and pM = (1 a + aprojR (TPM) (TPm), have that for any n. E N, =P. Ipn.
In Appendix I, show that in the Gaussian setting the bidirectional procedure (4.1) does not accumulate error when the vector field is approximated, while the unidirectional one (8) does.
Vector field parameterisation. Contrary to existing procedures (Shi al., 2023; Peluchetti, 2023; Liu, 2022), we do not parameterise v² and of using two separate networks. Instead, we consider an additional input E {0, }} such that vo ss v² and vp -) se v². This allows us to substantially reduce the number of parameters in the model. The conditioning Appendix K. Before stating our full algorithm in Algorithm I, we introduce a batched parametric version of (10). For ease of notation, write Interp, for the operation corresponding sampling from Qrjo, i.e.
in the network is detailed in
Interp, (Xo, X, Z) = -1 + XX + V EE(1 t)tZ.
(12)
We are now ready to introduce the batched parametric version of (10). For given batch of inputs X and X timesteps t ~ Un0 1)) and X, Interp, (Xo, X,, Z) with Z ~ NO Id) @ we compute the empirical forward and backward losses
1
(00;1,,,,,,,,
Ive im1
- (xi - x!) /(1 -1))1
(13)
B
1
No (0, 8', x:) (X) x:) /t/ 10
(0,;1,Xo,xx, X.)
B
We present the resulting a-DSBM in Algorithm 1. Note that in this algorithm, we maintain an Exponential Moving Average (EMA) of model parameters, as is common in diffusion models (Nichol and Dhariwal, 2021). During the finetuning stage, when we generate samples use model's inputs, we then have choice of sampling using the EMA or non-EMA parameters. At test time, we always sample using the EMA parameters, as it is known to improve the visual quality (Song and Ermon, 2020). In Algorithm 1, we specify a E (0. 1] as a stepsize parameter. In practice, we use Adam (Kingma and Ba, 2015) for optimization, thus the choice of implicit and adaptive throughout the training. To emphasize the importance of the parameter a, we sweep over its value with an explicit solver SGD in a toy setting, see Appendix K.2. We refer to Appendix K for more details on our experimental setup.

---

Algorithm 1 a- a--iiffusion Schrödinger Bridge Matching
1: Input: datasets To and TI, entropic regularisation è, number of pretraining and finetuning steps N Ppretraining and N Ninetuning batch size B and half batch size B/2, EMA decay 7, initial parameters 0 and initial EMA parameters B = 0 a E (0, 1]
2: for n E , .Npretraining do
3: Sample (1) (To m1) )
4: Sample & ~ (n, 1)) and Z ~ NO0, Id) and compute X, Interp, (Xo, X1, Z)
5: Update 9 with a gradient step on _/ [ (tlib, x xx, + l (th+l:B, xx0 1b X x
6: Update EMA parameters: (EMA = Gmm + (1 - 7)
7: end for
8: for E {l,..., .it...,,Ninecuninng! J do
9: Sample (Xo,Xi) ~ (To & )
10: Sample X, solving forward SDE (11)-(fwd) with eggna (1, or no (1,-) starting from X.
11:
Sample X, solving backward SDE (11)-(bwd) with rggxu (0. or uo -) starting from X;
Sample ~ Uni 66 and Z (O, Id) and compute X Interp, (Xo: X Z²)
13:
Sample ~ Un, 1))80 06 and Z ~ (0,, Id) @b and compute X; Interp,- (Xo:, X1, Z*)
14:
Update 0 with a gradient step on _ (1", X, X?) + (tt,, Xo. X;) and stepsize a
15:
Update EMA parameters: grma = 70th + 7)0
16:
end for
17:
Output: (0, (PNA) parameters of the finetuned model
5 Related work
Solving Schrödinger Bridge problems. Schrödinger Bridges (Schrödinger, 1932) have been thoroughly studied through the lens of probability theory (Léonard, 2014) and stochastic control (Dai Pra, 1991; Chen et al., 2021). They recently found applications in generative modeling and related fields leveraging recent advances in diffusion models (De Bortoli et al., 2021; Vargas et al., 2021; Chen et al., 2022). Extensions of these methods to other machine learning problems and modalities were studied in (Shi et al., 2022; Thornton al., 2022; Liu al., 2022; Chen et al., 2023; Tamir et al., 2023). Shi et al. (2023); Peluchetti (2023) concurrently introduced the DSBM algorithm which relies on a new procedure called IMF, while the DSB algorithm introduced in (De Bortoli et al., 2021) is based on the standard Iterative Proportional Fitting (IPF) scheme. Neklyudov et al. (2023a,b); Liu et al. (2022) generalise DSBM arbitrary cost functions, albeit the expense of having to learn the reciprocal projection which is no longer given by a Brownian bridge. These new methodologies translate improved numerics when compared their IPF counterparts, but they remain reliant on alternating between the optimisation of two losses. Finally, we note that the Schrödinger Bridge flow and the Q-IMF procedure can be linked to the Sinkhorn flow recently introduced by Karimi al. (2024), see Appendix H.1 for a detailed discussion.
Sampling-free methodologies. Sampling- free methodologies have been proposed to solve OT related objectives. In (Liu et al., 2023a; Somnath et al., 2023; Diefenbacher et al., 2024; Cao et al., 2024), the authors perform one step of DSBM, i.e. only consider the pretraining stage of our algorithm. While the obtained bridge might enjoy transport properties, it does not solve an OT problem. In another line of work, Pooladian et al. (2023): Tong al. (2024a.b); Eyring al. (2024) have proposed simulation-free methods to minimise OT objectives. However, they target not the OT problem, but a minibatch version of it which coincides with OT only the limit of infinite batch size, see (Pooladian et al., 2023, Theorem 4.2). Other sampling-fre free methods to solve the Schrödinger Bridge problem include Kim al. (2024); Gushchin et al. (2024b) both of which rely on adversarial losses solve the OT problem. In (De Bortoli et al., 2021; Vargas et al., 2021; Liu et al., 2022; Shi et al., 2023; Peluchetti, 2023) the adversarial objective dropped and instead the procedure requires alternating objectives during training and is not sampling-free et al. (2024); Gushchin et al. (2024a) in which the Schrödinger Bridge potentials are parameterised with mixtures of Gaussians, allowing for fast training in small dimensions. Finally, recently Deng et al. (2024) introduced a variation on Schrödinger Bridge for generative modeling, which while still not sampling-free does not require learning forward process.
We also highlight the line of work of Korotin
7

---

0.8
1.0
Noralized Frobenius distance
Iterative @1K Base
iterative @2.5K Online
0.6
Iterative @5K
0.8
Covariance
0.4
0.6
0.5
Iterative @ 1K
0.2
iterative #2.5K
Online
Iterative @ 5K
Optimum
30k
50k
10k 20k 30k 40k 50k 60k 70k BOk 90k100k Iraining iterat ons
Training iterations
Figure 2: Evolution of the covariance during online and iterative DSBM finetuning for forward and backward networks. The finetuning starts after 10K steps of training a bridge matching model. For the iterative case, alternate between forward and backward updates with varying frequencies, ie. changing after 1K, 2.5K and 5K steps. Left: Gaussian with scalar covariance matrix. Right: Gaussian with full covariance matrix. We compute the normFrob between C, and its estimate using Bridge Matching (Base), a-DSBM (Online), and DSBM (Iterative with @ x training steps per model fit)
6 Experiments
In this section, illustrate the efficiency of o-DSBM different tasks. In Section 6.1, compare a-DSBM to DSBM in a Gaussian setting where the EOT coupling is tractable and show that a-DSBM recovers the solution faster than DSBM. In Section 6.2, we illustrate the scalability of method through a range of unpaired image translation experiments.
6.1 Gaussian case
We compare D-DS to DSBM in the Gaussian setting where = -N0,, = -0 and Q is associated with (VEBe)uejo.)) with vē 0.5. In this case, the BOT coupling (0, E.). with Σ, given by
oo o'ld) oo ofd
Σ,
where of (1/2) 2 E),
with 1d being a x d identity matrix. We consider d = 50, oo = 01 1, resulting in of ~0 0.88. To showcase the robustness of a-DSBM, we consider the initial coupling Po.1. where (X,, x,) Po.1.
Xn ~ N0, Id), X, ~n, and let poo = Po. Q In this setting, the base model, i.e. bridge matching, significantly underestimates the true covariance o², as shown in Section 6. Additionally, the figure illustrates that online finetuning approaches the true solution faster than the original iterative
DSBM finetuning. For the latter, we can set how often we alternate between updating the forward and backward networks, and as this frequency increases, the behaviour approaches that of the online finetuning.
Full covariance Gaussian case. Let æu N(Ho Eu), = -[(1, En) with Σ, Id 2 for i E (0. and Zo. Zj independent d matrices with unit Gaussian entries. We also set Nu = A41 = 0. We consider the Entropic Optimal Transport (EOT) with regularization = 0.5 and given by
LS C. C.
II = T=NNuu, E.), E. =
with C, /2,
Idd, with D, = /2 + o Id))//2 Let normFrob ||A - -A-Bl Fro Fro be the normalized Frobenius distance between matrices A and B. The results are presented in Section and confirm those presented in the original manuseript considering diagonal covariance.

---

intial samples
VZ+00
VZ+0225
VT=0.75
YZ=20
YT+5.0
base
0.6
0.00 0.14 0.22 0.37
0.61
1.00 1.55
2.72
4.46
7.30
V/
Figure 3: Left: FID and Mean Squared Distance (MSD) on EMNIST to MNIST translation before and after finetuning with different values of Right: AFHQ-64 samples after the finetuning. For both, we use a bidirectional model with online finetuning, More results are in Appendix K.3 and K.4.
6.2 Image datasets
Similarly to Shi et al. (2023), we apply our method to image translation problems, such as MNIST digits to EMNIST letters (LeCun and Cortes, 2010; Cohen al., 2017), Wild to Cat domains from the Animal Faces-HQ (AFHQ) dataset (Choi et al., 2020), downsampled to 64 x 64 and 256 x 256 resolutions and CelebA 64 64.
The whole training procedure can be framed as a two-stage process: first, we train base model on the true data samples, performing bridge matching (Peluchetti, 2021; Albergo and Vanden-Eijnden, 2023; Lipman et al., 2023; Liu al., 2023a), and then we finetune this model. We compare models that combine different vector field parameterisations (two networks vs. one bidirectional net), finetuning methods (iterative vs. online), and sample generation strategies during the finetuning stage.
Following the established practice (Choi et al., 2020), we evaluate our models using FID (Heusel al., 2017) for visual quality, and mean squared distance (MSD) LPIPS (Zhang et al., 2018) for alignment. It is important to note that for image translation tasks at hand, FID scores are not ideal, as FID was designed for natural RGB images, which not the case for MNIST, It is also not well suited for small sample sizes as it is the case with AFHQ, where the test set in each domain has fewer than 500 examples. Thus quantitative results in Table should be interpreted cautiously, and we recommend a visual inspection of samples to complement these quantitative measures, especially for AFHQ models. Samples from the models along with the training and evaluation protocols are given Appendix K.
Compared the iterative DSBM, our online finetuning a-DSBM reduces the number of tunable hyperparameters, i.c. inner and outer iterations, refresh rate and the size of the cache for storing generated samples. This simplifies implementation and makes the algorithm more practical. The primary remaining hyperparameter, the variance of Brownian motion E, requires careful tuning as it influences the trade-off between the visual quality and alignment, as was also observed in Shi et al. (2023). An appropriate needs to balance the two: setting too low results in poor visual quality, while high values of cause poorly aligned and oversmoothed samples. Figure illustrates how FID and MSD metrics vary with for the case of MNIST. Additionally, it demonstrates the impact of on the generated samples for the AFHQ-64 64 model.
We run a-DSBM on CelebA with image size 64 x 64 with o= 2.0. We do not change the training hyper-parameters compared AFHQ. Visual results are reported in Figure and Figure In Figure 5, we show the influence of during the pretraining. The visual quality of the transfer is much lower for 0 = 0 than for 2.0. The case Flow Matching). Given the poor quality of the samples, we do not perform finetuning with o = 0. In Figure 6, we compare the visual quality and alignment DSBM and a-DSBM after 4000 training steps, corresponding to two outer DSBM iterations. In this case DSBM is trained with a bidirectional network and both procedures consist of finetuning the pretrained model obtained with a = 2.0. We note that the alignment is better in the case of a- D-DSBM.
corresponds the first step of Rectified Flow (i.e.

---

EMNIST MNIST
AFHQ-64 Wild
Cat
Method
FID
MSD
FID
LPIPS
DSBM
10.59
0.375
Pretrained two-networks model
6.02
0.564
25.97
0.589
(a) iterative finetuning
5.25+0015
0.345=0.001
25.41
0.485
(b) online finetuning
4.28:10.07
0.000
28.752:t1.999
0..87
(c) online finetuming without EMA
4.23:10.171
0.361 =0.002
32.665
0.445
Pretrained bidirectional model
6.33
0.572
29.44
0.584
(d) online finetuning
4.39
0.387 =00003
26.579
0..02
(e) online finetuning without EMA
4.57
0.369
30.638.1.022
0.451 10.002
Table I: Results of image translation between EMNIST and MNIST, and AFHQ x4 between Wild and Cat domains. DSBM results are from Shi et al. (2023). Our reimplementation of DSBM corresponds to (a). For MNIST and AFHQ models, used and ş 0.75², respectively. Each finetuning run was done with 5 random seeds, and we report mean scores standard deviation.
(a) Cat
Wild
(b) Wild - Cat
Figure 4: Online DSBM transfer results on AFHQ 256x 256 dataset between Cat and Wild domains.
Top row-initial samples, bottom row-transferred samples.
Discussion
In this paper we have introduced a- e-Diffusion Schrödinger Bridge Matching (a-DSBM), a new methodology to solve Entropic Optimal Transport problems. a-DSBM is improved version of DSBM, which does not require training multiple DDM-type models. We have shown that non-parametric version of this method recovers the Schrödinger Bridge (SB). In addition, a-DSBM casier to implement than existing SB methodologies while exhibiting similar performance. We illustrated the efficiency of our algorithm
variety of unpaired transfer tasks.
While a-DSBM solves one of the most critical limitations of DBSM, namely the alternative optimisa- tion, several issues remain to be addressed in order for the method to scale comparably to generative DDMs. particular, the method is not sampling-free as during training requires sampling from the model from the previous iteration to obtain the training data for the current iteration. While it seems difficult to derive completely sampling-fres method to solve SB problems without resorting to the Minibatch OT approximation, there is still room for improvement.
Figure 5: Translation Female
Male on
Translation Female -> Male on CelebA. Left: output after finetuning with
Figure
CelebA. Left: pretraining with = 0. Right:
pretraining with a 2.0.
DSBM. Right: output after a-DSBM finetun-
ing.
10

---

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.08777.
Albergo, M. S. and Vanden-Fijnden, E. (2023). Building normalizing flows with stochastic inter- polants. International Conference on Learning Representations.
Ambrosio, L.,, Gigli, N., and Savaré, G. (2008). Gradient Flows in Metric Spaces and in the Space of Probability Measures. Lectures in Mathematics ETH Zürich. Birkhäuser Verlag, Basel, second edition.
Bauschke, H. H. and Kruk, G. (2004). ceelection-projection method for convex feasibility problems with obtuse cone. Journal of Optimization Theory and Applications, 120.3)5003-31.
Benamou, J.-D. and Brenier, (2000). A computational fluid mechanics solution to the Monge- Kantorovich mass transfer problem. Numerische Mathematik, 84(3):375-3935
Black, K., Janner, M., Du, Y., Kostrikov, I., and Levine, S. (2023). Training diffusion models with reinforcement learning. arXiv preprint arXiv:230S. 13301.
Brekelmans, R. and Neklyudov, K. (2023). On Schrödinger bridge matching and expectation maximization. NeurIPS 2023 Workshop Optimal Transport and Machine Learning.
Brock, A., Donahue, J., and Simonyan, K. (2019). Large scale GAN training for high fidelity natural image synthesis. In International Conference on Learning Representations.
Cao, Z., Wu, X., and Deng, L.-J. (2024). Neural shrödinger bridge matching for pansharpening. arXiv preprint 1141::204.11416.
Chen, T., Liu, G.-H., Tao, M., and Theodorou, E. (2023). Deep momentum multi-marginal Schrödinger bridge. Advances in Neural Information Processing Systems.
Chen, T., Liu, G.-H., and Theodorou, E. A. (2022). Likelihood training of Schrödinger bridge using forward-hackwerr SDEs theory. In International Conference on Learning Representations.
Chen, Y., Georgiou, T.T. and Pavon, M. (2021). Optimal transport in systems and control. Annual Review of Control, Robotics, and Autonomous Systems, 4.
Chen, Y.. Goldstein, M., Hua, M., Albergo, M. S., Boffi, N. M., and Hunden-Eijnden E. (2024). Probabilistic forecasting with stochastic interpolants and Föllmer processes. arXiv preprint arXiv:2403.13724.
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 on Computer Vision and Pattern Recognition.
Chong, M. J. and Forsyth, D. (2020). Effectively unbiased FID and inception score and where to find them. IEEE/CVF Conference on Computer Vision and Pattern Recognition.
Cohen, G., Afshar, S., Tapson, J., and van Schaik, A. (2017). EMNIST: an extension of MNIST to handwritten letters. arXiv preprint arXiv: 17X.::1702.05373.
Courant, R. and Hilbert, D. (2008). Methods of Mathematical Physics: Partial Differrntial Equations. John Wiley & Sons.
Cuturi, M. (2013). Sinkhom 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-339.
Daras, G., Dagan, Y., Dimakis, A., and Daskalakis, C. (2023a). Consistent diffusion models: Miti- gating sampling drift by learning to consistent. In Advances Neural Information Processing Systems.

---

Daras, G., Shah, K. Dagan, Y., Gollakota, A., Dimakis, A., and Klivans, A. (2023b). Ambient diffusion: Learning clean distributions from corrupted data, In Advances in Neural Information Processing Systems.
De Bortoli, V., Hutchinson, M., Wirnsberger, P., and Doucet, A. (2024). Target score matching. arXiv preprint arXiv:2402.08667.
De Bortoli, V., Liu, G.-H., Chen, T., Theodorou, E. A., and Nie, W. (2023). Augmented bridge matching, arXiv preprinl UrXiv:2311.06678
De Bortoli, V., Thornton, J., Heng, J., and Doucet, A. (2021). Diffusion Schrödinger bridge with applications to score-based generative modeling. In Advances in Neural Information Processing Systems.
De Lange, M., Aljundi, R., Masana, M., Parisot, S., Jia, X., Leonardis, A., Slabaugh, G., and Tuytelaars, T. (2021). A continual learning survey: Defying forgetting in classification tasks. IEEE Transactions on Puttern Analysis and Machine Intelligence, 44(7):3366-333
Deng, W., Luo, W., Tan, Y., Biloš, M., Chen, Y., Nevmyvaka, Y., and Chen, R. T. (2024). Variational Schrödinger diffusion models. arXiv preprint arXiv:2405.04795.
Dhariwal, P. and Nichol, A. (2021). Diffusion models beat GANs on image synthesis. In Advances in Neural Information Processing Systems.
Diefenbacher, S., Liu, G.-H., Mikuni, V., Nachman, B., and Nie, W. (2024). Improving generative model-based unfolding with Schrödinger bridges. Physical Review D, 109(7):076011.
Dupuis, P. and Ellis, R. S. (2011). A Weak Convergence Approach to the Theory of Large Deviations. John Wiley & Sons.
Eyring, L., Klein, D., Uscidda, T., Palla, G., Kilbertus, N., Akata, Z., and Theis, F. (2024). Un- balancedness in neural Monge maps improves unpaired domain translation. In International Conference on Learning Representations.
Fan, Y., Watkins, O., Du, Y., Liu, H., Ryu, M., Boutilier, C., Abbeel, P., Ghavamzadeh, M., Lee, K., and Lee, (2024). Reinforcement learning for fine-tuning text-to-image diffusion models. In Advances in Neural Information Processing Systems.
Ge, Z., Liu, S., Li, Z., Yoshie, O., and Sun, (2021). Ota: Optimal transport assignment for object detection. In IEEE/CVF Conference on Computer Vision and Pattern Recognition.
Genevay, A., Peyre, G., and Cuturi, M. (2018). Learning generative models with Sinkhorn divergences. In Artificial Intelligence and Statistics.
Gushchin, N., Kholkin, S., Burnaev, E., and Korotin, A. (2024a). Light and optimal schrödinger bridge matching. arXiv preprint arXiv:2402.03207.
Gushchin, N., Kolesov, A., Korotin, A., Vetrov, D. P., and Burnaev, E. (2024b). Entropic neural optimal transport via diffusion processes. In Advances Neural Information Processing Systems.
Heusel, M., Ramsauer, H., Unterthiner, T.. Nessler, B., and Hochreiter, S. (2017). GANs trained by two time-scale update rule converge local Nash equilibrium. In Advances in Neural Information Processing Systems.
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. In International Conference on Machine Learning.
Hu, E. J., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L., Chen, W., et al. (2021). LoRA: Low-rank adaptation of large language models. International Conference on Learning Repre- sentations.
12

---

Hudson, D. A.. Zoran, D., Malinowski, M., Lampinen, A. K., Jaegle. A., McClelland, J. L., Matthey, L., Hill, F., and Lerchner, A. (2023). Soda: Bottleneck diffusion models for representation learning. arXiv preprint arXiv:2311.1790)
Karimi, M. R., Hsieh, Y.-P., and Krause, A. (2024), Sinkhorn flow mirror flow: A continuous-time framework for generalizing the Sinkhorn algorithm. In International Conference on Artificial Intelligence and Statistics.
Karras, T., Aittala, M., Aila, T., and Laine, S. (2022). Elucidating the design space of diffusion-basel generative models. In Advances in Neural Information Processing Systems.
Kim, B., Kwon, G., Kim, K., and Ye, J. C. (2024). Unpaired image-to-image translation via neural schrödinger bridge. In International Conference on Learning Representations.
Kingma, D. P. and Ba, J. (2015). Adam: A method for stochastic optimization. In International Conference on Learning Representations.
Korotin, A., Gushchin, N., and Burnaev, E. (2024). Light Schrödinger bridge. In International Conference on Learning Representations
LeCun, Y. and Cortes, C. (2010). MNIST handwritten digit database.
Lec, K., Liu, H., Ryu, M., Watkins, O., Du, Y., Boutilier, C., Abbeel, P., Ghavamzadeh, M., and Gu, S. S. (2023). Aligning text-to-image models using human feedback. arXiv preprint MrXi::230.112192
Léonard, C. (2014). A survey of the Schrödinger problem and some of its connections with optimal transport. Discrete & Continuous Dynamical Systems-A, 3444)
Léonard, C., Roelyy, S., Zambrini, J.-C., al. (2014). Reciprocal processes. a measure-theoreticc point of view. Probability Surveys, 11:237-269.
Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., and Le, M. (2023). Flow matching for generative modeling. In International Conference on Learning Representations.
Liu, G.-H., Chen, T., So, O., and Theodorou, E. A. (2022). Deep generalized Schrödinger bridge. In Advances in Neural Information Processing Systems.
Liu, G.-H., Vahdat, A., Huang. D.-A., Theodorou, E. A., Nie, W., and Anandkumar, A. (2023a). 12sb: image-to-image Schrödinger bridge. In International Conference on Machine Learning.
Liu, Q (2022). Rectified flow: A marginal preserving approach to optimal transport. arXiv preprint arXiv:2209. 145577777.
Liu, X., Gong, C., and Liu, Q. (2023b). Flow straight and fast: Learning to generate and transfer data with rectified flow. In International Conference on Learning Representations.
Masip. S., Rodriguez, P., Tuytelaars, T., and van de Ven, G. M. (2023). Continual learning of diffusion models with generative distillation. arXiv preprint arXiv:2311.1400
Mnih, V., Kavukcuoglu, K., Silver, D.,. Rusu, A. A., Veness, J., Bellemare, M. G., Graves, A., Riedmiller, M., Fidjeland, A. K., Ostrovski, G., et al. (2015). Human-level control through deep reinforcement learning. Nature, 518(7540):529-533
Neal, R. M. and Hinton, G. E. (1998). A view of the EM algorithm that justifies incremental, sparse, and other variants. In Learning in Graphical Models, pages 355-368. Springer.
Neklyudov, K., Brekelmans, R., Severo, D., and Makhzani, A. (2023a). Action matching: Learning stochastic dynamics from samples. In International Conference on Machine Learning.
Neklyudov, K. Brekelmans, R., Tong, A., Atanackovic, L., Liu, Q., and Makhzani, A. (2023b). A com- putational framework for solving Wasserstein Lagrangian flows. arXiv preprint arXiv:2310.
Nichol, A. Q. and Dhariwal, P. (2021). Improved denoising diffusion probabilistic models. In International Conference Machine Learning.
13

---

Parisi, G. I., Kemker, R., Part, J. L., Kanan, C., and Wermter, S. (2019). Continual lifelong learning with neural networks: A review. Neural networks, 113:54-71.
Pavon, M., Trigila, G., and Tabak, E. G. (2021). The data-driven Schrödinger bridge. Communicaatonns on Pure and Applied Mathematics, 74:1545-1573.
Peluchetti, S. (2021). Non-denoising forward-time diffusions. https://openreeven wtw.pp/nnrriii....eefortrrrm. net/forum? id=oVfIKuhqfC.
Peluchetti, S. (2023). Diffusion bridge mixture transports, Schrödinger bridge problems and genera-
tive modeling. Journal of Machine Learning Research, 24:1-51.
Perez, E., Strub, F., de Vries, H., Dumoulin, V., and Courville, A. C. (2018). Film: Visual reasoning with general conditioning layer. In AAAI.
Peyré, G., Cuturi, M., et al. (2019). Computational optimal transport: With applications to data science. Foundations and T in Machine Learning, 11(5-6):55-6077.
Pooladian, A.-.,, Ben-Hamu, H., Domingo-Enrich, C., Amos, B., Lipman, Y., and Chen, R. T. (2023). Multisample flow matching: Straightening flows with minibatch couplings. In International Conference on Learning Representations.
Rafailov, R., Sharma, A., Mitchell, E., Manning, C. D., Ermon, S., and Finn, C. (2024). Direct preference optimization: Your language model is secretly reward model. In Advances in Neural Information Processing Systems.
Rolly, S. (2013). Reciprocal processes: stochastic analysis approach. In Modern Stochastics and Applications, pages 53-65 67. Springer.
Ronneberger, O.,
Fischer, P., and Brox, T. (2015). U-net: Convolutional networks for biomedical image segmentation. In Navab, N., Hornegger, J., Wells, W. M., and Frangi, A. F,, editors, Medical Image Computing and Computer-Assisted Intervention - MICCAI 2015, pages 234-244- 244,21, Cham. Springer International Publishing.
Schrödinger, (1932). Sur théorie relativiste de l'électron et l'interprétation de la mécanique quantique. Annales de l'Institut Henri Poincaré, 2(4):269-3) -10..
Shi, Y., De Bortoli, V., Campbell, A., and Doucet, A. (2023). Diffussion Schrödinger bridge matching. In Advances in Neural Information Processing Systems.
Shi, Y., De Bortoli, V., Deligiannidis, G., and Doucet, A. (2022). Conditional simulation using diffusion Schrödinger bridges. In Uncertainty in Artificial Intelligence.
Smith, J. S., Hsu, Y.C.., Zhang, L., Hua, T., Kira, Z., Shen, Y., and Jin, H. (2023). Continual diffusion: Continual customization of text-to--maae diffusion with c-lora. arXiv preprint (tf.:2304.0602..
Sommerfeld, M., Schrieber, J., Zemel, Y., and Munk, A. (2019). Optimal transport: Fast probabilistic approximation with exact solvers. Journal of Machine Learning Research, 20( 1-0)):1--3.
Somnath, V. R., Pariset, M., Hsieh, Y.-P., Martinez, M. R., Krause, A., and Bunne, C. (2023). Aligned diffusion Schrödinger bridges. In Uncertainty in Artificial Intelligence.
Song, J., Meng, C., and Ermon, S. (2021a). Denoising diffusion implicit models. In International Conference on Learning Representations.
Song, Y. and Ermon, S. (2020). Improved techniques for training score-based generative models, In Advances in Neural Information Processing Systems.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. (2021b). Score- based generative modeling through stochastic differential equations. In International Conference on Learning Representations.
Su, X., Song. J., Meng, C., and Ermon, S. (2023). Dual diffusion implicit bridges for image-to-image translation. In International Conference Learning Representations.
14

---

Tamir, E., Trapp, M., and Solin, A. (2023). Transport with support: Data cota-iondiiinall diffusion bridges. Transactions on Machine Learning Research.
Thorton,, 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., Fatras, K., Malkin, N., Huguet, G., Zhang, Y., Rector-Brooks, J., Wolf, G., and Bengio, Y. (2024a). Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research.
Tong, A., Malkin, N., Fatras, K., Atanackovic, L.,, Zhang, Y., Huguet, G., Wolf, G., and Bengio, Y. (2024b). Simulation-fire Schrödinger bridges via score and flow matching. In International Conference on Artificial Intelligence and Statistics.
Vargas, F., Padhy, S., Blessing, D., and Nüsken, (2024). Transport meets variational inference: Controlled Monte Carlo diffusions. In International Conference on Learning Representations.
Vargas, F., Thodoroff, P., Lamacraft, A., and Lawrence, N. (2021). Solving Schrödinger bridges via maximum likelihood. Entropy, 23(9):1134.
Yang. K., Tao, J., Lyu, J., Ge, C., Chen, J., Li, Q.. Shen, W., Zhu, X., and Li, X. (2023). Us- ing human feedback to fine-tune diffusion models without any reward model. arXiv preprint arXiv:2311.15 2.1..3331.
Zajac, M. Deja, K., Kuzina, A., Tomczak, J. M., Trzciński, T., Shkurti, F., and Mitoš, P. (2023). Exploring continual learning of diffusion models. arXiv preprint arXiv:2303.15342.
Zhang. R., Isola, P., Efros, A. A., Shechtman, F.,, and Wang, O. (2018). The unreasonable effectiveness of deep features as perceptual metric. In Proceedings of the IEEE conference on Computer Vision and Puttern Recognition, pages 586-595.
Zhou, T., Wang, W., Konukoglu, E., and Van Gool, L. (2022). Rethinking semantic segmentation: A prototype view. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition.
15

---

A Appendix organisation
The supplementary material is organised as follows. First in Appendix B, we analyze an Euclidean counterpart to the a-IMF sequence identified in Section 3 and the associated flow. In Appendix C, we show that the Markovian projection can be recovered as the parameterisation of the vector field that minimises the accumulation of errors, extending the results of Chen et al. (2024). Theoretical results are gathered in Appendix D. In particular in Appendix D.1 we show that the proposed non- parramemetricc method coincides with the a-IMF and prove the convergence of the a- a-IMF. In Appendix D.2, we show the connection between the non-parametric and the parametric updates. In Appendix E, we provide more background on DSBM and propose an extension of the DSBM methodology. Consistency losses similar to (Daras et al., 2023b; De Bortoli et al., 2024) are proposed in Appendix F. Model stitching procedures are described in Appendix G. We comment extended related work in Appendix H. In particular we draw connections with Sinkhorn flows (Karimi et al., 2024), Reinforcement Learning policies, Expectation-Ne MatiiiisanMiiiiiiiiinnnn schemes following (Brekelmans and Neklyudov, 2023) and comment on finetuning of diffusion models. In Appendix I, we investigate the accumulation of bias in Gaussian setting and compare woward-fowward and forward backwwrdd methods. In Appendix J, we derive the preconditioning of loss following the principles of (Karras et al., 2022) in the case of bridge matching. Additional results and experimental details are presented in Appendix K.
B Euclidean flow and iterative procedure
In this section, we study a simplified counterpart of the Schrödinger flow and of DSBM in a Euclidean setting. The goal of this section is draw some conclusions in the Euclidean case which also remain true empirically when analyzing the Schrödinger Bridge problem.
We consider the set A {(z, y) E IR2 y 2 x} and the set A2 = ((x, v) E IR2 y S 0). Loosely speaking, one can identify A with the reciprocal class R(Q) and A2 with the set of Markov path measures projm- In that case, we have that for any (z, E IR, proja, ((I, y))
(( + y)/2, (x + v) (z,y) & A and otherwise, proja, ((x, 3)) = g). In addition, we have that for any (z.y) E IR? projng (1z, y)) = (x, 0) if (x, y) & A and proja, ((z, v)) = (z. v) otherwise. We consider the following flow (Ze, Vo 12) given by
(x 3t) = proja, (projag (()) - (ot, ve).
Let (ro, Bo) & A and (ro, & A2. Denote T the explosion time of (xt, Vr), i.e. for any t 2 T we have that (Ii; yo) = oo, where R {oo} is the one-point compactification of R2. Finally, denote < T such that for any & E [0, T]. (It,, Me) & Ai and (It, M=) & A2. Then, we have
Ox M1) (-21/2, x2/2 y1).
Hence, we have that Ie = To exp[- for any t E [ and = zo exp[-t/2] + (In expl - z0)/2/- Therefore, we get that
T T too and we have that for any
t t 0
x2 xu cxp[-t/2], Yt = x4 + 22(00 x0.
Hence, ((r,, ()) converges exponentially fast to (0, 0) with rate 1 1/2.
We now investigate the rate of convergence of the alternate projection scheme, i.e. the Euclidean equivalent of DSBM. We define ([In, vn))neN such that for any n E N,
(Ix+1, v1+1) proja, (proja, (Z Bn))) (zn/2,0).
Hence, we get that Vn xn2 T and therefore ((In, Mn))) NNN converges exponentially fast to (0, 0). Note that this procedure corresponds to a discretisation of the flow ((z,, yr))NNN with stepsize a = 1.
More generally, we define for any a E (0, 1]. ((1a) v ) such that for any n E N,
(1 vh+1) = aprojA, (proja, (4 v))))) + (1 - a= v/) T1
Hence, we get that Vn xn2 -" and therefore ((Ir,) V/x) EEN converges exponentially fast to (0, 0).
can be shown that for any n E N, o. = 21 a a/2)" and in addition, we have that
n
0 = (1 - v6 + an
Σ
- a) -/
k=0
16

---

Therefore, we get that
Un = (1 2(1 - -0 -1- - a1)" a//2)
We now analyse the complexity of the different discretisations assuming that the cost discretising the flow with stepsize i.e. | S E, we get total cost O = (l) the terms that do not depend on log(1/ (1/e). Hence, if Ca is constant then the choice a 1 is the best possible one in the range a & (0, 1]. Otherwise, one has to consider the ratio C* (l//((/(111//))) a//)), where the lower is the better. The flow procedure and the iterative based one are presented Figure 1.
(0, 1] is C In that case in order reach the threshold value where we have neglected
Based on this simplified Euclidean experiment, we draw some conclusions which also remain true in our setting, see Appendix K for more experimental details. First, we have that different discretisations of the flow yield different convergence rates. Large stepsizes incur faster convergence. This suggests to choose schemes with a E (0, 1) assuming that C < C in that case. To draw a parallel with our setting, in the case of DSBM (case a 1), we need solve the projection subproblem each step which incurs a great cost. On the other hand, one step of the online algorithm only requires sampling once from the model and performing one gradient step.
1. However, if the cost of choosing a 1 is too high then might turn to alternative
C Minimisation of errors and Markovian projection
For given non-Markore Man-MMaoovinn (stochastic) interpolant process (see definition below), there exist an infinite number of Markov processes admitting the same marginals (Albergo and Hunden-Eijnden. 2023). In this section, when it is well-defined, we show that the Markovian projection corresponds to the process which minimises an error measure (defined further) in case one has access the oracle of x; + E X, = zz).
Stochastic Interpolant. We first start by recalling the framework of Albergo and Vanden -ijjennEiinden (2023). Consider a coupling II between and one builds a (stochastic) flow between and π using the following interpolation procedure
X, = Interp; X, Z) = X,
3,,, + ZZ
(X.,
II, Z ~ (d,
where 01 Bu = ?1 = and ag 31 1. This defines a Mon-Markovian process. We denote by xe the induced unconditional distribution of X, Let us now consider the Markov process (Xç) given by
= = d + 3,X, + ( X, X
E
X No:
(14)
where (B:) (Br)e((011 is d-dimensional Brownian motion and is an additional hyperparameter. It can then be shown that (E()) satisfies that X; ne for all t E [0. 1]; see e.g. (Albergo et al., 2023, Theorem 2.8, Corollary 20) Hence (Xf)ee(o.1) a (stochastic) flow mapping To onto _ Note that (xE)e in (14) can rewritten as
dx;- = (át/ox) + (BL Brare X1 + (ir - -nôr/or -7/ X, = X Ed
(15)
In the specific case where ag- t. 3r and 2 o / /
then (15) becomes
X, = Interp, (Xo, X1, Z) = 1) X, + XX1 + ao 1(1 - -)Z.
This corresponds to the marginal distribution of the bridge associated with (, B.) In this case, (15) becomes
dx E X.) ojx, = x + dB
for ef = (2r náz/az) /0)) = 20 In Proposition C.I, will show that this choice of (Ex)rejo.) optimal in some sense.
Consider (X 1 given by dx;
X, x X, +5dd,,
17

---

with X6 ~ To and where Z(t,z) x) is an approximation of E|Z = z]. We have the following result.
Proposition C.1 (Optimality and stochastic interpolant); Denote Pe, respectively i. the path
measures associated with (x))e10.1) and (Xt)ee(0.1) respectively. Consider {(e) = (L (P:). Let = argmin ((e). Then we have
(Ei) = 27/ch - -27/ 27////a
In particular, at =t, 8 and v fra/t(1111 - then E.. V200. The value (#7 (e corresponds to Markovian projection when it is well defined.
Proof. We have that for any
(ir -2dd://(-_///2))))1) KTP(PF (T=) {o² E
where A, |2 X E X, x and the expectation is w.r.t. IPs. We have that KL(P*1P = 0 if s", which concludes the proof.
Proposition C.I is related to (Chen et al., 2024, Section 3.4). Therein it is noticed that, in the case of Augmented Bridge matching (De Bortoli et al., 2023), the choice of & does not affect the joint distribution of (Xi, ( The authors then select (Ex) so as to minimise an approximation error. They show that, in that case, they recover the Föllmer process.
We now show that Proposition C.I can be further strengthened to establish that e* is also the optimal value if we interpolate between ng and Ti, my and g. for any [0. and n, the distribution of X,, Consider in this context for any s,t E [0, with t 2 8, 1//%/ 2 2 as the following interpolation model.
X, (at/a,)X, + (Bt a,3./a,)Xi V a?72/aZZ.
(16)
where X, ~ Ts, X2 ~ Ti and Z ~ (O,, Id). Assume that a₁ = 1-111,7, 3r = and 7 = o /t1/t(1-0 t)
for any E [0, 1] then (16) corresponds to the Brownian bridge associated with (o.Bi) with endpoints X, at time 8 and X, at time 1. We have the following proposition.
Proposition C.2 (Stochastic interpolant with intermediate time points): Define (Xi, (X given by
dx; (ds/or)XF,, +E (B(( Brdr/os) X, X
(17)
+E [Htt,e / a / (2 1 Xx X Er, X Taa
with Nt. =t
o////003 Then for any |s 1}. X and X, defined by (16) have the same /v =
distribution.
Proof.
We let s E [0, 1] and X1, X, E Rd. From (16), we have directly that for any E [s, 1]
dX, = [(::/a.)X. X, + (st - (r.3./o.)) X1 + Jr. Zjdt, where Z NO0, Id). In addition, rearranging (16), we also have that
X, = (a./a,) - (a.3:/au - 33) - (1, ) Hence, by combining these two expressions, we get that
(āc/ dX, [(x//a))X, + (ß. dn.3//w,)X, + (fr. nt. (ās/ (r)) It follows that (Xr,))ee(.11) given by
dx, (ár/a,) + (3, B,0r/ EX X, = Xt,l + (fr, /t, E/Z X, x,,), X ~ Tas
such that for any E [s, 1] the same distribution as X, defined by (16). Then, we conclude similarly to (Albergo et al., 2023, Theorem 2.8. Corollary 20)
D
18

---

We now consider the following approximate version of (17)
dx; = ( + (3. 3 IE
x1 X
X=
[,8
+ (ft, / / (2 Z(t, X + Et, X Ta.
Similarly to Proposition C.I we consider the best choice of E to minimise the interpolation cost.
Proposition C.3 (Optimality and stochastic interpolant): Let s E [0, Denote PF, respectively IP. the path measure associated with (ee(o)) respectively (Xt,))e Consider ((E) [LP (P)) Let e" = argmin, (.)) Then we have
(=1)² 2 - -272 2r//./
In particular, E* does not depend on s E [0, 1] and for every $1, 82 E [0, 1] with $1 I 82. we have that (Xi,s, )eE(s2.1] and (x [c coincide.
Proof. Similarly to Proposition C.1, we get first that for any 8, [0, 1] with S
i Yt, 27,,0 ar/an.
(18)
Second, have that for any s, f [0, with 8 5
tt st 2t,5 2//0% -2á,ar /
(19)
Third, we have that
n 0./// = /n// à,
(20)
Combining (18), (19) and (20), conclude.
D Theoretical results
In this section, we prove the main theoretical results of the paper. In Appendix D.I, we first prove the convergence of the a-IMF sequence, i.e. we prove Theorem 3.1. Second, show that the con-parametric updates (8) correspond to the a-IMF sequence, i.e. we prove Proposition 3.2. In Appendix D.2, we link the non-parametric updates the parametric updates.
D.1 Non-paraeetrie sequence and convergence
Let Q E P(C) be associated with (VeBe)eejo.)) where (Be) 0.11 is a d-dimensional Brownian motion and E > In this section, we abuse notation and denote P(C) the set of path measures which will consider Q P(C) associated
not necessarily probability path measures. In particular, with (VEB.) with Qo = Leb. In that case, the Leiblac divergence still well-defined and we refer to (Léonard, 2014) for more details. We recall that have defined (IPm "r" for any n E N and a E 1] by
P projom(
po++1 = -1 a)F" + oprojæ (proj (IP))
In addition, for any n E N, E [0, 1) and x E Rdd we have defined
(x) (x)) dn V Lo
where
CF(v,, = /
1
|v((xe)
F1 Iz
2
dFo, x1 (2)
(21)
2 (Rd)3
-
1
(x(x)))
T1 - It 1-t
dproj (P)
2
(Rdj3
We define (P.z )nEN associated with (27), where for any suitable vector field v, F, is associated with
dX, = (X,)dt + VEdB,,
19

---

where (Bo)re(o.1] is d-dimensional Brownian motion.
In order to rigorously prove Proposition D.1 detailed further, we introduce P2(C), such that IP E P2(C) if P E P(C) and for
/ (Rey2
+ o, < +o0.
Note that if P E P2(C) then we have that for any E [0, 1]
f | dprojx (P), +o0.
In addition, we recall that ø E (m) for μ P(RR) o R + Rd and
and
(d())))))(()) < +00. [ Rd
Finally, we define
A [(0,P) PEPI(C), & (P))).)
Then for any E [0,1), 1), we define L. A R given for any (v, P) E A2 by (21).
Proposition D.1 (Non-parametrie updates are a-IMF): Let E (0, 1]. (P",,P"")me as in (4), 6n and n² (1 a:)lPm + axroix (P*). Assume that for any n E N, Pun is well-defined. Then, for any n e N, Pu P
Proof. First, we have that for any t E [0, 1), U, A2 and EE (P,,) we have
+ + P) = Le(u,, PP) +e
(( 12093
In -== ydPoo. (ro, xn (xs
+ (62/22
10(x:))1
(R4)s
= L.(r, +e (0(x)), vx(2))
(3
-(/ mndproia (P) (1)) - xz ) - -1)d-)i))
()) (P),
+ (s2/2) / [[((=:))1 (2)))
Hence, we have that
Pe = (xx (Eproin (7) X1 X, = xz) - x2) /(1 - (Po))
(22)
Assume that for some n. E N we have that for any L E [0.1) and æa E Rd,, we have v/(x)) (Epe [X, X, = æı] xx) (11 - t). We are going to show that for any E [0, 1) and x, E Rd, we
have uf (It) = (Ep.-1 [ X, = = 21) (1 - t). For any t E [0, 1) and xt Rd, we denote
(xx) 5, (dprojæ ( )e/dp())
Since have that 6n a and n² (1 a) + aprojx (P"), we obtain for any [0, and It Rd
82 aldproiR (Pd)) (PP)),/(((1 .PP aproiR (z)))
(23)
so that
1 1 (xx) (1 a) / PPPPP aprojR (I")) (zz))
(24)
20

---

Therefore, combining (8) with (23), (24), (22), we get that for any t E [0, 1) and E R
of (xe)) = (1 -- (xx)) (xr)
+ 0f (x)) [Eprij (Pm)
[1 X x1] - xz) ((1-t) -t)
(1 6, (1x)) (Bpp [X|
X, = xt] xx) ((1 - 1)
+ 8% (axa) (p (Pm) [1
X, = =d] - m1) (11 1)
( 2idli 1/1 (z) (xx)) 2z /(1 /(1-t) -t)
(1 5, (xx))
Rd
+ 0g (m1) (fg x xidprojæ (P") )1p (11 (x)) xe ) /(1
(0) - w/)/(1 - -)d[(1 - 8ن (z, + 55 (2x))rooir (P)))(((21
2]d[(1 - a)) aproja (Pm) (111((21
Hence, we have that for any t E [0, 1) and E R, v.t (x,) (L [1 X, = z1l xz) /(L
t). Since, for any [0, 1) and De E Rd, ol (xx) (Eço [x X, = myl mm/ (1 - by definition,
we get that for any n E N, t E [0. 1) and 22 Rd, oz (2+) (Epp. [X X, = =+) mr) (11 t).
Using, Definition 2.2, we get that P.- projm (IP*), which concludes the proof.
Before stating our convergence theorem, we show the following result which is direct consequence of (Léonard, 2014, Theorem 2.12) and (Léonard et al., 2014, Theorem 2.14). We recall that the differential entropy of a probability measure π is given by
H n
=--==-/
-
log((d=//decb)
RR
if i admits a density with respect to the Lebesgue measure and +oc otherwise.
Lemma D.2 (Characterisation of Schrödinger Bridge): Recall that Q is associated with (VEBe)eejoa) and assume that Qo Leb. Let no, Ti P(Rd) such that
[ [a]12dm(()) < +00,
H(m)) < +00;
for i E (0, Let IP* such that P is Markov, F E R(Q), P To and P TT. Then IP* is the Schrödinger Bridge, i.e. the unique solution to (2).
Proof. First, we have that Q is equivalent to Leb @ Leb. Indeed, we have that for any 2o, 21 E Rdd
(dQo, @ Leb))) x1) = (2Tz) 4/2 expl- 1200 - x1 (2:)
Similarly, we have that for any E (0, 1) and a, E Rd,, Qo,1((-12x;) is equivalent to Leb @ Leb. Indeed, we have that for any E (0, 1) and To, Fı, 21 E Rd
(dRonn((():))((Lee)) @ Leb) 21) = (2Tet) d/2 expl- || 22l12/(2et))
x (2TE(1 t)) t) expi x1 112/(2(1 t))]
Hence, for any E (0, 1) and He E Rd, Qo (2s)) is equivalent to Qo Since P is Markov and T E R(Q) we get that there exist 99 and yp which are Lebesgue measurable such that for any C we have that
(dP" (dQ)()))(w) (i(w)) )p (w1).
(25)
Second we verify that the conditions (i)- of (Léonard, 2014, Theorem 2.12) are satisfied. First, @ is Markov and hence (i) is satisfied. Then, (ii) is satisfied since for any E (0, 1) and x; e R4
21

---

Q0,1je(-(ar) (1re)) is equivalent to Q0 We have that Qn = Q1 = Leb and (iii) is satisfied. We have that for any 2o. _
(dQo.1/((Leb @ (2o, x1) = -4/2 x xz11 -
2 (2Tz) ( cxp1- 1|= /s]]
Hence, (iv) is satisfied and we let A Rd addition, we have that for any Tu. 11 Rdd
+ Rt be given for any E R by A(x) = |</18/5. In
f expl expl-1|=0 /
x1) < to0.
Hence, (v) is satisfied and we let assumption (vi) and (vii) are satisfied. We conclude the proof upon using (Léonard, 2014, Theorem 2. 2b)) and (25).
R2 R+ given for any Rd by B(x) ||x||//s. By
We are now ready to state our main convergence result.
Proposition D.3 (Convergence a-IMF); Let (0, 1] and (IP (Pm))ew defined by (4). Under mild assumptions, we have that 10 Ip = IP*, where IP* is the solution of the Schrödinger Bridge problem (2).
Proof. Using the convexity of the Kullback- divergence with respect to its first argument (sec e.g. (Dupuis and Ellis, 201 1)). the data processing incquality (see e.g. (Ambrosio et al., 2008, Lemma 9.4.5)), the fact that the Schrödinger Bridge is Markov and in the reciprocal class of Q (see e.g. (Léonard, 2014, Theorem 22) and (Léonard et al., 2014, Theorem 3.2)), and the Pythagorean theorem for the Markovian projection (Shi et al., 2023, Lemma 6), we have that for any n E N (Tp(t=+ pP*) KT.((1 Tr oproja (
s (1 a: KL PT) + a a
(proj (P*
(1 KL P*) + aKLL
J,, IP,.))
s (1 a) a))KL(PM PT) + (KL (P*)) [IP))
$ (1 a) KL (T)) (P*)) - aKL (P"" I projm((P))
(26)
Therefore, we get that Hence, it follows that
aKL(f" |proo M (P²)) s KL(P"MP KL
Σ KT. (P projj (Tm))) s 2KT. (P9 +00
BEN
So we obtain lim KT KT |rrjj (Pm) = 0. In addition, using (26) we have that KL(P""IP*) S (L too for all E N. Using (Shi al., 2023, Lemma 6), we also get that (L (Pm) I[P)) < KL(P" < too for any E N. Hence both the sequences (M (nE") and (P") EEN = (proj ))Ee are relatively compact in P(C). Let PP E P(C) be an
adherent point to the sequence (TPm) NEN and y N + N increasing such that lim P.P(n)) = P. Similarly, let Φ: + N increasing such that (ó(n))nEN
subsequence of (((()))EN such that lim proj (fPm) = P., with IP' and adherent point to the sequence (proj.n (P*))) Using the lower semi-continuity of the Kullback-Leeille divergence in both arguments (Dupuis and Ellis, 2011 1), we get that
PL
lim inf (I(0)())) proj 4-00000
0.
Since the set of Markov measures and the set of reciprocal measures w.r.t. Q are both closed, we have that IP is Markov and in the reciprocal class of Q. Since we also have that P. = wo and P: = W2, we get that P = p* using Appendix D.I. Since every adherent point of (fPn) ) is PA, we have that lim, oo fmm- = P*. Similarlyy using that lim KL( |proj (fP)) = 0 and again the lower semi-continuity of the Bellback-Leibbes divergence in both arguments, we get that every adherent point of ((P)) is JP. Hence, we have that lim pa PP. which concludes the proof.
D
22

---

D.2 From parametric to non- non-parametric.
In this section, we show that the parametric updates considered in (9) are preconditioned version of the non con-aaametrri updates considered in (8). We first recall the non con-prametric loss
!. Lo(ve, P)dt
1 [. S |cller)) 2
I1 3t ||
(uu, P)
1 1 (p)o)R(P))) (P),, x1 ddf
and the parametric loss
31 -1- "²1" 2 Jadxad (l((z))
1 2/2 J
L(O, =
1- dproiRic) (1,,
The non-parametric sequence (v")nen is given by (27), i.e. we have for any n E N, L E [0, 1] and x E Rd
(x) (x) ón V Li(e, (r))
(27)
Similarly the sequence of parametric updates is given for any n E N, [0, and R by
pn+1 = On - a V (0,, ).
We recall that Ipor. is a stop gradient version of IP a We are going to show that on average the parametric algorithm yields direction of descent for the non-parametric loss. We assume that the set of parameters θ is an open subset of IRP for some E N. For any E [0, 1] and x E IR we assume that 0 + vf (x) is twice continuously differentiable and denote Dovf (z) E Rdd its Jacobian and Di (x) its Hessian. For any 0 O, we denote
hg = VoL(O,P").
We show the following result.
Proposition D.4 (Velocity field parametric update): Assume that there exists C such that for any 0 E 2 and x E R
['1- DD33) ashe (no, ho)ds C,
(28)
where hg = V., We have that for any n E N, t E [0, 1) and x e Rd
8n
(x) v0- (x))
aDo On (x) So' /
Van PPn. (+)d o(a),
where μ² = (1 am + aP,on.
Proof. First, have that for any μ P(C)
VoI.(0,P')) - [// Dou (v (v)( (x1 xx)/(1 s))dproiR (P )=, (1),
Therefore, using (22), we get that
VOL(O,PM) - ((¹ / (v. (x..))
Let f E θ and denote 0' = 0 a VoL(O, IPO), Using Taylor expansion, we get that for any 0 O, we have that
8' (x)= (x(x (D) /'' / (xu((xx) T V Cu, (xx)du (x,)ds
+ aa² f (1- S asshe (2)) hojds.
23

---

Since the functional gradient is not applied on the second coordinate, we can drop the stop gradient operator and therefore we have for any 0 θ
of (x) (x) aDouf (x) [//o (xx) Vu /
+ + [ (1 D - (ho, ho)ds.
Combining this result with (28), conclude the proof.
The corresponding update on the velocity field is given for any n E N, te e [0, 1] and r e IRd by
(x) -aDu (x) L. Po. (G) +(a))
We immediately have the following corollary.
Proposition D.5 (Parametric direction of descent): For any E N, if
LC [
(o)d 40,
then we have
[ (V La (o. S. a a-
lim
(())
E
Background material on DSBM and extensions
In this section we recall some basics on Markovian and reciprocal projections in Appendix E.1. We explain the link between the concept of iterative refinement and Schrödinger Bridges in Appendix E.2. Then, we briefly present Diffusion Schrödinger Bridge Matching (DSBM) (Shi al., 2023) in Appendix E.3 and propose some new extensions in Appendix E.4.
E.1 Markoy and reciprocal projections in practice
In this section, recall the definition of the reciprocal and Markov projection. We provide more details on how these different projections can be performed and illustrate them on simple examples.
Markov projection. First, recall the definition of the Markovian projection.
Definition E.1 (Markov projection): Assume that Q induced by (VEBt) for e > 0. Then, when it is well-defieed, for any P. R(Q), the Markovian projection M proj M is the path measure induced by the diffusion
dX; = of (X VEdB,,
oi (x,, (ETT) X1
X, = x.] - x1) / (1-t).
X ~ Pa.
In Figure 7 and Figure 8, we illustrate the effect of the Markovian projection, following the example of (Liu, 2022). We consider two distributions To and πı such that
1 -= N 2d) + + ([(-2,)) क = ===NN0
N(2. + [ (1,22) Id).
To =
In Figure 7, we display samples from the distributions mo and क well trajectories from the path measure IP = (To @ π ) Practically, this means that we sample X, and X₃ I independently and then consider Brownian bridge between Xo and X,. The SDE associated with the Brownian bridge with scale ε > 0 is given for any & E [0, 1] by
dx,- = (X1 x1)((1 - t)dt + VEdB,.
(29)
Note that the measure P (To X oo. is in the reciprocal class, i.c. I R(Q).

---

Stochastic process
22
Figure 7: Samples from the original distributions To (left) and Ti (right) are shown in red, while sample paths from IP (To @ Ti Poo, shown in blue.
Next in Figure 8, we display samples from the distributions To and क as well as trajectories from the path measure F (P)) Note that in Figure 8, contrary to Figure 7, observe less crossings between the trajectories. Indeed in the limit case where - 0 the Markov measures IP* is an ODE with regular coefficients and therefore admits unique solution for every starting point in the space no crossing is possible. In particular, note that most of the trajectories starting from the upper-left Gaussian end at the upper-right Gaussian. Similarly, most of the trajectories starting from the lower-left Gaussian end at the lower-right Gaussian.
2
Stochastic process
Data
2
Figure 8: Samples from the original distributions are shown in red, while sample paths from M Mr are shown in blue.
In practice, computing the Markov projection involves finding the optimal drift of. This optimal drift the minimizer of regression problem, see (Shi et al., 2023) for more details. Hence, computing the Markovian projection requires training neural network to define a vector field.
Reciprocal projection.
First, we recall the definition the reciprocal projection.
Definition E.2 (Reciprocal projection): P E P(C) is in the reciprocal class R(Q) of Po.1 We define the reciprocal projection of P(C) IP* proj (P) = Po. We will write projæ instead of projR(g)) simplify notation.
To sample from P++ = proj (P), we only need to sample (Xo, X1) Po., and then to sample from the Brownian bridge conditioned on (Xn, X1). This means that in order to sample X P.. we only need to sample (Xo, X1) ~ Po.1 and then compute
X; (1 +X + Vet(1 t)Z,
(30)
25

---

with Z ~ (O, In particular, sampling from I = projR (IP) does not require training any neural network. However, practice, in order to obtain samples (Xo, X1) ~ we have that IP associated with an SDE and therefore obtaining (Xo, X1) requires unrolling the SDE associated with P. In Algorithm 1, the measure IP is associated with an SDE with parametric drift e..
In Figure 9, we continue study of the example of (Liu, 2022) that we used to explain the concept of Markovian projection. We consider the path measure M obtained as the Markov projection of P (Tu @ Ti 200 In Figure 9, display samples from the distributions To and πa as well as trajectories from the path measure P = Mn. Q In order to sample from I' we first sample Xx, X1) ~ M0.1 This involves unrolling the SDE associated with M. Once we have access to samples (Xo, X1), draw trajectories from the Brownian bridge following the SDE (29). We can also sample from any time without having to unroll the SDE (29) by simply sampling from (30).
This is what is done in Algorithm 1.
1-
Stochastic process
Data
-1
Figure 9: Samples from the original distributions shown in red, while sample paths from F projæ (M) are shown in blue.
E.2 Iterative refinement and Schrödinger Bridge
The Schrödinger Bridge problem (2) can be solved leveraging techniques from diffusion models and bridge matching. De Bortoli al. (2021); Vargas et al. (2021) consider an alternating projection algorithm, corresponding to a dynamic version of the celebrated Sinkhorn algorithm. Peluchetti (2023); Shi et (2023) introduce the Iterative Markovian Fitting procedure which corresponds to perform an alternating projection algorithm on the class of Markov processes and the reciprocal class of the Brownian motion. can be shown that the solution of this iterative algorithm converges to the Schrödinger Bridge under mild assumptions, see (Peluchetti. 2023; Shi et al., 2023). We highlight that the case where e + then DSBM is equivalent to the Rectified Flow algorithm (Liu al., 2023b). One of the main limitation of those previously introduced procedures which provably converge to the solution of the Schrödinger Bridge problem is that they rely on these expensive iterative solvers and requires to consider two networks, one parameterising the forward process To - I and one parameterising the backward I - To.
E.3 Diffusion Schrödinger Bridge Matching
Diffusion Schrödinger Bridge Matching corresponds to the practical implementation of the Iterative Markovian Fitting procedure proposed in Shi et al. (2023); Peluchetti (2023). The IMF procedure alternates between projecting on the Markov class M and the reciprocal class Rq. In what follows, denote M p22++1 . and II p2nn E R(Q). We also recall that Q is (rescaled) Brownian motion associated with (ooBe)rejo.1) and that therefore sampling from Q (+|2o, X1) corresponds to sampling from
dx, (z) X.)((1 t)dt oodB
X 20.
We recall that the main computational bottleneck of the DSBM lies in the approximation of the Markovian projection. Indeed, using (Shi et al., 2023, Definition 1, Proposition 2), have that
26

---

M = M (IT) is associated with the process
dx;- ( x x,) t)dt + addB,
X no-
We also have using (Shi al., 2023, Proposition 2) that M' can be approximated using M given by dx.- = vg- (1, x,) + addB, En,, (II(X] - X.) - - vo, Jdd,,
X ~ To:
(31)
0- argmingee [
(32)
where {uo } a parametric family of functions, usually given by a neural network.
Hence, since we can approximate proiR(q) (M) and proj we can approximate the IMF proce- dure. This is the DSBM algorithm introduced in Shi ct al. (2023); Peluchetti (2023). We describe the first few iterations. Let IT Mg. 1jo, where N1g To. Iīς #1. Learn M. x pr given by (31) with vg- given by (32). Next, sample from II2 projR(Q)) (Mf+) = My. Qjo.1 by sampling from M and reconstructing the bridge Q Upon iterating the previous procedure, we obtain a sequence (II", M )EN. To mitigate the bias accumulation problem caused by approximating only the forward process, we alternate between a forward Markovian projection and a backward Markovian projection. We give more details on the advantage of using forward- backward parame- terisation instead of a forward in Appendix I. This procedure is valid using (Shi et al., 2023, Proposition 9). The optimal backward process is approximated with
dY, us. (1 4, Yj)dt + oodB,,
Yo T
(33)
[
Φ" argmin
Eno.. (Xlo -Xo X,)/t us d.).
(34)
We recall the full DSBM algorithm in Algorithm 2.
Algorithm 2 Diffusion Schrödinger Bridge Matching
1: Input: Joint distribution IT. tractable bridge Qjo,1 number of outer iterations N N.
2: Let II 11, Q
3: for {O...... NN
N - 1} do
4: Learn Ug- using (34) with II = 112",
5: Let M be given by (33).
6:
Let 1[ = 2 Q 0.11
7:
Learn ug. using (5) with 11 = 112
8: Let M2 be given by (31).
9: Let 2 M Q,
10: end for
11: Output: ugr, U
E.4
Reflection-projection extension
First, we consider reflection- method similar to the one investigated in Bauschke and Kruk (2004), We recall that the DSBM algorithm is associated with a sequence (P"")en such that for any E N
pn++1/2 (roj
pP++1
projR(o) (1mn+1/22..
In a dnleccion-projecttion scheme, one of the projection is replaced by a reflection. As noted in Bauschke and Kruk (2004), this can yield faster convergence rates in practice. We consider the sequence (IP") ne"") such that for any N
p=+1/2 1/2 proj M (P"),
p=+1
(r)j 12/2+1122_ pT).
(35)
In what follows, we make the assumption that 2p-+1/22 P a probability measure, even though is not clear if this path measure is non-negative. However, even by making this strong assumption, we show that we can recover DSBM in Algorithm 4. By considering a relaxation of the reflection- projection scheme.
27

---

First, note that for any N. Pl is Markov, see De Bortoli et al. (2023) for instance. Hence, we assume that In is associated with
dx; = Xn) + oodB,,
X ~ To.
(36)
7++1 assuming that we can sample from In and 1p++1/22 Estimating pp++1, First, we compute 62
Since put1 is in the reciprocal class, have that pu+1 is associated with
dX,
(Bpn-1 1|0.t
[X] Xx, Xo] - ())(111)de + oodB,
We refer to De Bortoli et al. (2023) for a proof this fact. Hence, using (35), we have that
argmin, ('' / xd
x+1
20) - (x1 x1)/(1 - +) d U,1,1 (2o, xt, 21)
0,0
/
[ JRU 10,o(t, 20, (2] xr) - ))] ar dQ (0
argmin,2
% S (t,o0t,z,,,00) - (x1 z1)/(1 dF dQ (2t 20,
(37)
Next, we turn to the estimation of [Pn+3/2
Estimating 3n22 Next, we assume that for any n E N, p/+ is associated with
dx; of (X,)dt + oodB
Xo ~ To
Using (36), we have that of a+1 given by
u argmin, S/' See [lve(t,a)) not (t, zt 20)1 ?dP (zo,
We note also that using (Shi et al., 2023, Proposition 2) and the fact P" is in the reciprocal class of Q also have that
/ I [lee(t,ze)) - (x1 - (11)/11- +)]
of m+ argmin,,2
x1)
L/' S [1((t,2)) - -0x s1)/(1 (d,,)
(38)
Hence, assuming that we can sample from Pa and pn++1/2 then we can estimate a.7 and of wn+1 i.e. sample from pa+1 and pn44/22 Note that the losses (37) and (38) only differ by the conditioning with respect to the initial condition To and therefore the optimisation can be conducted in parallel.
We are now able to propose the following projection-reflectioo algorithm, see Algorithm 3.
28

---

Algorithm 3 Reflection Diffusion Schrödinger Bridge Matching
1: Input: Vector field and conditional vector field of and v0,0- noise level dg and associated bridge
Q number of outer iterations N & N, batch size
2: for n E NN N
1] do
3: while not converged do
4: Sample X
5: Sample X 1 using dx!: 1i (X B:: dt + aodB;
6: Sample | using dx 1: = (
1:7
I: _ 1: )t + dodB;
7:
I Σ B
R= |
(x
x1_X"11 -
E k=
- ||
1-t
I-t
8:
Lo = fo |> k= 10,0 (X xb) 1-t
xixxi
12 (1/2) Σ (X ( x ||
9:
a+1 Gradientstep(4 L)
10:
n+l = Gradientstep( Co)
end while
12:
end for
Output: of N+1 N+1
2.0
Note that in Algorithm to Algorithm 2, one can construct forward backward extension to alleviate some of the bias accumulation problems. Finally, we can interpolate between DSBM and this new reflection algorithm and DSBM by introducing an hyperparameter 2 0 and consider the following extension given in Algorithm 4
only consider the optimisation of a forward process but similarly
Algorithm 4 Reflection Diffusion Schrödinger Bridge Matching
1: Input: Vector field and conditional vector field of and eY,,- noise level and associated bridge Q number of outer iterations N N, batch size
2: for {0, NV N
1} do
3: while not converged do
4: Sample X1:
To
@B
5: Sample X 1:/ using dx! 1.B = of (X. IB )d + JadB.
6: Sample 1 1:B using dx 1:B = = (x:
1:B _x] 1:B dt + oodB:
7:
fo' + B
k= |/ (X)
x 1|2- Σイ= | (X,)
x -x 11
8:
Lo = S P |0,0 (X) Xb)
x
1-t 1 a EB=
|1,,, (X. Xb)) - 13
9:
o+ Gradientstep(4. C)
10:
0,0
Gradientstep( Co)
11: end while
12:
end for
13:
Output: 1/ 200
N+1 N+1
Using different values of a 2 0 Algorithm 4, we recover different existing algorithms. If = 1, we recover DSBM Shi et al. (2023). Finally, if a 1/2, recover the reflection algorithm Algorithm 3.
F
Consistency in Schrödinger Bridge
The idea of training both the forward and the backward jointly was mentioned in (Shi et al., 2023, Section G). However, was still assumed that, while being trained jointly, the forward and backward vector fields were obtained using an argmin operation, see (Shi et al., 2023, Equation (43). (44)). In addition, in (Shi et al., 2023, Section G) a consistency loss was proposed in order to enforce that the forward and backward processes match, see (Shi et al., 2023, Equation (49)). In this section, we leverage new results from Daras et al. (2023a); De Bortoli et (2024) in order to enforce the internal consistency of the model.
29

---

First, note that for any (Xt) (Xt)te(a.1 associated with P E R(Q) we have for any 0 0 St S
that
t- to
+ 1 -t
X
X
X + t0,t, Z,
Z I(0,
t --to to
tı -
where
(1 (0)(61 4 to
to.1,f
Let Pt be the density of X₁ with respect to the Lebesgue measure, we have that for any 0 1 to Sts S 1 and E Rd
Pa(x))
(.. to.
2/2 -d/2 exp
(
||2
41 to It,
11 stol?
Peo (Ft,)Pt, (2t, )d. dret
20 20,00
Using the change of variable Tto - Ita + De and 24, + 21, + xs we get that for any 0 $ to S <t S and x E Rdd
Vlog
( (Itt) log Pt, (to,)) (Tta Tt, (x))d
(39)
JRAxR4
This identity for the score has already been presented in bridge matching context in (De Bortoli al., 2024, Section 3.3). Let P E R(Q) then we have that ( is such that for any
Ft. see (Shi et al., 2023, Proposition 2). We have that for any / E [0, 1]. (7(x) + (2) of V (xr). Combining this result and (39), this suggests considering the
[0. 1]. p
following consistency loss
(tons.(to.dd1)) Bl[lva (1, 1, x,) + no 0,
(40)
(to: 1, X,) 20 to. xto) vo(1, 1, X,,,, vo(1 -1, X, )13]-
Similarly (13), we can consider an empirical version of (40).
G Model stitching
In Algorithm 1, the finetuming stage requires a pretrained bridge matching model interpolating between To and πi (lines 2-7). However, for large datasets with complex distributions To and πı e.g. ImageNet, training this bridge model from scratch can be computationally expensive. To improve efficiency, we can leverage existing diffusion models targeting To and æT. Specifically, we assume access to generative models transferring between (0, Id) and To, and between (O, Id) and TT In the rest of this section, we show how one can adapt Algorithm I to this setting. We then comment on the link between the proposed algorithm and Dual Diffusion Implicit Bridges (Su et al., 2023).
Setting. For simplicity, assume that we have two pretrained diffusion models for To and T.. We describe our procedure for To. Consider a forward process of the form X, = Xo + o,Z, with (0, Id), where ae hyperparameter. Note that we could have considered an interpolant of the form X, X, + ag Z instead, see Song et al. (2021a) for instance.
We assume that the model X; X Z is associated with the forward diffusion model
dX; gd
(41)
where assume that 9. 2 for all E [0, 1]. Note that we have that for any t j0, o? So' os In particular, we have that for any 8, E [] with <
X, X, + oo oZ
Z ~ (O,
Our goal is to solve the following Entropic Optimal Transport problem
{// (
I argminneP(RRa))
c(æ, v)dl v) cH(II) HI = Tp, II, m1
(42)
where E >> is some entropic regularisation. We assume that æ > 0 is fixed and assume that there exists [0, such that of =/2 We now consider dynamic version of (42) with
1p* argminpe (KL())))) Po no. Pr m1).
(43)
30

---

where Q is associated with (o,) (41). Note that contrary to the setting presented in the main paper, here we do not consider the integration between time O and but between time 0 and t'. It can be shown that for any E [0, t'. (X+)cefo.r) associated with Q e is given by
X, Interp, (Xo, Xx, Z) (1-
oi df
oi Xx
of Z, of
Xo +
? of
or
z Z~ N0,,0d.
Solving (43) is equivalent to solving (42). We now propose an algorithm to solve (43). corresponds to the finetuning stage of Algorithm 1 with specific initialisation, similar to DSBM-IPF in Shi et al. (2023).
By Ug. we denote DDM model associated with π.:
dx, = ctl X,)dt + g,dB,, Xo Id,,
X, ~M.
(44)
Similarly, denotes a diffusion model associated with Tn:
dY, = (t, ; Y dt + g.dB,, Yo ~o,,(0,11) Y1 ~ no
(45)
In analogy to Equation (11), the two equations above correspond to the forward and backward SDEs. For given batch of imputs X. and X|:B, timesteps t (ni =1)@8, and interpolations X and X7, we compute the empirical forward and backward losses as the following modification of Equation (13):
1
Σ
Us (t' x
(xi - x /0:l?
B
B
F00tt,Xo; X
Σ
Dog (t', X
(X X / / - -
Algorithm 5 corresponds an online version of DSBM-IPF (Shi et al., 2023) with the initialisation given by two generative models. Algorithm 5, we finetune the trained vector fields to solve the interpolation task. At inference time, the SDE associated with vector field ey interpolates between T + To. while the SDE associated with the vector field Us interpolates between To - ø1-
Algorithm 5 a- a-Diffusion Schrödinger Bridge Matching for DDM finetuning
1: Input: datasets Ty and T1, number finetuning steps Ninetuning. batch size B, DDM parameters
④ and 0.
do 3: Sample (X,, ~ (ro & ) ~ Unif([O, 1J), Z ~ NNO,
2: for E
4: Sample X;, by solving (44) starting from X.
5: Sample X3, by solving (45) starting from X;
6: Sample 1 Uni t7) Zt [0, Id) and compute X Interp,, (Xu, Xx,,Xt, Z*)
7: Sample 1a ~ Uni ))) Z ~ (00 Id @B Id)@B and compute X; = IItterp,- (X,, X2, Z*)
8: Update 0 with gradient step {+ 1f, Xo, X
9: Update á with gradient step on (d: 8 X
end for
11:
Output: 0, e parameters of the finetuned models
Our model stitching approach is related to Dual Diffusion Implicit Bridges (DDIB) (Su et al., 2023), which uses pretrained diffusion models, but without further finetuning. As highlighted in Shi et al. (2023), DDIB is inferior to DSBM in terms of quality and alignement of the samples.
31

---

H Extended related work
We highlight links between our proposed flow and Sinkhorn flows in Appendix H.1. We draw connection between our practical approach and Reinforcement Learning in Appendix H.2. We discuss how a-IMF related (incremental) Expectation- Expectation-Maxxiisattoo in Appendix H.3. Finally, we discuss how our algorithm can be seen
instance of continual learning in Appendix H.5.
H.1 Links with Sinkhorn flow
In this section, we discuss the links between our approach and the Sinkhorn flow introduced by Karimi et al. (2024). We start by recalling how Sinkhorn flows are defined and then discuss how they are related to our approach.
Sinnkhorn and Sinkhorn flows. We first consider the static EOT problem and recall the Sinkhorn procedure, also called Iterative Proportional Fitting. We define a sequence of coupling (II", ).EN. Le. for any n E II E ( Rd). We let I7 Q0,1 and we consider for any n E
11 = arrgiin(KLQ) II") II E P(R x IR,,, 1o =ohh
(46)
11+1 (KL I1") 11 E P(R RR), 111 m
In Karimi et al. (2024), the authors generalise (46) by introducing an extra hyperparameter E (0, and defining
IT" = argmin (KL [TT) II P(R x Rd), II, = =ihh
(47)
1 argmin(7KL(IT II") + (1 > IIn). II P(R x Rd), HI = Tob
Using (Karimi et. 2024, Lemma 2), have that for any y E (0, 1], any and any 2o. 21 R
(dīi (di v1) exp fxp (2r) + 9 (x1));
(48)
with 12 oº =- 0 and for any N. 7 (0. 1] and xι Rd
9 (T)) = (x1) /dz)
(49)
In addition, using (Karimi et al., 2024, Equation (9)) we have that for any N. (0, 1] and
To E Rd
log (f
(2)) =
- (1/
Rd
When letting 7 + 0, (48) and (49) suggest to consider for any 2 0, 20. 21 E RRd
(dif" Idd' = explff + g²
II" = argmin (LL II') II II, = =11h
where for any 0, Rdd
0.f%(zo) = / log(dīi; /dm (zi (dll((11aaa
(z1) = Log(dii;/dm))(a)
Comparison with Schrödinger Bridge flows. In order to compare our approach with the one of Karimi et al. (2024), we start by rewriting the 7-Sinkhorn algorithm defined by (47). To do so, we introduce the projection on the measures with fixed marginal.
Definition H.1 (Projection on marginals): Let 11 P(Rd x Rd) and To E P(Rd), we define proju. (II) as follows
projo. (II argmin(KK((i II) Ne P(R R Hi, ma}.
Similarly, for any II E ((R x Rd)) and Tı E P(RR), we define proji.s 1. (II) as follows
proji.m (II) Drgmin/KL./e II) ñee P(R Rd) i1 m
32

---

7-Sinkhom
7-IMF
Loss function
KL(UI p (II) KL ITT
So Eproircos (2) [11 (X,,
x, ||a]de
i-t
regularisation
So Ju (x))()))))))) )
Update
Implicit
Explicit
Table 2: Comparison between y-Sinkhorn and 1-IMF
With these definitions, have that for any n N, I1m proju, (proj1.z, (II")), with (II") "eeN the original Sinkhorn sequence defined by (46). Similarly, we have that the original Iterative Markovian Fitting (IMF) sequence (P"*),EN as defined in (4) with = 1 satisfies for any N, pin+1 (C)) (projy((P)) The analogy between the Sinkhorn iterates and the IMF sequence was already highlighted in Shi et al. (2023); Peluchetti (2023) and further studied in Brekelmans and Neklyudov (2023). We know show that similarly, we can draw an analogy between the sequences defined in (4) with E (0,1) and the sequences obtained in -Sinkhorn. To do so, start by introducing for any II, II ((R x R
CIPF (TI, MT KI. projj, (IT),
RIPF (TI, (II, KT. (m).
RIPF(
With this notation, we can now rewrite (47) for any n E N as
IT") II P(R x Rd), To = mo)- m ( " (TI, ITT) + (1-1)/7)
Now, we are going to see that (50) linked with the discretisation of the path measure flow described in (4). Recall that for any suitable u, we define the path measure P, associated with
dX. = uz (X, + V
X
define
I(v,P)) [ C(v,, P)dt (aa.Pttt(('/ [/ x pa(xx)
It 1 1-t
(P)
Similarly, for any y P(C), we define
Ra (P..P.P)) = f - (xx))
Next, we define the sequence of path measures (P")nen such that for any n E N
pa+1 argmin(C((PPFr) PP) RL (LI, P*) P= Pr. for some v}.
(50)
Now, if we denote (D"))en the sequence such that for any n E N, J = P.... then we have that for any n E N, EE [0, 1] and x E Rdd
( (x)))) 5v nr (1 (())tt
(51)
Recall that (P*) given by (4) is associated with (v") nEE such that for any n E N, t [0, and Rdd
x
(x) (( (z)).
(52)
see Proposition 3.2. Therefore, the only difference between (52) and (51) is that (52) is an explicit update whereas (51) is implicit update. We summarise the differences between 7-Sinkhorn and the discretisation we introduce in Table 2.
H.2 Links with Reinforcement Learning
In this section, draw connection between Algorithm 1 and self-play in Reinforcement Learning. In particular, we introduce a generalisation of Algorithm 1 which uses the concept of replay buffer commonly used in Reinforcement learning, Mnih al. (2015) for instance.
We first present a generalisation of Algorithm called Replay Buffer Diffusion Schrödinger Bridge Matching Algorithm 6. We define buffer B collection of samples {XX) x where
33

---

N E N is the size of the buffer equipped with two functions Add and Sample. We have that (R2d) - (R22)) where Q is a probability space. In practice Add takes a random number (the function can be stochastic), any number of proposed samples well as the current buffer. As an output Add returns the updated buffer. We also define Sample Ω x N x (R2d)N - LREN (RR This function takes a random number (the function can be stochastic), a natural number & representing the number of samples to return as well as the current buffer. As output Sample returns batch of k samples from the buffer.
Add : Q x LkkN (
Algorithm Replay Buffer Diffusion Schrödinger Bridge Matching
1: Input: To, TI, & (entropic regularisation), N Vpretraining (number of pretraining steps), Ninetuning (number of finetuning steps), B (batch size), y (EMA parameter), (initial parameters), Bfwd (forward buffer), Bowd (backward buffer)
2: 0=
3: for E ,O,..., N prretraining ) do
4: Sample (X]- X (To @ 8B L Un, 1]). Z ~ ~N(0 Id)o
5: Compute X Interp, (X 1:
XL.B Z1:B) using (12)
6: Update ! with gradient step on f 0 (1 7)0 +00
7:
end for
Vinneuuiigg 9: if n = Ofeeereehh then
8: for ( do
10:
Sample [X]-2 1:1 Yd ~ (To @ Tí 0 ,t t~ (n, z ~ (d),
11: Sample [X: Y!:B) using (11) with initialisation (X):B, Y):B)
12: Bfwdd Add((X): 1:B x xI= Bfwd)
Bbwdd Add((YP, Y! b
14:
end
if
[X
X (B x
Sample B, Bfwd)
16:
(Y): Y 1B = Sample B, Bowd)
17:
Compute X 11:B = Interp, (X) 1:B
x 1:B Z1:B) using (12) 1
18:
Compute YI-B I-t = Interp, (YI:B Y.. L.B zz: using (12)
19:
Update 0 with gradient step on 6B 0 (1 7) + 70
20:
end for
21:
Output: (0, 0) parameters of the finetuned model
In Algorithm 6, we allow for more flexibility than the online procedure by leveraging the concept of replay buffer originally introduced in Reinforcement Learning Mnih et al. (2015). The concept of replay buffer has been used previously in Schrödinger Bridge works, with the notion of cache where every Arefrrsh steps cache is emptied and filled with new samples, Areerrshh 1, N B for both
Blwd and Biwd we have that for any w Q and (X x1:B) E (R22)
Add(w: (X) 1:B x1:B) )) (X] 1:B X
Sample(w, B, (X xI:B) (x
This means that the Add simply fills the buffer with the new samples while Sample just return the whole current buffer. In that case recover Algorithm 1. For more general update rules, the replay buffers 13wdd and Bbwd allow us to collect previous samples and therefore to keep a memory of the past experiences. In future work, we plan to investigate popular choice in experience replay and their impact on the performance of Algorithm 6.
H.3 Links with Expectation Maximisation
In this section, we make a connection between DSBM and the Expectation Maximisation (EM) algo- rithm, and show that the discretisation of the Schrödinger Flow proposed in Algorithm corresponds to some incremental version of an idealised algorithm, as discussed in Neal and Hinton (1998). We would like to emphasize that the link between the EM algorithm and Diffusion Schrödinger Bridge based methodologies was already highlighted by Vargas et al. (2024); Brekelmans and Neklyudov (2023). Below, we follow the framework of Brekelmans and Neklyudov (2023) and recall the following definitions.
34

---

Definition H.2 (Projections and maximisations): Let be a subset of P(C). Then, for any
IP E PCC)), when it is well-defined, we define its E-projection on A as Ip = argmin KL,(Q PP). Similarly, for any IP E P(C), when it is well-defined, we define its M-projection on A as P
geg (P Q).
In Brekelmans and Neklyudov (2023), the authors choose M-projection because this corresponds to the Maximisation step in EM algorithm while the E-projection corresponds the expectation step in the EM algorithm. In Brekelmans and Neklyudov (2023), the authors highlight that the Iterative Expectation-Expeetaaiiee
procedure, i.e. the alternating projections are both E- E-projections. In contrast, the Iterative Markovian Fitting procedure is a Maximisation- Maximisation procedure, i.e. the alternating projections are both M M-projections. In particular, we can define the following sequence of path measures (IP")ne, where for any n E N we have
Proportional Fitting procedure
pn+1/22 = argmin KL(P" [P),
pa+1 = argmitl KL(Pp+ +1/2 [P))
In addition, we have that
pu+1/2 P Putt,
1 n-1 ((r
pn+1 arg KL [P),
since we have that P P Hence, online procedure Algorithm 1, which corresponds to the discretisation of the flow path measures (3) can be rewritten as
p=+1/22 P 12 2 Gradientstep(((r, PP), pa+1 argminge KL(P,, of [P).
Therefore, proposed algorithm can be seen incremental version of the Maximisation- Maximisation algorithm associated with DSBM instead of an incremental version of the Expectation- Maximisation algorithm discussed in (Neal and Hinton, 1998).
H.4 Links with finetuning of diffusion models
Algorithm I can be seen as a method to finetune bridge matching. Finetuning of diffusion models and flow matching procedures is an active research area. Most of the existing methodologies optimise for an external cost after a pretraining phase. These procedures rely on Reinforcement Learning strategies (Lee et al., 2023; Black et al., 2023; Fan et al., 2024). Recently Direct Preference optimisation (DPO) (Rafailov al., 2024) has been applied to the finetuning of diffusion models in (Yang al., 2023; Rafailov et al., 2024). Our approach departs from these works as the objective we minimise is given by the EOT cost. However all of these approaches involve some level of self-play, i.e. are not simulation free.
H.S Links with continual learning
Continual learning develops techniques train models when the dataset changes during the training, usually to solve different tasks De Lange et al. 22 Parisi et al. (2019); Zajac et al. (2023). In the context of diffusion models, continual learning has been investigated in Masip al. (2023); Zajac et al. (2023): Smith et al. (2023). In (Masip et al., 2023), the authors consider a weighted loss between diffusion model loss and distillation loss which ensures some consistency between the model being trained and the previous task model. Similarly to our approach this distillation loss is not simulation-Fi but, contrary our loss, the clean samples are not obtained by unrolling the diffusion model but by applying one-step prediction operator. In (Zajæc et al., 2023), consider different replay buffer techniques train continual diffusion models and observe that experience replay with a small coefficient can bring improvements. Finally, in (Smith et al., 2023), the authors consider the continual training of -ox-tto-mmaee diffusion model with LoRA (Hu al., 2021).
I Forward-Forward, Forward-Backwar and accumulation of error
In this section, we investigate how error accumulates in the context of DSBM. In practice, we observe similar conclusions in the case of the online version of DSBM. We compare two methods: one which only trains forward model and one which trains forward and a backward model.
35

---

In what follows, we assumee that %o = Ti (o), we also assumee that Q is associated with (v2..) We recall that for any / E [0, 11, we have that
X =1- + x 2 10l1-ty2 t)Z,
Z ~~ I0,,O
We going consider approximate schemes implement IMF.
Forward-forward First, consider the following sequence of path measures (IP" )nEK. We set pd = (To @ mT) Q,11 For any n E N, we define pp2++2 = P Q i.e. p2n+2 = projæ (p2n+1). In addition, define p2=+1 proj5 (IP2n) such that p20+1 associated with (Xx))E(U,1) where for any & [0, 1]
dx, [(Brgp
X1 x x,) () + et 2dB,
Xo ~ n0.
(53)
with IR Recall that if we define |2n+1 proj M (F2n) we have that for any E [0, 1), pan-1 is associated with (Xr))e(o.1) where for any / E [0,
dx.
(Epan X1 X.] - X,) t)dt + 22dB
Hence, projM corresponds to making an error of order x + Ex on the estimated velocity field. Doing so, now longer have that for any E N, PT Tj. In what follows, how the error accumulates for the sequence Po,1-
going to show
Before stating Proposition L.I, we introduce f R² + such that for any co,0, C1,1, Co,1 > 0 and t e [0, 1]
f(co,0 C0,1, [-(1 t)coo 10,1 (1--2)ep 21) 22]]
/[(1 t), + t 2t(1 t)co,1 + 2t(1 - t)].
s)s + 2st. Finally, we define We define F(Co.0; C,,1, 10 [(u,,, Co,
feor(eu,: 0,11 = (011
well
faar (0,,0, exply expp)_ co,,, 1)](1 +2 L'
Proposition L1 (Forward-Forward updates): For any n E N, we have that p2n+1 N(0, 2m) where
0,1
ld
a+ ld c
g
1 ld
and for any N
n+1 fvar (1, Cj 10,, c,,1
0, Lror (1, Ci e).
Proof. Let (Fo,1)QQo, where Fo,1 is a Gaussian random variable with zero mean and covariance matrix V E g2dx2d such that
Σ ld co ld co. cı,lld)
where Td is the d-dimensional identity matrix and we assume that Co C1,1 > We denote P (PP). We have that P, a Gaussian random variable with zero mean. We now compute its covariance matrix. First, we have that
E XX (1 xXXXXXX X [11 t)co,1 ter ld.
36

---

We also have that
, = (1 - -01
+ + + 26(1 t)ld
= =[(1- (1.t))+(o.1+2((( - 10,,,11111,111- + - t)lld
= =[1- -1- =11-t+ + 2t(1 1)co,1]Id.
Therefore, get that for any / [0, 1] and 24 Rdd
E2 X1 X _, ([(1 f) + tet,1/[1 & tl + 2t(1 t)co,1)2z.
Hence, we have that for any t E 00 1] and E Rd
Ep [1 X, zi|
([(1 t)co, tr,,1/111 - + Ec, 21(1 (1- t)co,11 1)2tz
=([( t) + 10, - - -
Pa,1
2t(1-t)c..//11-t t. + 2t(1 t)co,1l)z
-[[1 ()(1 - -2))o. + +1(1
- - Pca, 21(1 1)ca,
= =(1 -(1-t) 11-t ([(1 - -2t)co,1 + - - + 2t - (1)to)),,)
So follows that
BE X1 X, x1 x1) -)
= ([(1 - 20. + 111,1
-12 1/ Pei, + 20(1 - 1)c.,1))x.
(54)
Note that if we set co.1 = e² and C₁,1 1, we recover (Shi et al., 2023, Lemma 13) with = 2. Denote Ip+ = p Combining (54) and (53) we get that IP* is associated with (X1)se(011) such that for any E [0, 1] we have
dx, = [[(( 2() t0p1 - 1/111 - + 1Ra,1 + 21(1 ()o,a)) + X + V2d.
Hence, we get that
X, =xxp] G Co,1, (I1, +o++(22 (2/ onp
z,
where Z ~ (d) is independent from X, and for any t E [0, 1]. C0.1 C1.1, we have
C( 2 / [(1 2,, te,, - 11/11 - 1² fa,,1 21(1 100. 2st.
In addition, we define
Jcor (c., = =xc.)=)))).=([((...............
Ivar (Co (1,, =xpp=((,,,,,,,,,,... +2 /. expl-G((,co,,,
Hence, we have that
X = Geov (Co.1, 61,1,)Id
E X Ivar (Co,1, C1,1, ElId
Therefore, since for any E N, we have that p2n++ = projj (P22) and p22+22 p2n+ define (i,)ee such that for any n E N
we
Epan X = Id,
Epan X2 = =.,
Note that for any n E N. we have that
Epanti XxxX]] 001 I+1Id,
Epan+1 [ [X
cT,
n+IId.
Since p0 = (mo @ xT) Q we get that c0,1 = and C1,1 = 1. We have that for any n E
c = yoow (ci,1 of,1
n+1 G
Ivar (var(c (T,,)))
which concludes the proof.
37

---

Forward-backward. Next, we consider the following sequences of path measures (P NEN and (IP***))xx We set p0,* 10 (To @ n Q For any N, we define p2n+2, 0,1
@ a,11
and p2=+22 = 0..1 p2 Q n.1, ic. 2, projæ (F2n+ 1+) and p2n+2. = projæ (p2n+1.))). In
addition, we define p2n+1, = projjM (I22.+) such that for any E [0, 1], 12 is associated with (Xx)1eju.1) where
dx; ((Egpn 1 X,] (1) t) + eX, Vdd
Xa no.
with & E IR. Similarly, we define p2n+1* proj'i (Tp@n,)) such that for any E [0, 1]. p2=+1. is associated with (Y1-t)rejo.1) where
dY.- = {(Egan. Xo ojs
Y.] - Y.) (1 + a v2dB,, Yo ~ mT
Proposition 1.2 "Forward-Backwart updates): For any n E N, we have that 10.1 122++1, (0,, Id) and Po.tl.
(N,00 En+I+Td) where
gatl>=
ld
1.1 2+1+'1d
gnth+=
0.0 1+l
Itl atli
00,
0,1
I+1
1,
U
Id
and for any N
[,1
1+,,>
S (0000 1, 0,,1, E),
a+
Dit foor (0,00 1, ntt Co,1,E), E)
ç0,1
a+1+ fvar (1, (0,1,))
0,00
a+1, Soow (1, 11,1 0).
0,1
The proof is similar to the one of Proposition I.I.
Proof. Let P (Fo, Q where Fo,1 is Gaussian random variable with zero mean and covariance matrix D E p2dx2d such that
cooolld Io
co.ldd) Id
Σ
where Id is the d-dimensional identity matrix and co.1, Co.0 > 0. We denote IP* = proj We have that Pipe Gaussian random variable with zero mean. We now compute its covariance matrix. First, we have that
EXXXX]] = (1 x
[[1 t) + t]ld.
We also have that
x = (1 (1-1))E[Xox,] + -
+ [ + 2/(1 1)Td
= =(0 + 2f(1 - t) + 2f(1-))]4
= 12 - +1-11)2c0.0 + 21(1 -t)co,
Therefore, we get that for any t E [0 and E Rddd
Er, X1 X, x, ([(1 c) + f2/ - 12 + +1 - 1)2co.0 + 26(1 t))
Hence, we have that for any e [0, 1] and x; E Rd
Ep, X X, xx] - It - ([(1 ()co,1 + 01/(21 12 + (1 (1-1) (1-1)²com 1)0cc.. co.0 21(1 1)co,.] 1)
= =([1 - t) + - 2t + f2 (1 - c) - 2t(1 - c))
/(21 12 + (1 - ()-1)2c0.0 + 21(1 - l)co,1)).;
= =([11 11 (1) -2t)c.1 - c) t(1 t)l/ (2t - f2 + (1 t) + 2t(1 - 1))
= =(1 (1-+)(((1 2t)co,1 t)co.o tl/12t t² + (1 - t)²coo 2t(1 t)co,
38

---

Finally, we have that for any E [0, 1] and 2t E R
(Br, [11 X mz] xxt -)
= ([(1 c0. - -(1 - t) - - -), 0 + 2t(1 - t) (55)
Note that if we set C ²² and Co,0 1, we recover (Shi et al., 2023, Lemma 13) with
Denote IP- proj jM (P). Combining (55) and (53) we get that IP is associated with (X.))e(0.1) such
that for any E [0, 1] we have
dx, {{{(1 co. 1)0,00 0/(21 1/121-112111111111 12+ () +21 - +000)-))),))))))XX.)) +e) 2dd,
Hence, we get that
X, 2 Co.1 Co.o. +2 / expl-H(s,,) c0,1 00,0 E)lds exp Co.1, Co.o /
where Z ~ Id) is independent from X, and for any E [0, 1]. co c1,1, > we have
(t, =) 2 = /,
- -1
co.0 - 1 2Et.
In addition, we define
goo (eor(eo,1, co,0, 8) = exple (1,, co,
Ivar (co,1 Co.0. =) = e1n 1+ + f.. I. expl- Ht, co.1 Co,1,
Hence, we have that
E[XoXi] = Jcov :
E|X,X]| = Svar (co,1, co,0, s)Id.
(56)
Remember that p0,> p0,t (no @ 1. In addition, for any N. we have 22n+2, = +2o+ and p2n+2, = Fo.l p Q0,1 iee p2n+2, = projæ (p2n+l,e) and
p2n+2+ =projæ (p2n-1,*). In addition, we also have p2n+1,+ proj and p2n+1e
projj(((P2m))) We also define (ch,1, n6.7) (ii)) such that for any NEM N
Epao. X
Id,
Epan. X
I...d.
Finally, we define (c o NEN such that for any N
Epa.. X = II,
Epan. X, = c0,0
Using this definition and (56) we get that for any n EN
0
n+I+ Seov ( 21 n.t
(0,1 00,0 E), n.
G
n+l,
Ivar Co.1 co,0 E),
( nht n+
601 Jeov
(c) (111)).
2).
0,0
u+l.
Ivar ( :..1:)). e).
n, n Co In addition, we have that Co.1
= 0 and G1i
Coo =
0,+
1. This concludes the proof.
Proposition I.I and Proposition I.2. we derive the sequences corresponding to the evolution of the variance and the covariance throughout the DSBM iterations in forward- forward mode or sequences for different values of ε > 0. We recall that corresponds to the eтror made in the Markov projection, iee. proj. is replaced by proj in the forward-forward mode and proj is replaced mode. First, if we consider the perfect scenario, ie. e = 0, then we observe that both the forward-forwan mode and the forward-backwe satisfy that Epz.. X,XX] = Id, see Figure 10 and Figure 11. Additionally, we can show that in the mode (v2 1)Id, see Figure 10 and Figure 11. However, as increases the behavior between the forward-forwand sequence and the forward- harward-backward sequence significantly differs. More precisely, the error explodes as & increases along the DSBM iteration mode, the error remains
Error accumulation.
mode. In what follows, we showcase the behavior of these
by projem and proj in the forward-backward
mode
perfect scenario, i.e. 0, then both the forward-forward mode and the forward-backwer
satisfy that lim E,2.. [X
for the forward-forwand mode. On the contrary, in the
bounded along the DSBM iterations, see Figure 10 and Figure 11.
39

---

109644
Figure 10: Evolution of (1) Egan [ Id nEN in log-space along DSBM iterations (x-axis). Different curves correspond different values of E, i.c. the larger e the larger the error in the Markovian projection. Left: evolution in the forward-forward mode. Right: evolution in the forward- backward mode.
377
377
Td EEN log-space along DSBM iterations (x-axis). Different curves correspond to different values of i.. the larger e the larger the error in the Markovian projection. Left: evolution in the backward mode.
Figure 11: Evolution of Egan [X
mode. Right: evolution in the forward-
J Preconditioning of the loss function
In this section, we provide details on the scaling of the loss function we implement when training our online version of DSBM. We adapt the method of (Karras al., 2022, Appendix B.2) the case of bridge matching. We only present our derivations in the case of the forward training of the online version of DSBM, i.e. (9). The preconditioning of the loss described in this setting can be readily extended to the forward-backward backward loss we consider in practice, i.e. the parametric version of (10).
We consider the following objective function for any E [0,
e, A XeEr (c(x,) + GG,. X 1-1
1111
(57)
We also define for any E [0, 1] and T, E Rd. of (xx)) (nn((eearr) cts. Hence, is an input scaling function, ef is an output scaling function and of is a skip-connection function. During the training of the online version of DSBM, PP will be given by pn, where Pn P P, where the sequence (0,, neN is given by (9). Here, we apply the principles of Karras et al. (2022) to the case where
(To @ Ti 1 i.e. initialisation of the sequence. In what follows, we Em Xx || d. Note that our considerations can be generalised to of d. We also have that
assume that Ezo Xpp [2]
E,, [|| X | and d and E
X, (1 X), + x + /et(1 t)z,
Z ~~ ~(0,1d
(58)
Using (58), we have that for any E [0, 1]
Ex | = (1 - 2 ||| | |2 + E
| [13] + el(11 l)d
[[1 -)2 + t + Et( 1 t)
We set of so that ||||xX,1|11 d for every t E 10. 1]. Hence, we have that for any t E [0. 1]
= - + - + t).
40

---

Figure 12: E E [0, 10).
From left to right (EEo))) (G)eeo.1 and ((e))) for different values of
Next, we rewrite (57). For any t [0,1] we have that
4 = A, (c +
G x, - 112)
(ct)
|- c + --t
x // 1/
X, mn! (X,)))
X1 (1+(2+ x) 1/c/1191
//c is the target of the network in the regression loss. We are going to fix c and c such that i) B |2 = d, ii) cf is as small as possible in order not to minimise the error propagation made by the neural network. Using (58), we have that for any / [0, 1]
Hence, we get that for any E [0, 1]. T, = x (1+ x)
1-1
E... [|. ((= (1
[x, ||
- -2(1 () (X,, X,)]
(1 + c/ - 2(1 cg t))td
Hence, we get that for any t [0
((1 ( |X/ 1001-1222400000//111 -1²
We now minimise (c with respect to (1 + c( - 1)). We get that
1 1+ 11 () - 1/(E,, x,,l/////11//)))
Hence, we get that of = t/[[11- () that for any E [0, 1]
+ +f t))] - 1/(1 - t). With that choice, we get
(cg) = (1- (1- + et(1 ())) -
In Karras et al. (2022), the weighting function X is set so that the weight in front of the regression loss is equal to one for all times t [0, 1]. Hence, Karras et (2022) suggests set X = 1/(c))² However, in practice we observe better results by letting A = 1. This means that the effective weight given by 1/(c))². Therefore, for any 1 [0, 1] we have
(c)) = (1 +1+ (c 2)t(1 - 1))-
= ((= 2)t 1)/(1 + + 2)t(1 - -)), (c)) = (1 + (a 2)t(1 t))/(1 t).
K Experimental details
In this section, delve deeper into the specifies of each experiment, implementation details, and share additional results.
We consider two ways of parameterising the vector fields: in DSBM, we can use two separate neural networks to approximate the forward and backward vector fields, or we can use a single neural network that is conditioned on the direction. In the latter case, we do the conditioning in a similar fashion to how DDM's neural networks, U-Nets or MIPs, are conditioned on time embeddings. After all, if we work with continuous time variables t [0. 1], then the direction signal {0, 1}
41

---

can be thought of as a target time. Thus, we perform the same initial transformations on t and ie. computing sinusoidal embeddings followed by 2-layer MLP, and use the concatenated outputs in adaptive group normalisation layers (Dhariwal and Nichol, 2021; Hudson et al., 2023; Perez et al., 2018).
To optimise our networks, we use Adam (Kingma and Ba, 2015) with 8 = (0.9, 0.999). and we modify the gradients keep their global norm below 1.0. We re-initialise the optimiser's state when the finetuning phase starts.
All image samples in the paper generated using EMA parameters it has been known to increase the visual quality of resulting images (Song and Ermon, 2020). Sampling is also the integral part of DSBM's finetuming stage, both iterative and online. Here, we have two options: sample with EMA non-EMA parameters. The non-EMA sampling might be easier to implement, while EMA sampling stable training and slightly better quality, e.g. AFHQ samples in Figure 23 and Figure 24 for comparison.
results in
For every model used in the paper, provide hyperparameters in Table 3.
2D
Gaussian
MNIST
AFHQ-64
AFHQ-256
Channels/hidden units
256
256
64
128
128
Depth
3
3
2
4
4
Channels multiple
n/a
n/a
1, 2, 2
1,2,3,4
1, 1, 2, 2, 3, 4
Heads
n/a
n/a
n/a
Heads channels
n/a
n/a
n/a
64
64
Attention resolution
n/a
n/a
n/a
32, 16, 8 0.0
32, 16, 8
Dropout
0.0
0.0
0.1
0.0
Batch size
128
256
128
128
128
Pretraining iterations
50K
10K
100K
100K
100K
Finetuning iterations
150K
40K
150K
20K
20K
Pretraining learning rate
lc-4
le-4
lc-4
2c-4
2c-4
Finetuning learning rate
le-5
le-4
le-4
2e-4
2e-4
Pretraining warmup steps
n/a
n/a
n/a
5K
5K
EMA decay Parameters count
n/a
n/a
0.999
0.999
0.999
133.4K
371K
8.8M
194.4M
226.7M
Table 3: nyper-parameters describe only one of the two identical networks. Approximate parameters counts are given for bidirectional networks, except for the Gaussian case, where we only experimented with a 2-networks model.
for each model. Note that for 2-networks models, the architectural
K.1 2D Experiments
In addition to the experiments presented in the main text, we test our models in the simplest 2D data settings used in Tong al. (2024a) and Shi et al. (2023). Note, that now-dimensional datasets might not be the ideal showcase for a-DSBM given that one can successfully employ less computationally demanding techniques based minibatch-OT methods (Tong al., 2024b).
The results of our bidirectional model finetuned with online updates are given in Table During finetuning, generate samples using 100 Marran steps to solve the forward and backward SDEs. At test time, we solve the forward probability flow ODE (PF-ODE) given by:
dx, 2 [Lo (,), (0, X,) ] dt,
X, mo.
(59)
To evaluate model fit, compute 2- 2-Wasserstein distance between the true and generated samples (generated with 20 Euler steps). Additionally, we estimate path energy as a measure of trajectory simplicity: Ex ( ||v (t, [))]] where (1, X,) is the drift of PF-ODE in (59), and the integral approximated using 100 steps. We have made a deliberate effort to closely replicate the experimental setup of Shi et al. (2023) to ensure the comparability of our results. However, as
42

---

illustrated in Figure 13, 2- Wasserstein distance can be very noisy even with 10K samples in the test set. To mitigate this variance. we averaged the 2- 2-Wasserstein distance across five random sets of 10K samples per run, and then averaged these results across multiple runs. Despite these measures, we recommend a future redesign of these 2D experiments to facilitate more robust comparisons between methods.
Matlad
7. M.Woeemooin
Paf
-> muait
-+ scuenc
NV Rgusosians
Resanoiats
A -hecanee
NV -agansiars
mnson Spariia
DIBM-IMP"
CTon 203AY
0110390013
(1255211/40
WTI
127
1.7760005
15.10el
D-RSSM
6.168 40,011
021540891
02202-0047
194401266
1.49 40014
3,03 40,05
15,0840.190
$7,0640 590
Table 4: 2- Wasserstein distance and path energy for the 2D experiments. We report means standard deviations across random seeds. DSBM-IMF* and OT-CFM* results are copied from Shi al. (2023).
20.0
17.5
Prqquency
12.5
10.0
50-
14
1.5
16
177
18 19
2.0
2-Masserstein
Figure 13: A histogram of 2- 2-Wasserstein distances for the 'moons-> 8gaussians' task. These distances calculated between 10K samples from finetuned a-DSBM model and 8gaussians distribution, with both sets generated using 100 different random seeds. The wide spread of scores indicates that 2- 2-Wasserstein distance, even computed on 10K samples, may be ideal metric for evaluating model fit in this context.
K.2 Gaussian data
To parameterise the forward and backward drifts, we use a 2-layer MLP network with 256 hidden units. To process time variables, compute sinusoidal time embeddings, followed by 2-layer MLP with 256 hidden units and 50 output units. The resulting time embeddings are then concatenated with X,, the drift networks receive (00-dimensional input vectors.
For iterative DSBM finetuning, we perform 40K steps with varying number outer iterations, i.e. when we switch between training the forward and the backward networks. Alternating every 5K steps, corresponds to outer DSBM iteration. Similarly, changing the direction every IK steps, leads to 40 outer iterations.
We do not have cache dataloader like in the original DSBM implementation' thus we generate training samples on the fly by sampling either from the forward or the backward network. For this simple task, we also do not use EMA.
During training and evaluation, we use Euler Maruyama method with 100 equidistant time steps between 0 and 1. The covariance is evaluated using 10K samples.
Additional comparison o-DSBM vs OT bridge-matching. We consider the scalar Gaussian setting. We highlight the dependence of OT bridge matching (256/1 1/88) on the batch size, as the mini-batch OT coupling can be far from the true OT coupling as the batch size decreases. All experiments are run with similar compute and architecture/training hyperparameters,
Table 5.
Ablation of the hyperparameter a. Instead of letting be determined by Adam and adaptive for a-DSBM, we explicitly set it by using Stochastic Gradient Descent (SGD) with learning parameter ag see Figure 14. We also ran online a- a-DSBM with 10 but the training diverges in this case.
"https://githuu: http:///hhuu.ommyyyeeeeeeee com/yuyang--eeo shi/dsba- pytorch
43

---

Table 5: Comparison of OT methods
Method
Covariance
Optimal
0.882
o-DSBM
0.890
Bridge Matching
0.491
OT Bridge Matching (256)
0.853
Bridge Matching (16)
0.840
Bridge Matching (8)
0.824
Normalized Frobenius distance
1.5
Online @Adam
Online @le-3
Online @le-5
Online @le-2
Online @le-4
Base
1.0
0.5
10k 20k 30k 40k 50k 60k 70k 80k 90k100k
Training iterations
Figure 14: normFrob between C, and its estimate for a-DSBM with different values a.
K.3 MNIST + EMNIST transfer
We closely follow the setup of Shi al. (2023) and De Bortoli et al. (2021), and train the models to transfer between 10 EMNIST letters, A-E and a-c, and 10 MNIST digits (CC BY-ND 4.0 license). We use the same U-Net architecture with hyperparameters given in Table
For DSBM finetuning, we perform 30 outer iterations, i.e. alternating between training the forward and the backward networks, while each outer iteration network trained for 5000 steps. We do not have cache dataloader and generate training samples on the fly by sampling either from the forward or the backward network with EMA parameters.
During training and evaluation, we use Euler-Maruyama method with 30 equidistant time steps between 0 and 1. For evaluation, we compute FID based the whole MNIST training of 60000 examples and a set of 4000 samples that were initialised from each test image in the EMNIST dataset. MSD computed between 4000 initial EMNIST test examples and their corresponding MNIST samples.
In Figures 15-18, we provide forward and backward samples, i.e. EMNIST MNIST and MNIST - EMNIST, from models that differ in parameterisation, finetuning methods, and sampling strategy. For all the models above, used = 1. Figure 19 illustrated the behaviour of the samples when sweep over the e hyperparameter.
Pretraining bidirectional model on 4 v3 TPUs takes hour, while the online finetuning stage requires hours on 16 TPUs. The number of pretraining and finetuning steps chosen match the experimental setup of Shi et al. (2023).
K.4 AFHQ: Cat + Wild
consider the problem of image translation between Cat and Wild domains AFHQ (Choi et al. (2020); CC BY-NC 4.0 DEED licence) introduced by Shi et al. (2023). Each domain has approximately 5000 samples in the training set, and around 500 samples in the test set. We resize the original 512 x 512 images to 64 xx66 664 or 256 resolutions.
44

---

a. 1
a
8 8C BB 11111111111111111011001111111111
dDb8 4CL 8
2 P B
9 3 9
1 52
(a) Initial EMNIST letters (b) Bridge matching:
(c) Bridge matching DSBM finetuning: FID=5.25, MSD=0.345
(d) Bridge matching online finetuning:
FID=6.02, MSD=0.564
FID=4.28, MSD=0.368
Figure 15: EMNIST to MNIST transfer with 2-networks model.
6
4
4 20c.CCA2200 9 7
a a <Addce<<b(9015344 2tt.
1 0
3 2
/
S 7 7
3
e e 2
9
5
8
7
5
8
d 4Db8 A0bE&dCAcC 4
8
8
2
8
ebdE 8 BC2
9 9
5
59
(a) Initial EMNIST letters (b) Bridge matching:
(c) Bridge matching
(d) Bridge matching + online finetuning: FID=4.39, MSD==0.387
FID=6.33, MSD=0.572
+ online non-EMA finetuning: FID=4.57, MSD=00339
Figure 16: EMNIST to MNIST transfer with a bidirectional model.
[
4 7
8 3.
1
0 6 6
bcc
3 0 a 4 9
A
7 6 3
1
5 8
1
$ 4
4
3
4
&
Z
4
9
5
6
6 3
(a) Initial MNIST digits (b) Bridge matching:
(c) Bridge matching DSBM finetuning: FID=3.56, MSD=0.330
(d) Bridge matching
FID=7.50, MSD=0.553
online finetuning:
FID=3.67, MSD=0,357
Figure 17: MNIST EMNIST transfer with a 2-networks model.
45

---

4 7 8
6
1 3 3 4
C
|
6
3 0 a 4
7
103/112240
4
2
4 5 8 1 4
[ 8 2
0
8 49
6 4 /
7
4
9
9
6
(a) Initial MNIST digits
(h) Bridge matching:
(c) Bridge matching + ouline non-EMA
(d) Bridge matching
FID=7.97, MSD=0.572
+ online finetuning:
finetuning: FID=4.16, MSD=0.370
FID=3.97, MSD=0.392
Figure 18: MNIST to EMNIST transfer with a bidirectional model.
Our U-Net (Ronneberger et al., 2015) implementation based on Ho et al. (2020) with a few im- provements suggested in Dhariwal and Nichol (2021); Song et al. (2021b) such as rescaling of skip connections by 1/v2, using residual blocks from BigGAN (Brock al., 2019), and convolution-basee up- and downsampling. Hyperparameters parameterisation of the vector fields, we obtained slightly better results using EDM precondition- ing Karras et al. (2022), which we derive in Appendix J for the case of bridge matching. During training, we use horizontal flips as a way to augment the data.
given in Table 3. Compared to the straightforward
During training and evaluation, we use Euler-Man Maruuaaa method with 100 equidistant time steps between and 1. When evaluating the quality of Cat - Wild transfer, we compute FID based on the whole training set of 4576 examples in the Wild domain and a set of 480 samples that were initialised from test images in the Cat domain. LPIPS and MSD are computed between 480 initial Cat images and Wild samples from the model. The same procedure is followed when evaluating in the reverse direction from Wild to Cat. Given that train, and especially the test sets are small, the quantitative results for AFHQ are likely unreliable (Chong and Forsyth, 2020). In Figure 22 we provide samples from the models finetuned either with an iterative or an online method. While their FID scores are different, the samples look similar between the two models.
As we discussed in the main text, hyperparameter trades off the visual quality and alignment of the samples in the resulting transfer models. In Figure 20, we provide AFHQ 64 x 64 samples for pretrained and finetuned models with different values of e. In addition to its relation to EOT, from DDM perspective, can be seen the controlling factor of the noise schedule. As observed by Hoogeboom et al. (2023), noise schedules should be adjusted for different image sizes by shifting the noise schedule of some reference resolution where it is proven to be successful. In our case, if we find a good value of for 64 X 64 images, then a shifted æ for the 256 x 256 resolution can be computed
Thus, if we choose ve 0.75 for AFHQ-64, then for AFHQ-256, we can expect v = 3.0 to also work well. Samples from an AFHQ-256 model trained with ve 3.0 are given Figure 27.
as $256 = 864 6.
256
On 16 v3 TPUs, the bidirectional base and finetuned AFHQ-64 models take 4 and 14 hours to train, respectively. For AFHQ-256, the base model trains for hours, and finetuning takes additional 37 hours. While we did not experiment with varying pretraining and fine-tun fine-tuning iterations, these training times suggest that a longer pretraining stage followed by fewer mne-tuning steps may be desirable.
46

---

e 6 B
12-033 17w072 27-01H Ve in sumpoo
BB
e 6
ba
3 3 8
e
b
5
5
8
2
4444
8
8
8
6
1
2
6
vre V/-1.0 V2-081
8
2//4
6
8
2
P
6
66
6
5
4
2
/
4
4
39
6
65
5
9
4
49
6
5
9
4 9
(a) Base model
(b) Finetuned model
Figure 19: MNIST samples transferred from EMNIST letter inputs (top row) using base (pretrained) and fine-tuned models for different values of Low noise values result in poor sample quality, particularly in the base model, which finetuning cannot fully rectify. Conversely, excessively high restricts information passing from the inputs to the outputs, leading to poor alignment. Additionally, high increases blurriness due to increased noise levels, thus requiring more denoising steps.
47

---

(a) Base model
(b) Finetuned model
Figure 20: AFHQ 64 x 64 Wild Cat transfer results for different values of Ve bidirectional model before and after online finetuning. Low values of lead to poor sample quality in both base and finetuned models. Excessively high & values impede information passing from the inputs to outputs, resulting in poor alignment. High values of also increase blurriness due to noisier SDE trajectories, thus requiring more denoising steps during sampling.
48

---

W
2003
(a) Base model
(b) Finetuned model
Figure 21: AFHQ 64 x 64 Wild - Cat transfer results for varying number of function evaluations (equivalent to time discretisation steps in the Maruyama method) in bidirectional model with V 0.75, both before and after online finetuning. Post-finetuning, clearer images are achievable with fewer steps. This observation aligns with findings from Rectified Flows (Liu et al., 2023b).
49

---

(a) Iterative finetuning: Cat + Wild. FID-27.76, LPIPS=0.503, MSD=0.093
(b) Iterative finetuning: Wild - Cat.
FID-25.24, LPIPS=0.483, MSD=0.094
(c) Online finetuning: Cat + Wild. FID=32.12, LPIPS=0.503, MSD=0.097
(d) Online finetuning: Wild + Cat. FID=27.32, LPIPS=0.48S, MSD=0.116
Figure 22: Samples and metrics from a 2-networks model architecture finetuned with DSBM's iterative procedure vs online finetuning, Within each two rows, initial and transferred samples are on the top and bottom respectively.
50

---

(a) Forward: Cat + Wild
(b) Backward: Wild -> Cat
Figure 23: Uncurated samples for AFHQ 64 x 64 transfer in bidirectional model with online finetuning with non- non-EMA sampling and ve = 0.75. Within each two rows, initial and transferred samples are on the top and bottom respectively.
51

---

(a) Forward: Cat ++ Wild
(b) Backward: Wild -> Cat
Figure 24: Uncurated samples for AFHQ 64 x 64 transfer in bidirectional model with online finetuning and Vł = 0.75. Within each two rows, initial and transferred samples are on the top and bottom respectively.
52

---

(a) Forward: Cat -> Wild with inputs from Wild.
(b) Backward: Wild -> Cat with inputs from Cat.
Figure 25: Samples for AFHQ 64 64 transfer bidirectional models with online finetuning and different values of e. The models are only trained on Cat and Wild domains, To and T. respectively. Thus, the forward direction the models expect Cat samples as inputs at O, and transfer them to the Wild domain at / = 1. The reverse transfer holds in the backward direction. Here, we test the models' behaviour when inputs not come from the same distribution as during training: we feed Wild samples in the forward direction, and Cat samples in the backward, which is the opposite from what the models expect. Ideally, the model should leave these inputs unchanged, which it does to varying degrees depending e, variance of the Gaussian noise. As increase less information can pass from the input to the output, thus making them less alike.
53

---

(a) Forward: Cat Wild with inputs from Dog.
(b) Backward: Wild with inputs from Dog.
Figure 26: Samples for AFHQ 64 x 64 transfer in bidirectional model with online finetuning and VE = 2.0. The model is only trained on Cat and Wild domains, my and Tı, respectively. Thus, in the forward direction the model expects Cat samples inputs t 0, and transfers them to the Wild domain at = The reverse holds in the backward direction. Notably, the model generalises well to the unseen AFHQ Dog domain, often producing high-quality translations. These results come from model with v = 2.0, which is higher than our chosen default value of V = 0.75. Higher noise allows the model better deal with -ut-of-distribution inputs.
54

---

(a) Forward: Cat + Wild
(b) Backward: Wild -> Cat
Figure 27: Uncurated samples for AFHQ 256 x 256 transfer in bidirectional model with online finetuning and V = 3. Within each two rows, initial and transferred samples are on the top and bottom respectively.
55

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's s contributions and scope?
Question:
Answer: [Yes]
our main theoretical and experimental contributions are claimed in the abstract and demonstrated in the paper. We summarize our main contributions hereafter. Theoreti- cally, we identify new family of sequence of path measures related to the IMF algorithm, called a-IMF. We show that these sequences correspond non-parametric updates. We then introduce parametric update that corresponds to an online version of the DSBM algorithm. We show that our procedure retains the favorable properties of DSBM while not requiring the expensive repeated inner minimisation procedure of DSBM.
Justification:
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: The limitations are addressed in the discussion section. The main limitation of our algorithm is that is not sampling free methodology. future work, we would like to see how to mitigate the fact that our algorithm depends on some self play.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: All theoretical results are proven in the supplementary material, see Ap- pendix D.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Full experimental details are provided in Appendix K.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [No]
Justification: Due to IP restrictions, we cannot share the codebase used for this paper. However, we plan release some notebooks in order reproduce experiments in small scale setting.
6. Experimental Setting/eetails
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimiser, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Full experimental details are provided in Appendix K.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
56

---

Justification: All metrics are computed using multiple random seeds and error bars are provided.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [Yes]
Justification: Full details on the compute requirements are given in Appendix K.
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: cttp://./../...ci/iiiiioooeeeeee
Answer: [Yes]
After careful review of the NeurlPS Code of Ethics, we can ensure that the research presented in this paper conforms with the Code of Ethics in every respect. Indeed, we see no immediate safety, security, discrimination, surveillance, deception, harassment, environment, human rights or bias and fairness concerns to our work. In addition, we release details and documentation regarding the datasets and models used. We disclose essential details for reproducibility and have ensured that work is legally compliant.
Justification:
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
This paper addresses the problem of unpaired dataset translation and proposes an improvement to the DSBM methodology. As the current paper mostly theoretical and methodological we do not see immediate societal impact of this work and therefore do not discuss these issues. However, we acknowledge that large scale implementation of our algorithm might suffer from the same societal biases generative models. We hope to address the limitations of such models when turning to more experimental work.
Justification:
11. Safeguards
Question: Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g., pretrained language models, image generators, scraped datasets)?
Answer: [NA]
Justification: The paper poses no such risks.
12. Licenses existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We have referenced the license of the datasets we use and cite the original papers that produced the code packages and datasets that we use in that paper.
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
57

---

Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
58


Reviews:
[{"review_id": "0ya0gdQh3w", "reviewer": "Reviewer_ovZP", "summary": "The paper proposes a novel algorithm for mass transport problems, aiming to compute maps that transport one distribution to another. This paper introduces the Schrödinger Bridge Flow algorithm, a dynamic entropy-regularized version of OT, eliminating the need to train multiple DDM-type models. The algorithm discretizes a flow of path measures, ensuring the Schrödinger Bridge as the only stationary point. The paper demonstrates the algorithm's effectiveness on various unpaired data translation tasks, showcasing its potential to solve Entropic Optimal Transport problems with improved computational efficiency and practical implementation compared to existing methodologies.", "questions": "1.See Weaknesses\n2.I am curious about the complexity and cost of the training method proposed by the authors. Could the authors provide more detailed quantitative data, such as training time comparisons with related methods?", "limitations": "Although the authors proposed a novel model to address mass transport problems and provided some theoretical proofs, the method is simulation-based. I believe that the training cost is higher compared to simulation-free methods, and the model performance does not achieve state-of-the-art (SOTA) results. Additionally, there is a lack of direct comparison with existing SOTA methods.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1.This paper is eloquently written, and its ideas are easy to follow.\n2.While I did not strictly follow all the proofs provided by the authors, I found them to be generally detailed, clear, and easy to understand.\n3.The authors provide both quantitative and qualitative results on several datasets, along with a relatively detailed discussion and visualization. However, it is worth noting that the figures in the appendix appear to have low resolution, and the positioning of the legends is somewhat odd.\n4.I understand that the convergence of the parametric methodology may require several strong assumptions. Therefore, it is reasonable that the authors have not provided further proof of convergence.", "weaknesses": "1.Assumptions in Theorem 3.1: Theorem 3.1 only mentions \"under mild assumptions\" without specifying them in the main text. It is necessary to state these assumptions clearly in the main body of the paper and provide references to justify their rationality. While a more detailed explanation can be provided in the appendix, the main text should at least include a summary of the assumptions to ensure the reader understands the conditions under which the theorem holds.\n2.Experimental Results and Comparisons: The experimental results on some datasets are not particularly impressive, even though the authors claim that their method approximates OT maps more accurately than existing methods. The results in Table 1 do not show significant progress or state-of-the-art (SOTA) performance on certain datasets. Additionally, it is confusing that the authors explicitly mentioned Rectified Flow, Flow Matching, and OT-Flow Matching in the literature review, but did not include comparisons with these methods in the experiments. For instance, Rectified Flow has conducted experiments on mass transport, and other related methods, such as Denoising Diffusion Bridge Models from ICLR 2024, have tackled similar tasks. A direct comparison of the proposed method's performance with these existing methods would provide a clearer picture of its advantages and limitations.\n3.Discussion of Related Work: The paper could benefit from a more comprehensive discussion of related work, including diffusion models, flow matching, and the application of Wasserstein gradient flow-based models in mass transport tasks. This would help in better contextualizing the authors' innovations and contributions. By situating their work more clearly within the existing literature, the authors could highlight the unique aspects and potential advantages of their proposed method, enhancing the readers' understanding of its significance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments, we appreciate your acknowledgements of the paper’s merits.\n\n> It is necessary to state these assumptions in the main body of the paper and provide references to justify their rationality.\n\nDue to space limitations, we have detailed most of the technical assumptions for our theorems in the supplementary material rather than in the main paper. However, if our submission is accepted, we will be granted an additional page, which will allow us to include the assumptions as detailed below.\n\n**Theorem 3.1**\n\nThe assumptions for the result to hold are stated in Lemma D.2. We recall them here. Let $\\pi_0$ and $\\pi_1$ be the distributions at time 0 and 1. We require that $\\pi_0$, $\\pi_1$ have distributions w.r.t. Lebesgue measure, with finite entropy and finite 2nd moments. This finite 2nd moment assumption is quite weak. While the Lebesgue density hypothesis is commonly considered in the diffusion model literature, see e.g. [1,2].\n \n**Proposition 3.2**\n\nThis proposition is detailed in Proposition D.1.The sole assumption in this case is that $\\mathbb{P}_{v^n}$ is defined for all $n$. This means that the SDE with drift $v^n$ admits a weak solution. The existence of weak solutions of SDE is a well-known topic in the literature of SDE; see e.g. [3]. In practice this condition is satisfied since $v^n$ is replaced by a NN approximation, which satisfies the conditions for the SDE to admit a solution. \n\n**Proposition 4.1**\n\nThe proof is similar to the one of Proposition 3.2 and therefore holds under the assumption that $v^{n, \\rightarrow}$ and $v^{n, \\leftarrow}$ define solutions of SDEs. \n\n[1] Lee et al., Convergence of score-based generative modeling for general data distributions, 2023\n\n[2] Conforti et al.,  Score diffusion models without early stopping: finite Fisher information is all you need, 2023\n\n[3] Stroock et al., Multidimensional Diffusion Processes, 2006.\n\n> The experimental results on some datasets are not particularly impressive, even though the authors claim that their method approximates OT maps more accurately than existing methods\n\nWe want to emphasize that we do not claim our approach yields better results than DSBM (or any other method that recovers the EOT). There is no expectation of better results from our method compared to DSBM, as both converge to the same solution. Currently, we state “$\\alpha$-DSBM is easier to implement than existing SB methodologies while exhibiting similar performance” (l.313).\n\nThe main advantages of our methods are computational: 1) significantly simpler implementation, requiring only one fine-tuning stage and eliminating the need to alternate between two optimisation problems  2) utilization of only one NN (with the bidirectional implementation) 3) faster convergence, as shown in Fig. 2 (and in the Gaussian study in the one-page rebuttal). We hope this explanation clarifies the key computational contributions of our paper. We will revise our manuscript to further underline this point if it remains unclear. \n\n> It is confusing that the authors explicitly mentioned Rectified Flow (RF), Flow Matching (FM), and OT-FM in the literature review, but did not include comparisons with these methods. \n\nRF corresponds to DSBM for $\\varepsilon = 0$. Our $\\alpha$-DSBM algorithm thus provides an alternative online implementation of RF for $\\varepsilon = 0$ and results for this online implementation are provided in Fig 3 (first column of right figure).  For unpaired data translation, RF is not competitive and adding noise improves performance. However, for generative modeling, RF outperforms DSBM and $\\alpha$-DSBM using $\\epsilon>0$. We conjecture that when one marginal distribution is Gaussian, adding further noise to the interpolant hurts the generative capabilities. In addition, the fixed points of RF exhibit straight paths but are not necessarily solutions to the OT problem, see [2] for a counterexample.\n\nFM [3] can be seen as the first iteration of DSBM with $\\varepsilon = 0$. It consistently produces inferior results compared to RF, this is why we did not include it. \nOT-Flow Matching [4] is closer to our setting and, if the batch size has infinite size, then we recover the OT solution. In the one-page rebuttal we show that $\\alpha$-DSBM outperforms this method in the Gaussian case.\n\nFinally, we also highlight that DDBM [5] do not solve the (E)OT problem but requires paired data to train the model and therefore is not comparable to our approach.\n\n[1] Liu et al., Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow, 2023.\n\n[2] Liu, Rectified Flow: A Marginal Preserving Approach to Optimal Transport, 2022.\n\n[3] Lipman et al., Flow Matching for Generative Modeling, 2023.\n\n[4] Tong et al., Improving and generalizing flow-based generative models with minibatch optimal transport, 2024.\n\n[5] Zhou et al., Denoising Diffusion Bridge Models, 2024\n\n> The paper could benefit from a more comprehensive discussion of related work\n\nIn the revised manuscript, we have expanded the related work section in the appendix to offer a more comprehensive overview. Notably, we discuss a taxonomy graph of diffusion methods and their connection to OT, see the one-page rebuttal.\n\n> Could the authors provide more detailed quantitative data, such as training time comparisons with related methods?\n\nBidirectional online learning approximately halves the number of network parameters, thereby reducing GPU memory usage and wall-clock training time. For AFHQ-64, pretraining bidirectional and 2-network models takes 4 and 9 hours, while fine-tuning takes 14 and 19 hours, respectively. We use the same number of gradient updates for both models. The computational cost difference comes from the use of a bidirectional versus a 2-network model. Note that bidirectional models are incompatible with the original iterative DSBM and are suitable only for online finetuning. We have included these cost and complexity considerations in the appendix.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TFzdHXeFb3", "reviewer": "Reviewer_xpT6", "summary": "This work introduces $\\alpha$-DSBM, a new way of training DSBM-like models, which does not require a Markovian projection at each step and eliminates the need to train multiple models. The main advantage over previous DSBM-based approaches is that $\\alpha$-DSBM only needs to train a single model with a single model, thus exhibiting a more stable training procedure.\n\nThe authors thoroughly contextualize their work within related work and provide detailed theoretical derivations to motivate their approach. Empirically $\\alpha$-DSBM is validated by comparing to existing DSBM-based approaches across different unpaired image translation tasks as well as some toy data examples.", "questions": "- How does $\\alpha$-DSBM perform compared to bridge matching with mini-batch OT sampling empirically? While it is clear that these methods introduce significant errors because of the mini-batch approximation, it is unclear how this affects results empirically compared to $\\alpha$-DSBM. I think this is the most important additional competing work (apart from DSBM), as these methods share their overall goal.\n- Does it make sense to combine $\\alpha$-DSBM with mini-batch OT sampling for the initial pretraining? How would this impact $\\alpha$-DSBM?", "limitations": "- The main limitation of $\\alpha$-DSBM is that new data needs to be generated with the current model during its fine-tuning, making the procedure not simulation-free and, thus, more expensive. This is sufficiently addressed in the paper.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The new $\\alpha$-DSBM is introduced based on thorough theoretical motivation and derivations.\n- The authors manage to present and contextualize their method very clearly within the scope of related work across the different bridge matching approaches (Appendix E) and highlight connections accordingly (e.g. connection to the reflow procedure).\n- In general, the paper is well-written and, given the complexity of the topic, quite clear to follow.\n- Empirical validation includes important ablations giving further insights into the hyperparameter choice of e.g. $\\varepsilon$.", "weaknesses": "- While the experimental section of the paper thoroughly analyzes and compares DSBM and its different flavors, including the proposed method, a comparison with other competing methods could further support $\\alpha$-DSBM through empirical evidence. Specifically, a comparison to bridge/flow matching with mini-batch OT similar to [1, 2], a highly optimized CycleGAN as in [3], and adversarial-based OT methods like [4] could be good candidates for further comparisons. The authors also mention these as competing methods.\n\n\n[1] Alexander Tong and Kilian Fatras and Nikolay Malkin and Guillaume Huguet and Yanlei Zhang and Jarrid Rector-Brooks and Guy Wolf and Yoshua Bengio. \"# Improving and generalizing flow-based generative models with minibatch optimal transport\". IN TMLR 2024.\n\n[2] Luca Eyring and Dominik Klein and Théo Uscidda and Giovanni Palla and Niki Kilbertus and Zeynep Akata and Fabian Theis. \"Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation\". In ICLR 2024.\n\n[3] Dmitrii Torbunov and Yi Huang and Huan-Hsin Tseng and Haiwang Yu and Jin Huang and Shinjae Yoo and Meifeng Lin and Brett Viren and Yihui Ren. \"UVCGAN v2: An Improved Cycle-Consistent GAN for Unpaired Image-to-Image Translation\". In Arxiv 2023.\n\n[4] Beomsu Kim and Gihyun Kwon and Kwanyoung Kim and Jong Chul Ye. \"Unpaired Image-to-Image Translation via Neural Schrödinger Bridge\". In ICLR 2024.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive comments and feedback.\n\n> While the experimental section of the paper thoroughly analyzes and compares DSBM and its different flavors, including the proposed method, a comparison with other competing methods could further support-DSBM through empirical evidence. \n\n> How does 𝛼-DSBM perform compared to bridge matching with mini-batch OT sampling empirically?\n\nWe thank the reviewer for their suggestion to evaluate our approach against competing methods. While we agree that comparisons with [2,3,4] would be valuable, we have chosen, in this rebuttal, to focus on [1] which was also mentioned by reviewer ovZP and is most similar to our approach. However, we will mention [2,3,4] as possible alternatives with similar goals in our extended related work section. Our comparison with OT-bridge matching in the Gaussian setting is reported in the attached one-page rebuttal. We will include a comparison to [2,4] in the final version. \n\n[1] Tong et al., Improving and generalizing flow-based generative models with minibatch optimal transport, TMLR, 2024.\n\n[2] Eyring et al., Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation, ICLR, 2024.\n\n[3] Torbunov et al., UVCGAN v2: An Improved Cycle-Consistent GAN for Unpaired Image-to-Image Translation, arXiv:2303.16280, 2023\n\n[4] Kim et al., Unpaired Image-to-Image Translation via Neural Schrödinger Bridge, ICLR, 2024.\n\n> Does it make sense to combine 𝛼-DSBM with mini-batch OT sampling for the initial pretraining? How would this impact 𝛼-DSBM?\n\nThis is indeed doable and is a great suggestion. This would require changing our pretraining strategy to include the mini-batch OT sampling during the training. Our guess is that the convergence of $\\alpha$-DSBM (and even DSBM) would be faster. To support this result, we refer to the results of [1,2] which show that the convergence of the Sinkhorn algorithm depends on the Kullback-Leibler divergence between the original guess (here the output of the (E)OT-FM methodology) and the target Schr\\”odinger Bridge. \n\n[1] Bernton, Schrödinger Bridge Samplers, arXiv:1912.13170, 2019.\n\n[2] L\\’eger, A gradient descent perspective on Sinkhorn, Applied Mathematics & Optimization, 2021.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "k3mGcalKhF", "reviewer": "Reviewer_Ti9G", "summary": "The authors consider developing a new algorithm for a Schrödinger Bridge problem of translating between two probability distributions. Motivated by the fact that current Schrödinger Bridge (SB) approaches either use mini-batch optimal transport techniques or require training diffusion at every iteration, the authors propose discretizing the Schrödinger Bridge Flow (SB Flow) of path measures that converge to SB. They call this discretization $\\alpha$-Iterative Markovian Fitting ($\\alpha$-IMF) procedure. Using the fact that Markovian path measures characterizing the considered flow can be parametrized by diffusion vector field $v$, authors propose a non-parametric scheme for updating $v_n$ corresponding to $n$ discretization step of SB Flow and prove its correctness. To implement the procedure in practice, authors prove a parametric counterpart to the non-parametric update scheme that allows for optimization over a vector field $v_{\\theta}$ parametrized as a neural network called $\\alpha$ Diffusion Schrödinger Bridge Matching ($\\alpha$-DSBM). Furthermore, the authors computationally validate the proposed algorithm on both vector and image state spaces. The authors compare their $\\alpha$-DSBM algorithm to other DSBM versions [3] in the Gaussian setting. Image state space experiments were held on image-to-image translation tasks, i.e., MNIST $\\leftrightarrow$ EMNIST, Wild $\\leftrightarrow$ Cat on AFHQ 64/256, with quantitative evaluation and comparison with the original DSBM algorithm [3].", "questions": "- Since the AFHQ dataset is quite small, experiments on unpaired male-to-female translation on the Celeba [5] dataset could provide a more meaningful comparison (as it was in DSBM[3]).\n- How do DSBM and $\\alpha$-DSBM compare in terms of computational budget? Could authors provide a quantitative comparison?\n- How does variation of $\\alpha$ affect speed of convergence and deviation from the SB solution?\n- Can authors provide experimental results of $\\alpha$-DSBM on more complex Gaussian distributions or multimodal low dimensional data, i.e., gaussian mixtures? \n- Is it possible to extend the proposed approach to GSBM [4], where, as far as my understanding goes, authors change the reciprocal projection step compared to DSBM [3]?\n- Can this method with small $\\epsilon$ can be applied to the generation problem? This could drastically reduce the computational budget of such and be applied in modern text-to-image generative models [2] similar to the work [1]. It would be great to see such experimental results. \n- Is it possible to perform *parametric* updates for Sinkhorn flow in a similar to $\\alpha$-DSBM way?\n\nOverall, the paper looks very promising. I will consider increasing my score if the authors address some of the listed weaknesses and questions, especially related to the same computational budget comparison, and evaluate their method on one of the listed more complex setups with the known ground truth distributions.\n\n[1] Liu, Xingchao, et al. \"Instaflow: One step is enough for high-quality diffusion-based text-to-image generation.\" The Twelfth International Conference on Learning Representations. 2023.\n\n[2] Esser, Patrick, et al. \"Scaling rectified flow transformers for high-resolution image synthesis.\" Forty-first International Conference on Machine Learning. 2024.\n\n[3] Shi, Yuyang, et al. \"Diffusion Schrödinger bridge matching.\" Advances in Neural Information Processing Systems 36 (2024).\n\n[4] Liu, Guan-Horng, et al. \"Generalized Schrödinger Bridge Matching.\" The Twelfth International Conference on Learning Representations.\n\n[5] Liu Z. et al. Deep learning face attributes in the wild //Proceedings of the IEEE international conference on computer vision. – 2015. – С. 3730-3738.\n\n[6] Hicham Janati, Boris Muzellec, Gabriel Peyré, and Marco Cuturi. Entropic optimal transport between unbalanced gaussian measures has a closed form. Advances in neural information processing systems, 33:10468–10479, 2020.\n\n[7] Gushchin N. et al. Building the bridge of schrödinger: A continuous entropic optimal transport benchmark //Advances in Neural Information Processing Systems. – 2023. – Т. 36. – С. 18932-18963.", "limitations": "The authors adequately addressed the limitations and potential negative societal impact of their work.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The overall theory is novel and interesting. The authors prove that the proposed $\\alpha$-IMF procedure with non-parametric updates converges to the solution of SB.\n- The parametric counterpart to $\\alpha$-IMF ($\\alpha$-DSBM), which can be used with neural networks, is proposed and empirically verified.\n- The possibility of online $\\alpha$-DSBM finetuning seems beneficial in terms of computational resources compared to regular DSBM iterative training.\n- A sufficient comparison of the proposed $\\alpha$-DSBM with the previously known DSBM algorithm [3] is presented for a wide set of hyperparameter $\\epsilon$.\n- The Appendix presents an extensive additional theoretical review of the proposed concept and its connections to existing research.\n- The study on the parameterization of diffusion in both directions by one neural network with varied conditional vectors is proposed and tested.", "weaknesses": "- The proposed $\\alpha$-DSBM algorithm [3] is, in some sense, an advancement of the previously known DSBM algorithm [3]. It is expected to be more efficient because it removes the necessity to learn a distinct model at each iteration. However, the paper lacks a study of image quality boost given the same computational budget. Thus, it is not clear whether the authors achieved the goal of developing a more efficient algorithm.\n- The convergence of parametric $\\alpha$-DSBM is not proved. While it may be a hard theoretical problem to prove, at least a more extensive empirical evaluation of the setup with the known ground truth solution would be beneficial. The presented setup with a scalar covariance matrix is too simple. It would be a more solid argument to evaluate the Gaussian setup with a full covariance matrix [6] or use a mixture benchmark [7] and compare it with other neural network methods. \n- The FID metric is similar for DSBM and $\\alpha$-DSBM for the AFHQ experiment in Table 1 considering measured standard deviation. Thus again highlighting the lack of study with the same computational budget to justify that $\\alpha$-DSBM is more efficient.\nThe chosen AFHQ dataset may not be the best choice considering the small test set, which may introduce bias in measuring the FID metric. In turn, the visual comparison is hard since most images produced by DSBM and $\\alpha$-DSBM and given in Figure 19 are similar, and it is hard to deduce which algorithm is better. A bigger dataset such as Celeba [5] may help to solve this issue. At least a bigger dataset would give a good estimate of FID metrics. \n- It seems that some hyperparameters for DSBM AFHQ experiments are missing. The number of pre-training iterations, grad updates per IMF iteration, and IMF iterations are not specified. What is the training time and computational budget? \n- There is no study on *cornerstone* hyperparameter $\\alpha$. What is the value of $\\alpha$ used in the experiments? Is it equal to the learning rate of the diffusion model parameters?\n- The absence of code may cause trouble with reproducing the presented results.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments and thoughtful questions, we are glad you enjoyed the paper.\n\n> The paper lacks a study of image quality boost given the same computational budget. \n\n> It would be a more solid argument to evaluate the Gaussian setup with a full covariance matrix [6].\n\n> A bigger dataset such as Celeba may help to solve this issue.\n\nWe appreciate the reviewer's suggestions regarding our experimental setup. Concerning the image quality improvement given the same computational budget, we'd like to emphasize the inherent challenge in designing appropriate metrics for evaluating OT-based algorithms, as we are considering two distinct metrics: 1) Image quality (FID score): This roughly quantifies the alignment between the marginals $\\pi_0$ and $\\pi_1$ and the output distributions of the sampling process, 2) Alignment between samples from $\\pi_0$ and $\\pi_1$: This quantifies the minimization of the OT cost. Our objective is to minimize both the FID score and the alignment cost, necessitating the tracking of both these metrics. Nevertheless, we have conducted experiments comparing DSBM and $\\alpha$-DSBM with equivalent computational costs, see the one-page supplementary.\n\nAdditionally, we have now evaluated our methodology in a complex Gaussian setting using full covariance matrices, similar to the approach in [1]; see one-page additional rebuttal. The findings from this full covariance matrix setting agree with our earlier conclusions from the scalar setting.\n\nFinally, we have also evaluated our method on CelebA and refer to the one-page supplementary.\n\n[1] Janati et al., “Entropic optimal transport between unbalanced gaussian measures has a closed form”, 2020.\n\n[2] Chen et al., “Gradient flow in parameter space is equivalent to linear interpolation in output space”, 2024\n\n> The number of pre-training iterations, grad updates per IMF iteration, and IMF iterations are not specified. What is the training time and computational budget?\n\nWe have added the missing hyperparameters for the training of DSBM in the case of the AFHQ experiment to the revised manuscript. The number of pretraining iterations is fixed to 100 000. The number of gradient updates per IMF iteration is 500 and the number of IMF iterations is 40.\n\n> There is no study on cornerstone hyperparameter 𝛼. What is the value of 𝛼 used in the experiments? Is it equal to the learning rate of the diffusion model parameters?\n\nWe agree that we should have provided more discussion about the choice of parameter $\\alpha$.  Currently, we say (l.218) “In Algorithm 1, we specify $\\alpha \\in (0,1]$ as a stepsize parameter. In practice, we use Adam for optimization, thus the choice of $\\alpha$ is implicit and adaptive throughout the training.” So indeed, $\\alpha$ is linked to the learning rate and if one were to use SGD and not Adam, $\\alpha$ would be exactly the learning rate. To complement this answer and also fully answer reviewer KqVg, we train a model with SGD and a sweep on the values of $\\alpha$. The results are reported in the one-page rebuttal. We found  that setting explicitly the value of $\\alpha$ in SGD can yield similar results as letting $\\alpha$ adaptive using Adam. However, for too large values of $\\alpha$ the algorithm diverges.\n\n> The absence of code may cause trouble with reproducing the presented results.\n\nWe are working on releasing the notebooks to reproduce our experiments on Gaussian data. Unfortunately, due to confidentiality agreements and intellectual property concerns, we are unable to open source the full source code. We have included a code snippet of our main training loop implementing our online methodology ($\\alpha$-DSBM) in the appendix of the revised paper. \n\n> Is it possible to extend the proposed approach to GSBM\n\nWe thank the reviewer for this remark. Indeed our online approach can be used to improve the computational efficiency of every method derived from DSBM. For example Generalized SBM [1] would also benefit from the techniques used in $\\alpha$-DSBM. We have added a comment on this in the revised manuscript. \n\n> Can this method with small epsilon can be applied to the generation problem? [...] [2] similar to the work [1].\n\nYou are correct: $\\alpha$-DSBM can be used with $\\varepsilon$ small (or even $\\varepsilon = 0$, even though in that case, our theoretical framework does not guarantee the convergence to the solution of the OT problem). Hence, our technique can potentially enhance the performance of the 2-Reflow of [1] in a generative model context. However, we emphasize that in [2], there is no finetuning stage and only a flow-matching pretraining (i.e. this text2img model only corresponds to the pretraining of our model). As the pretraining and finetuning of such models is long and expensive, we could not conduct such experiments during the rebuttal period, however we plan to include such text2img experiments in the final version of the paper.\n\n[1] Liu et al., InstaFlow: One Step is Enough for High-Quality Diffusion-Based Text-to-Image Generation, 2024.\n\n[2] Esser et al.,  Scaling Rectified Flow Transformers for High-Resolution Image Synthesis, 2024.\n\n> Is it possible to perform parametric updates for Sinkhorn flow in a similar to $\\alpha$-DSBM way?\n\nThe applicability of a methodology similar to $\\alpha$-DSBM to Sinkhorn-flow [1] is not straightforward. First, $\\gamma$-Sinkhorn is an inherently static algorithm. In order to apply a methodology closer to $\\alpha$-DSBM, we would need to develop a dynamic version of $\\gamma$-Sinkhorn. This would correspond to extending DSB [2] to an online setting. It is unclear whether the objective of $\\gamma$-Sinkhorn, see Lemma 1 in [1], can be modified to yield a tractable loss for an hypothetical $\\gamma$-DSB version.\n\n[1] Karimi et al., Sinkhorn Flow: A Continuous-Time Framework for Understanding and Generalizing the Sinkhorn Algorithm, 2023.\n\n[2] De Bortoli et al., Diffusion Schrödinger Bridge with Applications to Score-Based Generative Modeling, 2021.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WNE9SdmjzH", "reviewer": "Reviewer_KqVg", "summary": "This paper proposes $\\alpha$-IMF, which is an incremental version of IMF method in DSBM paper. Specifically, this work demonstrate the convergence properties of $\\alpha$-IMF and implement it through an online learning method. Here, $\\alpha$ is implicitly reflected. Moreover, the functional flow interpretation provides a basis for online learning, allowing for online learning as opposed to the iterative approach of DSBM. From an engineering perspective, the authors suggest parameterizing forward and backward control simultaneously by adding a forward/backward indicator (0 or 1) as an input variable to the neural network. The method shows feasible results on toy (Gaussian) data and image-to-image translation tasks.", "questions": "- If we perform this algorithm with full gradient descent algorithm, is it possible to explicitly control $\\alpha$?\n- Can we develop exact algorithm that reflect $\\alpha$ explicitly? If so, can we obtain some experimental results that compares $\\alpha$?", "limitations": "The authors adequately addressed the limitations.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "-\tThis work is theoretically fruitful. The introduction of $\\alpha$-IMF leaded to a more expansive theoretical framework. This flow-based perspective enables online learning. \n-\tReduced the sensitivity to hyperparameters such as the iteration number, the number of phase in DSBM.\n-\tExpands the interpretation of the Schrödinger bridge by demonstrating connections with reinforcement learning (RL) and Sinkhorn flow.", "weaknesses": "- The experiment on practical data is insufficient and the performance improvement is incremental. The paper only presents real-world I2I on Cat <-> Wild. Moreover, the performance of bidirectional online learning improves only incrementally. Totally, it needs more evaluation on other kinds of dataset, e.g. Male <-> Female / Handbags <-> Shoes. \n- It would be beneficial to quantify the advantages of bidirectional learning or online learning over DSBM. For instance, comparing the actual wall-clock time or GPU memory usage would be useful. \n- The current algorithm reflects $\\alpha$ implicitly, and so there is no analysis provided on $\\alpha$. If there were any intuition or insights into the role of $\\alpha$ (step size) or if an algorithm could be developed to explicitly control $\\alpha$, it would offer more valuable insights.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments, we appreciate your acknowledgements of the paper’s merits.\n\n> The experiment on practical data is insufficient and the performance improvement is incremental. The paper only presents real-world I2I on Cat <-> Wild. Moreover, the performance of bidirectional online learning improves only incrementally. [...] For instance, comparing the actual wall-clock time or GPU memory usage would be useful.\n\nBidirectional online learning approximately halves the number of network parameters, thereby reducing GPU memory usage and wall-clock training time. In the case of AFHQ-64, pretraining bidirectional and 2-network models takes 4 and 9 hours, while finetuning takes 14 and 19 hours, respectively. We use the same number of gradient updates for both models. The computational cost difference comes from the use of a bidirectional versus a 2-network model. However, note that bidirectional models are incompatible with the original iterative DSBM and are suitable only for online finetuning.\nA fair comparison of DSBM and α-DSBM in terms of the convergence speed, is then possible with a 2-networks model. It is best illustrated by the Gaussian example where there is a single evaluation metric.\n\n> The current algorithm reflects 𝛼 implicitly, and so there is no analysis provided on 𝛼. If there were any intuition or insights into the role of 𝛼(step size) or if an algorithm could be developed to explicitly control 𝛼, it would offer more valuable insights.\n\nIf we perform this algorithm with full gradient descent algorithm, is it possible to explicitly control 𝛼?\nCan we develop exact algorithm that reflect 𝛼 explicitly? If so, can we obtain some experimental results that compares 𝛼?\n\nWe agree with the reviewer that parameter $\\alpha$ merits further discussion. In our current manuscript, we state in  (l.218) “In Algorithm 1, we specify $\\alpha \\in (0,1]$ as a stepsize parameter. In practice, we  use Adam (Kingma and Ba, 2015) for optimization, thus the choice of $\\alpha$ is implicit and adaptive throughout the training.” So, $\\alpha$ is linked to the learning rate: in a Stochastic Gradient Update (SGD), $\\alpha$ would correspond to the learning rate. To address this point comprehensively (and also answers reviewer Ti9G), we train a model with SGD and a sweep on the values of $\\alpha$; see additional one-page rebuttal. We found out that setting explicitly the value of $\\alpha$ in SGD can yield similar results as letting $\\alpha$ adaptive using Adam. However, for too large values of $\\alpha$ the algorithm diverges.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We sincerely thank all reviewers for their valuable time and insightful feedback. We appreciate the thoughtful questions and the overall positive response.
We have provided detailed responses to each reviewer's comments. In summary, the key feedback points we have received are as follows:
* **Additional experiments**: as requested by all reviewers, additional experiments have been conducted to better illustrate the efficiency of the proposed method. In the one-page rebuttal, we have conducted additional experiments regarding the choice of $\alpha$, comparisons with OT bridge matching and experiments on CelebA. 
* **Better support for the parametric case**: as suggested by reviewer TiG9 we have conducted an improved analysis of the parametric case. Following the recommendation of reviewer ovZP we have also included the assumptions our main theorems rely on in the main body of the paper.
* **Effect of the parameter alpha**: Reviewer TiG9 and reviewer KqVg have asked for more details regarding the choice of the parameter alpha. We have conducted additional experimental analysis of the effect of this parameter and report them in the one-page rebuttal.

Attached to this rebuttal is the additional one-page PDF with additional results and tables.

Meta-review: This paper introduces a novel approach to solving Entropic Optimal Transport problems via the Schrödinger Bridge framework, demonstrating theoretical convergence and empirical improvements in image translation tasks. The reviewers generally found the work to be strong, with two reviewers (xpT6 and Ti9G) giving it an accept score of 7, showing clear support for the paper.

While Reviewer ovZP had concerns about the clarity of assumptions and comparisons with SOTA methods, these were sufficiently addressed during the discussion, leading to an increased score of 6. Reviewer KqVg also maintained a score of 6, acknowledging the incremental contributions but remaining positive overall.
Given the overall positive scores, particularly strong positive attitudes from two of the reviewers, I recommend accepting this paper.
Decision: Accept (spotlight)