BM$^2$: Coupled Schrödinger Bridge Matching
Authors: Stefano Peluchetti
Venue: Accepted by TMLR ()

Abstract: A Schrödinger bridge establishes a dynamic transport map between two target distributions via a reference process, simultaneously solving an associated entropic optimal transport problem.
We consider the setting where samples from the target distributions are available, and the reference diffusion process admits tractable dynamics.
We thus introduce Coupled Bridge Matching (BM$^2$), a simple \emph{non-iterative} approach for learning Schrödinger bridges with neural networks.
A preliminary theoretical analysis of the convergence properties of BM$^2$ is carried out, supported by numerical experiments that demonstrate the effectiveness of our proposal.

Published in Transactions on Machine Learning Research (12/2024)
BM Coupled Schrödinger Bridge Matching
Stefano Peluchetti
stepelu@iskkano.dt
Sakana AI
Reviewed on OpenReview: https: openreview. net/ forun? id= Mg
Abstract
Schrödinger bridge establishes dynamic transport map between two target distributions via reference process, simultancously solving an associated entropic optimal transport problem. We consider the setting where samples from the target distributions are available, and the reference diffusion process admits tractable dynamics. We thus introduce Coupled Bridge Matching (BM²), simple non-iterative approach for learning Schrödinger bridges with neural networks. A preliminary theoretical analysis of the convergence propertics of BM is carried out, supported by numerical experiments that demonstrate the effectiveness of our proposal.
1 Introduction
The Schrödinger bridge problem seeks process, the Schrödinger bridge, with prescribed initial and terminal distributions, such that the distribution of the Schrödinger bridge minimizes the Kullback-Leibler (KL. divergence to the distribution of reference process. Schrödinger bridges play a central role in measure transport theory (Marzouk et al., 2016). Notably, it is known that the initial-terminal distribution of a Schrödinger bridge provides solution to a corresponding entropic optimal transport problem (Peyré Cuturi, 2020). Schrödinger bridges thus provide an effective framework for finding an alignment. between samples from two target distributions. Furthermore, diffusion-basset generative models (Ho et al., 2020; Song al., 2021) can be interpreted solving trivial instances of the Schrödinger bridge problem (Peluchetti, 2023). Consequently, Schrödinger bridges offer more general approach to contemporary generative applications.
We consider the sctting where samples are readily available from both target distributions, and where the reference process is a diffusion process solution to a stochastic differential equation (SDE). We thus introduce Coupled Bridge Matching (BM?), a novel methodology aimed at computing the Schrödinger bridge given the reference SDE and samples from the two marginal distributions of interest. BM builds upon Bridge Matching (BM), introduced by Peluchetti (2021). Our approach advances recent contributions by Peluchetti (2023); Shi et al. (2023) by removing the need to solve sequence of optimization problems. A neural network is employed to jointly learn forward drift function and backward drift function corresponding to the forward and backward dynamics of a Schrödinger bridge. BM2 achieves several key desiderata:
non-iterative: training is conducted through standard stochastic gradient descent within a single optimization loop;
(ii) exact: the idcalized version of BM2 yields the target Schrödinger bridge without approximations: the only sources of error involved in its practical implementation are the neural network approximation error and the discretization error due to sampling the learned SDE;
(iii) efficient gradient: the gradient of the loss function with respect to neural network parameters depends solcly on few random variables sampled at the current optimization step;
Peluchetti (2021) used the term "Diffusion Bridge Mixture- Mixcure-Matching Transport" (DBMT), but we follow Shi al. (2023) in using the sleeker nomenclature "Bridge Matching" for this transport.

---

Published in Transactions on Machine Learning Research (12/2024)
(iv) simple loss: the loss function avoids derivative terms with respect to neural network inputs and does not impose hard constraints (such as conservative vector field requirements) on the neural network approximator.
These features collectively enhance the efficiency and applicability of BM2 in solving Schrödinger bridge problems. Training is robust, as it does not depend on hyperparameters that are lypically challenging to set. without time-consuming pilot runs, such as the number of training steps per optimization iteration (i) or the level of approximation (ii). Moreover, the memory requirements are modest due to (iii). Finally, the implementation is straightforward (i, iv), as illustrated in Algorithms 1 and 2 and in the annotated PyTorch code of Listing 1.
Content: This paper is structured as follows. In Section 2, we formally introduce the Schrödinger bridge problem with associated reference process dynamics. Section reviews Bridge Matching, while Section 4 introduces Coupled Bridge Matching, discussing its theoretical properties and implementation aspects. Numerical experiments are presented in Section 5, followed by a discussion of related works in Scction 6. Section 7 concludes the paper. For clarily, more general formulation of BM2 is deferred to Appendix A, all proofs to Appendix B, an additional numerical experiment to Appendix C, and code listings to Appendix D.
Notation and Assumptions: To enhance accessibility, we refrain from discussing the more technical aspects related to the Schrödinger bridge problem in its path measure formulation. The excellent treaties of Léonard (2014b;a) and Bortoli et al. (2021, Appendices D, H) already serve this goal. We denote distributions with upperease letters and their corresponding (Lebesgue) densitics with lowercasc letters. All stochastic processes considered are d-dimensional, continuous, and defined on the unit time interval [0, 1]. For a stochastic process X with distríbution P (denoted X ~ P), use subscripts to specify marginal distributions, joint distributions, and conditional distributions of P. Pr: marginal distribution of X, at time t, with density Pt; Po.1: mntiallterminal joint distribution of X1); Plo: distribution of X given its initial value Xo- Superscripts indicate a distribution P's dependency on another distribution Z, as in p², or a sequence of distributions, as in p()), 2 1. For a d-dimensional distribution Qu. we define the stochastic process mixture distribution QoPio as: (Qo E === P E (200) From a generative perspective, X~ QoPo is obtained by sampling Xo ~ Qo and then X ~ Po conditionally on Xo. The marginal- conditional decomposition of P over its initial value is thus P PoPio- Similarly, for a dxd-dimensional joint distribution Qu,1, define Q0,1Po, such that X ~ Q0 Plo. is obtained by sampling X1) ~ Q0 and then X ~ P (- X1) conditionally on Xo and X,. Time is always indexed on a common forward timescale, on which all stochastic processes' distributions are defined. The dynamics of a diffusion process X X~ can be formulated in both forward and backward time directions, through corresponding forward and backward SDEs. In backward SDEs, decreases from to (di negative), which is denoted by L E [1, All Brownian motions are independent. Unless otherwise noted, each diffusion process is a Markov diffusion process which is (weak) solution to an associated SDE.
2 Problem Setting
2.1
Schrödinger Bridges and Entropic Optimal Transport
For two target d-dimensional distributions Vo and v and a reference stochastic process distribution R, the dynamic Schrödinger bridge (SB) problem seeks to find
SWo.Vi.. = argmin min KL R), PEP(No.V..)
(1)
where KL(- -) is the KL divergence and P(Do, D1) is the class of distributions of stochastic processes having initial distribution Vo and terminal distribution V, We narrow down (1) to the case where It is the distribution of a diffusion process. In this case, under suitable conditions (Léonard, 2014b), (1) admits a unique solution which is also a diffusion process. From this point forward, Vo, Vi and are considered fixed. For brevity, we will thus denote the Schrödinger bridge W..R simply as S, and apply the same notation convention to any distribution dependent on these variables.
2

---

