Tree-Based Diffusion Schrödinger Bridge with Applications to Wasserstein Barycenters
Authors: Maxence Noble
Valentin De Bortoli
Arnaud Doucet
Alain Durmus
Venue: NeurIPS 2023 spotlight (2023)
Keywords: Schrödinger bridge
optimal transport
diffusion model
Wasserstein barycenter

Abstract: Multi-marginal Optimal Transport (mOT), a generalization of OT, aims at minimizing the integral of a cost function with respect to a distribution with some prescribed marginals. In this paper, we consider an entropic version of mOT
  with a tree-structured quadratic cost, i.e., a function that can be written as
  a sum of pairwise cost functions between the nodes of a tree. To address this
  problem, we develop Tree-based Diffusion Schr\"odinger Bridge (TreeDSB), an
  extension of the Diffusion Schr\"odinger Bridge (DSB) algorithm. TreeDSB
  corresponds to a dynamic and continuous state-space counterpart of the
  multimarginal Sinkhorn algorithm. A notable use case of our methodology is to
  compute Wasserstein barycenters which can be recast as the solution of a mOT
  problem on a star-shaped tree. We demonstrate that our methodology can be applied in high-dimensional settings such as image interpolation and
  Bayesian fusion.

Tree-Based Diffusion Schrödinger Bridge with Applications to Wasserstein Barycenters
Maxence Noble*
Valentin De Bortoli
CMAP, CNRS, École polytechnique, Institut Polytechnique de Paris, 91 120 Palaiseau, France
Computer Science Department, ENS, CNRS, PSL University
Arnaud Doucet Department of Statistics, University of Oxford, UK
Alain Oliviero Durmus
CMAP, CNRS, École polytechnique, Institut Polytechnique de Paris, 91 12 Palaiseau, France
Abstract
Multi-marginal Optimal Transport (mOT), generalization of OT, aims at min- imizing the integral of a cost function with respect to a distribution with some prescribed marginals. In this paper, we consider an entropic version of mOT with a tree-structured quadratic cost, i.e., a function that can be written as a sum of pairwise cost functions between the nodes of a tree. To address this problem, we develop Tree-based Diffusion Schrödinger Bridge (TreeDSB), an extension of the Diffusion Schrödinger Bridge (DSB) algorithm. TreeDSB corresponds to a dynamic and continuous state-space counterpart of the multi-marginal Sinkhorn algorithm. notable use case of our methodology is to compute Wasserstein barycenters which can be recast as the solution of a mOT problem on a star-shaped tree. We demonstrate that our methodology can be applied nigh-dimensional settings such as image interpolation and Bayesian fusion.
Introduction
In the last decade, computational Optimal Transport (OT) has shown great success with applications in various fields such as biology (Schiebinger et al., 2019; Bunne et al., 2022), shape correspondence (Su et al., 2015; Feydy et al., 2017; Eisenberger et al., 2020), control theory (Bayraktar et al., 2018; Acciaio et al., 2019) and computer vision (Schmitz et al., 2018; Carion et al., 2020). While OT commonly seeks at computing the transport plan that minimizes the cost of moving between two distributions, can naturally be extended to the multi-marginal setting (mOT) when considering several distributions. This extension of OT has notably been studied in quantum chemistry (Cotar al., 2013), clustering (Cuturi & Doucet, 2014) and statistical inference (Srivastava et al., 2018). In particular, popular application in unsupervised learning of mOT with Euclidean cost consists in computing the Wasserstein barycenter of a set of probability distributions (Agueh Carlier, 2011; Benamou et al., 2015; Livarez-Estebaan et al., 2016; Peyré et al., 2019).
Interior point methods can be used to solve OT and mOT problems but they come with computational challenges (Pele & Werman, 2009). In order mitigate these limitations, one often considers entropic regularization of OT, known as Entropic OT (EOT). This regularized formulation can be efficiently solved discrete state-spaces using the celebrated Sinkhorn algorithm (Cuturi, 2013; Knight, 2008; Sinkhorn & Knopp, 1967), which admits a continuous state-space counterpart referred
*Corresponding author. Contact at: maxeenee.oble--ouuillooeee pww.lllooeooolyyccceuee...ee
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