Published in Transactions on Machine Learning Research (12/2024)
The forward and backward dynamics of X ~ S are given by:
Xo Vo, dX, [X, t)dt + odW, t [0,
(S)
(5)
MI,,, dX, = (X, + odW, E [1,0],
for the SB-optimal drift functions Us These functions are related to the Schrödinger potentials (Léonard, 2014b) and are not analytically available aside from very specific choices of Mo. 1 and R.
We assume that Ro,1 admits density ro.1- Once S is obtained, the solution to the static Schrödinger bridge problem given by So
Sa, = arg argmin min (L(Co,1 |I Ro,1);
Co.1 Co.
(2)
arg mig Co. Co.1
(o Xo)] - HC..,)).
In (2), C( V1) denotes the class of dxd-dimensional joint distributions with marginal distributions Vo and I1 commonly referred to as the class of couplings of Vo and Vi, and IH,(),,) -cc..l- log (X,, is the entropy of Co.1-
The entropic optimal transport (EOT) solution for the cost function k(zo, m1) and regularization level & is given by:
Eo.1 = arg min
Co, Co.1 [x( Xo)] cH,,))
(3)
Thus, for each choice of R in (2), So,1 solves corresponding problem (3). As in the following, when R is associated to (R), 1 solves the EOT problem (3) for the Euclidean cost n(2),, zo) 1/21120 2 I|? and regularization level o².
We refer to Peyré & Cuturi (2020); Léonard (2014b); Gushchin et al. (2023) for related background material from complementary perspectives.
2.2 Reference Dynamics
We focus on the case where R is the distribution of a scaled Brownian motion:
Xo ~ Vo. dX2 = odW, E [0,
(R)
with σ > 0. Our approach is not limited to the choice of SDE (R), BM² readily extends to the broader class of reference SDEs examined in Poluchetti (2023). The main requirement for the applicability of BM² is the analytical availability of (4, 5) for the chosen reference SDE. We address the case, commonly employed in generative applications, of d.x, = /B,W,, for a schedule 3g explicitly Appendix A, and refer the reader to Peluchetti (2021; 2023) for the general setting. As our developments are orthogonal the specific choice of reference process, we focus the simplest case for explanatory reasons.
We collect here various results concerning (R) that will be utilized in the following:
Ryjl (ooo) . NNo, o² d²1), Rejo, (- x1) NN 1) x21, 02 1)); poi (It, 11 == o² Vx, (
(4)
(5)
I1 It.
(6)
1 t.
Voi (Iz, L ro) : à² V,, log logrgp(r:
30 It
(7)
1
xo(1 1) + 211 Ft 1(1 (1)
You ( x0, = 2 Vz, ( 21)
(8)
Conditioning X ~ R on the endpoints Xo xo, X1 = I1 results in the diffusion bridge distribution R with associated forward and backward SDEs:
Xo 3o: dX, poi (Xe, 1, ddld odW,, E [0, 1].
(Ro.1)
3

---

Published in Transactions on Machine Learning Research (12/2024)
Xi æ1, dx, -um (X,, odW, I [1,0].
(Ro..)
3 Bridge Matching (BM)
We succinetly review Bridge Matching, and refer to Peluchetti (2021; 2023); Shi et al. (2023) for more details. BM takes as input a joint distribution Q0,1 with marginal distributions Qu. Q1 and SDE, (R). Firstly, a stochastic process 17Q.,1 is constructed as a mixture of diffusion bridges 1),,),, such that the endpoints (1) of X 11 are distributed according to Q0.1. This process, which is a mixture of diffusion processes, is not itself a diffusion process in general (Jamison, 1974). However, we can obtain a marginal- matching diffusion process with distribution M for which M = 11,0,1 0 s t S Consequently, X ~ M is a diffusion process for which Xo ~ Qo and X1 ~ Q1, i.. it defines a dynamic transport from Qu to Q1-
Concretely, let 17 Q0,1 Ro..1 The BM transport based on Qa,1 with distribution M realized by
Xo ~ Qo
dX, = pm Ao.1 (X,, odw,, [0,1],
(M)
E B MooI Xx, 1, [x)
(M
X1 ~ Q1,
dX,- u (Xe, d) + odW, te [1,0]
E En (vo1 Xo)
and satisfies M = 11
fo. 5
As conditional expectations are mean squared error minimizers, suitable training objectives for the drift Qa, and 2'm
functions um
Qa. are derived from
Qo.1 arg min
m(X,, dtt
(9)
17.
2
1 2
/ (X,,,, Xo) v(X,, t) 11
mm
00. arg min n1.
(10)
by replacing each integral with an expectation over uniform time t ~~ 2(0 1), and then approximating both expectations with Monte Carlo estimators. While we will rely exclusively on (9, 10) in the experiments of Section 5, pm Qo.1 and um Qa can be inferred from paths X ~ I also by performing maximum likelihood estimation or by employing a drift matching estimator (Liu et al., 2022; Peluchetti, 2023).
We conclude this section by reviewing prior BM results relevant for BM2. Define:
P = {d-dimensional, contimuuss, stochastic processcs on [0. 1]}
R (P PIP P= Ro,1 Po. for some Pa.
M (P P P is (Markov) diffusion process),
S = (P E P P is Schrödinger bridge for some target marginal distributions} = R R.M,
where the equivalence established by Jamison (1975) under appropriate assumptions. We additionally define the following restrictions: P (W. = (P E P Po Vo), P(: Wn = (P E P P. = v,), P(Vo, 1) {P PI Po Vo and P 1. Restrictions M, and C employ the same notation.
For Q E P, it is instructive to view BM map between distributions defined by the composition of two projections: Q Re uIQ... MY MQo. Here, the reciprocal projection Rp P - R projects Q onto the reciprocal class R. while the Markovian projection Mp : R + M projects ITQ..1 onto the class of diffusion processes, see Shi et al. (2023). follows that if R, then P = Rp(P), and P E,,, then p Mp(P). Consequently, if P PES, S, P = (Mp (P) for the BM map (Mp Rp), and conversely if P (Mp Rp) then S.
3.1 Iterated Bridge Matching (I-BM) and Diffusion Iterative Proportional Fitting (DIPF)
In the dynamic setting, Peluchetti (2023): Shi al. (2023) demonstrate that, under suitable conditions, iterative application of the BM procedure to initial coupling Ca E C(V,, V,) results in convergence
4

---

Published in Transactions on Machine Learning Research (12/2024)
toward S. Specifically, defining T(0) == M and 7(0 == M' for 1 1, it holds that KLL((() (| +
as + co. In practical applications, the independent initial coupling given by the product distribution Ca.1 = 1 frequently employed.
In the static setting, the classical procedure employed in solving problems (2, is known by several names: the Sinkhorn algorithm (Peyré & Cuturi, 2020), the Iterated Proportional Fitting (IPF) procedure (Ruschendorf, 1995), or Fortet iterations (Fortet, 1940). The iterates are given by D = VoRioo D = v K D = VoD and so on. At each iteration, one of the target marginal distributions is replaced while the remaining conditional distribution is kept fixed. Alternatively, one can start from 1 Roj following the same logic. Under suitable conditions (Ruschendorf, 1995), KL convergence (L || D - 0 is established. The key insight of Bortoli et al. (2021); Vargas et al. (2021) is that it is possible to extend the IPF iterations to the dynamic setting. In this case, the IPF itcrations are implemented by learning the time reversal of a diffusion process at each iteration. We refer to the resulting training algorithun, as proposed by Bortoli et al. (2021), as Diffusion Iterative Proportional Fitting (DIPF). Bortoli al. (2021) establishes the convergence properties of the DIPF iterates, see their Propositions 4, and Section 3.5.
ijo:
4 Coupled BM (BM²)
As starting point for the derivation of BM2, consider the system of equations
H Kaa VoMM Ko, k = V M Hó.1
(11)
whose variables are diffusion distributions Ho.,, Koo,, and H', K', That is, H Kaa is obtained as the BM transport based on K₁,1 conditioned to have initial distribution Vo, while K is obtained as the BM transport based on H6,1 conditioned to have terminal distribution Vi. Equivalently, (11) is expressed as
Xa Vos dX, pm
Ko.1 (X,, l)dl + odW, LE
(H*6.))
Ho. X1 ~ L dX,- = UU (Xe, + odW, te [1.0].
(K
All of Jon Um, MM are defined in Section 3. System (11) defines an update step (H', K') (11) (HKo,, K K We are interested in the fixed points of such updates, i.e. (H', K*) such that (H'', K") (11) (H² K'). It holds that H' = K' = S is fixed point to (11). As S E S(Yo, 11), S 11 S0,0, see the review at the end of Section 3. Consequently, VoM50 = VoSpo = S and V VIM. So.1 = v1S = S. In this case, the SB-optimal
drifts Ne and v, of (S, 5) respectively replace pmm Kh. assumption that H' = K', equivalently that (H.,,,, K are the time reversal of each other, this fixed point unique. Let G= H' = K", we have G = VoM Go.1 G.M Moos M MG and Go = Vo, G1 1,, thus G = SVV., 12) S. We have shown the following:
and Um Ho.1 in (Hhb,, R...)) Under the additional
Lemma (Fixed points of (11)). Under suitable conditions (Léonard, 2014a), the updates (H', K') (1) (H*,,, (,%,,,,H,,,, parametrized by diffusion process distributions, admit H' = K' = S as fixed point. If H' K', this fized point is unique.
K²,1 = n4, and Um Ho. = U., (11) has reached an equilibrium. The updates (H', K") (11) (o K When pm
Ke. Ho.. i.e. by minimizing the losses (9, 10), where Qa,1 is respectively equal to K6,1 and Ho.1- Our proposal, BM2, follows from replacing the complete minimization of (9, 10) with partial and stochastic minimization of (9, 10) through stochastic gradient descent.
are realized through the computation of the drifts fmm. and Vm
More precisely, consider the forward and backward SDEs with distributions F(O) and B(0):
Xo Vo. dx. (x(( adW,, EE [0, 1),
(F(O)
5

---

Published in Transactions on Machine Learning Research (12/2024)
X1 v1 dx, -xx,,,xXX,,,0)d + adW,, L [1,0].
(B)))
(x,(Xe,t,)) 0) and us (x,,,, t, 0)
drift functions to be learned, which are implemented through neural network with parameters 0. Let 0' represent the values of 0 at given step during training, and define the losses
[/// (X,, X1) (X, 0,
Ly(000)= = E
12
(
(0:'0') = E
(o 2 / (X,,, - 2%6 (X,, ())l
[2
[
(12)
L(0;0) 6*) = (0; 8¹) + (0:6'')
At each optimization step, BM attempts to minimize L(0; 0¹ in 0 via step of stochastic gradient descent, starting from 0 = 6' and koeping 6' fixed, resulting in 6". The subsequent optimization step employs 0' <- 8". The complete training objective is presented in Algorithm 1, where sg() refers to the stop-gradient operator L(0; 0¹) is minimized in the first arguments only and discretize() represents a generic SDE discretization scheme. For completeness, we outline the standard SGD training loop in Algorithm 2, where sgdstep() refers an update step via a generic gradient descent optimizer.
It should be noted that merely performing coupled drift matching of F and B, wherein F learns the drift consistent with paths from B and vice versa, does not yield the Schrödinger bridge a fixed point (Bortoli et al., 2021). The introduction of the mixing process II is crucial in ensuring this property. Moreover, L(0; must be minimized only with respect to its first argument: the application of the stop-gradient operator sg() is not an efficiency consideration but necessary component.
Algorithm BM training loss computation
outputs: 1(0): sampled loss value
inputs: 0: current parameters
1: function loss(0)
2: fo Vo
D. Marginal sampling
3:
far,-.., fi ~ sg(discretize(fo, Mt, 0)))
b Discretization of (F(0))
4:
bi ~ V
P Marginal sampling
bI-Ar...., bo/bb sg(discretize(D, At, ub(.,- 0))
Discretization of (B)
t t~ 1)
Time sampling
Tf. Rejo, (fo,
Bridge sampling (5)
8:
Tbe R. bi)
Bridge sampling (5)
(0) + 1/21|m01 (Tb,, b1) (Tb t, 2
BM based on Bo,1 (6, 9)
10:
(0) + ( t, uD(, 0)
BM based on Fo,1 (7, 10)
11:
1(6) - (+( + (f)
12:
return 1(0)
Algorithm BM
training loop
outputs: 8": trained parameters
inputs: 0% initial parameters
1: function train(0")
2:
0 7-00 00
while not converged do
1(0) loss(()
Sample loss with Algorithm b Perform SGD step
0 + sgdstep((, Vol(0))
return 0
4.1
Implementation Aspects
The following aspects not presented in Listing but impacts the performance of BM2

---

Published in Transactions on Machine Learning Research (12/2024)
Path Caching: in Bortoli et al. (2021); Shi et al. (2023), to enhance efficiency, cache the initial and terminal endpoints of the paths sampled in lines 3 and 5 of Algorithm 1, and periodically refresh the cache during training. Notably, it is unnecessary to cache entire paths; only the endpoints are required for bridge sampling, which is advantageous from a memory perspective. Bridge sampling offers the additional benefit of increased sample diversity: for cached (fixed) endpoints, the samples corresponding to lines and 8 differ at. each step.
we utilize single neural network to parametrize both (z 0, and t,( As the training process is not iterative, it is unnecessary to introduce multiple neural networks (or parameters), one for each iteration.
Model:
Sampling EMA: as in Ho et al. (2020); Song et al. (2021), to improve the stability of training we apply the Exponential Moving Averaging (EMA) to the parameters employed in path sampling in lines 3 and of Algorithm
Loss Singularities: the losses of lines 9 and 10 of Algorithm diverge for t and - 0 respectively. Singularities of these kind are common to scalable losses for generative diffusion models. In our numerical experiments we simply restrict sampling of t to U(c, 1 c) for a small c >> 0. More sophisticated alternatives involve either employing the dynamics of Appendix A for an appropriate scheduling Sr, or learning terminal- value predictors in place of drift terms, recovering the latter through (6, 7)
Two-Stage Training: a siguificant challenge in early training is the inouulationnifferente mismatch. To achieve reliable results, the drift functions μ (Z t, 0) and vs(x,, t, must be accurately learned in regions where the corresponding SDEs (F(0)) and (B (0)) will be simulated. However, the processes F(0) and B(0) typically differ at initialization. Because the drift of (F(0) is inferred from samples of (B))) (and vice versa), the approximation quality can be poor; see Peluchetti (2023, Sections 2.3 and 6.2) for detailed analysis of this issue within the DIPF framework. To address this challenge, the BM transport based on the independent coupling V 1 can be learned in both directions in first stage. By construction, the mismatch. Morcover, at convergence, the processes F(f)
BM transport circunvents the simulation-inferon and (0) are the same. Subsequently, in the second stage, the BM2 transport can be learned employing the first-stage solution as initialization.
Berward-Bacwwarr Consistency: the mutual time reversal relationship between (F(O)) and (Z())), i.e., the equivalence of F(0) and (00), can be encouraged leveraging the diffusion time-reversal result of Anderson (1982), yielding the additional consistency loss term:
1 /./ [1r((X,,,,) vs(Xe,t,0) nX,, Xo, di).
E ()))(0;0'== f Ro.1 10 (1 Ro. (0'))
(13)
where 0' denotes an independent copy of 0 (implemented via the stop-gradient operator), and yo1 (ar, to, 21) is defined in (8). Indeed, for a mixture process ll =
the following relationship holds:
o² v log n Co.1 (a,)= = E hoi (Xe, t, Xo, X1) = =x,. o
Although (13) shares similarities with the consistency loss proposed by Shi et al. (2023, page 33), (13) ((."'))) When a single neural work
utilizes the dynamic mixture 1 (TT Ja.: + (TB(0")..) is employed to parametrize both (r,,,, t, 0) and vo(2,, t, 0), estimating (13) in addition to (12) results in a negligible computational overhead per SGD step.
4.2 Convergence Properties
At each training step, BM performs a partial and stochastic minimization of the loss L[0; 0') from (12) with respect to 0, where L(;; 0') is defined by an expectation over a distribution dependent on 0', yielding d". Subsequently, 6' is updated to match 8", and the process advances the next training step. The alternation between expectation and maximization steps bears resemblance to the classical Expectation Expecation-Naxiiizatio (FM) algorithm (Dempster et al., 1977).
7

---

Published in Transactions on Machine Learning Research (12/2024)
4.2.1 Complete Minimization
We start by establishing in Theorem that the version of BM² where L(0; is fully minimized at cach training step recovers the I-BM and DIPF iterations for two specific initialization choices of (F(0), (()).. The prior convergence results of Bortoli et al. (2021); Shi et al. (2023); Peluchetti (2023) (see the review of Section 3.1) toward S thus apply.
To facilitate the presentation of the convergence results in this section, we introduce, with a slight abuse of notation, the following functional versions of the losses (12):
(1((t1 = E
12 J.
M (
II 0,
2 /
(14)
Lo(vii mby) = E
(X,,,, -X -
IT
[( ubi ply, v16) == (g( + (o((s))))))
In (14) we identify My,0h th with F, B, and u'y, vi with B' (the remaining quantities defining F, B, F', B' are fixed). We will use Ly vf)), Lo(ve; ((y) and L, 0"), (s(00)) 01) interchangeably. We are now ready state our first convergence result.
Theorem 1 (Complete BM Iterations). Consider the SDEs (F(0), B(8)), with initial drifts mj 26 and
(0) (0)
corresponding distributions p()), B()), For each i 2 1, let (u) (6) ,0b ) arg (in) Lu,, u; .tb (-11),,
(1-1)
resulting in the distribution iterates F(i, B(0, We distinguish two cases:
(0) = u = 0: both the iterates F "0), (i) Ps
and the iterates B F are equiva- lent to the DIPF iterates, started respectively from the forward and from the backward time direction;
(ti) [
(0) = pm , (0) = Um
Co.1
Con1 for some C C'EC(VO 11): F() = -B( = =1(0 for each 2 0 where I() are
the I-BM iterales.
4.2.2 Partial Minimization
In the EM algorithm it suffices to perform partial maximization steps. A partial result for the setting where [0;; 0') partially minimized with respect to at each step stated Theorem 2, which is based on Lemma 1 and Lemma 2.
Lemma 2 (Loss Interpretation). holds that
KT((Bo || Vo) + (1y(11; (TT Ro.1 P) + C(B) = RL(MP 8o.1 [] + C2(B),
(15)
KL(Fi || v1) Lo(vo: F11
Fo.1 B) + Di(F)
Fo.1 || BB (M
D2),),
for Ci (B),C2(B) independent of F, (F(F), D2(F) independent of B, with 1 CB s Cz(B) and 0 s (F) S (P)...
The losses (uy; vz) and [Ls (1y) are easily amenable to optimization in their first arguments, as seen in Algorithm Lemma relates these losses to more interpretable KL divergences between distributions. By (15), a decrease of Ly vb) due to change in fz corresponds equivalent decreases of KL(M Ro.1 || F) for fixed vb; B. Thus, partial minimization of ( 06) brings F closer to B the BM transport based B and the result of complete minimization step, by means of reverse KL minimization. Symmetric considerations apply to Lo(ve; (1)) as function of its first argument. Putting this result and Lemma 1 logether yields Theorem 2.
Theorem 2 (Partial BM² Iterations). At each optimization step, decreases of L0 0') and Lo(0; 0') in 0 correspond to equivalent decreases of IKL (8') F(0)) and (L(MFo, :o. B(0).. If the losses (8; 8') and (0; 0¹) cannot be decreased in 0, i.e., at optimalily, and if F(() B()), then F(0) B(0) S.
8

---

Published in Transactions on Machine Learning Research (12/2024)
4.2.3 Infinitesimal Minimization
We conclude our Cheoretical investigation by relating our proposal to the work of Karimi el al. (2023), which introduces continuous variant of the IPF procedure. In IPF, the two target marginal distributions are replaced sequentially, one at a time. Each step corresponds to solving a static Schrödinger half-bridge problem (Léonard, 2014a), where in (2), C(Vo, V1) is replaced by either C(VD.) or C(., 11). The approach proposed by Karimi al. (2023) retains either the even or odd steps of the TPF scheme while substituting the alternate steps with partial minimizations of the corresponding half-bridee bridge problems. In the limit of infinitesimally small improvements, this yields a dynamical system for the evolution of the iterates over continuous algorithmic time.
We demonstrate that similar result can be obtained for a modified version of BM², where forward KL divergences are minimized instead of reverse KL divergences. The resulting dynamical system is a symmelrized version of the one obtained by Karimi et al. (2023). Let F", represent the current state in the optimization process. We consider partial minimization of KL(F || MM,,,, instead of KL(M || F), in F and a partial minimization of (L(B) M%,.), instead of (ML( B), in B. As in Karimi et al. (2023), partial minimization is formulated as
F(A) = arg min XKIL((F) M B6.1) + A)KKL(F)) F*), FEM(Vo.)
(16)
B(A) arg min AB M Fo.)) + (1 - A)KL(B B3), BEM(
where À E [0, controls the extent of the minimization. We begin by establishing two stability results: the updates (F', B') (16) (F(A), preserve both R and S.
Lemma 3 (R-stability of p(A),B(A) B(A)).
E R, then ((A),OC E R for each X E [0,
Lemma 4 (S-stability of F(A), B(A).. If F", B' 8, then M(A,BO E S for each [0, 1].
Provided that the initial values F', B' E S, Lemma 4 establishes that the iterates defined by the updates (F", BB) (26) (FWW),() always remain in S. It is straightforward to ensure that F", B' E S at initialization by setting the corresponding drifts to zero: 1'y,0 0, which we will assume henceforth. As M n,, B' and F F', (16) can be reformulated in simpler terms:
F(A) = arg min AKL B') + (1 () rEM(Vo:)
mig
(17)
( = arg min AKL B (1) (1 B'). BEM(-0.)
By Lemma it suffices solve (17) in the static setting,
(A) arg min AKL(Fu,1 || Bb,1) + +1 (11 F).1), Fo. (C()))
(18)
(A arg min AKIL (Bu,1 || Fo.1) + +11+01100 A)KLL B.
The dynamic solutions are then recovered by F(22) ((A) R
1o.1
We assume that Fo, By, Vo, 1 admit densitics. By calculus of variations, the solution to (18) is given by ((A)) (1)) = (x (x) (o0), and (20. = b/11 (00l/1))//(00) where f(A (21 x and 60 (xo(x1) Sop (2o|21) ( ) (2ol=)) The IPF iterations are recovered when = 1. Instead, taking the limit X 0 and applying Bayes theorem, obtain the evolution of
(1o(=1/00) ( (x1 (1

---

Published in Transactions on Machine Learning Research (12/2024)
log [10 (21 (20) and log 6/11 (20) a function of algorithmic time [0, oo) through the dynamical system
d log fll (21 (ar)
$1 (11 (20)
log
(L (20) |l bn (201)1)01((1)) L (0,
dl
0011 (ar (20 12) (x1)
(19)
d log aif (2ol21) log
o0 (1) (11)
(IL(1) (<o|>1) |I (1 LE 10,
dl
1/0 (ar (r (o)o)))))
In (19), (TL( || denotes the generalized KT. divergence between unnormalized densities, is the case here for the second arguments, and the initial conditions $1j0 (21l and 80 (2ol21)) are determined by (F(0), (0)) with null drift terms. (19) can be contrasted with Karimi et al. (2023, Equation (13)). In Appendix C.1 we report simple numerical application of (19) to the Gaussian setting, which recovers S.
5 Numerical Experiments
To evaluate the performance of BM² on EOT problems, we utilize the benchmark developed by Gushchin et al. (2023). For the reference process (R), this benchmark provides pairs of target distributions Vn. v with analytical EOT solution So,1 and analytical SB-optimal drift function μıs. We focus on the mixtures benchmark, which consist of a centered Gaussian distribution as So = Vo and mixture of 5 Gaussian distributions for Siju- S v is not mixture of Gaussian distributions, but has 5 distinct modes. The benchmark is constructed for dimensions d {2, 16, 64, 128} and entropic regularization parameters 1,
For each fully trained method, characterized by a stochastic process distribution P and forward drift function Up, we assess performance using two evaluation metrics:
KL((S)) P) where, by Girsanov theorem (Oksendal, 2013),
KL P) = S 20 /. (1ns(Xe,)) () (X,
(20)
cBW3-UVP(S.., where
100
cBW3-UVP(S.., Po.1)
(5
(21)
[
BW2..)) is the squared Bures-Wassersternn Wasserstein distance, i.c. the squared Wasserstein-2 distance between (assumed) multivariate Gaussian distributions (Dowson & Landau, 1982), and Vs(X,] is the variance of X1 Sj.
We focus on the divergence IKL((S) || P), rather than KL( || S), low KLL(S P) more accurately indicates that P approximates S effectively across the entire support of S. The data-processing inequality implies that (IL(S,,1 Po.1) < KL((S) || P). The cBW3-UVP(... metric, introduced by Gushchin et al. (2023), is a normalized and conditional extension of the standard BW?(-,,) distance. Results for evaluation metrics (20) and (21) are summarized Table and Table 2, respectively.
In our benchmarking, we compare BM against the I-BM and DIPF methods (Section 3.1). Each experiment is repeated five times, including both model training and metric evaluation, to obtain uncertainty quantification. We use 000 Monte Carlo samples to estimate (20, 21). For simplicity, we employ the Euler-Maruyama scheme (Kloeden & Platen, 1992) with 200 discretization steps (Ai = 0.005) in all path sampling procedures. Each method undergoes 50, 000 SGD training steps with a batch size 1, 000, settings similar to those used by Gushchin et al. (2023), enabling qualitative comparison of our results with theirs. We use the AdamW optimizer with a learning rate of 10-4 and hyperparameters:
(0.9. 0.999), e 10-8 wd 0.01, where wd. denotes weight decay. Time is sampled as ~ U(e,) e) for 0.0025.
10

---

Published in Transactions on Machine Learning Research (12/2024)
6=0.1
6=1
==10
Method
d=2
d= 16
d=64
d= 128
d=2
d=16 16
d= 64
d=128
d=2
d= 16
d= 64
d= 128
BM
0.01
0.20
1.03
3.06
0.01
0.11
1.43
8.29
0.11
2.25
13.13
40.46
0.01
0.02
0.07
0.16
0.00
0.00
0.03
0.36
0.01
0.04
0.13
0.49
BM;
0.43
3.76
39.55
127,2 1.4
0.04
0.13
5.36
18.66
0.15
2.61
13.78
43.12
0.09
0.46
1.9
0.01
0.03
0.35
0.73
0.00
0.05
0.24
1.43
I-BM
0.03
0.20
1.24
5.70
0.01
0.16
1.94
7.79
0.16
4.09
17.17
19.17
0.01
0.02
0.04
0.42
0.00
0.01
0.04
0.07
0.00
0.03
0.21
0.55
DIPF
0.59
2.39
7.93
34.77
0.23
1.21
13.13 0.79
36.51
0.81
28.25
113.8 7:2
345.8
0.14
0.05
1:23
0.82
0.06
0.18
1.05
0.06
2.12
6.1
Table 1: Monte Carlo estimate of KL((S || P) function of e and d, standard deviation in gray.
For BM2; we employ a single feedforward neural network with layers of width 768 and ReLU activation, resulting approximately million parameters. We initialize the neural network parameters such that the resulting initial forward and backward drift functions evaluate to zero everywhere, choice that has proven effective in experiments. As mentioned in Section 4.1, implement path caching and exponential moving average for parameters used in path sampling. The cache contains 5, 000 initial-terminal values from both (F(0)) and (B()), refreshed every 200 training steps. We omit reporting the results for both the two-stage training procedure and the consistency loss described in Section 4.1, neither approach demonstrated consistent performance improvements across the considered benchmark.
For I-BM and DIPF, each outer loop iteration comprises 5, 000 SGD steps, totaling 10 outer loop (algorithmic) iterations. Following best practices (Bortoli et al., 2021; Shi et al., 2023), we alternate time directions over iterations for both algorithms. Each method employs two separate neural networks for forward and backward time directions, maintaining total parameter count close to million, matching BM model size. As with BM2, we implement path caching (for DIPF, entire discretized paths are cached) and EMA for sampling, with EMA decay rate of 0.99.
We also consider BM,, variant of BM that learns Schrödinger bridges for Vo, Vi across multiple values. This amortized version leverages 2M non-iterative nature. At each optimization step, is sampled from (00.1,44 4) and utilized in discretizing SDEs (F(O), (()) (lines 3 and 5 of Algorithm 1) and in bridge sampling (lines 7 and 8 of Algorithm 1). The neural network implementing drift functions (r(2, 1, O) and us(2, 4, 0) is modified to accept as an additional imput, resulting in conditional drift functions p (I, t, 0, o) and un(2, 1, 0,. Path caching is adjusted to store a values corresponding to cached paths.
In Table 2, we additionally include three baselines. EOT: sampling from the EOT solution, accounting for the bias due Monte Carlo estimation. SB(discr): sampling from the SB solution via the SB-optimal drift Ps: additionally accounting for Euler Maruyama scheme discretization error. & sampling from the independent coupling. Results for additional discretization intervals are reported in Appendix C.2, and we illustrate in Figure 1 the evolution of (21) during training for representative benchmark setting.
UUUVP(SS..L Po.1)
20
10
BM²
DIPF
(SGD steps 10,000)
Figure 1: Evolution of the metric (21) for d = 64, over SGD steps for BM2, 1-BM and DIPF.
11

---

Published in Transactions on Machine Learning Research (12/2024)
6.0.1
E=1
6= 10
Method
d=2
d= 16
d= 64
dd=28
d=2
d=16
d=64
d= 128
d=2
d=16 16
d= 64
d 128
EOT
0.02
0.05
0.34
0.91
0.09
0.17
0.43
1.14
0.12
0.18
0.23
0.38
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.00
SB(discr.)
0.04
0.07
0.35
0.92
0.10
0.17
0.45
1.18
0.12
1.15
5.38
10.48
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.01
0.01
195.8
186.3
162.6
145.1
136.1
127.6
113.0
93.61
8.07
4.88
4.22
4.45
0.14
BM
0.73
4.64
6.84
8,28
0.14
0.41
1.72
8.30
0.14
2.30
41.14
264.4
0.40
0.58
0.59
0.62
0.03
0.04
0.09
1:17
0.011
0.04
1.99
7.0
BM²
8.38
16.06
44.15
83.84
0.20
2.61
25.89
641.76
0.14
2.57
58.76
323.0
0.84
2.06
0:08
I-BM
1.07
4.25
7.19
16.63 63
0.20
0.53
2.20
7.79
0.14
5.21
135.8
578.7
0.50
0.68
0.28
2.07
0.00
0.04
0.35
0.79
0.02
0.11
1.3
9.8
DIPF
7.82
15.30
20.12
29.36
1.66
5.98
13.1
28.86 3.52
0.69
6.85
72.63 0.70
226.1
2:51
1.53
1.02
Table 2: Monte Carlo estimate of cBW3-UVP(SS.,, Po.1) function of e and d, standard deviation in gray.
We now discuss the results presented in Tables and 2. BM2 demonstrates superior overall performance across dimensions and entropic regularization settings in both metrics. I-BM also shows good performance, particularly in comparison the DIPF procedure, which aligns with the findings of Shi et al. (2023).
As expected, the performance of all methods deteriorates as the number of dimensions increases. This is because the metric(20) scales linearly with the nunber of dimensions, assuming constant error rate in estimating each component of the true drift μt. Similar considerations apply to the metric (21).
While BM3 4 exhibits a performance gap compared to BM2, it yields reasonable results in low-dimensional settings (d= = 2, 16). This gap may be due to increased pressure on model capacity or the need to normalize loss levels across a values. All methods perform poorly in the high regularization setting (e 10), especially in high dimensions (d = 64, 128), which we include for completeness. It should be noted that, in such cases, sampling from the independent coupling (a trivial solution) is preferable to sampling from the SB-optimal SDE for the chosen discrelization interval.
6 Related Works
Relevant works that, like BM2, address the dynamic Schrödinger bridge problem (1) include:
Schrödinger Bridge Flow: in concurrent work, De Bortoli et al. (2024) propose non-iterative methodology called o-DSBM (Diffusion Schrödinger Bridge Matching). a-DSBM is optimally implemented backward SDE approach, in which case a-DSBM's formulation and training objective align with our proposal (12, F(0), (0)), see De Bortoli et al. (2024, Equations (10) and (11)). The formulation of a- D-DS begins by establishing probability flow in the space of path probability measures, whose discretization yields a-MMF (Tterative Markovian Fitting). Under mild conditions, two theoretical results are established: (i) the a-IMF iterates converge to the SB solution, and (ii) the non-parametric updates of the functional loss (14), obtained through functional gradient descent with respect to the drift functions, recover the a- a-IMF iterates. The practical implementation, a-DSBM, thus replaces non-parametric drift. functions with parametric neural network approximators, and employs standard stochastic gradient descent on the parametric loss (12). This theoretical framework provides strong convergence guarantees for a-DSBM and paves the way to further theoretical developments. Furthermore, De Bortoli et al. (2024) demonstrate the method's effectiveness through extensive numerical experiments on high-dimensional computer vision problems, complementing our synthetic benchmark results.
through forward-backward
I-BM and DIPF: closely related to BM2 are the iterative, sample-based DIPF (Bortoli et al., 2021; Vargas et al., 2021) and I-BM (Shi et al., 2023; Peluchetti, 2023) procodures, which do not satisfy desiderata (i). Built on similar bridge matching principles, BM2 can be viewed as a modification of I-BM that employs a single optimization loop, resulting in simpler algorithm that have empirically shown to be competitive.
12

---

Published in Transactions on Machine Learning Research (12/2024)
Forward-Backward SB SDE: SB problem. Both approaches employ loss functions that require divergence computations (violating desiderata (iv)) and the use of two distinct neural networks. The first method is iterative, resembling DIPF (violating desiderata (i)), while the second method involves differentiating through entire discretized paths, resulting in high memory consumption (violating desiderata (iii))
Chen et al. (2022) propose two training algorithms addressing the dynamic
The subsequent works concentrate on solving the static Schrödinger bridge (2), or EOT (3); problem. Once this is achieved, solutions to the dynamic problem are trivially obtained through the standard decomposition S = o.. Although these works differ in nature and objectives, include them here due to their shared characteristic with BM the non-iterative nature of the algorithm.
Light SB: in two notable works, Korotin et al. (2023) and Gushchin et al. (2024) propose non-iterative, sample-based based EOT solvers for the Euclidean cost function, ie., for the specific choice of reference dynamics (R). Korotin ct al. (2023) introduccs an approximation to (an adjusted version of) the Schrödinger potential for Vi via mixture of Gaussian distributions, resulting in a mixture of Gaussian distributions approximation to Sijo- Gushchin et al. (2024) builds upon this approximation and introduces an additional sample- based training objective that takes as input any coupling Co, E C(Vo, V1), whereas Korotin et al. (2023) requires the independent coupling Vo While also non-iterative, the proposals of Korotin et al. (2023); Gushchin et al. (2024) differ from BM in two key aspects: (a) they learn a solution in the static set ting instead of the dynamic one, and (b) they employ mixture of Gaussian distributions approximations, rather than neural network approximators for the drift functions. Consequently, these methods may face challenges in scaling to modern generative ML applications. Light SB, in both variants, demonstrates strong performance in the benchmark presented in Section 5: the results of (Gushchin et al., 2024, Table 1) and (Korotin et al., 2023, Table 2) are directly comparable with the results of Table 2. However, it is worth noting that this benchmark is particularly well-suited for Light SB, as acknowledged by its authors, since each benchmark target So, is constructed such that Sijo is itself a mixture of 5 Gaussian distributions,
7 Conclusions
In this work we introduced Coupled Bridge Matching (BM²), a novel approsch for learning Schrödinger bridges from samples. BM builds on the principles of Bridge Matching while addressing key limitations of existing iterative methods. Our approach offers several advantages, including a simple single-loop optimization procodure, exactness in the idcalized sotting, modest memory requirements, and a straightforward loss function. The numerical experiments demonstrate that BM is competitive with and often outperforms existing iterative diffusion- based methods like I-BM and DIPF across various dimensions and entropic regularization settings.
On the thcorctical front, there is substantial room for improvement. Firstly, while bearing some resemblance to the standard convergence result for the EM algorithm, Theorem 2 lacks a quantity analogous to the likelihood being maximized in the FM algorithm. It remains unclear whether decreases in KL((M/..1(0)) || F(0)) and KLL(M Fo. (0') (()) can be linked to decreases in (F(( [| S) and (L(((() || S). Secondly, the requirement that F(0)= (()), equivalently that (F(0)) and (B(0) are time-reversals of each other, appears unnecessary. Notably, all numerical simulations conducted do not explicitly enforce this condition, which emerges naturally during the training process. Thirdly, it would be valuable to study problem (16) where reverse KL divergences are partially minimized, aligning more closely with the BM2 algorithm. In this scenario, Lemma 4 no longer holds, and it may be necessary to impose a corresponding additional constraint to maintain tractable analytical computations. The attractors of (19), and of a corresponding dynamical system arising from reverse KL, minimization, can be investigated to assess further convergence properties of BM2.
On the empirical front, the applications of BM2 in contemporary generative machine learning tasks remain unexplored. Given the promising results from previous studies employing Bridge Matching, such as those by Liu et al. (2023) and Somnath et al. (2023), it is anticipated that BM2 could be effectively applied to various domains, including image generation, audio synthesis, and molecular design. Future work could investigate the scalability and performance of BM in these domains.
13

---

Published in Transactions on Machine Learning Research (12/2024)
References
Brian D.O, Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313-3200 1982. ISSN 03044149.
Valentin De Bortoli, James Thornoon, Jeremy Heng, and Arnaud Doucct. Diffusion Schrödinger Bridge with Applications to Score-Based Generative Modeling, In Thirty-Fig Fifth Conference on Neural Information Processing Systems, 2021.
Tianrong Chen, Guan-Horng Liu, and Evangelos Theodorou. Likelihood Training of Schrödinger Bridge using Forward- Backward SDEs Theory. In International Conference on Learning Representations, 2022.
Valentin De Bortoli, Iryna Korshunova, Andriy Mnih, and Arnaud Doucet. Schrödinger Bridge Flow for Unpaired Data Translation, September 2024.
P. Dempster, N. M. Laird, and D. B. Rubin. Maximum Likelihood from Incomplete Data Via the EM Algorithm. Journal of the Royal Statistical Society: Series B (Methodological) 39(1):1-22, 1977. ISSN 0035-9246.
D. C Dowson and B. V Landau. The Fréchet distance between multivariate normal distributions. Journal of Multivariate Analysis, 12(3):450-455, 1982. ISSN 0047-259X.
Robert Fortet. Résolution d'un systeme d'équations de M. Schrödinger. J. Moth. Pure Appl. IX, 1:83-1000000 1940.
Nikita Gushchin, Alexander Kolesov, Petr Mokrov, Polina Karpikova, Andrei Spiridonov, Evgeny Burnaev, and Alexander Korotin. Building the Bridge of Schrödinger: Continuous Entropic Optimal Transport, Benchmark. In Thirty-Seventh Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2023.
Nikita Gushchin, Sergei Kholkin, Evgeny Burnaev, and Alexander Korotin. Light and Optimal Schrödinger Bridge Matching. In Forty-First International Conference on Machine Learning, 2024.
Jonathan Ho, Ajay Jain, and Picter Abbeel. Denoising Diffusion Probabilistic Models. arXiv, 2020.
Benton Jamison. Reciprocal processes. Zeitschrift für Wadhschheiniccheeisshheeiee und Verwandte Gebiete, 30 (1):65-86, 1974. ISSN 1432-2064.
Benton Jamison. The Markov processes of Schrödinger. Zeitschrift für Wadrschehnlichkettsshcerr und Verwandle Gebiele, 32(4):323 331, 1975. ISSN 1432-2064.
Mohammad Reza Karimi, Ya-Ping Hsieh, and Andreas Krause, Sinkhorn Flow: Continuous Time Framework for Understanding and Generalizing the Sinkhorn Algorithm, 2023.
Peter F. Kloeden and Eckhard Platen. Numerical Solution of Stochastic Differential Equations. Springer Berlin Heidelberg, 1992.
Alexander Korotin, Nikita Gushchin, and Evgeny Burnaev. Light Schrödinger Bridge. In The Twelfth International Conference on Learning Representations, 2023.
Guan-Horng Liu, Arash Vahdat, De-An Huang, Evangelos Theodorou, Weili Nie, and Anima Anandkumar. I2SB: Image-to-Image Schrödinger Bridge. In Proceedings of the 40th International Conference on Machine Learning, pp. 22042 22062. PMLR, 2023.
Xingchao Liu, Lemeng Wu, Mao Ye, and Qiang Liu. Let Build Bridges: Understanding and Extending Diffusion Generative Models. In NeurIPS 2022 Workshop on Score-Base Methods, 2022.
Christian Léonard. Some Propertics of Path Mcasures. In Catherine Donati-Martin, Antoine Lejay, and Alain Rouault (eds.), Séminaire de Probabilitss XLVI, Lecture Notes in Mathematics, pp. 207-230. Springer International Publishing, 2014a.
14

---

Published in Transactions on Machine Learning Research (12/2024)
Christian Léonard. A survey of the Schrödinger problem and some of its connections with optimal transport. Discrete & Continuous Dynamical Systems, 34(4):1533, 2014b.
Anton Mallasto, Augusto Gerolin, and Hà Quang Minh. Entropy-regularized 2- W-ssssesstein distance between Gaussian measures. Information Geometry, 5(1):289 323, 2022. ISSN 2511-249X.
Youssef Marzouk, Tarek Moselhy, Matthew Parno, and Alessio Spantini. Sampling via Measure Transport: An Introduction. In Roger Ghanem, David Higdon, and Houman Owhadi (eds.), Handbook of Uncertainty Quantification, pp. 1-41. Springer International Publishing, 2016.
Stefano Peluchetti. Non-Denoising Forward-Time Time Diffusions. 2021.
Stefano Peluchetti. Diffusion Bridge Mixture Transports, Schrödinger Bridge Problems and Generative Modcling. Journal of Machine Learning Research, 24(374):1-51, 2023. ISSN 1533-7928.
Gabriel Peyré and Marco Cuturi. Computational Optimal Transport. 2020.
Ludger Ruschendorf. Convergence of the Iterative Proportional Fitting Procedure. The Annals of Statistics, 23(4):1100-1174, 1995. ISSN 0090-5364.
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion Schrödinger Bridge Matching. In Shirth-Seventh Conference on Neural Information Processing Systems, 2023.
Vignesh Ram Somnath, Matteo Pariset, Ya-Ping Hsich, Maria Rodriguez Martinez, Andreas Krause, and Charlotte Bunne. Aligned diffusion Schrödinger bridges. In Robin Evans and Πlya Shpitser (eds.), Proceedings of the Thirty-Ninth Conference on Uncertainty in Artificial Intelligence, volume 216 of Proceedings of Machine Learning Research, pp. 1985-1995. PMLR, 31 Jul-04 Aug 2023.
Yang Song, Jascha Sohl-Dickstein, Diedorik Kingma, Abhishck Kumar, Stefano Ermon, and Ben Poole. Score-Based Generative Modeling through Stochastic Differential Equations. International Conference on Learning Representations, 2021.
Francisco Vargas, Pierre Thodoroff, Austen Lamacraft, and Neil Lawrence. Solving Schrödinger Bridges via Maximum Likelihood. Entropy, 23(9):1134, 2021.
B. K. Øksendal. Stochastic Differential Equations: An Introduction with Applications. Universitext. Springer, 6th ed., 6th corrected printing edition, 2013.
15

---

Published in Transactions on Machine Learning Research (12/2024)
A Additional Dynamics
In this section we consider simple extension to the dynamics of Section 2.2, and refer the reader to Peluchetti (2021; 2023) for the more general case. Here, consider the case where the reference distribution R is given by the solution to:
Xo ~ Vo. /.dWW E [0,
(22)
with o 2 O, 3: : [0, R strictly positive and continuous. With best == 3, 5 3, is chosen such that bu:1 = 1, to disentangle the contribution of 3 from the contribution of a. Indeed, under these conditions, Be defines time-warping: if X2 is the solution to (R), then Xto. has the same distribution as the solution to (22). Consequently, the solutions to (2) and (3) are independent of Br.
When employing (22), the definitions in Section 2.2 are replaced as follows:
Raj(- (mo) N (x o bu:)),
(23)
R: (0,,(_) (1) = N (moben + 21 bub o bu:bba);
(24)
Be
(11, B² V., log (01
(x) -xt),
(25)
Br
uon (n,,,, x0) == o² Be Vx, log (ogrooalal =
(Io - xzl,
(26)
bot
ro (0, == o²3, Vx, log (1 = Br
(wobest + 2 2ibbot zz).
(27)
bn:bb.1
B Proofs
Theorem 1 (Complete BM Iteaations). Consider the SDEs (F(0), ((0)), with initial drifls My 51 and
(0) (0)
corresponding distributions (()) For each i 2 1, let (ul)) (2) .U6 _.,)) arg (i L(u, us pig ./b ") resulling in the distribution ilerales p(i, B(), We dislinguish two cases:
(i-1) 1)
(0) = ub = 0: both the ilerates F()) B(1), F(2),. ..... and the ilerates B(0) F()) B(2)... are equiva- (i) Pi lent to the DIPF iterates, started respectively from the forward and from the backward time direction;
(ti) P (0) = um , 2b Um
Co. (0) Co.1 for some C'E E CCVV, 11): F = -B = 160) for each i 2 0 where I() are
the I-BM iterates.
Proof. Define Q associated with
(
X1 11 dX, odW,, te [1.0].
which is not the time reversal of (R), but Rio.1 = Q
Firstly, consider the case of initial mill drifts: Ms (0)
Up (0) 0, corresponding to F(0) Ro = V Rio R
F Rio.1 E and B(0) = 11Q = 12Qn R = B Rio, E S. As B(0) = B = M
we have
FO) = LoM (
10 = 1/ (0) = 1o Rio,1 E S. As F(O) = 11 (0)
1 M Po. B() = v M plo: = 4 1
( (00 1 0j1 R S. By induction, F()) V Bi 1|f Roo,1 S and B v F(i Rio.1 E S, 1 1. We now tjo
construct two forward- backward sequences. For the sequence F(0) B(1) p(2) ...., we have P.l Vollo (0)
B = v, 0j1 1 = Vo 1 which are the static IPF iterates: one marginal gets replaced at a time keeping the conditional distribution fixed. In the same way, for (0), B(2) we have B V 2 (u)
F = 1. ijo _ = F(1 which are again the static IPF iterates (for the backward formulation of the dynamic SB problem, i.e. via Q o as reference measure instead of Rio, and switched marginal distributions).
16

---

Published in Transactions on Machine Learning Research (12/2024)
As each pair F(), B6) is of the form F() Fo....... Ro,1, B6 Bo.1 R1o,1 we also recover the dynamic DIPF iterates.
(0) Secondly, consider H and v/ both corresponding to the BM transport based on the given coupling:
(0)
((0) M F(0) B(0) 7(0), Then, looking separately at either of the sequences F(9, i 2 1, and B(),
2 we obtain that F() B 1(0) i 1.
Lemma 2 (Loss Interpretation). holds that
KL(Bb] || Vo) + ( IKI Bo.1 F) + C(B) KL(M || [ C2(B), (L(Fi || 1 - Lo(eb: (II
(15)
Ro.1 B) + Di(F) = Po. || B) + D2(F),
for CBB (B)) independent of F, Di(F), D2(F) independent of B, with S Ci C2(B) and 0 <
(F(F) & (FF).
Proof- We consider only ( the arguments for Lo(vs; py) symmetric. By Girsanov Theorem (Oksendal, 2013) and by the marginal--ondiitonaal decomposition of Kullback-Leibler divergences we have
2 ! [|as(X,t)) Pm
KL(M Bo.1 || F) = (B
Mo) +
E IT P...
Bo.1 (X,, (X,,
KI Ro.1 F) KIL Bu
Mo)
+
E
(Xe, 11 Bo.1 (X,, (x))l 2
II 00.1
2
2 ./
1
KL([IT Bo.1 [[F) KL Bo
Vo)
E
(X,,(X,,t) Moz (X,,,, x1)1l
0o.1
= KL Vo) + Ly(m5:
where 11T
Bo.1 (Xt,t, = Ro. [no] (X,, t, X,) Xol, pm
Bo.1 (Xt,t) = Pn [uOI 4, and F is
distribution of the non-Markov diffusion solution to the auxiliary SDE
Xo Mo. dX: (X,( poi (X,, X1) + _ (X,, t, Xo)]dt adWe, E [0.
(F)
By the lower property of conditional expectations and by the conditional Jensen inequality follows that BL. B
F) Po FF)
E
(X,,)) Moo (X,, t, (l,r(Xe,t) 1 Bo.1 (X,, t, Xt)
0o.1
2 20
E
|| (X,,t, 1, x1) 11 Bo. (X,, Xo) = ( 2 0.
2
By the Pythagorean property of the BM transport (Liu el al., 2022; Peluchetti, 2023)
Bo. || F) KL(M || F) (I Bo.) = K(B) 0.
Taking CB K(B) completes the proof.
0
Lemma 3 (R-stability of P(A), B(A). Tf F', B' E R,, then F(A), B(A) E R for each X E [0, 1].
Proof. By the marginal-condittonal decomposition of Kullback-Leible divergences
[F || B') = (Fo. Bb.)) +
[K]L B/o,)),
Fa.1
KIL(F ( F*) = (L Fo.)) + E [KL((Fjo, || F%,,)), Fa..
17

---

Published in Transactions on Machine Learning Research (12/2024)
and Blo,1 Fl,.1 R, hence
p(A) = arg min (Fo, || 1),1) + (1 1)o,L(Fo, Fo.1) + E [KL || Roo,));
PEP(
Fo,1
B(A) arg min (Bo, F + (1 (Bo, || B),1) ()
E
[K |l Ro,,)),
Bo,
and thus 17601 / 10.1
RI.,, which completes the proof.
|0,1
Lemma 4 (S-stability of F(A), B()). If F", B' E S, then F(A)B(()) E S for each E [0, 1].
Proof. In view of Lemia 3, we have to verify that 0.1 Bo solve the BOT problems (3) for some marginal distributions if 1 1 do. For simplicity, we assume that all of 0.1 Bo,1 F6,1 Bo, admits positive densitics on Rd and that vo and v admits positive densities on Rd. The steps of this proof carry over the more general theoretic setting.
We know that ful (ro, =1) = vo0())) (m)) and bo (20, 21) / (10l=)) where f1j (a (20) x (1)0 (i (1 (10(011/0))=-4 (=x))=>
and / (zo (2)) x Soj (2o ( (1 (- (see Section 4.2). On the other hand
(0,, exp fob (zo) + of (x1) m(20,x1)) (0, = exp f (Io) (x1)
e x((0,,11)
e
for the Schrödinger potentials? of (mo), of (x1) and (20) + of (x1) (Léonard, 2014a). It follows by direct computation that ful (Fo, x1) and 60y1 (2o, 21) satisfy:
fol (I), = exp 06 (z)) (2o) + of - x( x1)
60 (00. =exp
exp { db (20) + (x1) x(10,T1)
for some other Schrödinger potentials (xu),, of and (zo), (x1).
C Additional Results
C.1 Infinitesimal Minimization, Gaussian Case
Consider the one-dimensional case 1, with target Gaussian marginal distributions Vn = N af) and V1 =N([,, of), and a reference diffusion distribution R associated with (R). In this setting, the solution to the static Schrödinger bridge problem (2) is known analytically and is given by bivariate Gaussian distribution (Mallasto et al., 2022).
We hypothesize that conditional Gaussian densities for f10(31 (z) (2)) and bi (2o|z1) solve (19). Specifically, we propose F10 = (N(A/ zo + of) and ! = ~ + at, o where Af, af,af A/, at R and of of E Rso are algorithmic-tinm time dependent scalar parameters. By construction, F/") NVo,, (z) and BSTh -N for cach 1 2 0. Substituting these expressions for ( (2o) and b (zolzı) into (19) yields a six-dimensional ODE system in the parameters. The initial conditions are A A/ = 1, of = aB = 0, of = og = o², corresponding to initial null drift terms for (F(O), (0)), as discussed in Section 4.2.
To numerically solve the ODE and determine the values of Af of of Af, af, of over E [0, LJ for some L >>0 O, we evaluate (19) for three different pairs of (zo, 21). This provides sufficient constraints identify the
We formulate the potential with respect the Lebesgue measure Rd.
18

---

Published in Transactions on Machine Learning Research (12/2024)
2
1
0
F(1) [Xi]
por [X1]
Cpt) (oo. X1]
Figure 2: Algorithmic-time evolution of Ees [X1], V Po [11], Cpa [Xo, X1), compared with Es[X]], Vs[Xi],
Cs(Xa, X1) as dashed gray lines.
paramcters. Subscquently, we verify that the proposed functional forms for f100 (21 (mo) and bij (20) indeed solve (19).
Figure 2 illustrates the evolution of F [Xi], V Foo [X1], and Cru [Xo, Xi] over algorithmic time 1. These quantities represent the mean and variance of X1 and the covariance between Xo and X₁ according to F(),, respectively. The corresponding values Es [X]], [X[], and Cs [x, for the static Schrödinger bridge solution So,1 from Mallasto et al. (2022) are depicted as dashed gray lines, demonstrating convergence.
We examine the scenario where fo -2, μ1 2, 01
C.2 Results for Additional Discretization Intervals
In Table 3 we report the results for metric (21) obtained by considering the BM2 I-BM and DIPF methodologies for different discretization intervals At = 1/TT where 7' is the number of time-steps. For all methods we employ the same number of time-steps at training and inference (testing) time. We recall that in Scction 5 200 time-steps have been employed to produce the results of Tables 1 and 2, and that we rely exclusively on the Fuler- Maruyama discretization scheme.
e= 0-0.1
==
10
Method
d=2 2
d= 16
d
64
d= 128
d=2
d= 16
d= 64
d= 128
d=2
d= 16
d= 64
d 128
SB(100)
0.08
0.10
0.37
0.91
0.10
0.19
0.49
1.25
0.12
2.96
14.86
28.90
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.02
0.03
SB(50)
0.23
0.22
0.44
1.00
0.15
0.21
0.61
1.49
0.13
7.24
37 27.24
72.34
0.00
0.80
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.01
0.04
(100)
1.17
5.08
7.12
7.94
0.14
0.41
1.71
8.91
0.15
4.06
55.10
295.82
0.79
0.22
0.71
0.78
0.D4
0.03
0.14
0.35
0.02
0.03
1.13
10.60
(50)
0.80
4.59
6.46
8.47
0.21
0.48
1.85
9.57
0.16
8.17
78.76
202.22
0.52
0.877
0.73
0.09
0.015
0.10
0.27
0.01
0.04
1.85
41.38
DIPF(100)
7.28
15.72
16.44
26.69
1.43
5.96
11.98
22.60
1.14
7.36
85.42
212.19
4.39
1,31
0.39
0.71
U.45
0.312
U.822
0.86
U.07
0.26
2.64
1.29
DIPF(50)
6.90
12.37
14.98
26.91
1.44
5.09
10.73
20.20
2.18
11.32
75.80
226.81
3.01
1.40
2.36
2.71
0.33
0.665
0.29
U.0G
0.35
1.02
0.99
I-BM(100)
0.98
4.54
6.35
13.60
0.18
0.51
2.09
7.30
0.15
5.95
113.91
494.32
0.222
0.988
0.84
1.23
0.04
0.05
0.23
1.06
0.02
0.10
0.93
3.59
L-BM(50)
1.58
4.21
6.35
13.75
0.28
0.60
1.96
7.31
0.16
9.14
114.40 0.89
432,20
0.61
0.20
1.01
2.93
0.08
0.04
0.25
1.24
0.022
0.07
5.82
Table 3: Results analogous to Table 2 but for varying discretization intervals At = 1/T, where the number of time-steps T' is indicated in parentheses after each method name.
19

---

Published in Transactions on Machine Learning Research (12/2024)
D Python Code
dimenslons: B: batch; D: data; T: time_atepa + 1
required: (bpple_Obbtcch_in,, device), sample_ (batch_dim, device), Fnd_drift_fa(k, t), bud_drift _nnfftfff(x,, t) import torch as th
sampling from Reuu (6): (B, D), (B, D), (B,), O (B, D)
def sample, (x x_1, t, sigma):
D x_O.shapp
mean_t (1 t(..., Nonel)
t(..., None] x_1 (B, D)
var_t sigma=t2 t(..., None] (1 t(..., None])
(B, D)
z_t th.dandn (ikea__
(B, D)
11
x_t mean_t ch.aqrt(var_t) z_t x_t
(B,
13
14 fwd target (6): (B, D), (B, D), (B,) -> (B, D)
15 def fud_te [ad_target (x_t, x_1, t):
18 return (x_1 x_t) / (1 t(..., None]) # (B, D)
17
18 fwd EN target (7): (B, D), (B, D), (B,) >> (B, D)
19 def bud, tarttarget (x_t, x_0, t):
20 return (x_0 x_t) / t(.... None] # (B, D)
21
22 Buler-Maruyama discretization scheme: fn t), (B, D), (T), 0 -> (B, D)
28 def discretization(drife
initial_valug times, sigma):
D initial_valuus sal.vaall.espe...
times times(.... None] (-1 B)
(T,
x_prev_t initial_value
# (B,
for prev_t, in zip(times(:-1), cimea[1:]]. dt prev_t drift_t drift_fn(e feiit__f((x_pe__tt prev_t) drift_part_t drift_t dt[..., None]
(B), (B)
28
(B)
28
(B,
30
(B,
31
epa_t th. 1ik
(B,
32
diffueion_part.is (eigma th. (dt))) None] epe_t x_t - x_prev_t drift_part_t diffueion,part__3
(B,
33
(B,
31
x_prev_t- return x_t
(B,
36
30
ST BM loss computation: fn(b, d), fn(b, d) fn(x, t), fn(x, t), O. O. O. 0 def sample_losse lompl__osssssmplo_00 sample_1, fud_ ard_dfiff_fn, bud_drift_fn, batch_din, time_steps, sigma, device):
38
# sample from the target marginals:
f_o 1ample_0bbttchh dim, device)
# (B, D)
b_1 cemple__bbtcchddte
device)
# (B, D)
# sample according to current (F(0)) and (B())):
42
43
trd_times th, ti.nins (0.0, 1.0, time_steps 1, device=device)
# [0, 1/time_stepa 1/time_steps,
1]
44
bud_times th. th.nins (1.0, 0.0, time_steps 1, device=device)
# [1,
45
f_1 discretization fiffffddddff_ffeeeeeeeee f_O, fud_times, sigma) detach()
# (B, D)
b_o discretization Abaerettiaatioobbbdddffff._.._.
b_1, bud_times sigma) detach((
# (B, D)
# sample time and mixture processes based on (0) and (o)
t rand (a.ran (banchbb___ddi. ,im,), device=device)
# (B)
pi_f_t sample_ br wridge(f_o) f_1, signa) pi_b_t brnple_riddeeb__, b_1, signa)
# (B,
# (B, D)
# define regression targets and model predictions:
target_f_t fwd, hnd_targee((p1.. target_b_t bad_ mad_target(pp.. (_t_f_t, f_0, t)
b_1, t)
# (B,
53
# (B,
54
prediction f_t ved_drife__tnpi__b_ prediction_b. bxd_drift_ffmpiifft t)
(B,
58
(B,
58
# compute loss:
57
losa th. (assunttttrte___.. prediction f_t)=+2, dim=1) 2 1osa_b, th. (uwanntttttttb....aaaaaa predictiionb_ttxe dim=1) # (B)
# (B)
58
50
loss_t th. mhlsmm loss_b_t)
# ()
60
return losa_t
Listing Basic implementation of BM2 loss computation (Algorithm 1) in Porcrh.
20


Reviews:
[{"review_id": "ZcJzrVqK0v", "reviewer": "Reviewer_ezyt", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n- Well-motivated problem and solution. The paper is well-written with a comprehensive literature review and background sections.\n- The derivation of the new loss is based on rigorous theoretical argument. The proofs in the Appendix regarding convergence analysis appear to be correct.  \n- Some simple numerical schemes showing the method's competitive performance against Bridge Matching.\n\nWeaknesses: (mostly on the experimental side)\n\n- The experiments lack some of the recent baselines -- I was surprised to see a detailed discussion of these baselines in Section 6 (for example the strong baseline lightSB), but not inside the numerical experiments itself. \n- While the bridge estimation tasks and the KL/BW2-UVP metrics show competitive performance, it is unclear on the benefit of the new loss in the training process itself. It would be welcome if the author could show a simple plot demonstrating normalized training loss of BM^2 compared with BM/DIPF.\n - In a training framework that requires a single neural network to learn both forward and backward drifts, it is intriguing to see how well BM^2 can scale beyond learning bridges between Gaussian/mixture of Gaussians; especially on the stability of the training. My opinion is that the current set of experiments is not large enough to make a good conclusion. However, this is a minor concern, provided that the paper's goal is to introduce a new training regime for diffusion SB and analyze it theoretically.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduced a new approach to approximating a solution of Schrodinger bridge problem (SB) with neural networks. SB is a problem that is well developed theoretically in control and measure transport theory, but it is well-known that there does not exist a closed-form solution for general setting. The bridge matching (BM) framework (Peluchetti 2021; Shi et al. 2023) offered a new scalable solver for the SB problem, but is still considered to be costly as training requires solving an iteration of the optimization problem (a process called iterated bridge matching) that updates two neural networks sequentially (forward and backward drifts). This work extended the BM framework and propose a a new training scheme that is simpler: coupled both the forward and backward drift matching into a single loss. Some theoretical investigations are shown to analyze the equivalence of BM^2 with previous frameworks, along with empirical demonstrations to show efficiency of the new training loss.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See some of the questions on the experimental side in the Weaknesses section.", "broader_impact_concerns": "No.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "HASM0Db1ed", "reviewer": "Reviewer_Jqq8", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- Extremely clear and effective presentation.\n- Excellent contextualization within the existing literature and comparison to related methods.\n- The introduction of the coupled system is natural, resembling the structure of the EM algorithm.\n- Includes I-BM and DIPF as special cases when complete minimization is performed at each step.\n- Enables a non-iterative approach, with a single optimization loop instead of the outer and inner loops used in I-BM or DIPF. This partial minimization has strong theoretical backing in Theorem 2, suggesting promising effectiveness.\n- Demonstrates superior experimental results.\n- Offers further theoretical research opportunities, as highlighted in Section 4.2.3.\n\nWeaknesses:\n\n- Not really, maybe lacking higher dimensional experiments, to show improved generation quality or training speed.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces a novel approach for learning Schrödinger bridges. $BM^2$ addresses the problem of Schrodinger BM transport between two distributions, leveraging a reference diffusion process. Using a coupled system of forward and backward transport equations, it eliminates the need for iterative methods by employing a single optimization loop, using stochastic gradient descent, significantly simplifying the training process. The method is exact in its idealized form, with only neural network approximation and discretization introducing practical errors. The paper provides a preliminary theoretical analysis of $BM^2$ convergence properties, supported by simple numerical experiments that demonstrate its effectiveness against similar methods.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Questions and Requested Changes:\n- What about incorporating a consistency loss by enforcing $F$ and $\\overleftarrow{B}$ to be time-reversals of each other, as done in [1]? While you mention in the conclusion that this is not necessary, could it not potentially speed up the training process?\n- Contrary to DIPF and I-BM, the structure of $BM^2$ removes the problem of the initialisation choice for $\\mu_f^{(0)}, v_b^{(0)}$ (which is the choice of initial joint distribution for DIPF, for instance). Can you confirm it is true?\n- In the experiments, what value of $\\sigma$ is used for $BM^2$? What is the EMA rate?\n- What is the rationale behind the $BM^2_{\\sigma}$ model? What $\\sigma$ is used for sampling? Could this lead to a more efficient model for deterministic ODEs as $\\sigma \\to 0$? Is it valid to set $\\sigma = 0$ for sampling?\n- Regarding the neural network architecture, you mention using a single network for $\\mu_f$ and $v_b$. Does this imply two output channels? How is $\\sigma$ incorporated into the architecture— through positional embedding?\n- While the paper does not focus on transfer or generation on higher-dimensional datasets (like MNIST or CIFAR10) as stated in the conclusion, it would be beneficial to include some small experiments in this area and compare them to I-BM and DIPF. This would likely enhance the paper's appeal to a broader audience.\n\n\n[1] Yuyang Shi, Valentin De Bortoli, Andrew Campbell, Arnaud Doucet (2023) -- Diffusion Schrödinger Bridge Matching", "broader_impact_concerns": "Sufficiently addressed.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Tx64P0fW7s", "reviewer": "Reviewer_BMAP", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "*Strenghts*\n1. *Joint learning*. Unlike Iterative Bridge Matching (I-BM) and Diffusion Schrödinger Bridge (DSB), where the forward and backward processes are learned separately in an alternating manner (training the forward process first, then the backward process, and repeating the cycle), $BM^2$ learns these two processes jointly. By coupling the forward and backward dynamics during training, $BM^2$ reduces overall approximation errors. This coupling allows the two processes to mutually inform and constrain each other, effectively acting as checks on one another, leading to more consistent and accurate results.\n2. *Single optimization loop.* Unlike I-BM and DSB, which require an iterative refinement process where each iteration updates the drift functions based on progressively refined estimates of the true solution, $BM^2$ bypasses this need for iteration. Instead, $BM^2$ learns the optimal drift functions in a single optimization loop by directly training a neural network to find the right parameters, using stochastic gradient descent (SGD). This process only involves Monte Carlo sampling of the initial and terminal states, without the need to iteratively re-compute or refine the drift at each step.\n\n*Weaknesses*\n1. *Convergence guarantees.* The paper offers theoretical analysis, but its theoretical convergence is not as rigorously guaranteed as in iterative methods like I-BM and DSB. Specifically, Theorem 2 in the paper suggests that if $BM^2$ reaches a point where further reduction in the loss function is not possible and if the forward and backward processes are the same, then the solution matches the true Schrödinger bridge. This gives some assurance of convergence. However, this result doesn't provide a complete convergence proof in the way iterative methods do. The connection between reductions in KL divergence/loss functions during training and overall convergence to the true Schrödinger bridge remains unclear. The guarantee is more empirical and relies on specific conditions being met during the optimization process, such as the neural network’s ability to accurately approximate the drift functions and the avoidance of issues like local minima.\n2.  *Practical limitations*. In practice, the method’s performance could be limited by the discretization of the underlying stochastic differential equations (SDEs) and the quality of neural network approximations. These limitations are acknowledged but not thoroughly quantified in the paper.\n3. *Computational efficiency*. While $BM^2$ is theoretically more efficient due to its non-iterative design, in practice, its reliance on neural networks and stochastic gradient descent (SGD) introduces potential computational challenges. First, $BM^2$ is subject to common neural network limitations, such as overfitting, vanishing or exploding gradients, and sensitivity to network architecture. If the neural network fails to accurately model the dynamics of the underlying stochastic differential equation (SDE), $BM^2$'s performance may suffer. Additionally, $BM^2$’s training involves learning drift functions for both forward and backward processes simultaneously, which can become computationally demanding—especially as the problem’s dimensionality increases or the neural network architecture becomes more complex.\n3. *Empirical breadth*. The empirical results are promising but focus heavily on specific settings.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper presents a new method called Coupled Bridge Matching ($BM^2$), which addresses the Schrödinger Bridge problem and provides theoretical analysis and empirical results to support the effectiveness of this approach. The core idea is to jointly and simultaneously (rather than indipendently) learn the forward and the backward dynamics of the Schrödinger Bridge and to train a single set of neural networks in one optimization loop via stochastic gradient descent (SGD).", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "While the paper makes significant strides, a few aspects could benefit from further exploration. Look at the *Weaknesses* section.", "broader_impact_concerns": "There are no concerns on the ethical implications of the work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: All the reviewers and I found the paper highly interesting. The only notable weakness lies in the experimental section, which focuses solely on simple toy examples. Nevertheless, we agree that the paper introduces sufficiently novel ideas to make its contribution highly valuable.
Decision: Accept (poster)