to as the Iterative Proportional Firring (IPF) procedure (Fortet, 1940; Kullback, 1968; Ruschendorf, 1995). In the case of a quadratic cost, EOT is equivalent to the static formulation of the Schrödinger Bridge (SB) problem (Schrödinger, 1932). Given a reference diffusion with finite time horizon T and two probability measures, solving SB amounts to finding the closest diffusion to the reference (in terms of Kullback-Leibler divergence on path spaces) with the given marginals at times t = 0 and = T' This framework naturally arises in stochastic control (Dai Pra, 1991) where one aims at controlling the marginal distribution of a stochastic process at a fixed time. Recently, De Bortoli et al. (2021) introduced Diffusion Schrödinger Bridge (DSB), an approximation of a dynamic version of the IPF scheme on path spaces, see also Vargas al. (2021); Chen al. (2022). This methodology leverages advances in the field of denoising diffusion models (Song et al., 2021; Ho et al., 2020) in order to derive a scalable and efficient scheme to solve SB, and thus EOT.
Similarly to OT, mOT admits an entropic regularization (EmOT), which can be solved via a multi- marginal generalization of Sinkhor//PF algorithm (Benamou et al., 2015; Marino & Gerolin, 2020). Recently, Haasler et al. (2021) proposed an extension of the static SB problem in discrete state- space any multi-marginal tree-based setting. They notably made the correspondence between this formulation and EmOT, when the cost function writes as the sum of interaction energies onto the given tree structure, and introduced an efficient version of Sinkhorn algorithm to solve it.
Motivations and contributions. In this work, we investigate the continuous and dynamic coun- terpart of the tree-based framework from Haasler et al. (2021). To be more specific, we present an extension of the static SB formulation in continuous state-space to any multi-marginal tree-based setting, referred to as TreeSB. Then, we establish the equivalence between TreeSB and a formulation a (quadratic) tree-structured cost function, analogously to Haasler et al. (2021).
of EmOT relying Inspired by DSB, we develop TreeDSB. a dynamic counterpart of the multi maatinal IPF (mIPF) to solve it, by operating on path spaces and using score- based diffusion techniques. To bridge gaps in literature, we prove the convergence of mIPF iterations in a non-compact setting under mild assumptions, by extending results IPF convergence (Ruschendorf, 1995). Finally, we illustrate approach on examples of Wasserstein barycenters from statistical inference and image processing.
Although our approach can be applied to any tree, we focus star-shaped frees. In this setting, we show that TreeSB reduces to regularized Wasserstein barycenter problem. Our method comes with several benefits compared to existing works. First, it is out-of-samplle i.e., it does not require re-runming the full procedure when given a new data point. Second, our formulation of the Wasserstein barycenter problem obtained from TreeSB allows us to avoid numerical issues of having to choose the regularization too small, see Section 5. Finally, to the best of our knowledge, this is the first methodology to extend ideas from diffusion- based models to the computation of Wasserstein barycenters. In particular, we believe that the idea of iterative refinement, iee. solving the dynamic counterpart of a static problem, plays a key role in the efficiency and scalability of the method.
Notation. For any measurable space (X, X), we denote by (X the space of probability measures defined on (X, X). Unless specified, X' is defined as the Borel sets on X. For any E E N, let glt)- = P (RR))); we denote (1) by P. Assume that X = (Rd))² for some & E N. For any E X and any m,n {O..... f) such that m S n. let Fman (Im, Im,+1, ............... ..)).
Let Leb be the Lebesgue measure. For any non-negative function f X + R4, such that Ix fdLeb < +00, define H(ʃ) = fx log, fdLeb E (-oc, For any distribution E (X),, we define the entropy of μ as H(w) = (f(du/dLeb) if < Leb and H(m) = too otherwise. For any two arbitrary measures μ and v defined on (X,X), define the Kullback-Leible divergence between μ and as - -Ix du + Ix du i < v and KT( v) = +00 otherwise. For any T 0. we denote by C(10, T)) the space of continuous functions from T to Rd. For any path measure Pe E (C(O,, 27, )), we denote by Ext((P) E !P! the coupling between the extremal distributions of P. i.e., (P) = Por.. Note that, for given coupling No E g)) there may exist several path measures I verifying Ext(P) = To.r. For any undirected tree T = (V, E) with vertices V and edges E. we denote by {e, z') (or {u', (.) the undirected edge between v E V and v' e V, it exists. Given r E V, we denote by T, (V,E,) the directed version of rooted in r, where the directed edges E, are uniquely defined from the edges E, see Appendix B for further details. In this case, the edge linking 8 E V to v' E V T, denoted by (v, Finally, for any integers (n, K) E N x N we define mod(K) as the the remainder of the Euclidean division of n by K.
KL((//>= = Jx
2

---

2 Background and setting
Multi-marginal optimal transport. Let l E N Given cost function e: (124)++1 + R. a subset S c {0, .., and family of probability measures {uilies E Ølll, mOT consists in solving n arg min { c(zo:)dd ι gp(t+)) ni 14 Vi E S},
(mOT)
where my is the i-th marginal of n., i.e., Ty (A) = T (p))) for any A E (RR)), with proj, : zoot + Jz. Given some weights (I... E (IR+), the Wasserstein barycenter between the measures fuites is given by To in (mOT), in the case where S {1.....,) {} and c(zo.e) Σ -i= w.||zo [ol||ooxx,) =, (Peyré et al., 2019). In particular, when wi = 1/2, the distribution To can be regarded as the Fréchet mean (Karcher, 2014) of the measures {14)ess for the Wasserstein distance of order 2. Similarly to OT, (mOT) can be relaxed using the following entropic regularization
π arg min { (dr <TI() : x E gl(+1), n. = U4 Vi E ():
(EmOT)
where e > is a hyperparameter and v an arbitrary measure defined on ((I2))+1, B((dd)(+)).
Link with Schrödinger Bridge. We first recall the relationship between Schrödinger Bridge and EOT. Given T > 0, Q a (reference) path measure, i.e., Q E (C([0, 2] IR)) and two measures Po: M1 E (RR)). solving the SB problem amounts to finding the path measure defined by
P argmin(KLLPP QQ PE (CC( T) R),, F. μo, Pr =1
(SB)
If Q is associated with a Stochastic Differential Equation (SDE)?, of the form dX, = a +t dB,, with 2 0, then it can be shown, see (Léonard, 2014, Proposition 1) that Por verifies
Po. ( [L(() (o.r) E 9p( To = po, T1= #1)
(static-SB)
This is called the static formulation of SB. solving EOT with quadratic cost and regularization ε = 2 sinh(a"))) /a if > 0, s 27 if a = 0. Moreover, since 1 Porr T. where Q fo.r the measure Q conditioned on initial and terminal conditions, solving the dynamic problem (SB) is equivalent to solving (static-SB).
be shown that solving (static-SB) is equivalent to
Similarly, (EmOT) can be easily rewritten in static multi-marginal SB fashion
π argmin(KL(r]r")
go(tt), Ti p4 Vi E S
(mSB-like)
with /dLeb) x
//](dd/))///////////// (dd//d///eee /dLeb) (xo:e) where n² is the reference measure.
Diffusion Schrödinger Bridge. Recently, De Bortoli et al. (2021) introduced Diffusion Schrödinger Bridge (DSB), a numerical scheme to solve (SB). approximates the iterates of dynamic version
of the Iterative Proportional Fitring (IPF) scheme (Sinkhorn & Knopp. 1967: Knight, 2008; Peyré et al., 2019; Cuturi Doucet, 2014), which can be described as follows: consider sequence of path measures (P")) such that p Q and for any n E N
p2n+1 argmin (P(I))) Pr = m1}.
p2n+2 argminng (KL : Po = po)
This procedure alternatively projects between the measures with fixed initial distribution and the ones with fixed terminal distribution. For the first iteration, we get that [Pl μ4 @ Q Assuming that
Q is given by dX,, f((X, )dl + dB,, with f [0,T] Rd - Rdd then pl is associated with the time-reversal of this SDE initialized at μı, The time-reversal of an SDE has been derived under mild assumptions on the drift and diffusion coefficients (Haussmann & Pardoux, 1986; Cattiaux et al., 2021). In this case, we have (tejo)) ~ IPl wwth Yo ~ #1 and
dY, = (-IT- + llogrr: (Y:))dd dB:,
where p. is the density of P. w.r.t. the Lebesgue measure. The score V log pa is estimated using score matching techniques (Hyvärinen, 2005; Vincent, 2011). The first iterate of DSB, Pl, corresponds to denoising diffusion model (Ho et al., 2020; Song et al., 2021). DSB iterates further and not only parameterizes the backward process but also the forward process. can therefore be seen as refinement of diffusion models drawing bridge between generative modeling and optimal transport.
We refer Appendix C for details solutions of SDEs and associated measures.
3

---

Tree-based framework. Consider an undirected tree T EV, with vertices V and edges E. such that V is identified with {O..... .., f]. Inspired by Haasler et al. (2021), we restrict our study of
(EmOT), to the case where the cost function e is the tree-structured quadratic cost derived from T
c(Fo.t) " (usn))EE We
F,, 11
(1)
where Wa, is a weight on the edge {u, o'), which links to e' (and v' to v). Furthermore, in Haasler et al. (2021), we choose S. i.e., the set of vertices of T with constrained marginals, to coincide with the leaves of T. This framework recovers important applications. from Wasserstein barycenters to Wasserstein propagation, see Solomon et al. (2014, 2015). We emphasize that it differs from an OT problem defined on the space of graphs (Chen et al., 2016). Here, each node represents a probability measure (observed or to be inferred) and each edge represents coupling between two distributions.
We consider an arbitrary vertex r' V and choose in (EmOT) such that (dv/ (LLc) (x,(), where ypr is a density defined on Rd. Due to the form of and the reference measure x0 in (mSB-like) is therefore a probability distribution which factorizes along T, = (V, E,), the directed version of T rooted in We refer to Appendix B for more details on the notion of directed trees. In this setting, (EmOT) is equivalent to the tree-based problem
π argmin (KL((
E (VV), = Vi E 5}
(TreeSB)
with TO = 8 (e,,'EE, Tppe+
(2)
where (.
xx) N(x,, (2w,.rr Ia) and op < Leb with density fr. In a manner akin to Haasler et al. (2021), we thus establish, in continuous state-space, the correspondence between (TreeSB), static tree-based version of SB, and version of EmOT with tree-structured cost (1). In our work, we make the following assumption on the constrained marginals {u:hes.
AO. For any i E S, 14 < Leb and H(14) < co.
In what follows, we define K as the number of leaves of T denoting S = fio..... -x-1). and define the horizon times T = (/ for any {v, r/} E For any in E S, we will denote by T (V, Ex) the directed version of T rooted in the leaf ig. In the next section, we present our dynamic method to solve (TreeSB), called Tree-based Diffusion Schrödinger Bridge.
3 Tree-based Diffusion Schrödinger Bridge
In this section, we present a method to solve (TreeSB) in the case where E S, i.e., r is a leaf of T. We refer to Appendix E for the extension to the case where r E V.S. Without loss of generality, see Appendix E, we assume that ix- and choose 9r / such that = = Pix-
Dynamic approach to mIPF. In order to approximate solutions of (TreeSB), we consider the malti-marginal extension of the IPF algorithm, denoted by mIPF. Namely, we define sequence of probability distributions (e") such that for any E N
(mIPF) where k; = (n- 1) (od(K) and (k,, + 1) is identified with n mod(K). We define a mIPF cycle as sequence of K consecutive mIPF updates. In particular, each marginal constraint is considered exactly once during one mIPF cycle. In practical setting, our main aim is to sample from the (mIPF) iterates at the lowest cost. Although these updates can be made explicit, see Marino & Gerolin (2020) for instance, direct sampling unfeasible in practice when d is large. To overcome this limitation, we suggest to compute these iterates in a dynamic fashion with equivalent path measures.
n n1+1_ = argmin(KT.(TT//e)
: i E (((V), Tix, = t1
Since T² factorizes along T, see (2), one can show that the iterates of (mIPF) also factorize along T, see Section 4. Since these iterates all have a constrained marginal, we obtain the following decomposition for any π E N: म" Pix, ((v,r')E,, Te where E,.. denotes the set of edges of the directed tree T Tk..- Then, our approach consists in computing dynamic iterates, Le., path measures, along the edges of T that coincide on their extremal times with the static iterates (æ1 )eE.. Namely, for any E N, for any edge (v, v') E Ek., we define path measure Fn (iv.r') E (CC( T,,, R")) such that Ex ( -,, where Extt (PR,,
((.ww) stands for the joint distribution of Pn (w,w") at times Using the tree-based form of the (mIPF) iterates, we can thus sample from n" by (i) following the directed edges of TR.. (ii) diffusing along them the corresponding path measures (IP (eee.,)) and (ii) picking the samples on the vertices. When T is a bridge-shaped tree (2 vertices, edge), simply reduces to the dynamic reformulation of the IPF scheme. In what follows, we explain how to obtain our dynamic sequence.
O and Te,e" In particular, it comes that The IP P (r,,"),T.

---

Definition of the dynamic iterates. We first compute the iterate Fo, corresponding to the dynamic version of x defined (2), in Proposition 1. Then, we build the following iterates by recursion on E I and prove their well-posedness in Proposition 2.
Proposition 1. Let T K-1 = Ex)),, the directed tree associated with T= (V,E) and rool i Then, for any (v, o') E Ek-1, there exists PQ (v,,e) E (P(C([O, T,,w(, with Ext([P( (c,07) = Th (eowrr and such that PP (w.e")o is the distribution of (Br))e(D.Tu..)) recalling that Tu,e = E e2(,
Before deriving the dynamic counterpart of the (mIPF) iterates, we introduce several definitions. For any path measure P. we denote by p the time-reversal of IP. For any directed tree and any vertex e of this tree, p(v) refers to the (unique) parent of v, and e(v) to the unique child of when it exists, see Appendix B for more details.
N. Assume that we have de- fined the sequence of our dynamic iterates (P (eoe (ono)))o.")EEE....... mnnn up to stage n..
Let
Edges Ex Edges Ex+1 Path P
Consider the path P, = {(u,, 1141)) in the di- rected tree Th, such that 01 lk. and 2441 +1,++1 In particular, for any (u. r') E E1, ei- ther (v',v) E Pa (u,
Eg., P.. This is illustrated in Figure I when V = 2,,1,2,3,4)
S {2, 3, 4). ig = and f 4: in this case,
P P= [33, 1), (1, 0), (0, 4)} and (1, 2) is the only edge common to Eg and 1
Figure 1: Illustration of the change of root in toy tree with 5 vertices.
Consider now the directed tree TR..,. We define the (n + 1)-th iterate of our dynamic sequence by recursion the edges of this tree, following the breadth-firg order. In this order, (ik.+1. c(tR, +1)) (UJ-1, vy) is the first edge considered.
First, we define P(ot (1))1 Mian+o @ (P ()) In the case of bridge-shaped tree, this is
exactly the (n + 1)-th update described in DSB. Then, for any (v,v') E EE..+MM(UJ+1,00
(a) either (u, v²) E ER., \P,, and we define pptt1 pa+11
(u,e") (ple),,) 3 Pn (v,o")oo
(b) or (v', v) Pn. and we define pan+1
(w.1') =pa+1
(Tro).)).Trco). @ (o)
Proposition 2. Consider the sequence of dynamic iterates defined by (a) and (b). Then, for any n E N and any (u, o') E Ex.,, (lo,0') E P (C((,, IR))) and we have ExI(P (P = Towyy
(w,er))
Proposition 2 highlights the equivalence between the (mIPF) iterates and our dynamic iterates. These path measures are defined iteratively, by following the updates (a) and (b) along the edges of T. The key observation here is that the computation of each dynamic iterate reduces to sequence of updates (b) on a path linking two leaves of T. We emphasize that our iterates could be similarly obtained by directly considering dynamic formulation of (TreeSB) and introducing the formalism of deterministic time branching processes. We leave the study of this problem for future work. We now get into the details of our practical implementation, which relies on ccore-based methods.
Approximation of the dynamic iterates. The time-reversal operated in the update (b) can be computed explicitly, see Haussmann & Pardoux (1986) for instance. Indeed, assuming that (((",w) associated with dX,, fe,er,e (X, + dB, with X. The then, under mild conditions, its time-reversal ..)" is associated with dY, {-fr-r,," + V logPor,..T- Y,)dt + dB, where por... is the density of (e).,)). w.r.t. the Lebesgue measure. The score V log Po- can then be approximated using score-matching techniques (Hyvärinen, 2005; Vin- cent, 2011) which are now ubiquitous in diffusion models (Song et al., 2021) and used in DSB De Bortoli et al. (2021). Therefore, at iteration (n + 1), the update (b) is similar to the one of DSB for each edge on the path joining ix, and ik.+1- In practice, we parameterize the drifts ft,e.v' for any {u, o') E with neural networks fi,.,,. .. and use the mean- matching loss introduced by De Bortoli al. (2021). Note that doing so, obtain 2/E| neural networks. The whole procedure consisting in computing our dynamic iterates using the DSB framework is called Tree-based Diffusion Schrödinger Bridge (TreeDSB) and is summarized in Algorithm 1.
with Yo n:
5

---

end for The algorithm is initialized with fe. for all {u, o') E. This corresponds to Brownian motion dy- namics when sampling at the first iter- ation of TreeDSB, see Proposition 1. Note that in Algorithm 1, when we sample from P (ry,,)+1)* we update f..0,44iin, which will be used to sam- ple from Pa+1
Algorithm I TreeDSB (Training)
1: Input: T = (V,E), E), (ui)ies, 10 (U N E E
2: for n. = 0, 0,..., N do
3: Let kn = (n
(K)(K)
4: Get path between iR.. and 8k. +1, Pn = 10
5: while not converged do
6:
for j I,...,J J do
7:
Sample from P 2, (Buler-Maruyama)
8:
Compute mean matching loss ((0g51107))
(Oy11:09) in the next itera- tions. In order to sample from the dynamics P #0,,,)+13" we consider its Euler-Maruyama discretization, see Appendix F for more details. We de- scribe the different steps of the algo- rithm in the case of a toy example below, see Figure 2 for an illustration.
9:
O1j41.9; & Gradient Step(E(0n)411n))
10: Update f1.0,)1::5
11: end for
12: end while
13:
14:
Output: (0.,or)
TreeDSB on toy tree. We consider a star-shaped tree with three leaves denoted {1, 3} and its central node [0]. Following (2), we define x0 with r = 3 and y (duu//delt) During the first iteration of TreeDSB, T is rooted at vertex 3 and we compute samples from the forward path in order to learn the buckward path {(1. 3) In the next iteration, we re-root the tree T at vertex 1 and consider the forward path PI= ((1, 0), (0, 2)}. where the edges (1, 0) and (0, 2) are respectively given by the first iteration and the initialisation. This highlights that TreeDSB does not require to update the whole tree. The following iterations are done similarly. At each iteration n E N. we sample from xr by first sampling from Ms. at leaf ik,, and then following the parameterized SDEs on the directed edges of Tk..
Po = (33 0), (0, 1)} with Brownian motions, see Proposition
T0
ITg
50
JT
222
12
ITS
Figure 2: Illustration of one mIPF cycle solved by TreeDSB for a toy star-shaped tree. At each iteration, method learns the backward stochastic process (dotted arrows) that goes from the target leaf (green-circled), corresponding to the constrained marginal, to the current root of the tree (red-circled) by using samples from the forward stochastic process (solid arrows).
4 Theoretical properties of mIPF
In this section, we study some of the theoretical properties of the static iterates (π²) meN, that are equivalent our dynamic iterates according Proposition 2. In the case where the cost function c is bounded in (EmOT), results of convergence of (mIPF) exist (Marino & Gerolin, 2020; Carlier, 2022). However, our setting does not satisfy their assumptions, since our transport cost quadratic and the measures are defined on Rdd In what follows, we provide the first non non-quantitative convergence results for (mIPF)
con-compact setting.
For the rest of the section, we consider a static formulation of the multi-marginal Schrödinger bridge problem which is more general than (TreeSB), defined
π² argmin(LL((Tm (T) : π go(+1), ग. = 14, viiiie e S}
(static-mSB)
where S c {o....,
f), x E P, {uilies E 1l5 We consider the following set of assumptions.
A1. There exists family of measures (Vilie (iEeeeo... defined on (Rd BB(R)) such that TO < Di i
with density = d / i'00 vi) and h4 & vs with density r; = du//d/ for any i E
A2. { E (++11) KL(TT T < 0C, Tr = p4 Vi E 5} # 0.
A3. There exists family of probability measures {is (o...... such that TO ~ #9, where
T Øies Me ;e. Ug.

---

In particular, (static-mSB) recovers (TreeSB) by considering v; Leb for any i E {O..... f and h(Zg:e) = (z,(), exp(=c(zo:e)E] A1. We detail in Appendix D how A2 and A3 can be met in (TreeSB). Under these assumptions, the multi-marginal Schrödinger Bridge exists.
Proposition 3. Assume Al and A2. Then, there exists unique solution x² (static-mSB). In addition, assume A3. Then, there exists family {16 Fies of measurable functions 1/2 IR IR such that
(dx*/dx* = explet Es 15]
-9-a.s
In order to establish the existence and uniqueness result of Proposition 3, extend results from Nutz (2021) to the multi-marginal setting. A consequence of Proposition 3 is that the iterates of (mIPF) can be described using potentials.
Corollary 4. Assume AI, A2 and A3. Let meN be the sequence given by (mIPF). Then, for any
N with k., (n -1 1) mod(K) () and qn E N such that n qn K kn + 1, there exists a family
of measurable functions (1/9++1
Unnn+1 +1
iK .} such that
R., (dT") /dT = exple)
1 j=0
mutt (II,) j=k,+1 1 (Ii,)] x
In the tree-based setting, Corollary 4 explains why the (mIPF) iterations preserve the tree-based Markovian nature of n", We now prove that the marginal n² converges to μu for any E S, as n goes to infinity, i.e., we have marginal convergence the leaves T.
Proposition Assume Al and A2. Let (T")nen be the sequence given by (mIPF). Then, we have 1vllTV for any E S.
lim,-+o -!
The previous result does not ensure the convergence (œ") nen to the solution to (static-mSB). In particular, Proposition 5 does not provide the convergence of the marginals on the nodes E Vİ,, which is key compute regularized Wasserstein barycenters with TreeDSB. Relying on additional assumptions, we now derive the convergence of (mIPF).
A4. iee (ma) C (T²* closed.
A5. There exist E (0, oo) such that exp(1e), + S ē, for any n E N, any k E
2].
These assumptions can be seen as multi-marginal extensions of the ones of Ruschendorf (1995), see Appendix D for discussion and examples.
Proposition 6. Assume AI, A2, A3, A4 and A5. Let (ee be the sequence given by (mIPF). Then, we have lim,->oc ||== n l| O, where x² is given in Proposition 3.
To the best of our knowledge, Proposition 6 is the first convergence result of (mIPF) without assuming that the space is compact or that the cost is bounded. We highlight that traditional techniques to prove the convergence of IPF cannot be easily extended to the multi-marginal setting as pointed by Carlier (2022). In the case of bounded cost, quantitative results exist (Marino & Gerolin, 2020; Carlier, 2022). We leave the study of such results in the unbounded cost setting for future work.
5 Application to Wasserstein barycenters
Although Algorithm can be applied to trees T with fixed marginals the leaves, case of particular interest is star-shaped trees, i.e., trees with a central node, denoted by index 0, and such that S = f1,..., ... !) (see Figure 2 for illustration with é 3). In this section, we draw link between (TreeSB) and regularized Wasserstein barycenters. We recall the definition of the Wasserstein distance of order 2 with e-entropic regularization between and (Peyré et al., 2019, Chapter 4)
(w, v) = in ||~1 =oll x1) - =H(=) क E (2),, no = 14, क = =v}
(3)
In this work, we consider the (Ls, (e (Chizat, 2023) defined as follows
Wasserstein-2 barycenter problem
arg min Ei
W 2.6/mo (b) m4) + (e- ())H(() 9)
(regWB)
where (wi)vE(1....e) E (0, + The following proposition shows the equivalence between the barycenter problem (reg WB) and the multi-marginal Schrödinger bridge problem (TreeSB) over In particular, it allows us to use TreeDSB to estimate the solution μ of (regWB).
7

---

Proposition 7. Eet > O. Assume AO. Also assume that T is a star-shaped tree with central node indexed by 0, and that the reference measure of (TreeSB) defined in (2) verifies = i and 9r- = djpii-1/P //Mee-//LLeb>> > O. Under A2, (regWB) has unique solution The where n' solves (TreeSB).
The proof of this result is postponed to Appendix D. More generally, we show in Appendix D that, for any tree T, (TreeSB) is equivalent regularized version of the Wasserstein propagation problem (Solomon et al., 2014, 2015). Moreover, we present in Appendix E an extension of Proposition 7 in the case where the chosen root is not a leaf of T. We finally emphasize that the formulation of (reg BB) leads to a minimization of the entropy of the barycenter. In particular, this allows us to choose e reasonably large TreeDSB, which is stability advantage compared other regularized methods which do not consider this further regularization.
6 Related work
Diffusion Schrödinger Bridge. Schrödinger Bridges (Schrödinger, 1932) have been extensively studied using tools from stochastic control and probability theory (Léonard, 2014; Dai Pra, 1991; Chen et al., 2021). More recently, algorithms were proposed to efficiently approximate such bridges in the context of machine learning. In particular, De Bortoli et al. (2021) proposed DSB while Vargas al. (2021); Chen al. (2022) developed related algorithms. In Chen et al. (2023), the authors study a multi-marginal version of DSB in a linear tree-based setting, where the set of observed nodes is the whole of vertices. However, contrary to our setting, Chen et al. (2023) introduced a momentum variable. This allows for smoother trajectories which are desirable for single-cell trajectories applications and correspond to some spline interpolation in the space of probability measures (Chen et al., 2018). A general framework for tree-based static Schrödinger Bridges on discrete sate-spaces was given in Haasler et al. (2021). In this work, we extend their formulation to a dynamic and continuous setting, see Appendix D for more a thorough comparison.
Wasserstein barycenters. The notion of Wasserstein barycenter was first introduced in Rabin al. (2012) and then later studied in Agueh & Carlier (2011). The algorithms to solve this problem can be split into two families: the in-sample based approaches and the parametric ones. In-sample approaches require access to all the measures 14, which are assumed to be empirical measures (Cuturi & Doucet, 2014; Benamou et al., 2015; Solomon et al., 2015). Related to this class of algorithms is the semi-discrete approach, which aims at computing a Wasserstein barycenter between continuous distribution but rely on a discretization of the barycenter (Claici et al., 2018; Staib et al., 2017: Mi et al., 2020). Most recent approaches do not rely on a discrete representation of the barycenter, but instead parameterize it using neural networks. These approaches can be further split into two categories. First, measure-basee optimization approaches parameterize the measures using a neural network. This is the case of Cohen et al. (2020), where the barycenter is given by generative model, which is then optimized Fan et al. (2020) introduce an optimization procedure which relies on a min-max-min problem using the framework of Makkuva et al. (2020). More recently, Korotin et al. (2022) considered fixed point-based algorithm introduced in Álvarez-Esteban et al. (2016) to update a generative model parametrizing the barycenter. On the one hand, potential-buat -basee methods rely on a dual formulation of the barycenter. Korotin et al. (2021) parameterized the dual potentials using Input Convex Neural Network and considered regularizing losses imposing conjugacy and congruency. On the other hand, Li et al. (2020) consider dual version of the regularized Wasserstein barycenter problem contrary to other works. Our approach applied to start-shaped trees also approximates regularized Wasserstein barycenter. However, contrary Li et al. (2020), we do not consider a parameterization of the potentials in the static setting but instead, parameterize the drift of an associated dynamic formulation using Schrödinger bridges. To the best of our knowledge TreeDSB is the first approach leveraging DSB-like algorithms to compute Wasserstein barycenters.
7 Experiments
In experiments', we illustrate the performance of TreeDSB to compute entropic regularized Wasserstein barycenters for various tasks We choose to compare our method with state-of-the-art regularized algorithms: fast free-support Wasserstein barycenter (fsWB) (Cuturi & Doucet, 2014), and continuous regularized Wasserstein barycenter (crWB) (Li et al., 2020). In all of our settings, we consider star-shaped tree with K leaves and edge weights that are equal to 1/ K, resulting in a
Coode available https: / httg://tth.cccceeeeeooe/e//ee//e wtt.://fwwwuuuuoooooooooooorrrrr schrodinger----eeeeeere bridge.
8

---

sequential training procedure over 2K neural networks. The initial diffusion is always a Brownian motion parameterized as explained in Proposition I. Hence, the time horizon on each edge is defined by 7 = Ke/2. The order of the leaves is randomly shuffled between the mIPF cycles. We consider 50 steps for the time discretization on [0, T). We refer Appendix G for details on the choice of the schedule, the architecture of the neural networks and the settings of our experiments.
Synthetic two dimensional datasets. We first il- lustrate TreeDSB in synthetic two dimensional setting. We consider three different datasets Swiss- roll (vertex O, starting node r), Circle (vertex 2) and Moons (vertex 3) and compute their Wasser- stein barycenter (vertex 1) by running TreeDSB for 50 mIPF cycles with we show the estimated densities of the datasets on the leaves of the tree (we emphasize that the dis- tributions plotted on each leaf are generated from the central barycenter measure). Figure 4, we observe the consistency between the barycenters generated from the different leaves. In Appendix G, present additional results for this setting.
mmt cace 50, from ventee To
0.1. In Figure
HiF cyce: 56 from vemte
cytte: 50. from unitee
14
Figure 3: Estimated densities the leaves.
miPF cycle: 50, from vertex
mIPF cycle: 50, from vertex
miPF cycle: 50, from vertex to
10
5
5
0
-10
-15
-15
-10
-10
Figure 4: From left to right: barycenter estimated from the leaves Swiss-roll, Circle and Moons.
Synthetic Gaussian datasets. Next, we consider three independent Gaussian distributions with zero mean and random non-diagonal covariance matrices whose conditional number is less than 10, following Fan et al. (2020). In this case, the non-regularge reguuaguur barycenter can exactly computed. To evaluate the performance of the algorithms,
the Bures- Wasserstein Unexplained Variance Percentage (UVP), following (Korotin et al., 2021, Section 5). Given a target distribution jt* E g and some approximation μ. E 3, we define
BW (VP = 100 2B Var(x*)%
where (W2 W (N()), N(Ejxr) Cov(n")).
Method
d 16
d 64
= 128
256
fsWB (Cuturi Doucet, 2014)
0.06
2.8640.06
11.1240.05
14.47
17.41
crWB (I.i al., 2020)
0.02+4.01
1.52
11.41
5.7540.02
18.27 10.54
Tree DSB
0.
1.07
1.39+0.07
1.92+
2.62+0.07
Table 1: Gaussian setting: comparison with the regularized methods crWB and fsWB.
In this setting, choose μ² to be the non-regularized barycenter and assess the dependency w.r.t.
the dimension of the algorithms using the BW? BW?-UVP metric. In Table 1, we compare ourselves with the regularized methods Li al. (2020) (T.2-reg. equal to 10-4) and Cuturi & Doucet (2014). We run TreeDSB for 10 mIPF cycles with = 0.1. Bold numbers represent the best values up to statistical significance. While Li et al. (2020) and Cuturi & Doucet (2014) enjoy better performance in low dimensions (d 2), TreeDSB outperforms these methods as the dimension increases.
MNIST Wasserstein barycenter. We then turn to an image experiment using MNIST dataset (LeCun, 1998). Here, image is not considered as 2D-dimensional distribution as in Cuturi &
Doucet (2014) and et (2020), but as a sample from a high-dimensional probability measure (d 784). We aim at computing Wasserstein barycenter between the digits 2,4 and 6. To do so,

---

run TreeDSB for 10 mIPF cycles with r that corresponds to the digit 6 and E = 0.5. In Figure 5, we display samples from the estimated marginals on the leaves, to assess the reconstruction of the digits 2, 4 and 6, and samples from the barycenter, obtained by diffusing from the leaf corresponding to the digit Our results prove the scalability of TSeeDDB to the high-dimensional setting, compared to state-of-the-art regularized methods. Additional results on MNIST dataset are given in Appendix G.
2
a
2
a
4
4
4
4
4
6
6
6
6
0
2
Q
4
6
2
2
2
a
2
4
4
4
4
4
6
6
6 6 6 6
6
6
4
4
4
U
2
2
2
2
2
4
4
4
4
6 6
6 6
6
4
4
4
2
2
2
2
2
4
4
4
4
6
6
6
4
4
9
4
U
2
2
2
2
4
4
4
4
4
6
6
6
6
0
6
Figure 5: Samples from the estimated MNIST 2-4-6 marginals and from their Wasserstein barycenter.
Subset posterior aggregation. Finally. evaluate Tre in the context of Bayesian fusion (Srivastava et al., 2018), also called posterior aggregation. Given a Bayesian model and a dataset parti- tioned into several shards, this task aims at recovering the full data posterior dis- tribution from the posterior distributions computed on each shard.
Method
Without het.
With het.
IsWB (Cuturi & Doucet, 2014)
12.95
44.43 23.06
crWB (Li et al., 2020)
20.66
Tree DSB
8.69
8.90
Bayesian fusion setting: comparison with the regularized methods crWB and fsWB.
Table 2:
In particular, it has been proved that the barycenter of the subdataset posteriors is close to the full data posterior under mild assumptions (Srivastava et al., 2018). Here, we consider a logistic regression model applied to the wine dataset* (d = 42) and proceed follows. We first split this dataset into 3 subsets, with or without heterogeneity, and estimate the posterior parameters on each shard. Then, draw samples from the obtained logistic distributions define μ4, M42 M4s. Then, compute the Wasserstein barycenter these measures, and compare it to the posterior computed on the full dataset. As in the synthetic Gaussian experiment, we run Tre for 10 mIPF cycles e 0.1 and we compare ourselves with Li et al. (2020) (L2-reg. equal to 10-1) and Cuturi & Doucet (2014). We evaluate the methods using the BW3-UVP metric, where μ² is the estimated full data posterior, and report the results in Table 2. In both settings, we observe that our method outperforms existing regularized methods compute Wasserstein barycenters.
Limitations. One of the main limitation of entropic regularized OT approach is that their behavior is usually badly conditioned as E 0. In our setting. we observe that if e, equivalently T, is too low then the algorithm becomes less stable as the training of the models slows down. In the future, plan to mitigate this issue by incorporating fixed point techniques like the one used in Korotin et al. (2022). Finally, since our algorithm is based on DSB (De Bortoli et al., 2021), it suffers from the same limitations. In particular, training different neural networks iteratively incurs some bias in the SDE which is harmful for large number of mIPF iterations.
Discussion
In this paper, we introduced Tree-based Diffusion Schrödinger Bridge (TreeDSB) scalable scheme to approximate solutions of entropic-regularized multi-marginal Optimal Transport (mOT) prob- lems. Our methodology leverages tools from the diffusion model literature and extends Diffusion Schrödinger Bridge (De Bortoli et al., 2021). 1). In particular, it approximates the iterates of the multi- marginal Iterative Proportional Fitting (mIPF) algorithm, for which we prove its convergence under mild assumptions. We illustrate the efficiency of TreeDSB for image processing and Bayesian fusion, using the link between mOT and Wasserstein barycenters. In future work, would like study quantitative convergence bounds for mIPF in the unbounded cost setting. Another line of work would be scale TreeDSB to higher dimensional problems building on recent developments in the diffusion model and flow matching community (Lipman et al., 2023; Peluchetti, 2023; Shi et al., 2023).
'https:// lattps://archiee.... iis.e..c........ www.iicssi.......ooo http.//mm//aaattt///imnma
10

---

Acknowledgments
We thank James Thorton for the DSB codebase" and useful discussions. AD acknowledges support from the Lagrange Mathematics and Computing Research Center. AD and MN would like to thank the Isaac Newton Institute for Mathematical Sciences for support and hospitality during the programme The mathematical and statistical foundation of future data-da driae engineering when work on this paper was undertaken. MN acknowledges funding from the grant SCAI (ANR- SEER-19-CHIA. CHIA-0002).
References
Acciaio, B., Backhoff-Veraguas, Veraguas, J., and Carmona, R. Extended mean field control problems: stochas- tic maximum principle and transport perspective. SIAM journal on Control and Optimization, 57 (6):3666-3693, 2019.
Agueh, M. and Carlier, G. Barycenters in the Wasserstein space. SIAM Journal Mathematical Analysis, 43(2):904 2011.
Alseren- P. C., Del Barrio, E., Cuesta-Albertos, J., and Matrán, C. A fixed-point approach to barycenters in Wasserstein space. Journal of Mathematical Analysis and Applications, 441()) 744-762, 2016.
Bayraktar, E., Cox, A. M., and Stoev, Y. Martingale optimal transport with stopping. SIAM Journal on Control and Optimization, 56( 36(1):417-433, 2018.
Benamou, J.-D., Carlier, G., Cuturi, M., Nenna, L., and Peyré, G. Iterative Bregman projections for regularized transportation problems. SIAM Journal on Scientific Computing, 37(2)A1111-A113) 2015.
Bunne, C., Papaxanthos, L., Krause, A., and Cuturi, M. Proximal optimal transport modeling of population dynamics. In International Conference Artificial Intelligence and Statistics, pp. 6511-6528. PMLR, 2022.
Carion, N., Massa, F., Synnaeve, G., Usunier, N., Kirillov, A., and Zagoruyko, S. End-to-end object detection with transformers. In Computer Vision-ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part 16, pp. 213-229. Springer, 2020.
Carlier, G. On the linear convergence of the multimarginal Sinkhorn algorithm. SIAM Journal on Optimization, 32(2):786-794 2022.
Cattiaux, P., Conforti, G., Gentil, I., and Léonard, C. Time reversal of diffusion processes under finite entropy condition. arXiv preprint arXiv:2104.07708, 2021.
Chen, T., Liu, G.-H., and Theodorou, E. A. Likelihood training of Schrödinger bridge using forward- backward SDEs theory. International Conference on Learning Representations, 2022.
Chen, T., Liu, G.-H., Tao, M.. and Theodorou, E. A. Deep momentum multi-margine Schrödinger bridge. 2023. doi: 10.48550/ARXIV. 110..0155..
Chen, Y., Georgiou, T., Pavon, M., and Tannenbaum, A. Robust transport over networks. IEEE Transactions Automatic Control, 62(9):4675-4682. 2016.
Chen, Y., Conforti, G., and Georgiou, T. T. Measure-valued spline curves: An optimal transport viewpoint. SIAM Journal on Mathematical Analysis, 506))594775598,, 2018.
Chen, Y., Georgiou, T., and Pavon, M. Optimal transport in systems and control. Annual Review of Control, Robotics, and Autonomous Systems, 4, 2021.
Chizat, L. Doubly regularized entropic Wasserstein barycenters. arXiv preprint 1rX1v:2303.11844, 2023.
Claici, S., Chien, E., and Solomon, Stochastic Wasserstein barycenters. International Conference Machine Learning, pp. 999-1008, PMLR, 2018.
11

---

Cohen, S., Arbel, M., and Deisenroth, M. P. Estimating barycenters of measures in high dimensions. arXiv preprint arXiv:2007.07105, 2020.
Cotar, C., Friesecke, G., and Klüppelberg, C. Density functional theory and optimal transportation with Coulomb cost. Communications on Pure and Applied Mathematics, 66(4):548-58) 2013.
Csiszár, I. I-divergence geometry of probability distributions and minimization problems. The Annals of Probability, pp. 146-158, 1975.
Cuturi, M. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in Neural Information Processing Systems, 26, 2013.
Cuturi, M. and Doucet, A. Fast computation of Wasserstein barycenters. In International Conference on Machine Learning, pp. 685-693. PMLR, 2014.
Dai Pra, P. A stochastic control approach to reciprocal diffusion processes. Applied Mathematics and Optimization, 23(1):313 1991.
De Bortoli, V., Thornton, J., Heng. J., and Doucet, A. Diffusion Schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems, 34: 2021. 17695-17709, 2021.
Dupuis, P. and Ellis, R. S. A weak convergence approach to the theory of large deviations. John Wiley & Sons, 2011.
Eisenberger, M., Toker, A., Leal-Taixé, L., and Cremers, D. Deep shells: Unsupervised shape correspondence with optimal transport. Advances in Neural Information Processing Systems, 33: 10491-10502, 2020.
Fan, J., Taghvaei, A., and Chen, Y. Scalable computations of Wasserstein barycenter via input convex neural networks. arXiv preprint arXiv:2007. 04462, 2020.
Feydy, J., Charlier, B., Vialard, F.-X., and Peyré, G. Optimal transport for diffeomorphic registration. In Medical Image Computing and Computer Assisted Intervention- MICCAI 2017: 20th Interna- tional Conference, Quebec City, QC, Canada, September 11-13, 2017, Proceedings, Part I 20, pp. 291-299. Springer, 2017.
Flamary, R., Courty, N., Gramfort, A., Alaya, M. Z., Boisbunon, A., Chambon, S., Chapel, L.. Corenflos, A., Fatras, K., Fournier, N., et al. POT: Python optimal transport. The Journal of Machine Learning Research, 22(1):357133578 2021.
Fortet, R. Résolution d' un système d' ''quttions de M. Schrödinger. Journal de Mathématiques Pures et Appliquées, 19(1-)):83-105, 1940.
Haasler, I., Ringh, A., Chen, Y., and Karlsson, J. Multimarginal Optimal Transport with a tree- structured cost and the Schrödinger bridge problem. SIAM Journal Control and Optimization, 59(4))2428-2455) 2021.
Haussmann, U. G. and Pardoux, E. Time reversal of diffusions. The Annals of Probability, pp. 1188-1205, 1986.
Ho. J., Jain, A., and Abbeel, Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851, 2020.
Hyvärinen, A. Estimation of non-non normmazzdl statistical models by score matching. Journal of Machine Learning Research. 6(4), 2005.
Karcher, Riemannian center of mass and so called Karcher mean. arXiv preprint axxiv:1077.2000 2014.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. arXiv preprint arXiv: 1e..1.41.9688 2014.
Knight, A. The Knopp algorithm: convergence and applications. SIAM Journal on Matrix Analysis and Applications, 30(1):261-275. 2008.
12

---

Kober, H. A theorem on Banach spaces. Compositio Mathematica, 1940.
Koller, D. and Friedman, N. Probabilistic Graphical Models: Principles and Techniques. MIT press, 2009.
Korotin, A., Li, L.,, Solomon, J., and Burnaev, E. Continuous Wasserstein-2 2 barycenter estimation without minimax optimization. arXiv preprint arXiv:2102.01752, 2021.
Korotin, A., Egiazarian, V., Li, L., and Burnaev, E. Wasserstein iterative networks for barycenter
estimation. arXiv preprint arXiv:2201.12245. 2022.
Kullback, S. Probability densities with given marginals. The Annals of Mathematical Statistics, 39 (4):1236-1243, 1968.
LeCun, Y. The MNIST database of handwritten digits. http://yann. lecun. com/exdb/mmist// 1998.
Léonard, C. A survey of the Schrödinger problem and some of its connections with optimal transport. Discrete Continuous Dynamical Systems-A, 34(4):1533-1574 2014.
Li, L., Genevay, A., Yurochkin, M., and Solomon, J. M. Continuous regularized Wasserstein barycenters. Advances in Neural Information Processing Systems, 33:17755-17765 2020.
Lipman, Y., Chen, R. T., Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. International Conference Learning Representations, 2023.
Luo, Z.-Q. and Tseng, P. On the convergence of the coordinate descent method for convex differen- tiable minimization. Journal of Optimization Theory and Applications, 72(1):7-35, 1992.
Makkuva, A., Taghvaei, A., Oh, S., and Lee, J. Optimal transport mapping via input convex neural networks. In International Conference Machine Learning, pp. 6672-6681. PMLR, 2020.
Marino, S. D. and Gerolin, A. An optimal transport approach for the Schrödinger bridge problem and convergence of Sinkhorn algorithm. Journal of Scientific Computing. 85(2):27, 2020.
Mi, L., Yu, T., Bento, J., Zhang, W., Li, B., and Wang, Y. Variational Wasserstein barycenters for geometric clustering. arXiv preprint arXiv:2002 10543, 2020.
Minsker, S.. Srivastava, S., Lin, L., and Dunson, D. Scalable and robust Bayesian inference via the median posterior. In International Conference on Machine Learning, pp. 1656-1664. PMLR, 2014.
Nichol, A. Q. and Dhariwal, P. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, pp. 1162-177, 2021.
Nutz, M. Introduction to entropic optimal transport, 2021.
Pele, O. and Werman, M. Fast and robust Earth mover's distances. In 2009 IEEE 12th International Conference Computer Vision, pp. 460-467. IEEE, 2009.
Peluchetti, S. Diffusion bridge mixture transports, Schrödinger bridge problems and generative modeling. arXiv preprint arXiv:2304.00917. 2023.
Peyré, G., Cuturi, M., et al. Computational optimal transport: With applications to data science. Foundations and Trends@ in Machine Learning, ((656))355-6677 2019.
Rabin, J., Peyré, G., Delon, J., and Bernot, M. Wasserstein barycenter and its application to texture mixing. In Scale Space and Variational Methods in Computer Vision: Third International Conference, SSVM 2011, Ein-Gedi, Israel, May 29-June 2, 2011, Revised Selected Papers 3, pp. 435-446. Springer, 2012.
Ruschendorf, Convergence the Iterative Proportional Fitting procedure. The Annals of Statistics, pp. 1160-1174, 1995.
Schiebinger, G., Shu, J., Tabaka, M., Cleary, B., Subramanian, V., Solomon, A., Gould, J., Liu, S., Lin, S., Berube, P., al. Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell, (76(4))928-944 2019.
13

---

Schmitz, M. A., Heitz, M.. Bonneel, N., Ngole, F., Cocurjolly, D., Cuturi, M., Peyré, G. and Starek, J.-L. Wasserstein dictionary learning: Optimal transport-based unsupervised nonlinear dictionary learning. SIAM Journal on Imaging Sciences, (1)):6436778, 2018.
Schrödinger, E. Sur la théorie relativiste de I' l'électron l'interprétation de la mécanique quantique. Annales de I'Institut Henri Poincaré, 2(4): -10,, 1932.
Shi, Y., De Bortoli, V., Campbell, A., and Doucet, A. Diffusion Schrödinger bridge matching. arXiv preprint arXiv:2303.16833 2023.
Sinkhorn, R. and Knopp. P. Concerning nonnegative matrices and doubly stochastic matrices. Pacific Journal of Mathematics, 21(2):3433348, 1967.
Solomon, J., Rustamov, R., Guibas, L., and Butscher, A. Wasserstein propagation for semi coni-supervised learning. International Conference on Machine Learning, pp. 306-314. PMLR, 2014.
Solomon, J., De Goes, F,, Peyré, G., Cuturi, M., Butscher, A., Nguyen, A., Du, T., and Guibas, Convolutional Wasserstein distances: Efficient optimal transportation geometric domains. ACM Transactions on Graphics (ToG), 34(4):1-11, 2015.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score-based gener- ative modeling through stochastic differential equations. International Conference on Learning Representations, 2021.
Srivastava, S., Li, C., and Dunson, D. B. Scalable Bayes via barycenter in Wasserstein space. The Journal of Machine Learning Research, 19(1):312 346, 2018.
Staib, M., Claici, S., Solomon, J. M., and Jegelka, S. Parallel streaming Wasserstein barycenters. Advances in Neural Information Processing Systems, 30, 2017.
Stroock, D. W. and Varadhan, S. S. Multidimensional diffusion processes, volume 233. Springer Science & Business Media, 1997.
Su, Z., Wang, Y., Shi, R., Zeng, W., Sun, Luo, F,, and Gu, X. Optimal mass transport for shape matching and comparison. IEEE Transactions on Pallern Analysis and Machine Intelligence, 37 (1 (11):2246-2259, 2015.
Valiente, G. Algorithms on Trees and Graphs, volume 112. Springer, 2002.
Vargas, F., Thodoroff, P., Lamacraft, A., and Lawrence, N. Solving Schrödinger bridges via maximum likelihood. Entropy, 23(9):1133 2021.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, Attention is all you need. Advances in Neural Information Processing Systems, 30, 2017.
Vincent, P. A connection between score matching and denoising autoencoders. Neural computation, 23(7):1661-1674, 2011.
14

---

Appendix organization
First, additional notation is introduced in Appendix A. Then, we briefty recall some notions on undirected and directed trees in Appendix B. Similarly, martingale problems are introduced in Appendix C. The proofs of the main manuscript and additional theoretical results on Tree Schrödinger Bridges are given in Appendix D. Additional details on our consideration of the tree-based static SB problem are described in Appendix E. Details on the implementation of TreeDSB are given in Appendix F and the experiments are investigated in Appendix G.
A Additional notation
For any finite set E, we equivalently refer to the cardinal of E as card(E) or EE. Let (X, be a measurable space. For any x E (Rdje+1 and any m E {O....,)) e). let z-m = - 1, ,., For any family of measures (v..se defined on (X, and any i E (o...., {). let v- 2 je (0......)) U3. Let 7 = (i1,-.., 1n) C (1,...., and μ E 301) such that μ < Leb. We define 1 = [I,....eNNI and denote it by (ii...., { where q 8- q. We denote the marginal of μ along by pt,, i.e,, fr E gole) and we have for any A E (((()))..
nur (A) = 1x I1 ozz (Az)dz. In addition, note that fr & Leb. We denote the conditional distribution of μt given / by (1r(-))), i.e., (1)(1))) E g0 x (124)s and we have for any y E (R4)² and any A ( (RR))), pi (Alg) = ((x)//1(()) IT-,
83) I[ Oz.g, (Ay)) Remark that for any y E (R*)*, (11)) < Leb. For any subset J c F with card(J) = qs we also define Myj(-1-) E goli) x (R")1 such that for any y E (op(-te) = {p((11)) For a col- lection of functions (fiheer, with I C [I....,n] and n E N such that f; : Rdd + R, we define Bieif: (R")" I such that for any x = (Z1, Zn) E (Rdy", @ze)(()) =
B Introduction to trees
An undirected graph T = EV, with vertices V and edges E, is said to be an undirected free if it is acyclic and connected ( (Valiente, 2002, Definition 1. 1.19.). In particular, we have card(E)- = (V) - 1. The undirected edge between two nodes ii and v2 is similarly denoted by v2) or {v2, ). We say that T' = (V',E) ) is a sub-tree of T if T' is an undirected tree with vertices V V and edges E' C E. For any vertex V E V. we define the set of its neighbours N. as the set of vertices E V such that {e, o'} } E E. The integer card(N.)) is referred to as the degree of v. The vertices with degree I are called leaves, and we denote the set of leaves by VL C V. The (unique) path in T between two vertices v and v' is the sequence of two-by-two distinct edges {10, 11+1)); (with 2 1) such that e Uk+1 for any is (1....,P) n) such that k = 0 mod(2), 81
Undirected tree.
and Un+1 = -V. This path can be seen as a linear sub-tree of T. and we define n as the length of this path. We say that T is weighted if there exists map w E R4; in this case, w w((u1, 22)), or equivalently (f2,, v13) (also denoted by Woy.22 or Weg.01 is called the weight of the edge {v), )2). The tree T is said be rooted in E V if defines a partial ordering ST, V V such that for any U1, 02 E V, 01 ST. U/ if the node vn lies on the unique path between r' and 12.
Directed tree. Consider a directed graph T, EV, rooted in r E V. Any directed edge E E, from Uj E V to E V is denoted by (v), 122. T, is a said to be directed free rooted in r if (i) the underlying undirected graph T = (V, E) is an undirected tree rooted in and (ii) any (v1, 22) E E, is directed according to the partial ordering ST.r, {v1, E E and ta ST. U2. For any vertices (v, r') E V x V such that v ST,r v', the (unique) path in T, from v to v', denoted by pathy, (v, v'), defined as the directed version of the path in between and v' (viewed as sub-tree of T), which is rooted in v. We say that T, is weighted, if T is weighted and the edges of T, have the same weights as the corresponding undirected edges of T. For any (11, v2) E E,, we denote this weight by We, .ag- We say that T, is the (unique) directed version of T rooted in r. It is endowed with a canonical vertex numbering Ç V + (O...., card (V- - 1}. corresponding to a depth-first traversal of its nodes, starting from the root (Valiente, 2002, Definition 1.). This numbering is consistent
with the partial ordering on T, i..., 01 ST,r U2. ((v:) S ((12), and satisfies ((r) = 0. In the rest of the paper, we will write in an equivalent manner ((v).
15

---

For any vertices 02) E E x E such that ST. pathr, (#1, 02) corresponds to the ordered of edges in E, which define the ordered path between two vertices 8j and U2. For any vertex E V, we define:
(a) the set of children C, as the set of vertices v' V such that (u, Ey. In particular, for any E Vr., the set of leaves, one has C, = 0.
(b) its parent as the unique vertex p(v) such that (p(r), E E,, if v o r (the parent of the root is defined).
Note that N, = C, and, for any vertex E Vifr), N, = {p(u)) UC.
Definition (Tree-structured directed Probabilistic Graphical Model (PGM)). Consider a directed tree T (V, The directed PGM induced by т, (Koller & Friedman, 2009, Definition 3.4.).
denoted by TT,, is the family of distributions E gplVV which have Markovian factorization along T. i.e.,
- (IV) : T. (1,O))EE Tollo
Lemma 9. Consider an undirected tree T= (V. EV V. Let T' be a sub-tree of with vertices V such that r' E V'. Denote by Tú the directed version T' T' rooted in r'. Then, for any π E PT we have TV E PT.
Let (r, r') E V x V. We denote by T, EV,E,), respectively T,, = (V, Er), the directed version T rooted in r, respectively r'. We define the paths Pr. = patli,, ,, (r,r') E, and Pr.,r = patht, (r', E, It easy to see that
Proof.
(a) E Prr,r
(b) P Pr.r' [(12, 01): (v1, U2) Pr.r).
(c) Pr. [(02)) 01): (V1,
Pr..
Let ST, First note that for any (V1, U2) E E,., we have by Bayes decomposition Te₃ Tozles
Toy Toyl = Tor.e Then it comes
π Tr (VI.2)EEE Tzglo1
T. (e1.v))EP... Tuzlu1
(e1,2))EE, /Pr..
Tuglun
= (r2.)))P,,., Tezlen
(1..22)EE,AAP... Tez(m1
n. (V1i22)EP, Fagio (e1.,))EE//PP,,
T. (01,2))EE Txylor
and therefore, we have PT.--
Let T' be sub-tree of with vertices V" such that r' V'. First note that E,, C Ep,. Using the previous computation, we have for any A B ((IRR))V))
TV (A) = J[RO) vI The (xxr) X (or (n),E, Tezles (Feg (re,) ) II Se,, (Aor)dx
S(RO vI-Mi (Nr (Arr 2 (1,,, Tezlei (Ang(an,)) 2 (v,,,,),,,, Tezlos (Zog )aayy [Tp (OLNA)EE, Toalo)
which proves that my PT,
Discretized undirected tree. Let N 2 1. Consider an undirected T (V,E) with weights w. We say that T(V) (V(N), E(N) is a -ddiscretized version of T if it an undirected tree with weights (N) such that
(a) V(N) = V
{ot)
eEE, ke N1..,,
16

---

(b) E(M) = USEE Uk= {fre 1) with the convention that the vertices V.N and oe are defined such that {eg v =,
(c) Enepathy(oser) 1/we (N) = 1/W,,r if (0, EE
Remark that the leaves of T(N) are exactly the original leaves of T and that T(1) T. The non- uniqueness of T()) comes from the freedom of choice on the weights of its edges.
Discretized directed tree. Let N 2 1. Consider a directed tree T, = (V, E,) rooted in r E V with weights w. We say that T (V(M), E(7) is N-discretized version of T, if it is the directed version of T()) rooted in where ((N) is N-discretized version of the underlying undirected tree of T,-
C Background on martingale problems
In this section, we introduce the background on Stochastic Differential Equations (SDEs) and weak solutions of SDEs following the framework of (Stroock & Varadhan, 1997, Section 10.1, page 249). We recall that C (Rd) is the space of infinitely differentiable real-valued functions which vanish at infinity. In addition, we have that S4 is the space of d x d, symmetric, non-negative matrices.
Definition 10. Let T > 0 T +0, (0. R St and b (0,T) Rd R. locally bounded measurable functions. We define the infinitesimal generator, A, given for any f E C (RRd). t e [0, T) and IRd by
A(f(x) (be(2)), Vf(=)) + ( (V)(()))
(4)
We say that probability measure TP satisfies the martingale problem for for any E [0, T) and I E C (Rd), we have that (J(X;) - So (X.)))x.) is P--mattngale.
In the main document, see Section 2, say that "a path measure P is associated with dX, b(l x.)dl+o((,, X,)dB, with (B,)220 d-dimensional Brownian motion" if P solves the martingale problem associated with A given by (4). Unless specified, we always assume that such a path measure exists and is unique. Below, we recall the following theorem, see (Stroock & Varadhan, 1997,
Theorem 10.2.2), which gives sufficient conditions for the existence and uniqueness of solutions the martingale problem.
Theorem 11. Assume that for any x Rdd we have
(n, ao (s,x)0) 0 IRd, |0| 1, s [0,2]} > 0,
lim
(s. se [0,TT]
y=
In addition, assume that there exists C > 0 such that for any x E IRdd
supfl (t, x)]l s S [o, + sup((x, b(t,))) s [), S C(1+ ||x Then, there exists unique solution to the martingale problem with initialization 2o E Rd.
D Theoretical results on Tree Schrödinger Bridges
We respectively provide in Appendix D.I, Appendix D.2 and Appendix D.3 the proofs of the results of the main manuscript presented in Section 3, Section 4 and Section 5. Finally, make detailed comparison between our setting and the framework Haasler et al. (2021) 1) in Appendix D.4. In the rest of this section, consider undirected tree T = (V.E), E). where V + 1, and some subset which we denote by S = {Kk-). We define 5 = ViS.
5 C
D.1 Proofs Section 3
Proposition I is straightforward to obtain by combining the definition of the Brownian motion with the definition of =0 given in (2). The following lemma details the recursion relation between the (mIPF) iterates, which is key prove Proposition 2.
17

---

Lemma 12. Let (T") be the sequence given by (mIPF). Let n E N, kn (n 1) mod(K) kn + mo)(K). Denote by TR. respectively +1 with edges E the directed version of T rooted in ig.,, respectively in t1, We have:
(i) x² E TRa
(ii) x mnt1 Poin
" "v'l"" In particular, for any (v, (') Exc++1 " X (v.o"))E
ma
Proof. We show the result (i) by recursion on n E N, and will deduce (ii) from the proof. Using (2), first have Tº E TT,, where is chosen as k-- see Section 3. Thus, obtain the result (i) at step n = O. Assume now that n" E PT,, for some n E N.
Consider the paths p.. - patht (ik. +1)+1) and pn-1 path +u (1.,,11 kk.) Note that these two paths have the same length, denoted by J. and contain the same vertices, denoted by Vn- Let gp((+)) such that (x(("()"))))))))))))))))))())))_ +00. We have the following decomposition
KL(+|=") = (L(TV, (v.)) Saxeys KL(TVV Ii,, (iw,,,dv,,,, (IvV.)
Hence, the (n + 1)-th iterate of (mIPF) is given by n mht 80 =i.. with (gg TV,)
n+1
go(++)) Tin,t = Nin.ti
Since Th Te,, we have (i) VV_ X (EE, (px "bll and (ii) mV. E PP by Lemma 9,
where P Pn+1 is viewed as a directed tree rooted in i +1. Defining Vx+1 = Vn +ik we thus 2 (vve)))PP4 "V'le"
have mV.. Tha -V where Vi
Let E (J+1) such that Tih.11 Maa, and KL( +00. Similarly to the previous computation, we have the following decomposition
KL (.,) KL(TTe,,- 1 ) + Jgg (L( V 1i )d (Fiznti) 11) + Sge (I(Tin 1n I Vntilben )dpleent to (Dien = =KL(Mie, +1 t
Therefore, we obtain
n+1 Mients 2 AV Piuts (v.."))Pat "v'l"" Vv
Noting that Ex., \Pn Ex \Pn41 and recalling that nー TV
n+ TVV. it finally comes
x n++
-
T (v.*")EE
Tol
Pien+a
Tol
(5)
(u,"")PP44
(1,0"))E
Therefore, t1
P Te which achieves the recursion for (i), and obtain (ii) by (5).
Hence, Lemma 12 shows that the (mIPF) iterates admit a Markovian factorization on T, and can be defined recursively using the edges of T. We now provide the proof of Proposition
We will prove this result by recursion on E N. Observe that the initial- isation is directly given by Proposition 1. Assume now that the result Proposition stands for some E N. Let kin (n 1) mod(K), hn + 1 n mo)(K). Denote by TR,, with edges Ex,, respectively Tx,+1 with edges Ex,+1+ the directed version of rooted in iR,, respectively in +1. For any vertex of TA,+1. we define p(v) as the (unique) parent of and c(v) as the unique child when it exists. Consider the (n + 1)-th dynamic iterate defined by (a) and (b), ie., (P(ttt)) ((.t)))) (E,,)) To prove that this iterate has the properties stated in Proposition 2, we proceed by recursion on the edges of T following the bread-first order in +I,+1 In this order, the edge (ie.+1, +1). the first be considered, Remark that c(ik. is well defined since ik. +1 is leaf T.
Proof of Proposition 2.
Here, denote Teli,,, +1) by T. By construction, we have 1pa+1 (ix, offee (1)) Mik, &
(TPM (1 By recursion assumption, Pn (c(ie,
E
9 (C((0, 7), Rd)) since
18

---

(c(fk.-1) (kx+t) E Eks. Then, Pu+1
(1 is well defined path measure on [0, By
definition of the (mIPF) sequence, we have Mite 11
T
"x By recursion assumption, we also
n+
have that Ext(T
(e(Hkm p))ik. (1)) = Tolinn 1).s.. 11" Hence, comes that (I (c(is, 1) 11) 71o
Felfnn 1 12 (((k,,+t)llde,,. where the last equality comes from Lemma 12. Finally, we obtain (Panf
that Ext (fhm 11:0(fen 11)) Tiee 11:0(fka 12)" which proves the initialisation.
Assume now that p=++1 is well defined and has the right properties, up to some edge in TR. Consider the following edge, denoted by (v. vf) E +1, in the breadth-firs, order. By edge recursion, we have that Ext(P (p(e).,) T n++ p(r), and thus P ()(()):))TT = Define the path
Pn. -athh (iN,, +1). Then, we face two cases.
(i) Either (v, v') E Ev,, \Pn. Then, we have by (a) that p=+1 p=+1 (n,w²) (p(r).
P (u,or)o = n.. 2 PR (r,o"))0
In particular, Pht) is well defined path measure on [0, To,o.. Since (v, v') Ek.. (x,e)))
n+1 where the last equality Th T2p0r by recursion assumption. In particular, PR (n.e"):T.. Tplp = Tolp
comes from Lemma 12. We thus have Ext(( (u.er) -
(ii) Or (v',,) Pn. Then, we have by (b) that
pnttl (w,nr) = p (T()).)).T...)) & (P (o To @ (P (P)
In particular, pa+1 (w,"t) is a well defined path measure on [0, To,... Here, (u',u) E Ex,, and thus, Ext(PP (.) = The by recursion assumption. In particular, (TP (or, F...,.1 = x. "wlpe a n+1 where the last equality comes from Lemma 12. We thus have Ext (P = "n,ef"
o'ju
This achieves the recursion.
D.2 Proofs of Section 4
Remark on assumption A1. Although AI is not needed to establish the result of Proposition 3, Corollary 4 and Proposition 5, it is however crucial in the proof of convergence of (mIPF) stated in Proposition 6. Nevertheless, we choose to keep Al as an assumption in the statement of every theoretical result presented in Section for sake of clarity.
Additional definitions. We define the set Ps -ies Pi. where P. = { (++1)) pi). i.., Ps is the set of all probability measures E (lt+1) which verify
S fi( = Sea Se(x;; ((x:))
for any family of bounded measurable functions {fihes E C(Rd, ((RA,R)K Since Rdd is separable, there exists dense family of functions (ff JKEN with f (p:) for any l N and any E S, such that Ps if and only if
Saay--1 1 (2x)) = Ige (1:))(a(().)
or equivalently, upon centering J
Jasyr (x:))d((Z0)) 0.
In the rest the section, we consider such family (f+ JEEN
For any n E N we also define Pn Nies Pn, where P. for E g(++1)
fresyes (xx))d = 0, Vk [1.....m]] In particular, we have
S = OnEN- Ps.
(6)
Finally, (static-mSB) can be rewritten as
π argmin [L( π x0) m Ps)
(7)
19

---

Proof of Proposition and Corollary 4. In this part of the section, we present an extension of the theoretical results from Nutz (2021) to the multi-marginn setting. We first present two technical results, Lemma 13 and Lemma 14, which are respectively adapted from (Nutz, 2021, Lemma 2..).)
and (Nutz, 2021, Lemma 2.11.).
Lemma 13. Let (Kg)jes: be family of probability measures defined on (IR., B((B)). We define ØiES us JES Pj. Let A E 8 m0= B(R)) such that (A = Then, for 7 -almost any
r E A, there exists a family of sets {Xm Jmmo (12) such that
(a) P4 (x() = 1 for any i E S, and Pi [x((9)) = 1 for any j E
(b) A A T Im== Xm) satisfies xt E A and
Ln 1 Lmas F
22) A., Vx E A", Vm {O....,f)
Proof. Consider such set A. We define for any m E { ff the set
Xm = fu E Rd #= mm (A = 1}
where AW, = E (yo, ................................ Vm 1, u, Vm. VE-1) E A).
Take i S. Assume that (x:)
We recall that ~ # pr. Using Fubini's theorem and that SAz. _ for any 24 & X,, we have
(A) = SA d# (x- 60 (xse
= -Sa UAT: d# (x=:))ddp
fx. SA dF (x (o- (x) Jxy Un: dā (
< (x()))) + M4 (X:) =
which is absurd. Therefore, we obtain (x(),) = 1, and similarly, we have (ūz(X) = for any j E S For any y E (R24), any m E {O,..... l). we define the set
Av = (u E IRd (yo..... 9m- U, Um.... 9.--1 E A}.
Let i S. We have by Fubini's 's theorem
1 = (A) = JA das();) @ d
= SIzuYe SAp- (2x )dr
= =STI, x: f.aa-- du (d d## (x-1),
myti
where the last equality comes from the fact that |4(X;) for any E S, (15(()) for any E S
and that _ ies i4 X jjes My. Consequently, there exists measurable set A_; C I [ mati
X, such
that the following properties hold: (a) (A()) for any A-,, (b) #0 (A_+) 1. Similarly, this result holds for any j E 5 Le., there exists a measurable set A_; C m== X, such that the
mff
following properties hold: (a) fú, (A})= = for any y E A-j. (b) #0, (A-,) 1. We consider such sets (A m I for the rest of the proof and finally define the set
Ā
Am.
where Åm A- x {u E Ag, y A A_m) By definition, we have À C An m==n Xm. using the fact that Am CA for any E (O..... L}. In addition, for any E S, we get by Fubini's theorem n (A.) JA. dd.( @ d=
Sa_ _n, dd((a:))dF",(rr+)
(A_+) 11,
and similarly, we get 7 =²(A) for any j E Se. We can deduce that =²( = since (A S Em (AC,) = 0.
20

---

Let A. In particular, x* A. We define the set A for any m E (O..... Cj. We now establish the result of Lemma 13.
m==0 A X where X Xm Xmnim
We first prove (a). Let E S. Since x² E A, we have x" E A, and therefore r-, E A- By definition of A- we obtain that u4:AA (AAA 1 and thus,
pa (({])))) S (x() + pi((A,- = 0,
which gives (x.) = 1, and similarly, we have jág (x9)) = 1 for any j E 5c.
We now prove (b). Let m E {O....,f}. .}. Since x' E Å C A, we get o E Am Using that A C u Xm. we get x² E A Let x E A We denote xm = (Z)y.... 2 Em. ; x7). We need to show that x E A and rm E III X3 - =I nAa-"). First, since zÿ 2j or
3j for any {O,...,(), and E A and x= A we get that for any j E {0,..., .]. f), 2. Xj- Similarly, for any E l , 11. x E Aj Therefore, we get that xmh 11 (X, n nA)-))
Since Im E Am A (because E II;=11 (X, n Aj ")), we get that A, which concludes the proof.
n
Lemma 14. Let A C (RR)++) For any E {o....,fy, f), we denote X =projm (A9). We make the following assumptions.
(a) Assume there exists xt E A such that for any E A for any m ,0,...,ee. {} we have (z%..... Im Tma I mmt1 x7) A
x (-00, +00) (b) Assume there exists a family of functions {sp InEN AE(O.... with si,
such that for any E N and any k (O....,
: K 2}, we have Ph, (stty) =
(c) Denote (x) Ek- Pi (It) for any x E A Assume that for any x A F(x)
lim,- (x) exists and is such that F(x) E [-00,, with F(x*) E R.
Then, for any E for any XL, (x(x,) = lim (f(()) exists and is such that (x(x;) [-o0, +.0).
Proof. Consider A C (Rd)(+1 such that assumptions (a), (b) and (c) hold. Remark that we have F"(x*) Pik--1(2 (xix-.):
27) for any m E {o....,e} l). In Let x Ao. We denote xm (ai...... . 2m Fm 2 . particular, we have Th E Ao by assumption (a). Let define
Pia (xx)) -(x**)(((r') F(x)), Vk KO.....K- 2}. K-
Pix (xix ,) = F(a'K-1).
Using assumption (c), have pp((1s) 1-o0, too) for any S. Let ls (0...., K 2). We have by definition of F",
pi, (Fit) 1. (x'")
E -m= pi (x (3 1. (x*),
milk
where we used assumption (b) in the last equality. Since I's E A and x+ E Ao, we have by assumption (c),
lim oi (F.)) F(x*)) F(x)) = (x4)))
Furthermore, by combining the definition of F with assumption (b). we have
lim Pix (Iik-.) = F(xix 1) Pix-- (Iix-.), 2
which concludes the proof.
21

---

In what follows, before proving Proposition 3, we respectively show in Proposition 15 and Proposi- tion 16 how A2 and A3 can be satisfied in the case where x T,, in (2), that is
(I,O")E,, To'lo"
Proposition 15. Let T PT., Assume that n.. - ifr E Sorm! m. N(m,, J,Id), with my Rd and o, > Oir St. In addition, assume that for any (v, e') E Er, (-(00) = ( oe, Id) with J,,y' > Finally, assume that for any E S, Sga du; < to0 and H(44) +00. Then A2 is satisfied.
Proof. Let π = &isS/ Bies= with vi any Gaussian measure with positive definite covariance matrix. First, we have that
KL( m0)) E(o,"EE, Su (T Jd KL m9)
For any (v, 2²) E E,,, there exists Co.r" 2 0 such that
Jg KL I p Jdr. S Cr,o' - (((T.r)))))) + fgaxzd 1 - 12201/((20) JdTo & To (Iv:
I Co, - H(Tw)) + (1/62, Sau || 2 (xx) + (1/0?,
11 Sor |1 (Iwr < t00.
We conclude the proof upon remarking that KL.(Tr mp.) 400.
Proposition 16. Let Th E -T. Assume that nh = if E S or p = N(m,, I,Id), with m, E Rd and a, 0 r E St. In addition, assume that for any (u,v') E E,, (1),))))))) N(r J Id) with De,"" > O. Finally, assume that for any S, Pr admits a positive density w.r. the Lebesgue measure. Then A3 is satisfied.
Proof. We have that T² admits a positive density w.r. the Lebesgue measure. Letting #0 Øies μi jes ūy where ūy which admits positive density w.r.t. the Lebesgue measure for any j E 5" we get that # admits a positive density w.r.t. the Lebesgue measure and therefore T #0. which concludes the proof.
Using the preliminary results presented above, we now ready to prove Proposition 3.
Proof of Proposition 3. Assume A1 and A2. Since Ps is convex and closed in total-variation norm, there exists a probability distribution ग* solution to (7), or equivalently to (static-mSB), by using A2 with (Csiszár, 1975, Theorem 2.1.). Moreover, this solution is unique by strict convexity of KL(- 7
We now turn to the proof of existence of potentials defining (dr* "d by adapting the arguments (
of (Nutz, 2021, Section 2.3.). Define n" argmin (KL(T) #7) : ग E PP) } for any N Since {95) JEEN- C (p+++) is decreasing sequence of sets that are convex and closed in varia-vai norm such that (6) holds, we get from (Nutz, 2021, Proposition 1.17.) with A2 that
lim || 22000
x" l|rv 0,
or equivalently
lim [[dv" /dx0)/d>>))))))))))))))) (d** (dxx/dm")) (11) /dd" 200
= =0.
(8)
Following (Nutz, 2021, Example 1.18), there exists a family of bounded measurable functions (1p]) hene,essiiss with yp Rdd + E such that for any n N*
/d=") expl@t vp]
We consider such family (yp") Jnew" ies for the rest of the proof. By combining (8) and (9), we obtain, up to extraction,
(dx* /dm²) = lim exple explO p/1
aas..
(10)
22

---

We now define the following sets
A = {I E (Rd)e+1 1im + is p (xx) E [ +o0)),
B* - f= (R2) lirm es 1p (a:) -00 C
Using (10), we have (A*) = 1. Using A3, it comes F (A(A)) = 1. Moreover, we also get that x² (B*)) 1 by (10). Thus, it comes x xB()) > and # (B*)
using A3.
We then apply Lemma 13 to H and A = Ar. Since # (B")) > 0, it implies that there exists x² E B* and measurable set A B* verifying the properties (a) and (b). Following (Nutz, 2021, Corollary 2.12), we may assume without loss of generality in the statement of Lemma 13 that the sets X are measurable with IT X C A. In this case, we obtain that 14 (proj,(A9)
for any E E.
We now aim at applying Lemma 14 to the set A Remark that A directly satisfies assumption (a). For any n E N", consider the following transformation of the functions {4p hie
P - Sh. Pia )
Ak: E (O...., K -
Pix - Pik- + ER= 4P (1tt))
For any i S, we restrict pp, to X that the family { INEN" ,ess verifies assumption (b). Finally, since A C At and x* B", directly obtain assumption (c).
Therefore, Lemma may applied. It provides us with the family of functions {pi)ees defined by PP: : X [-00, +oo) with ş0: lim P a.sa.s. for any E S. Since (A(r)j(AAA)) 1 for any i E S, we may extend the functions sii to Rd In particular, we can find family of functions (vit hes with v/ Rd + [- +00) such that 1h. Pe a.sa.s Note that these functions are measurable as limits of measurable functions.
Since T #0 by A3, (10) turns into
(dn" /dm") exp D w3)]
aas...
(11)
Finally, we show that the functions v are p.-a.s. finite. Let i E S. Let us define A, {x, E Rd 1 (ax)) = -o Using (11), we obtain (dm* /dn"/d/")) x (RE))) = 0. Since ग. = /ii we have (A,) x* (A, (R SA ( (dr* /dr (d>)) 0
which gives the result.
We now turn to the proof of Corollary 4, which states that the iterates of (mIPF) can be expressed via potentials, in the same manner as the solution IT" (static-mSB).
Proof of Corollary 4. Assume AI, A2 and A3, We prove the result of this corollary by recursion on
n E N First take 1. this case, the first iteration of (mIPF) is multi-marginal SB problem of the form (static-mSB) where 5 = { with reference measure T0. Therefore, using A2 and A3, we can apply Proposition 3 and obtain existence of vb1 Rdd
+ R such that
(dw' /d exp(y 7 -
By taking vo for (1...., - 1}. we thus obtain the result at step m
Now assume that the result is verified for some E N with k, (n 1) mod(K) We define kn + 1 = n mo)(K)) and 4n E N as the quotient of the Euclidean division of n by K. In this case, the (n + 1)-th iteration of (mIPF) is multi-marginal SB problem of the form (static-mSB) where S = (ik.,+) with reference measure π². Using (13), we have that A2 is satisfied for this new (static-mSB) problem. Al I and A3 satisfied for this problem, given the form of n²" Therefore, can apply Proposition 3 and obtain existence of 2// 1k IRd + E such that
(d= "d>+)/dd") = explu/9++1 fxa+1 T -a.s.
(12)
By assumption, we have that n" < n Hence, we obtain n+1
< T and thus,
(dr*+1 /d>²) (dm "pm=+//d"")re (dm" /dx
n -
By combining (12) with the result of the recursion step n, we directly obtain the result at step + 1, which achieves the proof.
D
23

---

Proofs of Proposition 5 and Proposition 6. In this part of the section, we establish the proofs of results related to the convergence of (mIPF), respectively Proposition and Proposition 6, which can be seen as a natural extension of (Ruschendorf, 1995, Proposition 2. and (Ruschendorf, 1995, Theorem 3.1.).
Proof of Proposition 5. Under Al and A2, we obtain by Proposition 3 existence and uniqueness of solution to (static-mSB), which we denote by '². Since π Ps using recursively (Csiszár, 1975, Theorem 3.12.), the fact that {Ti = Msk π g vib is convex for any {O...., K - 1] and (mIPF), we obtain
KL(T"
| TT) = KL(T") x²) + Ei= KL(T' 1'
(13)
Therefore, we have Eis1 KI(T" | 1--1) < KT (T)) < 0o and thus,
+im KL(T*' | =
(14)
Let n E N with n > 2K, k E (0...., K - } and let 9n E N be the quotient of the Euclidean division of by K. We define ni 9, + k + with (nz - 1) k mod(K) if n S Otherwise, we set In n S 2K. particular, we have The = paa by definition of (mIPF). Therefore, we obtain || ITV < || - π ||V
(gn - 1)K + k + with (ng - 1) k mod K'). Note that we always have
|T" - xr" ---~""/r 1|V + - TH lTV
(triangle inequality)
(2KT((7" Th" 1==1)))/²
(2KT.(7"**) [=*)))/² (Pinsker's inequality)
where each term goes to 0 as +oo in the last inequality by (14), which achieves the proof.
We now turn to the proof Proposition 6, which requires several preliminary technical results. For the rest of this section, we define, for any N, On the quotient of the Euclidean division of n by K (in particular, qu = -))
Schrödinger equations. Under AI, A2 and A3, we know from Proposition 3 that the unique solution x to (static-mSB) can -a.s. written as (dx* /dx = exp ① Dies v47 where {v/ es are measurable potentials, referred to Schrödinger potentials. These functions are determined by the fixed-point Schrödinger equations
( (a.)) (z:)r((z4)/ J(way explEjes((1) ((0=
p4-a.s. Vie
which are obtained by marginalising n along its constrained marginals. This family of potentials is unique. Indeed, for any family of real numbers {Aig J (e(o......21 we have
(dn* /dm exp exp(Ates vil.
where Vip = 15 + dig for any k E {0.... .o...,,K K - 1} with di = Ain if E (O..... K 2] and X = =- -K-- Xix-
Remark on the initialisation (mIPF). Consider probability measure # (l(+1) of the form
(d# /d=") explA vi?1
(15)
where hies for any i S. Then, for any π E Ps, we have
his is family of measurable potentials with viP 12d - IR such that [Sa v.duul <
co
KL(m 9) (L #9) + StRayS Dies v. KL( #9) Eies Wadddp
Hence, -SSt equivalent to the multi-marginal SB problem
argmin(KL(7TH)))
E ge++1), मं Ps Vi SSE
We refer to (Peyré et al., 2019, Proposition 4.2) for the EOT counterpart of this result. This means that the solutions of the multi marlina Schrödinger Bridge problem invariant by multiplication of the reference measure by potentials on the fixed marginals. Consequently, the initialisation of the (mIPF) sequence may be chosen 7 instead of "
For sake of clarity, we now refer to the reference probability measure of (static-mSB) as or -- and to the initialisation of the (mIPF) iterates as x
24

---

Solving (mIPF) with potentials. To prove the convergence of the (mIPF) iterates to the solution Tt given by Proposition 3, we first rewrite these iterates with potentials, following the form of x*.
To do so, we recursively define the sequence of potentials {1b) InEN,iEs by
1 = =... = 4ik- = =0,
(16)
-i ( = log(rik- (fix-1)/ 1 Jeraye (xzo isc- (x-ix-)),,
and for any n E N and ... ,..,,,KKK K -
2b 1g++1 (x:.) logirs, (#rr)) SeReye exp exp(Ai vi 66n-1 -xi,) + Dm-tr W (Xim))
h( (x-s.))
(17)
recalling that yn is the quotient of the Euclidean division of n by K.
We now define the sequence of probability measures {-") JnEN by
dr /d = expl@x E=
the 1.++1
winm kn = (n - 1) mod(K)) n = qn K + kin + 1.
(18)
exple v..1 exply , and thus fna VSk dusk- [L -i.-,): ). Consequently, x0 can be chosen as the initialisation of (mIPF), following the previous remark, if we assume that [L(pix-1 aax-1) < co. In (TreeSB) with 1K-1, the latter assumption is directly verified since we choose -ik- = Mix-1'
In particular, have (dm
Let N, with knn (n 1) mod(K), k, + mod(K). Using (16) and (17), we get that iik. = pix,. Moreover, we have
dn /dx = explu 16 thn
2/
(19)
with the convention that wh- ix-1 O. In particular, we obtain that o "patt
n1+1
In conclusion, the sequence (x" JeEN defined in (18) verifies xn+1 = Mixn " for any n E N. By decomposition property of the Kullback-Leibler divergence, this sequence solves (mIPF) with initialisation n We consider such iterates in the following.
Since T.... Mix.," we have that
9h KL ("==)==---- = Sz (bb vie )dpie..
(20)
Before proving a multi-marginal counterpart to (Ruschendorf, 1995, Lemma 4.1), we state and prove the following result.
Proposition 17. Let No: xı two probability measures on Rdd such that To < my Then, denoting doo/dm the following assertions equivalent:
(a) KL( πz)= to0
(b) Sge log( (z)(ff(()))))o) < t00
(c) Sau (z) (())) < too
If one of these conditions is satisfied then Sge [log(f) dro & KT mı) + 2/c.
Proof. First, note that
Sge (lo((f) (r)())
S (ogg(((()))(()))
1/c
(21)
where we have used that for any E [0, 1], Iu log(u)) < 1/e. We have that (b) implies (c). Using the previous result we have that (c) implies (b). Hence (c) and (b) are equivalent.
addition, is clear that (b) implies (a). Finally (this is more of a convention), we have that KL( m))
Jge log(f (f (z)) (TT(()) for
< +oo. Using (21) this implies (c).
Finally, we have
((x) fgo [log( (z)) (x) for (z) 2 fer (in) ((
KL m1) + 2/c,
which concludes the proof.
25

---

We begin with the following lemma which controls the integral of the potentials uniformly w.r.t N. can be seen as the multi mariinaaa counterpart of (Ruschendorf, 1995, Lemma 4.1).
Lemma 18. Assume A4. There exist fcities (0, (0,+ such that for any function f (R4))t2 -
of the form - DiES St. we have
callfllu 2 [[All ViE.
(22)
For any n E N,, we have
(a) Eies Ju d/ S KL((T" a) < oc,
(b) S(ReYC (eey)n((iis 1े Dies "f)dd" s KL(T*) ()) <
(c) supneN Saa |,,11 du, oo, S.
Proof. First, we have that (22) is a direct consequence of (Kober, 1940, Theorem 1) and A4. Let us now prove (a). Using (20), have
Kn KL(TTM 1) = Eez Eizo KL (K* + KL((T)) 1
= EEO Eies Sa (wet1 - v/)dd + S (vi - Vixi )dui
Eies Ei=o Sp (bftt] vf)du (p Vix ddpix-
= - (20 - )dd + laa (ve Vixz_ )dpix-
Eies See 2 du: S KI((*) T)..
where the last inequality follows the proof of Proposition 5.
Since the first term in the inequality of (b) is equal to KT,(T* TTk), we obtain (b) using that KL((T' x S KL.(T") #) following the proof of Proposition 5.
Let us now prove (c). Since KT(x=* #) < oo, using Proposition 17, we have that Dies v E (+(")). From (b) and Proposition 17, we also get that Dies(v; - -.)) E (**),,, and thus
Jnedeee (ve + dx Co with Co O. Therefore, we have
[Raaer ① Dis v.r1 dr= S S(RUYerZ
① Dies v dr + S(Reee DDEEs(Tf- - dr s 2C.
Using (22), we conclude with A4 that for any E S, we have
SRl | dps I 2c,Co
which concludes the proof of (c).
The next lemma gives an explicit expression for (L(T" T). can be seen as the multi-marginal counterpart of (Ruschendorf, 1995, Lemma 4.2). Lemma 19. For any n E N, with kin = (n 1) mod(K), we have ddir.
KL( | F) = Sox ik..
+ - k., Sax viaat1 explu -b9 +
+ K-1
Jgs 1.. explu - wan 1d
Proof. Let n E M. with Kn = (n 1) mod(K). Using (18), we have
kT((T" (T) = SEgd w/gn +1 ddtiin
E=0 Ekan' kn-1
fee wgn "l + ΣK- f g0 dr p
(23)
Consider m E [kin + l,.., K - l}. Let mn be the closest integer to n such that m. > n and m (m, 1) mod(K). By (19), we have
d cxpl D ko 46 299h +1 +1
Using (19) recursively, we obtain
di exply Ug Idd
(24)
26

---

where we recall that "mm" Mim-
Consider now & E {0,...,ky kn - 1}. Let ln be the closest integer to n such that en n and
== (In 1) mod(K). By (19), we have
dr" = exp I j=k + Vi +1
① =
Ug
and using (19) recursively, we obtain
d = explitt
(25)
where we recall that mic = Mic- We conclude the proof upon combining (23), (24) and (25).
We are now ready to prove uniform integrubility result which the multi-marginal counterpart of (Ruschendorf, 1995, Lemma 4.4). Before stating Lemma 21, we prove the following well-known lemma. We recall that a sequence (Va)neN such that for any n E N, V, (,), uniformly integrable w.r.t. μ if (i) supneN Jg [ du +0 and (īi) for any 0, there exists K > 0 such that for any E N. JB(O.K | du S. S.
Lemma 20. Let f : IR + R, convex and non-decreasing on [A, +00) with A > 0 and lim,-++o (x)// too. Assume that supnen SR. ,l +00. Then, (Wn)en uni-
formly integrable w.r.t. μ..
Proof. Since f is convex, using Jensen's inequality, we get that sup I S[za Waldd/ < +oo and since limg-++oe fxx)/x = too we have supnen Jiz. n du < +00. Let E > O, there exists K > 0 such that for any x > K, x S <f(x)/ B with B supnnN J. f( < too. Therefore, we
have for any n E N
N Nnldd S (c/B) Biy f((aalddp se, BO,KY
which concludes the proof.
Lemma 21. Assume A4 and A5. Then, (exple) les WheN is uniformly integrable w.r.l. π.
Proof. is enough show that the sequence [f(expl+) f u + u log(u) is continuous, convex and such that lim [u)) u = +00, see Lemma 20. Let n N. We have
Sapll ies ) is bounded in L (TT), where
[ feeppl v ]d KK (T)
Wie dlix 2k= SEe thr explvn. - -V u
(Lemma 19)
=K-I Joa //7 djb,, + Σ k=o
K-2 SRe wh (exple - - wn++1 11
KL(TT 7) (e+ + Σ k=0 Saa with dui
(Lemma 18-(a), A5)
< KL(T") .) - (c+ Ek- suPnen Jao Ww dps, 00.
K-2
(Lemma 18-(c))
With the preliminary results stated above, we are now ready prove Proposition
Proof of Proposition 6. Using A4 and A5, have, by Lemma 21, uniform integrability of (exple) kes oo mffe in L (T). Therefore, the sequence (."K)net is relatively compact with respect to the weak topology of L denoted as the r-topology. We recall that lim, KL(
nk 0. This implies that = +nnN+ JEEN is also relatively r-compact. By trivial recursion, obtain that the sequences ( JNEN, where
[2..... K 1} also relatively 7-compact. Therefore, { meN is relatively r-compact and r-sequentially compact.
We consider an increasing function + N + such that (T²M " mmeq()) is r-convergent subsequence, and we denote by से its limit for this topology. particular, E Ps by Proposition We assume without loss of generality that (NN) c. KN.
27

---

Using the lower semi-continuige of the Kullback-Leibler divergence (Dupuis & Ellis, 2011, Lemma 1.4.3), we get
KL R) S lim inf KL(T" T) S lim sup (L(x"" (T)
Consider l E {O,..... - 2}. By (19), we have
dpig
dm. dmpRiell
dnn drnkez
explu v..
dm.
d d
and thus,
I|,s m. TV (1/2) Sae [d /dus d/d (1/2) Sg |exp[un. w 1/d
With Proposition 5, we obtain that fexplu? iit w converges to in L (pt.). In addition using the uniform integrability of {& JEEN and A5, we get
lim sup Sau v. explus. + lim supp--tuo Sar 4% duiz
K-2 We denote = Ke. Since KL [2) = Sge Vix Vix-. dusk-= + 2k=0 Jae why explitt, w "dyss by Lemma 19, we finally have
KL(@ #) S lim sup EE kk= Sax vi, duu. S KL( RT)
where the last inequality comes from Lemma 18.
Since म. 14, for any E S, using Proposition 5, we have स x² by uniqueness of n² Hence, x² is the only limit point {x") JeeN in the r-topology. In particular, KL(T" #) - KL(T" 7). Since Ps is convex, this last result implies |-" .T" rv O, the proof of Theorem 2.1 in Csiszár (1975).
We finish this section by highlighting that A5 is stronger than (Ruschendorf, 1995, B1). A natural extension of the latter assumption would consist of having guarantee on the (K 1) first potentials given by (17), presented below.
A6. There exist 0 < c < E such that for any l E {0,.... K 2), we have S exp(p S a
Under A6, (Ruschendorf, 1995, Lemma 4.3) can be adapted as written below. Lemma 22. Assume A6. Then, for any N"
(a) for any k E t0....,K K - 2}, there exists On,k E N such that
(c I exply%-1 vR1 S E- ((/6)om)((22)
(b) there exists Kn E such that
1/ek-1. (c/2)0mx-1(KK3 ((K-2) S exxplyn-1 - V. ( (K--)
a./c)
where InEN ke (O.......K is strictly increasing sequence that can be explicitly defined.
Proof. We prove the result by recursion on N Take n = 1. Let k E {O...., , - 2}. We define 01,4 = 0 and directly obtain (a) by AS since
V = 0. Let us prove (b). We have by (17)
Jmay exp ( K
vi [hdv--k-1
explu - -wtt v
Jendye exp k vi, [hdu-
Jaay exple E=o (w) - H vo Jide-ig
Jead exp k=d sa Ihrv-ix-11
R-2
Using (a) at rank n = 1, we have
K-2 1/cK-1 5 expll k=u ( - - 10111/66 K-1
28

---

and therefore, we obtain (b) by taking 01,K 0. Let us assume that the result is verified for some Nº. We have
J [x K 1 WT, Jidu
exple
1//0
J Teepp( 1i Vi Jadv
J expl " (17
x- 1ix 1+ K- 0/ "/dd to
eepp Vi 1Adv-
Using (a) and (b) rank n, we have
1/ck-2 (c/)) an, S scxp(+ {
K-2
K-2
(//
11/5 K- (0/c) (K-2) E ank
cx K-1 (c/0) < <xxplu
v ik- S k- -1, x
Therefore, we obtain
e (c/o)(K 2) Ex- and explp K-2 [ -リー + { Win
S (c/c) 2K Σイ and
e -c) and exply/ w s ē (0/c) -2) Σ an.
2)
2)
Now, we define an+1.0 Following the same steps as above, we recursively define
kk=1 an to obtain (a) for k 0. Consider now & {1...., K 2].
K-1
Σj=o Ontly
Ej'z+1 On..j'
On
which gives (a) at rank + 1. Let us now prove (b) at rank n + 1. We have
f expl x-2 w Jadv-1k-1
explu - v
I expl K-2 vi lhdw-ik-1
expl k twn -P H K- vi, Mide
I [epp or Ind-iix
Using (a) rank n + 1, we obtain
1/2K /////K Mk= amll <xxpl K-2 [1 1 H
< </ (K)c) Ef ant1,k
Eiioo
Therefore, by taking 0x+1KK-1 0 -K-- K-2 -k=0 Ontl,, obtain (b), which concludes the proof.
Li
Unfortunately, Lemma 22 only yields non-vacuous bounds in the case = 2. Indeed, when K > 2, the sequence fan, kneN "EENN..E(.......... (o..o......I leads to increase the bounds the quantities explv1? n- w.]. which motivates the use of A...
D.3 Proof of Section 5
For the rest of this section, consider the multi-marginal Schrödinger bridge problem given by (TreeSB) and establish in Proposition 24 the correspondence with the regularized Wasserstein propagation problem presented in Solomon et al. (2014, 2015). We first state a technical result.
Lemma 23. Let 0. Assume that no is given by (2), where E V is chosen arbitrarily. Then, for any π E P.., we have
<KL(T m²)
Eio,')E, (Wo, Br,, IIX X [13] eH(Tr,or) + Evev Evey card(C,)) + SKL(T, x*)
where recall that C, = (o' E V: (V. E.).
29

---

Proof. Since T E TH we obtain the following decomposition
KL x0))
= r [(,,))E Toljo T2 II [I,,, 78/po)
KL(TT, (T/ + "(,,,")E, J&a (L T To (2y)o()))
= -TT m E(o,ereE, JgaxR log oh d Ei.,'EE, Jaa (=le)()()))))
We finally obtain the result by using the definition of x and noticing that H(TT,w)) H(T,) for any (0,0") E,.
Ja (
Proposition 24. Let e > O and pro such that pa < Teb. Assume that xO is given by (2), where r E V is chosen arbitrarily, and that ypr = ddjo/dTa Also assume A2. Then, the set of marginals of the solution to (TreeSB) is exactly the solution to the entropic-regularized Wasserstein Propagation problem (Solomon et al., 2014, 2015) defined by
arg min E (1,0))E,, Wayy W W2.e 2.e/w,, Por) te EeEv Caryc card(C,)H((,) + =KLL Po): (WP)
{volvev E git1, = μ4, E S}
where we recall that C. {o' E E (v,t')
Proof. Assume that x0 is given by (2), where E V is chosen arbitrarily, and that yp, dpo/ dpo/dLeb. /Leb. In particular, we have x. = μo. Moreover, is clear that x verifies AI, and A3 by Proposition 16.
Let {NoJvev E ge++1 and (u(ow')) (ULLN")) (U,N"EE, (212))EE-), We define
F((uv}) -E(n"))E Wa, W2 G(ur,fu (v("0)p)) -E(,")E,,, (wo 11 - X xp]1 -
2.//wew (Vo, Ugt) Evey card( + +KKT(u po).
(m,)EE) H(v (o10") + +KKL(u) mO).
By definition of the regularized Wasserstein distance given in (3), we have for any (ur)uev goet1
F((vo)) = min(G((r, (plev())) (0,00 E g0l2), (exo)
vo, uf (u,1))
unr, V(U, V(u,' EE, (26)
In particular, we have F((me)) S G( (T. }} for any ((++)) We now prove the result of Proposition 24 in two steps denoted by Step 1 and Step 2.
Let us not assume A2 for now. In this case, we prove in Step 1a and Step 1.b that solving (WP) is equivalent to solving modified version of (TreeSB) given by
Step 1.
n² = argmin (TL E PT,, क. 14 Vi E S
(T,-TreeSB)
Remark that any solution (Te-TBe)BB solution (TreeSB), but the converse result may not be true.
Step La: (WP) + (T).TreeSB). Consider solution {Mr (ut)eev to (WP). For any (0,0²) E,, W2 2 (Ve, v.) is well defined and thus, there exists y(ow"") E I[,, ut) such that
p(u,e)) [| X (E/ w HH(T) A (((5.,)))
(27)
(1,0") Using the gluying lemma, build the probability measure n² Ve II 2 1 such that (i) x² E TT,, and (ii) t.t and (v,a") have the same distribution for any (v, v') E,.. In particular, we have गे p4 for any i E S.
30

---

Let us show now that x² is a solution to (T, (T,-TreeSB). Let π E Tr such that n; As for any i E S.
We have
eKL(TT (T)) G(Tr, (Tow))
(Lemma 23)
2 F((To))
2 F(1xT))
(definition of v²*)
= =Glf, (vlv)0())
(see (27))
=G(TT (m(nsor))
(definition of π
= <KL(TT' 1=)).
(Lemma 23)
Therefore, nT' is solution (TeeSBB))
Step 1.:: (T,-TreeSB) - (WP). Consider now a solution nt to (T,-TreeSB), Since x² E Pr,, have n² n, [1(,,0')E, Tel and मद = 14 for any E S.
Let us show that (ne solution to (WP). Let fuuhev E ge+ such that 1/ = fi for any i S.
Let (ulun")) be a family of probability measures such that ulua") E g012) 2 (e/o')
V (e,v') for any (v,,")= E E, Using the gluying lemma, we build the probability measure ग = 14 [[(,," Pojj such that (i)
T T, and (ii) Tu,e" and ((1,0") have the same distribution for any (v, v') E E,.. We have
<KI(T T = G( (f(ew))
(Lemma 23)
G(v,, G(umrpllee (vlore())
(definition of 1)
<KT(Tr* x7))
(definition of x**)
G( ( (m(()er)))
(Lemma 23)
By taking the infimum in the previous inequality over the families (u(v,0) (o, we obtain by (26) that
F((u+)) G( (m)me)))) F((xT)),
and therefore, eev solution to (WP).
We now assume A2. By Proposition there exists unique solution "* (++11) to (TreeSB) such that we -asa. have (dnT /dz")) = exp DDes w/], where {य7 hes are measurable potentials with d Rd + R. Since n E AT., we also have x* PT, ie., the potentials {vi hee do not modify the Markovian nature of x Therefore, n² also the unique solution to (T) Using the equivalence between (T, TTeeSBB and (WP) established in Step finally obtain the result of Proposition 24
Step 2.
In particular, Proposition directly derives from Proposition 24 by taking ig-1 and Ho Max-1
D.4 Comparison with Haasler et al. (2021)
In their work, Haasler al. (2021) study the static and discrete-state counterpart of our approach. Given state space X such that |X| n + 1 with E N, they establish correspondence between multi-marginal EOT with general tree-based cost and discrete-time multi-marginal static Schrödinger bridge, and provide an efficient method to solve these problems. In this section, we provide details on their framework and give precise comparison between our theory and their results.
To be coherent with the setting of Haasler et al. (2021), we adapt here some of our notation. Let us define ( IR (x+ For any g N the of probability measures X defined as Gai (M E Z (M, 1) 1}. We denote 3 = (1)), For any tensors M, P E Z the Kullback- Leibler divergence between M and P is defined as KL(M P) (M log M P 1) and the
31

---

entropy of M is defined as H(M) = -KL 1), where the operations are meant componentwise. In the rest of the section, we consider an undirected tree T EV) with M + 1 such that may be identified with (O.
Details on the results of Haasler et al. (2021), In their paper, the authors consider a cost tensor C. Z(++1) that factorizes along T, i.e., for any (io with for any {O...., C) je (O.....n3,, we have
Cj (OuT)EE Glue Jode
where ((u,e)) E Z(2) is cost matrix for transportation between the marginals at vertices v and v', see (Haasler et al., 2021, Eq. (3.1)). particular, this cost can be seen as the discrete counterpart of the tree-based cost introduced in (1) in the quadratic setting.
Given a subset S C V with [S] = K and set of marginals {uibes -. Haasler al. (2021) study the EOT problem associated to T, see (Haasler et al., 2021, Eq. (2.4)), which is given by
argmin{((C, M) - (M(M) M E go(tt)),roj,((( (M) μ4, Vi S}
(discrete-EmOT)
This problem may be solved with Sinkhorn algorithm (Cuturi, 2013; Knight, 2008; Sinkhorn & Knopp, 1967). for which the authors provide an efficient implementation adapted to the tree-based setting, see (Haasler et al., 2021, Algorithm 3.1). Moreover, they state the convergence of their method in (Haasler al., 2021, Theorem 3.5), as a direct consequence of the results presented in Luo & Tseng (1992).
In (Haasler et al., 2021, Section 4.2), is assumed that S corresponds the of the leaves of T, as we do, and it is shown an equivalence between Mescrete-EmOO] and the discrete-state static SB problem stated in (Haasler al., 2021, Eq 4.2), which is given by
argmin EE ((O,o')EE KL(M( (lla ( A(www))
(discrete- -reeSt
M((,00) E ,
Avoloeve E gelt1 " vy. M(r,' 1 Ver, V = P4, vi E S)
where T, = EV, is the directed version of T rooted in an arbitrary vertex r E S, and A(v.0")
exp(- (ap(-C(v.d)/e) (2) for any (v, e') E,.. Remark that A(v,²)) may not necessarily be transition probability matrix.
Finally, Haasler et al. (2021) provide two main numerical experiments. In (Haasler et al., 2021, Section 5.2), they consider a tree with 15 vertices, 14 edges and leaves, combined to the state-space X {0, 1] 1000050 and solve the corresponding (discrete-EmO1) problem for the quadratic cost. In (Haasler et al., 2021, Section 6). they apply their methodology to estimate ensemble flows on a hidden No, they consider tree T with internal vertices (modeling the distribution of N agents at time E {l...., T)), that are linearly linked, and such that each of these vertices is independently linked to S leaves of T (modeling observations at time & (1....,,)) 7). In this setting, the state space is given by X = (1,... 100} They solve the formulation (discrete- B where the reference measure is chosen random walk.
Markov chain. Given
Comparison with our results. We now establish remarks on the main differences berween our methodology and the work of Haasler et al. (2021).
First of all, the continuous state-space counterpart of (disc)) is given by arLmin(KL(T . St,,, T. Mi, Vi E S},
(28)
where 7 a reference measure which factorizes along Ty. In this case, To,e", T. and To u in (28) respectively correspond to the continuous version of M (v.e"), Vy and ((v,0') in (discrete-TreeSB) In contrast, our formulation of the multi-marginal Tree Schrödinger Bridge problem given in (TreeSB) is minimization problem over all probability measures E (e(++)) and is not restricted to the distributions that admit Markovian factorization along T as in (28). Hence, our framework may be considered more general. Remark that under A1, A2 and A3, Proposition 3 states that (TreeSB) admits a unique solution n² & n such that (dr /d=²) can be written with potentials. Then, n" E Pr. since T E PT,, and (TreeSB) is then equivalent to (28).
32

---

Furthermore, (EmOT) is more general than the continuous version of can recover by taking any measure V of the form (dv/dLeb) = expldies (Pi]] in (EmOT), where fis) is a family of potentials such that [g= p;djus < oo for any i E S. As a consequence, our setting allows us to choose the root r E V for the SB problem, whereas Haasler et al. (2021) only consider the case where E S. In the latter case, we establish in Appendix E that can be chosen arbitrarily, as stated by (Haasler et al., 2021, Corollary 4.3).
which we
Finally, TreeDSB deeply differs from the framework of Haasler et al. (2021) due its dynamic nature. Although we solve the same tree-based static SB problem (up to continuous/discrete state-space consideration), our approach consists in computing dynamic iterates (i.e., path measures) using diffusion- based methods instead of static iterates (i.e., distributions) using Sinkhorn algorithm. This paradigm is at the core of the DSB (De Bortoli et al., 2021) methodology, and offers an efficient approach to tackle high-dimensional settings, where Sinkhorn algorithm would fail.
Here, we present some advantages of the method proposed by Haasler et al. (2021) compared to ours. First, Haasler al. (2021) may choose any kind of tree-based cost in practice, while our methodology only holds for the quadratic cost. This limitation is shared with all approaches based on the DSB (De Bortoli al., 2021) methodology. Indeed, since the cost is determined by the reference path measure, we often choose quadratic costs associated with Brownian motions or processes. Moreover, Haasler et al. (2021) may consider various inhomogeneous (discrete) state spaces for the vertices of T, as presented in their numerical experiments. In our case, this approach is not compatible with our diffusion- based method. Finally, unlike Haasler et al. (2021), our method is not scalable with the number of vertices or edges in T due to computational limits. This limitation is common to all multi-marginal approaches which rely on neural networks to parameterize the potential and/or the distributions of the multi-marginal OT method, see Li et al. (2020); Fan et al. (2020); Korotin et al. (2022, 2021) for instance.
E Further results on TreeSB
Choice of the root r in (TreeBB). We recall that the reference measure T0 considered in (BreeSB),
which is defined in (2), verifies x0 E TT, for some fixed root V and x. < Leb with density sr- Moreover, we have T- (- x.) N( e/(2w.. )aa) for any v') E E,, and thus, x² is entirely determined by the choice of the root and the density on the corresponding vertex fr-
As presented in Appendix 1.1, we recall that (TreeSB) is equivalent to any multi-marginal Tree-SB problem with a reference measure # given by (15), i.e., H" writes as /d = exp ① ics w!], where {w/ his is family of measurable potentials with vif .dd - IR such that I.5ge a vif < co for any E S. In the case where r is chosen as a leaf of T, this result implies that (TreeSB) is unchanged if
(a) , = dv/dLeb where v E P is such that KL(H, (p) < oc, (b) r is replaced byr' r' E S, as long as H(HG) < co and H( < <0.
Therefore, under AO, the setting chosen in Section 3 is equivalent to any other setting where is arbitrarily chosen in S and p dv /dLb where KL (1)) < co.
Consider now the case where r St, i.e., is not leaf of T. Then, the choice of yo, can not be made arbitrarily anymore, since it determines further regularization on the r-th marginal of the solution to (Br In this setting, the sequence defined by (mIPF) is unchanged. Hence, TreeDSB proceeds in the same manner as presented in Section 3, except for the first iteration, which we detail now
Let us define P P= pathT;; (io, where T; (V, E..)) is the directed version of T rooted in ig. We recall that first iterate of (mIPF) is defined by
π = (KL x = E gil+t), Tio = pro
Following the proof of Lemma 12, it is clear that
0
0 1 = Mio
0 Tol
π
Mio
(0."")EP
(E, iP
(o,v')EEio
33

---

where we emphasize that P ((0,)') Es, (d), E.). Therefore, Proposition 2 still applies between r and in, by considering instead 1 practice, this means that the first iteration of TreeDSB consists in computing the time reversal of the path measures IPQ for any (v, v/) E P.
Extension of the regularized Wasserstein barycenter problem (reg Consider the regular- ized Wasserstein-2 barycenter problem defined as follows
2.e/w. F = arg min (2 w;W 2./ww (1, p4) + (etl(() + eKL(μ no) : p E P), (Ho-regWDD
where (W)) (1....f) E +00) and μo E P is a reference measure. This formulation admits a further regularization compared to (regWB), which tends to make μº closer Nao In particular, given a Wasserstein barycenter problem onto a star-shaped tree, the formulation (4o-regWB) may be more adapted than (regWB) if we have priori on the form of the regularized barycenter. In the case where fo = N(0, N(0,,fīaa) afta), letting do + oc, we recover the (le, (f - 1)e) Noubly-regularizze Wasserstein barycenter problem (regWB). In the same spirit as Proposition following result from Proposition 24, which proves that (Ho-reg WB) can be solved with TreeDSB.
we can derive the
and oo S such that pa & Leb. Assume AO. Also assume that T is star-shaped tree with central node indexed by O, and that the reference measure of (TreeSB) defined in (2) verifies
Proposition 25. Let
0 and fr duto dLeb 0. Under A2, (Ho-regWB) has a unique solution To. where x² is the solution to (TreeSB).
Below, we provide practical guidelines to parameterize μ4 when is chosen a Gaussian distribution.
Gaussian design of po in (eg-re Consider an undirected star-shaped tree T with K + 1 vertices and leaves , brought by Ho when it is Gaussian distribution, we set its mean Σ [[p]/ K and its diagonal (ov(p4))=
KJ. J. In order to incorporate the marginal constraints in the penalization
covariance matrix as a x (E wise and a positive hyperparameter. This choice of variance helps to correctly explore the state-space at the very first iteration of TreeDSB. which is key to ensure numerical stability. In this setting, (TreeSB) verifies A2 and A3, by Proposition 15 and Proposition 16. In particular, we use this approach for two of our experiments: synthetic Gaussian datasets and Bayesian fusion, see Appendix G.
where the inverse operation is component-
F Algorithmic techniques
Time discretization in TreeDSB. Denote kn (n 1on-1) mod(K) for any n E N. T EV) be a weighted undirected tree and consider the nulti-marginal Schrödinger bridge problem (TreeSB) associated to this tree. We recall that for any (u, r') E E, we define Tr, e/ (2w.,,))
Consider the path measures (TPm ( (V,,,, recursively defined by (a) and (b). By combining Proposition Proposition 2 and results on time reversal theory (Haussmann & Pardoux, 1986), we obtain by recursion that for any n E N, any (u, o') E E Ph (w.wr) is associated with a Stochastic
Differential Equation on [0, ] given by
dx,- fi (X + dB, X T
(29)
Let N N In order sample from the dynamics (29) at iteration n E N, we consider its Suler-Maruyane discretization on (N + 1) time steps,
Xm+1 Xm + Tm-1 1/ v (Xm) Vmm+1 2m+1, Xo T.. Em Yis and (7m m== E (0, oc) is a time schedule such that ΣM 7m To.e". This results approximating the path measure pm
(30)
where 7m ~ N(O. I2) for any m E (1,.... N.....,,, N}. fm
(w.wr) by the
joint distribution o. n..
(WWee) E (N++) defined by
n.N TH 2 OND T ((o"")mmllm N.N x (V,00))
... where "(v,"").m+ 1m (.(mm) = N(rm + Tm 15 Wf o, (xm), 7 Ta) for any m E N N n.N If N is chosen large enough, then (e. and Pp (o have approximately the same distribution
34

---

for any E NO,....Ny. Consequently, (1sgor) RR is naturally approximated by the joint distribution
Nu.N
(eooo) E (N(+1) defined by
R.NN (w,)'
B.N
m. m=
If N is chosen large enough, we obtain that
T R.N (o.o",,N 1|N (-1
(a YN IN-m w. ( + mm V Vlog Pe,,' Nm (WN-a)))
where Pv,ef, is the density of P. (v.,"),, w.t.t the Lebesgue measure.
Following the construction of our dynamic iterates, now explain how the sequence (T, InENN (E,, is recursively defined. Let n N, kn (n mod(K). Define the path P., = pathT km (knn: Then, for any (0,v)) E ERU+1.
1+IN (a) if (v, Es,, \Pn. then ((o,x")
ny
m== (m11"")
N-1 a.N
(b) if (u',v) E Pn. then =+1.N
(o,0")
T.
n+
N-1 m=0
" n.N
These computations may be obtained by considering the sequence given by (mIPF) to solve the multi-marginal Tree-SB problem associated to T(N) = (V(N), E(N)), the N-discretized version of (N)
(see Appendix B) with weights wem. =2mm/e, which is given by
π argmin ( : E plvin)) n. m-M.ViS
with π D.N
D.N R4 (E, "(UN'")1:N10''
To approximate the IPF recursion given by (a) and (b), we use on each edge of T the score-matching approach of De Bortoli et al. (2021), which avoids heavy computations of score approximations. The next proposition is direct adaptation of (De Bortoli et al., 2021, Proposition 3). Proposition 26. Assume that for any n. E N, any (v, v') E Ex,, with kn (n 1) mod(K), we have
D,N ((,,") )m (- (m)) oo (1m)) mmlaa)
Let n. E N. Consider the path P, patli ino (in,, +1) Let (v,w²) E E Define p" = n.N and In = N m - Then, ( Pa. we have
patt = argmingeLZ(ee..e
(31)
Eyn m+ [X (Xmx+1 + Fn mx.r' (Xmx) - Fm my (X
otherwise, we have pa+11 m,,,w"
Fn m,n,o"
In practice, we use two neural networks per edge (v, v'} E E, one for each possible direction of
the edge, such that (B m, x) F (x) and (e m, x) Fm ""..v...')) For any {v, o') E E, the parameter fi u is updated at iteration via the score matching loss defined by (31)
Proposition 26 (v, v') E pathT,,,, (in., (6, see Algorithm 1.
G Additional experimental results and details
The numerical experiments presented Section are obtained by our own Pytorch implementation, which is inspired from the code provided by De Bortoli et al. (2021). We first provide information on the general setting of our experiments in Appendix G. and then give details on each of them in Appendix G.2 along with additional results. We recall that a mIPF cycle is defined as subset of K consecutive iterations of (mIPF) and that the order of the leaves given by (io...., iK-13 randomly shuffled at each new mIPF cycle.
35

---

G.1 General experimental setup
Implementation of Algorithm in practice. mo)(K). Consider the path Pn. = pathT tn (dem. dR.) Assume that we are provided with a dataset Di...' which contains M samples from m Following Lines 7-9 in Algorithm 1, we apply processes (a) and (b) recursively on the edges (v, v²) E Pa.
Let E N. with k. = (n - 1) mo)(K). K, + 1
(a) Sampling step (Line 7). For any zo E Du, we sample from the diffusion trajectory (30) given by the Fuler Maruyama discretization of P.,," starting from zo. This gives us M x N trajectory samples. We then store the last iterate of each trajectory in a new dataset Dy', which thus approximates T.-. (b) Training step (Lines 8-9). In order to avoid heavy computation, we approximate the mean- matching loss (31) by unbiased estimator obtained by subsampling b elements from the full trajectories computed in the sampling process, see (De Bortoli et al., 2021, Eq. (97)-198)) Here, b refers to the batch-size parameter of the neural networks. Then, we perform gradient descent to optimize the parameter ,, which parameterizes the backward drift on the edge (u, v"))
To avoid any bias issue, the whole trajectories obtained at process (a) are refreshed at a certain frequency over the training iterations of the neural networks by once again simulating the diffusion (30). In our experiments, this refresh occurs each 500 iterations.
Setting of the time discretization. The number of time-steps N in the time discretization of the diffusions is chosen to be even and identical for each of the edges of the tree. Let {u, v'} E We now give details on the design of the time schedule {7k J related to the edge {u, e'), see Appendix F. Following De Bortoli et al. (2021), we choose this sequence to be invariant by time reversal and consider * to + (2k/N) %o) for any is {0,.... N/2} (the rest of the sequence being obtained by symmetry) where To is a free parameter and f determined by Σ 7 T Tw. In our experiments, we set . 50 and 7o 10 5
Sampling improvement. In our code, we implemented the corrector scheme of Song et al. (2021) and the probability based sampling approach detailed in (De Bortoli et al., 2021, Section H.3), but did not observe any significant improvement in our experiments using one of these techniques.
Choice of the architectures of the neural networks. In the case of the experiments related to synthetic datasets (two-dimensional toy datasets, Gaussian distributions) and to the subset posterior ag- gregation task, we implement the same architecture as presented in (De Bortoli et al., 2021, Figure 3). We refer to this model as "Basic Model" and detail it in Figure 6. the "Basic Model", the Positiona- |Encoding block applies the sine transform described in Vaswani et al. (2017), with output dimension equal to 32, and each MLP Block represents a Multilayer Perceptron Network. In particular, MLP- Block (1a) has shape (d. 128, max(256, 2d)). MLPBlock (1b) has shape (32, 128, max 2d)), and MLPBlock (2) has shape (2 x max(256, 2d), max(256, 2d), max(128, d), d). where d denotes the dimension of input data. We optimize the networks with ADAM (Kingma & Ba, 2014) with learning rate 10-4 and momentum 0.9. For each of the networks, we set the batch size to 4,096 and the number of iterations to 10,000 for the synthetic datasets and 15,000 for the subset posterior aggregation task. Our experiments ran on Intel Xeon CPU Gold 6230 20 cores @ 2. Ghz CPU.
MLPBBock (la)
Concatenate
MLPBlock (2)
Output
k
PositionalEncodingg
MLPBlock (Ib)
Figure 6: Architecture of the "Basic Model".
In the case of the experiments related to MNIST dataset, we use reduced UNET architecture based on Nichol & Dhariwal (2021), where we set the number of channels to 64 rather than 128. We implement an exponential moving average of network parameters across training iterations, with rate 0.999. We optimize the networks with ADAM (Kingma & Ba, 2014) with learning rate 10-1 and momentum 0.9. Finally, we set the batch size to 256 and the number of training iterations to 30,000. Our experiments ran using Nvidia A100.
36

---

Details on regularized state-of the art methods. We run the FB algorithm (Cuturi & Doucet, 2014) with the implementation provided by Flamary et al. (2021). For each experiment, we run 100 Sinkhorn iterations with 1500 samples for each dataset (iee, the maximum number of samples that it can generate) and set the regularization parameter e to its lowest value such that the algorithm is stable. Finally, for sake of fairness with our method, we initialise the barycenter measure with no when solving the problem (Ho-regWB) for synthetic Gaussian datasets and Bayesian fusion. To run the crWB algorithm (Li et al., 2020), we use the code provided by the authors. We consider the quadratic regularization, which is shown to be empirically more stable than entropic regularization. Following Fan et al. (2020), we choose the potential networks to be fully connected neural networks with hidden layers of shape 1max(128, 2d), max 128, 2d), max 128, 2d)). The activation functions are ReLu. We optimize the networks with ADAM (Kingma & Ba, 2014) with learning rate 10 for the subset posterior aggregation task and 10 3 for the Gaussian experiment. Finally, we set the batch size 4,096 and the number of training iterations to 50,000. We highlight that f&WB and crWB solve a regularized Wasserstein barycenter problem, which does not contain an additional penalization term on the entropy of the barycenter, contrary to TreeDSB.
G.2 Details the experiments
Synthetic Gaussian datasets. For each dimension that we consider, we generate three different triplets of random diagonal covariance matrices whose condition number is less than 10. We then run the algorithms on each triplet and aggregate the obtained results. The Gaussian datasets contain 1,500 samples for fsWB and 10,000 samples for crWB and TreeDSB. We run fsWB with the following settings (d. e) E {(2, 0.1), (16, 0.2), (64, 0.5), (128, 1.0), (256, 2.0)). We run TreeDSB for 10 mIPF cycles with regularization parameter
0.1, starting from the central node initialized to a Gaussian distribution μo chosen as detailed in Appendix FF with a = 1. Thus, we solve the regularized Wasserstein barycenter problem (Ho-reg (B), which contains an additional regularization with respect to 14n- This choice is justified, since the non-regularized barycenter is known to be a Gaussian distribution, and μ6 can be seen as an a priori for the regularized barycenter. For each of the three settings, we keep the best result among the 30 mIPF iterations. In this setting, TreeDSB and crWB have roughly the same training time.
Subset posterior aggregation. When considering a dataset splitted into several subdatasets, a common paradigm in bayesian inference consists in running Monte Carlo Markov Chain methods separately on these subdatasets, and then merge the obtained posteriors to recover the full posterior. The barycenter of these subdataset posteriors is proved to be close to the full data posterior under mild assumptions (Srivastava et al., 2018). In our setting, we consider the posterior aggregation problem for the logistic regression model associated to the wine dataset" (d (d=44 42) with 3 subdatasets. We consider here two splitting methods: (i) either, data is uniformly splitted between 3 subdatasets with respect to the label distribution, denoted by hiieehhomogeeeuus or (ii) data is splitted with some heterogeneity according a Dirichlet distribution whose parameter is randomly chosen, denoted by wine-her www.teeeeereeou......mom Following Korotin et al. (2021), we use the stochastic approximation trick so that the subset posterior samples do not vary consistently from the full posterior in covariance (Minsker et al., 2014). We implement the Unadjusted Langevin Algorithm (ULA) to sample from each subdataset posterior and from the full posterior. In each case, run ULA for 5.5- 10 iterations with a well chosen step-size, and obtain 9,900 samples after applying a burn-in in of order 10% and then thinning of size 500. We provide in Figure 7 some metrics which assess the quality of this sampling process. We recall that the the full posterior samples serve as ground truth in this experiment.
The results presented in Table 2 were computed as follows. For fsWB, we first subsample 1,500 samples out of the 9,900 samples from each posterior, and then run the algorithm with E = 0.5. We repeat three times this procedure and then aggregate the results. In the case of crWB and TreeDSB, we run the algorithms three times with various seeds. Similarly to the Gaussian setting, we run TreeDSB for 10 mIPF cycles with regularization parameter with a Gaussian distribution (A) chosen as detailed in Appendix F with a = 1, and thus solve the barycenter formulation (Ho-regWB). For each of the three settings, we keep the best result among the 30 IPF iterations. In this setting, TreeDSB and CB have roughly the same training time.
0.1. We start from the central node
37

---

ACF 0 thinning)
(Inclading thinning!
1.0
dene:
632x:
dott:
dala: 7
dell::2
0.8
data:
0.8
duta::
Aihremntaton
0.G
0.1
0.4
0.0
100
4.0
340
2033
500
400
500
Log
lag
MI1 pvnluodedd cen each trvin ditteset
NI. moskusted on each trwin dihenet
2.7
dats: full
Sato hd
682a:
50th
dala:
date
5 207
107
102
101
102
227
10
107
120
1.7
120
220
107
107
304
327
101
107
MC 100 oc
PIC horadon
Evaluation of the sampling process for hine-hhmoggeneus (left) and wine-heter seeeeheeeroogeeouu (right). We display the Autocorrelation function on 500 lags (above) and the evolution over the iterations of ULA of the negative likelliiddood (NLL) evaluated on each training dataset (below). In particular, the samples are decorrelated and the NLL has satisfying profile.
Figure 7:
Synthetic two-dimensional datasets. In this setting, we consider three different datasets (Swiss-roll, Circle and Moons) that each contain 10,000 samples. Since do not have priori on the shape of the barycenter between these datasets, we consider the regularized Wasserstein barycenter problem (regWB), i.e., is chosen as a leaf and corresponds one of the input datasets. We emphasize that this experiment is not intended to demonstrate the superiority of TreeDSB to compute 2D Wasserstein barycenters, but rather meant illustrate that (a) the marginals of the leaves are well recovered by the algorithm, see Figure 3, and that (b) the obtained barycenter is consistent when diffusing from the different leaves, see Figure persistently verified without difficulty. In this section, we rather aim at illustrating (b) by providing additional results which assess the quality of the barycenter obtained by TreeDSB with respect to the choice of the starting leaf and to the choice of the regularization parameter
all our experiments on 2D datasets, we observed that (a) was
To do so. we consider three different choices of regularization TreeDSB: (i) E 0.2 (50 mIPF cycles), see Figure 8, (ii) 0.1 (50 mIPF cycles), see Figure 9 and (iii) = 0.05 (60 mIPF cycles), see Figure 10. For each of these settings, we run TreeDSB with the starting Icaf chosen as Swiss-roll (first row), Circle (second row) or Moons (third row), and display the final barycenter obtained by diffusing from Swiss-roll (first column), Circle (second column) and Moons (third column). Note that the vertex 0 always corresponds to the starting leaf, the vertex I to the barycenter node and that Figure corresponds to the first row Figure 9.
We can make the following observations. First, the estimated barycenter is always coherent within each row, which assesses the convergence of our method. Then, for each value of e, the TreeDSB barycenter is rather consistent between the rows, ie., the choice of the starting Icaf does not have meaningful impact on our method. Finally, as expected, we observe that the support of the barycenter is less and less diffuse as long as E decreases.
38

---

miPF cycle: 5C, from vertex to 1
miFF cycle: 50, from vertex 2 to
miFF cycle: 50, from vartex to
15
15
15
10-
-19
-10
-10
-15
cycles
cycle:
sycle: 50,
15
10
10
-5
-5
-5
-10
-10
-15
-15
-15
-10
-5
10
15
-15
-10
-5
15
15
-15
-10
-5
miPP cycle: SC, from vertex 2 to
milF cycle: 50, from vertex 3 to
mill cycle: 50, from vertex to
15
15
15
10
5
0
-10
-10
-10
15
Figure 8: Estimated 2D barycenter obtained by TreeDSB with & = (50 mIPF cycles). First row:
starting from Swiss-roll. Second row: starting from Circle. Third row: starting from Moons.
cycle:
cycle:
cycle:
10
10
10
0
0
-5
-10
-10
-15
-15
-15
-10
-5
15
-15
-10
10
-15
-10
mIPE cycle: SC, from vertex 2 to
mIFF cycle: 50, from vertex 0 to
mIPF cycle 50, from vertex to
15-
18
5
55
-5
-5
-10
15
-10
miPi cycle: from vertex
mlFF 50, from vertex
mIPF 50, from vertex
15-
10
10-
0
-5
10
-00
-10
-15
Figure 9: Estimated 2D barycenter obtained by TreeDSB with 5 = 0.1 (50 mIPF cycles). First row:
starting from Swisy-roll. Second row: starting from Circle. Third row: starting from Moons.
39

---

miPF cycle: 62, from vertex to
miFF cycle: 60, from vertex 2 to
miFF cycle: 60, from vartex to
15
15
15
10-
0
-19
-10
-10
-15
cycles
cycle: 50, vertex
cycle: 60,
15
10
10
-5
-5
-5
-10
-10
-15
-15
-15
-10
-5
10
15
-15
-10
15
15
-15
-10
-5
miPP cycles 6C, from vertex to
miFF cycle:: 60, from vertex 3 to
miel cycle:: 60, from vertex to
15
15
15
10
5
0
-10
-10
-10
-10
15
Figure 10: Estimated 2D barycenter obtained by TreeDSB with & = 0.05 (60 mIPF cycles). First row: starting from Swiss-roll. Second row: starting from Circle. Third row: starting from Moons.
For purpose of illustration, we provide in Figure 11 the barycenter obtained by state-of-the-ar wwo-dimensional in-sample methods that are available POT library (Flamary al., 2021): non- regularized free-support Wasserstein barycenter (Cuturi & Doucet, 2014), (ii) pntropie-regularizee free-support Wasserstein barycenter (fsWB) with 0.5 (Cuturi Doucet, 2014) and (iii) entropic- regularized convolutional Wasserstein barycenter with 5.10- (Solomon et al., 2015), which is specifically designed for images. We notably observe that TreeDSB cannot capture the full complexity of the 2D barycenter compared to these methods. We infer that this gap comes from the dynamic nature of TreeDSB, since increasing the number of training iterations per IPF iteration or improving the complexity of the neural networks did not bring any significant change in results. Finally, recall that the methods (i), (ii) and (iii) do not scale well with dimension, and have to be completely again when data samples available, contrary to TreeDSB.
Support easct beryennter
Tre Support Serchem bacyerter
Conittitonall Shenrm denyentte in
70-
0-
Figure 11: Estimated 2D barycenter obtained by in-sample algorithms. From left to right: Cuturi & Doucet (2014) (non-regularizee), Cuturi & Doucet (2014) (regularized), Solomon et al. (2015).
40

---

MNIST Wasserstein barycenter. This setting can be qualified as high high-dimensional, since the data dimension is d 784. Here, each digit dataset contains 1,000 samples. As the wo-dimeensional setting, we do not have an a priori on the shape of the barycenter between MNIST digits, and thus consider the formulation (regWB), where the root r is chosen as a leaf. We propose below several experiments to assess the scalability of TreeDDB to this setting.
Digits 0 and 1. Figure 12. we report the results obtained by running TreeDSB on MNIST digits 0 and 1, for 15 mIPF cycles with 0.5, starting from the leaf MNIST-O. We display 25 samples from the reconstructed MNIST-O marginal (first column), from the reconstructed MNIST- marginal (fourth column), from the estimated barycenter by diffusing from MNIST-O (second column) and diffusing from MNIST-1 (third column). We notably observe that the digits are well recovered and that the barycenter samples are consistent. We draw the reader's s attention to the fact that TreeDSB showed numerical unstability with a regularization value lower than 0.5. For purpose of illustration, display in Figure 13 the Wasserstein barycenter obtained by non-regularized methods from Fan et al. (2020) and Korotin et al. (2021), and by the regularized approach from Li et al. (2020).
C
L
0
0
0
6
I
1
0
0
0
0
[
J
0
0
/
[
0
D
(
l
Figure 12: Tree DSB results for MNIST digits 0 and I, after 15 mIPF cycles with e = 0.5.
0
Figure 13: From left to right: Fan et al. (2020), Korotin et al. (2021) and Li et al. (2020).
Digits 2,4 and 6. In Figure 14, we report the results obtained by running TrecDSB on MNIST digits 2.4 and 6, for 10 mIPF cycles with E = 0.5. Here, we consider three settings which differ by the starting leaf r in the algorithm: MNIST-2 (first row), MNIST-4 (second row), or MNIST-6 (third row). For each of these settings, we display 30 samples from the estimated barycenter by diffusing from MNIST-2 (first column), diffusing from MNIST-4 (third column) and diffusing from MNIST-6 (third column). We notably observe a global consistency of the barycenter samples across the various settings. In Figure 15, report the results obtained by running TreeDSSB on MNIST digits 2,4 and 6, for 10 mIPF cycles with & = 0.2, starting from MNIST-6. We display 30 samples from the reconstructed marginals (first row), from the estimated barycenter (second row) by diffusing from MNIST-2 (first column), diffusing from MNIST-4 (second column) and diffusing from MNIST-6 (third column). As expected, we observe less noisy barycenter samples compared Figure 14, while still well recovering MNIST digits.
Digits 0,1 and 4. In Figure 16, we report the results obtained by running TreeDSB MNIST digits 0,1 and 4, for 10 mIPF cycles with æ = 0.5. We consider two settings which differ by the starting leaf r the algorithm: MNIST-O (second row) and MNIST-1 (first/third rows), for which we display samples from the reconstructed measures (first row). In Figure 17, we report the results obtained by running TreeDSB on MNIST digits 0,1 and 4, for 10 mIPF cycles with Æ 0.2. We consider two settings which differ by the starting leaf in the algorithm: MNIST-O (first/second row), for which display samples from the reconstructed measures (first row), and MNIST-1 (third row). For all of these settings, we display 30 samples from the estimated barycenter by diffusing from MNIST-0 (first column), diffusing from MNIST-1 (third column) and diffusing from MNIST-4 (third column). Similarly to the digits 2-4-6, we observe consistency within the barycenter samples, unconditionally to the starting leaf, and less noise as decreases. Note that the reconstructed MNIST digits are less truthful to the original datasets when is low.
41

---

4
4
4
4
0
1
4
4
4
4
4
4
4
4
4
4
9
4
4
4
4
0
4
4
u
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
Q
4 4
4
0
e
4
U 4
4
4
4
U
Q
4
Q
9
4
Q
Q
4
9
4
Q
L
L
4
Q
Q
Q
4
Q
J
4
2 L
8
9
6
θ
4
4
2
6
4
d
6
4
6
4
4
H
4
6
6
2
4
4
b
Q
4
4
6
6
6
d
Q
4
J
4
2
4
6
6
4
Figure 14: Tree DSB results for MNIST digits 2,4 and 6, after 10 mIPF cycles with 0.5. First row:
starting from MNIST-2. Second row: starting from MNIST-4. Third row: starting from MNIST-6.
2
2
4
8
2
2
4
4
4
4
4
4
6
0
6
2
2
2
2
3
4
4
4
4
4
4
6
6
6
6
2
2
2
3
2
2
4
4
4
6
6
4
2
2
2
8
2
4
4
6 6
6
6
2
2
2
8
2
2
7
4
4
4
J
6
6
6
6
6
4
4
7
P
4
4
d
4
0
4
4
4
4
u
q
7
u
2
2
9
4
0
4
4
4
4
Ø
U
L
4
4
4
J
y
0
9
o
o
0
4
Q
d
o
4
4
4
a
9
4
7
1
U
u
4
r
a
u
4
U
4
7
7
Figure 15: Tree DSB results for MNIST digits 2,4 and 6, after 10 mlPF cycles with 0.2, starting
from MNIST-6. First row: samples from the reconstructed marginals. Second row: samples from the
estimated barycenter.
42

---

4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
1
1
1
4
4
4
4
Figure 16: Tree DSB results for MNIST digits 0,1 and 4, after 10 mIPF cycles with s = 0.5. First row: samples from the reconstructed marginals, starting from MNIST-1. Second row: samples from the estimated barycenter, starting from MNIST-0, Third row: samples from the estimated barycenter, starting from MNIST-1.
43

---

0
0
0
0
0
0
0
D
0
0
0
d
9
O
0
0
J
0
X
I
Figure 17: Tree DSB results for MNIST digits 0,1 and 4, after mIPF cycles with s = 0.2. First row: samples from the reconstructed marginals, starting from MNIST-0, Second row: samples from the estimated barycenter, starting from MNIST-O. Third row: samples from the estimated barycenter, starting from MNIST-1.
44


Reviews:
[{"review_id": "NM6MpRxZfb", "reviewer": "Reviewer_2tnG", "summary": "Schröninger bridges (SB) can be understood as dynamical counterpart of the Entropic Optimal Transport problem. \n\nThis paper introduces the Tree Diffusion Schröninger Bridge (TDSB) model, which can be seen as a multi-marginal generalization of SB when the cost $c$ has a tree structure: given $\\ell$ marginals $\\pi_1,\\dots,\\pi_\\ell$, and a tree $T = (V,E)$ with $|V| = \\ell$ nodes, one has $c(x_1,\\dots,x_\\ell) = \\sum_{(v,v') \\in E} \\| x_v - x_{v'}\\|^2$. Let $S$ denote the leaves of $T$. \n\nIn that setting, the goal of the paper is to solve the TreeSB problem:\n$$ \\min_\\pi \\mathrm{KL}(\\pi | \\pi^0),\\quad \\pi_i = \\mu_i,\\ \\forall i \\in S. $$\nwhere the $\\mu_i$ are prescribed marginales, and $\\pi^0 = \\pi_r^0 \\bigotimes_{(v,v')} \\pi^0_{v,v'}$ is a reference measure that factorize along $T$, and $r \\in V$ is identified as a root of $T$, and the factors $\\pi^0_{v,v'}$ are chosen as normal distributions. \n\nIf $T$ has a star-structure, this problem is equivalent to the (doubly-regularized) Wasserstein barycenter problem (in sense that, at optimality, the marginal of $\\pi$ in the center of the star is the Wasserstein barycenter of the marginals on the leaves). \n\nThe contribution of the work is to solve this problem using an iterative algorithm along _paths_ of $T$ which can be summarized as a succession of $\\mathrm{KL}$ projections where only one of the marginals are constrained; if the tree is only made of two nodes (and one edge) we retrieve the Sinkhorn algorithm. In practice, the transition along a given edge (during one step of the algorithm) are performed by SDE of the form $\\dot X_t = f_\\theta(X_t) \\mathrm{d}t + \\mathrm{d} B_t$, where $B_t$ is the Brownian motion, and $f_\\theta$ is a drift that is learned (and parametrized by a neural network). This leads to an iterative scheme where one repetitively select path between two leaves and learns the drift to solve the KL projection along the path. \n\nFrom a theoretical standpoint, authors prove that (assuming the intermediate step converge, e.g. the drifts $f_\\theta$ end up being optimal) the proposed algorithm solves (asymptotically) the TDSB problem.  \n\nFrom a numerical perspective, they showcase the practicability of their approach on synthetic and simple benchmark datasets.", "questions": "- My main question, which can be developed (see below) would be _why should one use this approach rather than existing ones?_ (e.g. IPB, fsWb, etc.). Of course, the paper gives some insight, but:\n\n1. As said above, there is no insight on the computational burden of the proposed approach. \n2. While I find very interesting that the approach works on any tree-structured coast, the work only focuses on Wasserstein barycenter, somewhat limiting its impact (that is, one strength of this work is that it solves a more general problem than previous Wasserstein-barycenter-dedicated methods). Giving an example + experiment where a (non-stared) tree cost is involved would be insightful (it would give a situation where the proposed approach as no challenger, as far as I know!). ", "limitations": "authors have addressed the limitations of their approach. \n\nI do not identify ethical concern specific to this work. ", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- I find the developed approach quite elegant. It mixes different notions in optimal transport literature (multi-maginal OT, entropic OT/Schrödinger bridge, doubly-regularized Wasserstein barycenters, parametrization of \"transport maps\" (the drifts) by neural networks, etc.). It is nice to see all these approaches implemented in a _practical_ setting (which can be implemented). \n- Comparison with previous work is well presented.  \n- The convergence results of the proposed approach are strong.\n", "weaknesses": "## Numerical experiments \n\nMy main remark would be about the numerical experiments. They are quite limited in their comparison with state-of-the-art methods and rather acts like proof-of-concept. In details: \n- They do not showcase how the proposed approach solves new problems. Though reasonable, none of the results showcase feel \"groundbreaking\". This may just be a matter of exposition, but right now, loosely speaking, I am not \"impressed\" by the experimental results. For instance, the choice of synthetic datasets (Figure 4) is questionable in that the actual barycenter is not very intuitive and thus it is hard to visually tell whether the displayed output (Figure 5) is good or not; in addition, the \"consistency\" of the proposed approach is not striking (I admit that there is a decent similarity between the three outputs, but we cannot faithfully consider that the model did converge). \n- The comparison with SotA methods is incomplete; e.g. there is no running time involved. Given the (maybe only superficial) complexity of the proposed approach---such as the need for a possibly huge amount of NN (which are not very small according to Section G, as they are MLPs), one may be feared by a possibly huge running time (compared to simpler approaches like, say, the one of Cuturi and Doucet). If there is a computational tradeoff (better result for longer training time), it must be mentioned and discussed in details. \n- I would expect a comparison with the _Iterative Bregman Projection_ of (Benamou et al., 2015), given the \"similarity\" (in that they are iterative KL projections) with the proposed approach (and I would consider IPB as being \"more SotA\" than Cuturi&Doucet, for instance; for instance the IPB paper claims in Section 1.3 that \"IPB converges (...) orders of magnitude faster than Cuturi&Doucet\"). \n- A bit minor, but the description of the _subset posterior aggregation_ experiment is very minimalist. It is hard for the reader that is not already familiar with such experiment to have an idea of what is done there. ", "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. Below we provide a detailed answer to the reviewer’s concerns.\n\n> Why should one use this approach rather than existing ones?\n\nWe have provided a detailed answer to this point in the general response which we summarize here. First, contrary to previous out-of-sample methods, TreeDSB does not rely on a static formulation of EmOT but on a dynamic one, which allows us to leverage tools from diffusion models. We believe that the idea of solving a dynamic counterpart to the static problem plays a key role in the efficiency and scalability of the method. Additionally, the entropic regularization induced by our dynamic formulation is different from the one usually considered in static formulations and allows us to compute Wasserstein barycenters with low entropic regularization while avoiding the numerical issues of having to choose $\\varepsilon$ too small. This is well highlighted in our experiments on MNIST dataset, see Figure 6, where our estimated barycenter samples are close to the ones given by unregularized methods while the samples from [1] are of poor visual quality.\n\n> They do not showcase how the proposed approach solves new problems. \n\n> While I find very interesting [...] as far as I know!)\n\nWe have provided a detailed answer to this point in the general response which we summarize here. We chose the Wasserstein barycenter problem because there already exists strong baselines for this problem, and therefore, we can compare TreeDSB to these existing methods. We agree with the reviewer that it would be valuable to apply TreeDSB in a non star-shaped setting. For instance, in the case of single-cell genomics, it would be interesting to infer the distribution of features of cells at different time steps (temporal tracing) as well as the probability paths between treated and untreated states (effect of the treatment), see the general answer for details. We leave the application of TreeDSB to this framework for future work as it requires some domain knowledge of single-cell genomics and additional data collection.\n\n> I would expect a comparison [...] faster than Cuturi&Doucet\").\n\nWe would like to emphasize that the experiments in [2], where a comparison is made with [3], are only run on 2D data (see Figures 3 and 4 in [2]). For our general setting, we chose [3] as a baseline, following the comparison made in [4]. Nonetheless, we provide in the PDF a numerical result of this method in the 2D setting.\n\n> Though reasonable, [...] that the model did converge).\n\nWe emphasize that the only simple case where we have access to the true Wassertein barycenter is the Gaussian setting (for which we already show that TreeDSB performs well). We chose these 2D synthetic datasets as an illustration of our method, and aimed at showing that the marginals on the leaves are well recovered. In the PDF, we give improved results in this setting (which have better consistency on the barycenter node) with different values of $\\varepsilon$. We acknowledge that TreeDSB performs worse than in-sample methods in this setting, which can be explained by the fact that the diffusion-based setting suffers from compactness of the supports of the 2D distributions (see vertex 3). In this very specific case, it would be preferable to use [2] and [3] instead of TreeDSB. However, we highlight that [3] performs worse than TreeDSB in higher dimensional settings.\n\n> The comparison with SotA methods [...] the one of Cuturi and Doucet).\n\n> As said above, there is no insight on the computational burden of the proposed approach.\n\nWe have provided a detailed answer to this point in the general answer. We briefly summarize it. We highlight that the training time of TreeDSB should not be compared with the complexity of the algorithms from [2] and [3] which are in-sample methods. Instead, TreeDSB should be compared with the out-of-sample neural network-based methodologies such as [1]. In our experiments, we observed that the training time of TreeDSB and the algorithm from [1] were comparable. We acknowledge that the inference time of TreeDSB is however longer since it requires to sample from a SDE, a limitation common to all diffusion-model based techniques which could be addressed by using distillation methods.\n\n> If there is a computational tradeoff (better result for longer training time), it must be mentioned and discussed in details.\n\nThe computational trade-off of TreeDSB recovers the classic behaviour of OT methods: whenever $\\varepsilon$ is chosen too small (i.e., we expect to be closer to the true barycenter), the method takes more time to converge. This trade-off is well highlighted in the results of the 2D experiments in the PDF, where we consider 50 IPF cycles for $\\varepsilon=0.2, 0.1$ and 60 IPF cycles for $\\varepsilon=0.05$. As $\\varepsilon$ decreases, the obtained barycenter is less noisy, but requires more IPF iterations to converge. We will add these considerations in the section 'Experiments' in the main manuscript. \n\n> A bit minor [...] an idea of what is done there.\n\nWe briefly explain here how posterior aggregation works. In this setting, we consider some dataset, which is splitted into $\\ell$ subsets, and a Bayesian model. We assume that one posterior distribution is computed for each data subset. Then, posterior aggregation consists in computing the Wasserstein barycenter of these $\\ell$ posterior distributions, expecting that it recovers the posterior distribution from the whole dataset, see [5] for a theoretical justification. \n\n[1] Continuous regularized Wasserstein barycenters. Li et al.\n\n[2] Iterative Bregman projections for regularized transportation problems. Benamou et al.\n\n[3] Fast computation of Wasserstein barycenters. Cuturi and Doucet.\n\n[4] Scalable computations of Wasserstein barycenter via input convex neural networks. Fan et al.\n\n[5] Scalable Bayes via barycenter in Wasserstein space. Srivastava et al."}, {"review_id": "pe301fXlUf", "reviewer": "Reviewer_GsNC", "summary": "The authors consider the multi-marginal optimal transport (mOT) with entropic regularization and a tree-structured quadratic cost. The authors propose to address it by developing the Tree-based Diffusion Schrodinger Bridge (TreeDSB), i.e., dynamic and continuous state-space of the multimarginal Sinkhorn algorithm. The authors apply the proposed approach to compute Wasserstein barycenter, i.e., equivalent to mOT with star-shaped tree.", "questions": "+ It is not clear why one needs to address the dynamic and continuous state-space of the multimarginal Sinkhorn. It suffices to consider the multimarginal OT directly for Wasserstein barycenters and problems in the experiments. Could the authors comment on this?\n+ Although the authors give some discussion about the relation with Hassler et al. (2021) in the appendix, it is essential to clarify their difference in the main manuscript (e.g., in related work section) to emphasize the contribution. Could the authors summary their main differences? (since they are very closely similar).\n+ In line 110, $S$ is the set of probability measures on the “leaves” of $T$. However, in line 118, why $r \\in S$ (it seems that $r$ is the root of tree T as in line 113)? So, what does the tree $T$ look like?\n+ Please explain the Figure 2, its description is too short and hard to follow.\n+ As in line 157-158, there are $2|E|$ neural networks, how are these networks related?\n+ Kitagawa and Pass (2015) showed the equivalent between multimarginal OT and the Wasserstein barycenter. Could the authors comment the result in Proposition 8 and the claim in line 271-272 with those in Kitagawa and Pass (2015)?\n\nRef: Kitagawa, J. and Pass, B., 2015, September. The multi-marginal optimal partial transport problem. In Forum of Mathematics, Sigma (Vol. 3, p. e17). Cambridge University Press.\n+ In line 284, how does the choice of $\\epsilon$ affect the result, and how to choose it in applications?\n+ For results in Tables 1, 2, Figures 6, 7, it is better to also report the corresponding time consumption?\n+ It seems unclear how to choose $\\epsilon$ (hyperparameters) for the barycenter problem? Does it affect the result? Some hyperparamters seem missing, e.g., how many free supports are used for fsWB?\n+ For results in Figures 6 and 7, how to interpret these results? (especially Figure 7)\n+ in line 330, “a scalable scheme”? it is unclear whether the proposed approach is scalable from the proposed algorithm and experiments? Please elaborate the scalability of the proposed method.\n\nSome minor points:\n+ In line 2, what is $C([O, T], R^d)$? There is no explanation for the new notation $C( , )$\n+ in line 78, it is better to explain the Stochastic Differential Equation form, e.g., there is no explanation about several new notations, $X_t, a, B_t$\n+ Please explain the Figures. The description for each Figure is very short, it is unclear the meaning of each Figure, e.g., how to understand the Figures 2, 3?\n+ In line 125, what is the role of $K$? \n+ In Proposition 1, what is the relation between $T_{K-1}$ and $T$ (although there is some explanation in line 71-72, it is hard to follow it). It is better to consider to simplify the notations or explain the notations with more details.\n\n---\n\nThank you for the rebuttal.", "limitations": "The authors have discussed the limitation of the work. However, there is no discussion about the potential negative societal impact of their work.", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "+ The authors propose Tree-based Diffusion Schrodinger Bridge, a dynamic and continuous state-space counterpart of the multimarginal Sinkhorn algorithm.\n+ The authors illustrate some advantages of the proposed approach for Wasserstein barycenter.", "weaknesses": "+ The motivation of developing the dynamic and continuous state-space counterpart of the multimarginal Sinkhorn algorithm and experiments seem quite weak. For Wasserstein barycenters and problems illustrated in the experiments, it suffices to address the multimarginal OT (i.e., the static problem).\n+ Although there is some discussion in the appendix, it is essential that the submission is very closely related to (Hassler et al. 2021). It is an important part and should elaborate with details in the main manuscript to emphasize the contribution.\n+ The notations seem heavy and hard to follow. It is better to simplify the notations or give more explanation for the notations when they are introduced. At the present, it is hard to follow the meaning of the notation system.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our submission and for your constructive feedback.\n\n> The motivation [...] seem quite weak\n\nWe have provided a detailed answer to this point in the general response. First, TreeDSB introduces a new class of algorithms to solve EmOT using a dynamic framework. This formulation allows us to leverage tools from the diffusion models literature. We believe that the idea of iterative refinement, i.e. solving a dynamic counterpart to the static problem, plays a key role in the efficiency and scalability of the method.\n\n> closely related to (Hassler et al. 2021) [...] \n\nWe have also provided a detailed answer to this comment in the main response. We emphasize that the static problem is the same as the one from [6] (see “Our framework” paragraph in Section 2). However, one of our key contribution is to derive a **dynamic framework** to deal with this static problem as well as a practical algorithm to solve this new formulation (see Section 3). \n\n> What is $S$ ? [...]\n\nS is the set of leaves. As every leaf of the tree is associated with a marginal in a unique way, we make the slight abuse of notation to say that $S$ is the set of marginals of the leaves. We will correct this.\nThe tree $T$ corresponds to the graphical model of the prior joint distribution. The nodes are associated with the marginals while the edges are associated with the couplings. \n\n> The description for each Figure is very short, [...]\n\n**Figure 2** provides an illustration of how to bridge the edge iterates from Proposition 3 (i.e., the path measures $P^n_{(v,v')}$), which solve the dynamic Schrödinger Bridge (SB) problem, to the edge iterates of mIPF (i.e., the couplings $\\pi^n_{(v,v')}$), which solve the static SB problem. This is done by the operator $\\text{Ext}$ which takes as input a path measure $P$ (blue line in Fig. 2), and outputs the coupling $P\\_{0,T}$. In our case, we have $\\pi^n_{(v,v')}=\\text{Ext}(P^n_{(v,v')})$.\n\n**Figure 3** illustrates Proposition 3 with an example. Assume that $T$ is rooted in vertex 3 at iteration $k$ of our algorithm. The directed edges of $T$ are the red edges. Assume that we want to solve the dynamic formulation of mIPF at iteration $k+1$, where the marginal constraint is on vertex 4. Then, by Proposition 3, we simply have to operate **the time reversal of the path measures corresponding to the red edges that lie in the path between 3 and 4** (yellow dotted path). Once this is done, we can sample from the $(k+1)$-th iterate by following the red edges, starting from 4.\n\n> How are the NNs related to each other ?\n\nWe have two NNs per edge. Each NN has the same architecture.  This is similar to other neural-network based approaches [1,2,3,4].\n\n> Comparison with Kitagawa and Pass ?\n\nThank you for pointing us towards [7]. In [7] the authors consider a *partial* multi marginal OT problem while we consider an entropic regularized approach. As a result, our static formulation and the one from [7] are drastically different. \n\n> How does the reg affect results in 2D ?\n\nThe computational trade-off of TreeDSB recovers the classic behaviour of OT methods: whenever $\\varepsilon$ is chosen too small, the method requires more iterations to converge. This trade-off is illustrated in the results of the 2D experiments in the PDF, where we consider 50 IPF cycles for $\\varepsilon=0.1, 0.2$ and 60 IPF cycles for $\\varepsilon=0.05$. As $\\varepsilon$ decreases, the obtained barycenter is less noisy, but requires more IPF iterations to converge. We will add these considerations in the section 'Experiments' in the main manuscript.\n\n> Figs 6 and 7, how to interpret these results?\n\nIn Fig.6 we show non-regularized [2,3] and regularized [4] baselines for the barycenter computation. We also present the barycenter obtained with TreeDSB which is visually closer to the ground-truth computed with [2] than the competing method [4].\nIn Fig.7 we illustrate that TreeDSB not only reconstruct the barycenter distribution but also the distributions on the leaves.\n\n> in line 330, “a scalable scheme”? \n\n> report [...] time consumption?\n\nWe provide a detailed answer to this point in the general response. Compared to other out-of-sample competing  [1,2,3,4] methods, TreeDSB enjoys a similar training time and memory cost. In particular, the number of NNs also scales linearly with the number of edges. \n\n> What is $C([0,T], \\mathbb{R}^d)$ ?\n\n$C([0,T], \\mathbb{R}^d)$  denotes the continuous functions from [0,T] to $\\mathbb{R}^d$. We will add it to the notation section.\n\n> explain the SDE form [...]\n\nPath measures are defined at the beginning of Section 2. In line 78, the path measure $\\mathbb{Q}$ is associated with the process $(X_t)$, i.e. $(X_t) \\sim \\mathbb{Q}$. $(B_t)_{t \\geq 0}$ is a  $d$-dimensional Brownian motion. We will precise this in the new version of the paper.\n\n>  role of $K$?\n\nThis quantity is defined in line 115 “In what follows, we define K as the number of leaves of $T$”. To perform a full cycle of IPF, we have to solve iteratively on the whole set of the leaves which explains $\\mathrm{mod} (K)$. \n\n> relation between $T_{K-1}$ and $T$ ?\n\nT is the undirected tree used in the general formulation of SB. $T_v$, for some vertex v, is the directed version of T rooted in v, which is uniquely defined from T. The link between general trees and rooted trees is explained in our notation section.\n\n[1] Wasserstein iterative networks for barycenter estimation. Korotin et al.\n\n[2] Scalable computations of Wasserstein barycenter via input convex NNs. Fan et al.\n\n[3] Continuous Wasserstein-2 barycenter estimation without minimax optimization. Korotin et al.\n\n[4] Continuous regularized Wasserstein barycenters. Li et al.\n\n[5] Tackling the Generative Learning Trilemma with Denoising Diffusion GANs. Xiao et al.\n\n[6] Multimarginal optimal transport with a tree-structured cost and the Schrödinger bridge problem. Haasler et al.\n\n[7] The multi-marginal optimal partial transport problem. Kitagawa and Pass."}, {"review_id": "40OuYNUpdK", "reviewer": "Reviewer_BvM6", "summary": "This paper generalizes the diffusion Schrödinger bridge framework to multidimensional problems with tree-structured quadratic cost function. The method builds on representing each egde in the graph as two SDEs (one for each direction) parameterized by neural networks. These are used to propagate the information (from the known marginals) in the graph. Under a set of assumptions the method is shown to converge in TV.   ", "questions": "I do not see how the assumptions A1-A5 uses the neural nets. The neural nets typically have finite prespecified parameters. Can you explain how one can represent any measure (satisfying the assumptions), which are infinite dimensional objects. \n", "limitations": "The limitations are adequately addressed. ", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "It is a nice extension of the Diffusion Schrödinger Bridge framework, which allows for addressing a relatively large class of multimarginal optimal transport problems where each marginal is a probability measure. Also it is nice that convergence of the method is provided.    ", "weaknesses": "I think it would be nice with a little bit more information on the neural nets, about their structure and how they are used and updated. See the question below.", "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 acknowledgment of the paper’s merits.\n\n\n> I think it would be nice with a little bit more information on the neural nets, about their structure and how they are used and updated.\n\nAll information on the neural nets used in our experiments are given in Appendix G. The specific architecture of the neural network is given by a MultiLayer Perceptron (MLP) or a U-Net depending on the task (MLP for low dimensional data and U-Net for image data). The choice of these specific architectures follow from the ones made in DSB [1], which follow from the ones made in [2] in the case of the image data. We have added a sentence to refer to the supplementary in the main document. \nFinally, we briefly discuss what these neural networks model.  We recall that each edge {v,v’} of the tree is associated with 2 neural networks, and each one corresponds to a direction of this edge, either (v->v’) and (v’->v). Let {v,v’} be an edge. At iteration $(n+1)$ of TreeDSB, the neural net (v->v’) is updated if and only if the directed edge (v’,v) is in the path between the previous root chosen at iteration $n$ and the new root chosen at iteration $(n+1)$. This neural network models the drift of the SDE that describes the path measure $\\mathbb{P}^{n+1}_{(v’,v)}$. Through an entire IPF cycle, each network is updated only once, since each edge is visited once in both ways. Regarding how they are updated, the losses associated with these networks are the ones associated with the time-reversal in diffusion models. In particular, we have used the mean-matching loss from [1] to learn the time-reversal along each edge. We will include some details about the mean-matching loss as well as a complete pseudo-code of the training algorithm in the updated supplementary.\n\n> I do not see how the assumptions A1-A5 uses the neural nets. The neural nets typically have finite prespecified parameters. Can you explain how one can represent any measure (satisfying the assumptions), which are infinite dimensional objects.\n\nWe want to clarify that in the theoretical section of the paper (Section 4), we do not use the neural networks. This section is merely concerned with the static problem. In particular, we focus with the well-posedness of the mOT problem (in particular, we show that the solution of the mOT problem and the mIPF iterates can be obtained by updating dual potentials) and the convergence of mIPF procedure in that setting. The methodology introduced in our paper is an attempt at efficiently learning the mIPF iterates by leveraging a dynamic formulation and tools from the diffusion model literature. \nThe main motivation for establishing such results is that **we could not find any literature providing a proper mathematical background** for the treatment of multimarginal Schrodinger bridges. We see our contribution as a first step towards more theoretical results in the field of multimarginal Schrodinger bridges. In particular, in future work we would like to investigate extension of the theory to the full dynamic formulation and to provide convergence bounds of the mIPF for the (approximate) scheme based on TreeDSB.\n\n[1] Diffusion Schrödinger bridge with applications to score-based generative modeling. De Bortoli et al.\n\n[2] Score-Based Generative Modeling through Stochastic Differential Equations. Song et al."}, {"review_id": "KStMRmFtbt", "reviewer": "Reviewer_XAZX", "summary": "Authors develop a new method for computing multi-marginal optimal transport in the case where the cost function is tree structured. Their algorithm cleverly exploits the dynamic formulation of the problem. The correctness of the algorithm is thoroughly demonstrated, and benefits are shown over a number of toy examples as well as competitive baselines involving the computation of the Wasserstein Barycenter.", "questions": "1) What is the difference between your method and the one of Haessler et al?\n2)Could authors spell out applications where multimarginal OT inference in tree structured costs is important? Why is the focus only on the Barycenter?\n3)Authors are improving upon the competitive baselines such as Cuturi and Doucet (table 1). Why the method was expected to bring this benefit? is it related to the tree-ness or the fact that it implements a Schrodinger Bridge?. \n4)Likewise, why better results are obtained in data fussion? (Table 2). It seems to be implied that this is due to better performance in high dimensions. Why this should be the case? \n5)Line 175 seems to be an error in the Ext() statement.\n6)Type in line 305", "limitations": "n.a.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "It is a solid paper (theoretically and experimentally). Although the mathematical part is mostly technical and algorithm-related (but correct me if I am wrong), it is still quite sophisticated and praise deserving. Results section is compelling.", "weaknesses": "There are some weaknesses: first, the paper is slightly hard to read and while there are some helpful guides (as the first figure) at some points the paper gets helplessly hard to follow. Second, the paper should better explain what the contribution over the already existing methods (i.e Hasslear et al). It is not clear from the reading what this adds on top of that. Third, the paper is not compelling on its motivation. Why if the object of interest is a tree-based cost, all experiments focus on the barycenter problem? it gives the impression that the paper addresses a problem that doesn't exist in practice. Fourth, while the experimental results are super interesting, there is a lack of commentary as to why we should expect this.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback and their positive evaluation. \n\n> the paper is slightly hard to read [...]\n\nWhile we acknowledge that the current paper might be heavy with mathematical notation, we thought it was necessary to include them to precisely describe our TreeDSB methodology (Section 3) and our theoretical contributions (Section 4). However, we are happy to clarify the notation or to simplify our terminology if the reviewer has specific concerns regarding the presentation.\n\n> contribution over the already existing methods (i.e Haasler et al) [...]\n\n> What is the difference between your method and the one of Haasler et al? \n\nWe have also provided a detailed answer to this comment in the main response. We agree that the static problem  is the main as the one from [1], as we acknowledge in our “Our framework” paragraph in Section 2 of the current paper. However, one of our key methodological contributions is to derive a **dynamic framework** to deal with this static problem as well as a practical algorithm to solve this dynamic formulation (see Section 3). The introduction of the dynamic formulation is absolutely central to our work as it allows us to leverage tools from the diffusion model literature to design an efficient algorithm. Other minor differences between our static framework and the one of [1] are reported in the general answer.\n\n> Third, the paper is not compelling on its motivation.\n\nWe have provided a detailed answer to this point in the general response which we summarize here. First, TreeDSB introduces a new class of algorithms to solve EmOT. Contrary to variational-based or potential-based approaches (following the terminology of [2]), we do not rely on a static formulation of EmOT but on a dynamic one. This dynamic formulation allows us to leverage tools from the diffusion models literature. We believe that the idea of iterative refinement, i.e. solving a dynamic counterpart to the static problem, plays a key role in the efficiency and scalability of the method (see our general answer for an in-depth comment on the training-time, inference-time and memory costs properties of TreeDSB). We want to emphasize that if the static mOT is already available then there is no need to apply TreeDSB (or any other dynamic formulation). This is because the solutions of the dynamic mOT and the static mOT coincide up to marginalization. It is possible to view our work as a way to learn the static mOT solution by learning the dynamic one, the same way DSB [3] solves an entropic static OT problem using a dynamic formulation and leveraging tools from diffusion models.\n\n> Why if the object of interest is a tree-based cost, all experiments focus on the barycenter problem? it gives the impression that the paper addresses a problem that doesn't exist in practice.\n\n> Could authors spell out applications where multimarginal OT inference in tree structured costs is important? Why is the focus only on the Barycenter?\n\nWe have written a detailed answer to this point in the main response. Regardless, we provide summarize it in what follows. Our methodology is defined for any tree. However, we agree with the reviewers that so far we have only applied our method to the case of Wasserstein barycenters. We chose the Wasserstein barycenter computation example as there already exists strong baselines. We agree that it would be valuable to test our algorithm on other tree structures. Regarding the application of TreeDSB to other tree structured problems, we have described an application from single-cell genomics in the general answer. We briefly summarize this new experiment. The goal is to predict the effect of a treatment but also to follow the time evolution on some cell features. We obtain a tree structured problem which is not star shaped. \n\n> Authors are improving upon the competitive baselines such as Cuturi and Doucet (table 1). Why the method was expected to bring this benefit? is it related to the tree-ness or the fact that it implements a Schrodinger Bridge?.\n\n> Likewise, why better results are obtained in data fusion? It seems to be implied that this is due to better performance in high dimensions. Why this should be the case? \n\n> Fourth, while the experimental results are super interesting, there is a lack of commentary as to why we should expect this\n\nWe believe that the main reason behind the efficiency of the proposed approach in higher dimension is the dynamic formulation we introduce. Indeed by doing so we leverage efficient tools from the diffusion model literature which are known to scale with the dimension. We thank the reviewer for this question and have clarified this point in the revised version of the paper. We refer to the general answer for more details on the motivation and further discussions on the scalability of TreeDSB.\n\n> Line 175 seems to be an error in the Ext() statement\n\n> Typo in line 305\n\nWe have corrected all typos. \n\n[1] Multimarginal optimal transport with a tree-structured cost and the Schrödinger bridge problem. Haasler et al.\n\n[2] Wasserstein iterative networks for barycenter estimation. Korotin et al.\n\n[3] Diffusion Schrödinger bridge with applications to score-based generative modeling. De Bortoli et al.\n"}, {"review_id": "okb6IA7tIX", "reviewer": "Reviewer_fo2c", "summary": "The authors develop a novel method for approximating Wasserstein barycenters, which is based on diffusion models. This method combines ideas from the recently introduced Diffusion Schrodinger bridge and a framework for multi-marginal optimal transport with tree-structured costs.\nThus, although the paper focuses on the barycenter problem (corresponding to a star shaped tree), the proposed approach can be used for any multi-marginal optimal transport problem, where the cost is composed of pairwise terms that are structured as a tree.\nA diffusion Schrodinger bridge type algorithm is presented that approximates the multimarginal iterative proportional fitting (mIPF) algorithm. The authors present novel convergence results of mIPF, which in contrast to previous works neither assume that the space is compact nor that the cost is bounded.\nFinally, the application to Wasserstein barycenters is evaluated for several experiments.", "questions": "\nMy main question concerns the different update schemes mentioned in Section 3. The (mIPF) algorithm requires sampling from the whole tree. The authors argue that each (mIPF) update can be done on a path of the tree. In lines 160-173 it is described that one first samples along all edges of the path, and then learns the reverse dynamics for the path. In contrast, Algorithm 1 samples and learns the reverse dynamics for each edge on the path individually, and this is repeated iteratively along the path. That is, three different update schemes are mentioned:\na) sample along whole tree and learn reverse dynamics;\nb) sample along whole path and learn reverse dynamics;\nc) iterate over the edges on the path, and for each edge sample along each and learn reverse dynamics.\nPlease clarify why these are equivalent, especially the equivalence with scheme c) is not obvious to me.\n\n\nThe following list is ordered mostly chronological, and not according to relevance.\n\n1. In Section 2 the authors should make sure that the background material and setting are introduced in a clear and concise manner. This includes my following comments, but I suggest to carefully work through the whole section.\n- Are all the details on dynamic and static versions of Schrodinger bridges really needed? E.g., it is not clear what $\\Phi$ and $a$ are, and I believe they are not used anywhere in the paper anyways, so this comment could be skipped.\n- The link between (EmOT) and Schrodinger bridges is not well explained. A tree is defined in the second sentence, but it is only explained how it is related to (TreeSB) in the last sentence. Also you probably need a reference to equation (1) somewhere?\n- The final paragraph \"our framework\" could be used to give a much more helpful summary of the main idea of this work. I suggest to state more clearly the motivation, contribution, and main result, and tone down the technical details (which are hard to understand at this point).\n\n2. Proposition 3 should be presented in a self-sufficient way. It is not completely obvious here what $\\mathbb{P}^n$ and $\\pi^n$ denote.\n\n3. What is the convergence criterion in Algorithm 1, line 5?\n\n4. I feel that Section 4 does not really follow the flow of the paper, and I wonder why the authors don't stress the convergence result more as a main contribution of the paper. Also, I'm wondering if the result relies on any additional assumptions compared to previous results? Moreover, I'm a bit concerned about the fact that the one page of Section 4 relies on 10 pages of proofs, which are deferred to the appendix (and which I was not able to check). I find it problematic to have to hide 95% of the derivation for a strong theoretical result in the supplementary material. It might be preferable to not include this result in the paper due to space limitations, and publish it separately. This would also guarantee a thorough peer review of the result. In my opinion the paper makes a strong enough contribution without this Section. \n\n5. For the synthetic two dimensional dataset it would be interesting to also see the true (possibly regularized) Wasserstein barycenter for reference.\n\n6. For the MNIST application, I suggest to clarify in the main paper that you consider the problem in 784 dimensional space. This may not be clear, since images could be seen as distributions of 2-dimensional data, similar to the first numerical experiment.\n\n7. What do you mean by one IPF cycle, and how do you chose the number of IPF cycles in the different experiments?\n\n8. Since the paper presents methods for optimal transport problems with general tree structure, it would be very interesting to also include some experiments that are not a Wasserstein barycenter problem.\n\n9. Typos: \n43: hyperparameter and;\n112: Section 5. We;\n220 applied to trees", "limitations": "yes", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "I see the strengths of this paper in its originality and significance. The paper extends results for Diffusion Schrodinger bridges to a class of multi-marginal problems, including Wasserstein barycenters. Moreover, the paper extends results for tree structured multi-marginal optimal transport problems by developing a diffusion based method for these, which makes the problem scalable to high-dimensional settings.", "weaknesses": "The clarity of presentation should be improved in many parts of the paper. Please see comments below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thorough review of our work.\n\n> My main question concerns the different update schemes mentioned in Section 3. The (mIPF) algorithm requires sampling from the whole tree [...]\n\nWe clarify our training procedure. In Proposition 2 and Proposition 3, we draw a bridge between the mIPF static iterates (line 124) and the ones defined using a dynamic formulation. This connection is at the basis of TreeDSB. Contrary to the iterates of (mIPF), TreeDSB requires sampling trajectories at training and inference time. The description of the sampling procedure is presented in Figure 1. Our implementation follows the logic of this figure. We now provide some explanations regarding the example in Section 3 (lines 160-173). \n\nWe acknowledge that the first sentence in line 163 might be confusing in light of Figure 1. The goal was to highlight that we can sample from $\\pi^0$. However, to train the algorithm we only need to sample from a path (as highlighted in Figure 1). \n\nIn our example, the initial measure is rooted at node $3$. Since in the next iteration we consider a tree rooted at node 1, we only need to consider samples along the path $3 \\to 0 \\to 1$, see the sentence \"During [...]\" (line 165). Similarly, as before we do not need to fully sample from $\\pi^1$ to obtain $\\pi^2$, and therefore acknowledge that the sentence \"In order [...]\" (line 167) might also be confusing, as we only need to follow the path $1 \\to 0 \\to 2$.\n\nWe emphasize that during the training of the algorithm, we sample along paths and never sample from the full joint measure. So the option a) “sample along whole tree and learn reverse dynamics” is not used. Regarding the difference between b) and c), this is simply because each edge is parameterized by a different neural network and hence associated with its own loss function. The difference between b) and c) is only algorithmic.\n\nWe thank the reviewer for this question which allowed us to clarify our main example. We hope to have resolved the reviewer’s concerns regarding the implementation of our algorithm.\n\n> In Section 2 [...] background material and setting are introduced in a clear and concise manner.\n\n> 1. Are all the details on dynamic and static versions of Schrodinger bridges really needed? [...]\n\nWe would like to argue that the introduction of the dynamic and static Schrodinger Bridges are needed in order to explain the rationale behind Proposition 2 and Proposition 3. However, we acknowledge that the present introduction could be simplified (we will introduce the Brownian motion and not the Ornstein-Uhlenbeck).\n\n> The link between (EmOT) and Schrodinger bridges is not well explained [...]\n\nWe acknowledge that the current presentation is misleading as the EmOT is not necessarily defined for tree. In the revised version of the paper, we will make clear that the link between EmOT and Schrodinger Bridges is valid in any multimarginal setting. \n\n> I suggest to state more clearly the motivation, contribution, and main result [...]\n\nWe agree with the reviewer. In the revised version of the paper, we will highlight our contributions in a “Contributions” paragraph. We will relegate some technical details about mOT to the background section. \n\n> Proposition 3 should be presented in a self-sufficient way.\n\nIn the updated version of the paper, we will be more precise regarding the definition of $\\mathbb{P}^n$ (which is defined in Proposition 3). To clarify our results, we will merge Proposition 2 and Proposition 3.\n\n> I wonder why the authors don't stress the convergence result more as a main contribution of the paper.\n\nThe motivation behind this theoretical section was to provide a mathematical background to our method. We were surprised to find that the literature has several gaps when it comes to entropic mOT. Section 4 is an attempt at filling some of them. We are constrained by space but would be happy to provide more details regarding the theoretical results including a sketch of proof.\n\n> For the synthetic two dimensional dataset it would be interesting to also see the true (possibly regularized) Wasserstein barycenter for reference.\n\nWe have added this example to the one-page PDF associated with the general answer.\n\n> For the MNIST application, I suggest to clarify in the main paper that you consider the problem in 784 dimensional space.\n\nWe agree with the reviewer that it is not uncommon in the OT community to treat images as densities. In contrast, we work in a high dimensional setting and will clarify this point.\n\n> What is the convergence criterion in Algorithm 1, line 5?\n\nThe convergence criterion we choose is a fixed number of iterations. We stick with the common strategies in the training of neural network in deep learning which were used in [2].\n\n> What do you mean by one IPF cycle, and how do you chose the number of IPF cycles [...] ?\n\nBy IPF cycle, we mean that we have iterated on all possible roots. For example, in the case of the star-shaped tree with leaves 1,2,3 and central node 0, the first IPF cycle allows us to obtain $\\{\\pi^1, \\pi^2, \\pi^3\\}$. The number of IPF cycles (= number of IPF iterations/K) is chosen by hand like in other DSB approaches [2,3]. \n\n> [...] to also include some experiments that are not a Wasserstein barycenter problem.\n\nRegarding the application of TreeDSB to other tree structured problems, we have described an application from single-cell genomics in the general answer. The goal is to predict the effect of a treatment but also to follow the time evolution on some cell features. We obtain a tree-structured problem which is not star-shaped.\n\n> Typos\n\nWe have corrected all typos.\n\n[1] Multimarginal optimal transport with a tree-structured cost and the Schrödinger bridge problem. Haasler et al.\n\n[3] Diffusion Schrödinger bridge with applications to score-based generative modeling. De Bortoli et al.\n\n[4] Likelihood Training of Schrödinger Bridge using Forward-Backward SDEs Theory. Chen et al."}]

Author rebuttal:
We thank the reviewers for their insightful comments and are encouraged by their positive feedback regarding the proposition and soudness of our work. We provide detailed responses to each reviewer but summarize here their main feedback.

### 1. Motivation for the application to the Wasserstein Barycenter (WB) problem

First, we emphasize that our method is **out-of-sample** (OuS), similarly to recent non-regularized methods [1,2] or regularized ones [3]. This is in contrast to algorithms such as [4] or [5], which are **in-sample** (IS). Contrary to IS methods, OuS ones do not require re-running the full procedure when given a new datapoint.

Second, our WB formulation (line 225) is different from the one considered in the existing literature $\sum_{i=1}^\ell W^2_\{2,reg}(\mu, \mu_i)$, since it has an additional entropic term $\ell \varepsilon H(\mu)$. This term **penalizes the entropy of the barycenter** instead of promoting it, as done classically. Consequently, even for large values of $\varepsilon$, the entropic regularization on the barycenter may remain low [6], which allows us to avoid the numerical issues of having to choose $\varepsilon$ too small.

Finally, TreeDSB is, to the best of our knowledge, the **first methodology to extend ideas from diffusion-based models to the computation of WBs**. In particular, we believe that the idea of iterative refinement, i.e., solving a dynamic counterpart to the static problem plays a key role in the efficiency and scalability of the method.

### 2. Extension of TreeDSB to other trees

Although TreeDSB is defined for any tree, we agree with the reviewers that so far we have only applied it to star-shaped trees. We chose this specific example as there already exists strong baselines. We agree that it would be valuable to test our algorithm on other tree structures. In what follows, we briefly describe **the example of a tree-based setting in single-cell genomics** that we are currently investigating. 

Assume that we only have access to distributions of features of cells from a subset of the following nodes: $T$ (treated) or $U$ (untreated), at times $t_i$. This gives us a **non-star graph**, where we aim at modeling the edges $(T,t_i \to T, t_\{i+1})$, $(U,t_i \to U, t_\{i+1})$ (temporal tracing) and $(U,t_i \to T, t_i)$ (effect of the treatment), see [8]. As this application requires some prior domain knowledge in biology, we do not include it in the current work.

### 3. Comparison with [9]

Our main contribution is to find a **dynamic counterpart** to the static problem of [9]. This allows us to introduce ideas from Diffusion Schrödinger Bridge (DSB) [7], which are known to scale well with data dimension. We also highlight a few limitations of our framework compared to the one of [9], that are inherent to all existing dynamic versions of entropic OT [7,10]. First, our work is limited to the case where the cost is quadratic. Moreover, contrary to [9], the distributions on the nodes in our setting are required to be defined on the same state-space. This prevents us from evaluating our algorithm against [9] to estimate ensemble flows on a hidden Markov chain, see [9, section 6].

### 4. Running time

First, regarding the training time, we should not compare TreeDSB with IS procedures like [4] and [5], which do not rely on the use of neural networks. In contrast to IS methods, **OuS do not need to be rerun** when presented with a new sample, but require a  training procedure. For our experiments on Gaussian datasets and posterior aggregation, the training of TreeDSB required approximately 12 hours while the method of [3] required 8 hours. The memory cost of TreeDSB is similar to the one of competing methods.

Second, we acknowledge that at inference time, TreeDSB is slower than the competing one-shot methods [1,2,3], since it relies on a dynamic formulation. This is a well-identified bottleneck of diffusion models [11], that can be reduced through distillation procedures, see [12] for instance. In our experiments, we considered 50 time-steps per edge (3.5 s to infer 100 MNIST barycenter samples starting from a leaf on a A100 GPU).

### 5. Experiments in the PDF

In the PDF, we present **improved visual results** of the experiments conducted in the main paper (2D and MNIST datasets), which were obtained by slightly increasing the variance of $\mu_0$ (the initial barycenter measure). We also present new experimental results. In the 2D case, we illustrate the classic OT trade-off between the accuracy of the method and its convergence rate through the choice of the regularisation parameter ($\varepsilon=0.2, 0.1,0.05$) (Figure 1), and also present the results from SotA IS methods [4,5] (Figure 2). Finally, we present Wasserstein barycenters obtained for the MNIST experiment (Figure 3) with this new initialization.

[1] Scalable computations of Wasserstein barycenter via input convex neural networks. Fan et al.

[2] Continuous Wasserstein-2 barycenter estimation without minimax optimization. Korotin et al.

[3] Continuous regularized Wasserstein barycenters. Li et al.

[4]  Fast computation of Wasserstein barycenters. Cuturi and Doucet.

[5]  Iterative Bregman projections for regularized transportation problems. Benamou et al.

[6] Doubly regularized entropic Wasserstein barycenters. Chizat.

[7] Diffusion Schrödinger bridge with applications to score-based generative modeling. De Bortoli et al.

[8] Neural Unbalanced Optimal Transport via Cycle-Consistent Semi-couplings. Lübeck et al.

[9] Multimarginal optimal transport with a tree-structured cost and the Schrödinger bridge problem. Haasler et al.

[10] Likelihood training of Schrödinger bridge using forward backward sdes theory. Chen et al.

[11] Tackling the Generative Learning Trilemma with Denoising Diffusion GANs. Xiao et al.

[12] Knowledge Distillation in Iterative Generative Models for Improved Sampling Speed, Luhman et al.


Meta-review: All the reviews are positive. The paper is nicely written and proposes interesting theoretical and numerical contributions. The rebuttal included additional numerical results, which are convincing. It helped clarify several raised issues, in particular regarding computational complexity and connexions to closely related works. I recommend acceptance.
Decision: Accept (spotlight)