Schrödinger Bridge Matching for Tree-Structured Costs and Entropic Wasserstein Barycentres
Authors: Samuel Howard
Peter Potaptchik
George Deligiannidis
Venue: NeurIPS 2025 poster ()
Keywords: Schrödinger bridges
Optimal transport
Wasserstein barycentres
Flow-based generative models
Primary area: deep_learning
TL;DR: We extend the Iterative Markovian Fitting procedure for Schrödinger Bridge computation to the tree-structured setting, which notably includes entropic Wasserstein barycentres as a special case.

Abstract: Recent advances in flow-based generative modelling have provided scalable methods for computing the Schrödinger Bridge (SB) between distributions, a dynamic form of entropy-regularised Optimal Transport (OT) for the quadratic cost. The successful Iterative Markovian Fitting (IMF) procedure solves the SB problem via sequential bridge-matching steps, presenting an elegant and practical approach with many favourable properties over the more traditional Iterative Proportional Fitting (IPF) procedure. Beyond the standard setting, optimal transport can be generalised to the multi-marginal case in which the objective is to minimise a cost defined over several marginal distributions. Of particular importance are costs defined over a tree structure, from which Wasserstein barycentres can be recovered as a special case. In this work, we extend the IMF procedure to solve for the tree-structured SB problem. Our resulting algorithm inherits the many advantages of IMF over IPF approaches in the tree-based setting. In the case of Wasserstein barycentres, our approach can be viewed as extending the widely used fixed-point approach to use flow-based entropic OT solvers, while requiring only simple bridge-matching steps at each iteration.

Schrödinger Bridge Matching for Tree-Structured Costs and Entropic Wasserstein Barycentres
Samuel Howard Department of Statistics University of Oxford
Peter Potaptchik Department of Statistics University of Oxford
George Deligiannidis Department of Statistics University of Oxford
Abstract
Recent advances in flow-based generative modelling have provided scalable meth- ods for computing the Schrödinger Bridge (SB) between distributions, a dynamic form of entropy-regularised Optimal Transport (OT) for the quadratic cost. The successful Iterative Markovian Fitting (IMF) procedure solves the SB problem via sequential bridge-matching steps, presenting
elegant and practical approach with many favourable properties over the more traditional Iterative Proportional Fitting (IPF) procedure. Beyond the standard setting, optimal transport can be generalised to the multi-marginal case in which the objective is to minimise cost defined over several marginal distributions. Of particular importance are costs tree structure, from which Wasserstein barycentres be recov- ered as a special case. this work, we extend the IMF procedure to solve for the tree-structured SB problem. Our resulting algorithm inherits the many ad- vantages of IMF over IPF approaches in the tree-based setting. In the case Wasserstein barycentres,
defined
approach can be viewed extending the widely used fixed-point approach to use flow-hased entropic OT solvers, while requiring only simple bridge-matching steps each iteration. Our code is available https: SB_ LBB_MacchinggBBrr Barycentres.
//github. com/
Introduction
Transporting mass between distributions is ubiquitous problem with numerous applications in machine learning and beyond. Optimal Transport (OT) (Santambrogio. 2015; Peyré and Cuturi, 2019) provides principled approach for such problems, by seeking minimise the total cost of transportation according to a chosen cost function. Since the introduction of Sinkhorn's 's algorithm (Cuturi, 2013) and more recent neural approaches (Makkuva al., 2020), computational OT has seen great success across many domains such as biology (Schiebinger et al., 2019; Bunne et al., 2023), and extensively in machine learning (Genevay al., 2018: Cuturi et al., 2019: Corenflos et al., 2021). Recently, ideas from the powerful flow-based approaches that have revolutionised generative modelling (Song et al., 2021: Peluchetti, 2022; Lipman et al., 2023; Liu et al., 2023; Al- bergo and Ganden-Eijnden 2023) have been leveraged to solve the regular-ee dynamic OT problem, known as the Schrödinger Bridge (SB), Such approaches provide significant scalability ad- vantages, enabling approximation of OT maps between high-dimensional continuous datasets such as image data. Early flow-based SB solvers were based on the classical Iterative Proportional Fitting (IPF)) scheme (De Bortoli et al., 2021; Vargas et al., 2021; Chen et al., 2022), but such methods have since been superseded by those based on the Iterative Markovian Fitting (IMF) scheme (Shi et al., 2023; Peluchetti, 2023) due to its many superior properties.
Beyond the standard OT problem between two marginals, multi-marginal OT aims to find a joint coupling over multiple marginals while minimising the total cost. Tree-structured costs are often
Corresponding author: howard@stats. ox. ox.ac. uk
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

H
(a) Reciprocal process II: The Ys (x) are sampled from the current coupling IIs over S. Conditional on the Ys, points (x) al unknown marginals V are sampled Ys-is Qs-is. Brownian bridges are drawn along the edges between the samples.
(b) Markovianised process P: Vector fields are trained by bridge-matching along each edge. Sam- ples (x) from the next coupling IIs are obtained by simulating the resulting SDEs along the tree struc- ture, started at one of the known marginals.
Figure I:
The two stages of TreeIMF. On this tree, the marginals at the leaf vertices S (blue) are fixed. The marginal at the vertex V (red) not fixed, and can change during the procedure.
considered (Haasler et al., 2021), as they frequently arise in applications while also allowing for im- proved scalability by leveraging the tree structure. Of particular significance are star-shaped trees, as these correspond to the prominent Wasserstein barycentre problem (Agueh and Carlier, 201 I). The Wasserstein barycentre provides natural notion of 'average' for probability distributions, and is widely studied due to its importance in applications, including in Bayesian learning (Srivastava et al., 2018), clustering (Ye et al., 2017), and representation learning (Singh al., 2020) to name few. Computing barycentres is notoriously challenging (Altschuler and Boix-Adsera, 2022). Many suc- cessful approaches approximate the solution with finite set of points, but these in-sample methods struggle to scale well as dimensionality increases. Alternative methods aim to provide continuous approximations (Li et al., 2020) using neural parameterisations, and often require multi-level opti- misation procedures (Fan et al., 2021; Korotin et al., 2022). Of particular relevance diffusion-basee approach of Noble et al. (2023) which extends the IPF-based Diffusion Schrödinger Bridge (DSB) method of De Bortoli et al. (2021) the tree- SB setting, and is (to knowl- edge) currently the only neural ODE/SDE approach for barycentre computation.
work is the
One of the most successful and elegant approaches for Wasserstein-2 barycentre computation is the iterative fixed-point approach, which involves iteratively updating a candidate barycentre y by solv- ing for the OT map from to each marginal μu and updating according to the induced coupling. This approach was popularised by the seminal work of Álvarez-Esteban et al. (2016) and has formed the basis of many algorithmic developments, including machine learning (Korotin et al., 2022). has strong performance, and has been observed to converge quickly in only a few iterations (Lind- heim, 2023). However, the procedure requires solving a complete OT problem to each marginal at each iteration, which is expensive as solving even single OT problem be challenging.
Contributions In this work, we extend the IMF procedure to the tree-based SB problem (Haasler al., 2021). Our TreeDSBM algorithm provides an IMF counterpart to the TreeDSB method from Noble et al. (2023), closing clear gap in the existing literature (see Table 1) and translating the many benefits of IMF over IPF to the tree setting. For the specific case of Wasserstein barycentre computation, our algorithm can be viewed as extending the commonly used fixed-point approaches to the case of flow-based entropic OT solvers. In particular, we show that the iterations of the IMF scheme and the fixed-point barycentre solvers can be elegantly combined into a single iterative procedure, yielding a fixed-point-style algorithm that requires only inexpensive bridge-matching at each iteration. We demonstrate significantly improved empirical performance over TreeDSB, and show that flow-based barycentre solvers can offer competitive performance against existing algorithms for continuous Wasserstein-2 barycentre computation.
Table 1: Positioning of our TreeDSBM algorithm in the literature.
SB
TreeSB (Haasler et al., 2021)
IPF
DSB (De Bortoli et al., 2021)
TreeDSB (Noble et al., 2023)
IMF
DSBM (Shi al., 2023; Peluchetti, 2023)
TreeDSBM (Ours)

---

2 Background
2.1 Schrödinger bridges, optimal transport, and Wasserstein barycentres
We begin by reviewing the standard Schrödinger Bridge (SB) problem between two distributions, and its relation to Optimal Transport (OT). For notation used in the paper, see Appendix A. Given an initial and final measure AM and frr of population, the SB problem aims to identify the most likely intermediate dynamics of the population under the assumption that the movement is driven by stochastic reference process Q. The resulting dynamic SB problem is defined as
PSB
arg min (K)))) Po po, PT WT).
PEP(C)
The static SB problem instead considers only the coupling over the endpoints,
TTT arg min
[KT, (|Qo Th Moo 7 pT)-
(2)
Under mild assumptions, the two problems are equivalent; the dynamic SB solution can be expressed as mixture of bridges over the static SB coupling, psn 17 20or jo. (Léonard, 2013).
Connection to quadratic OT
The reference path measure Q is usually considered to be that of Brownian motion (e In this case, the static SB problem can be rewritten as
I arg min [E Xr (HI II Po, II pr):
(3)
which the entropy-regularised OT problem for the quadratic ground-cost c(wo, ØT) 1 =|1000
or 2|2 and e 2²T (see the Appendix for overview of OT). In the sequel, we assume that the reference process is a Brownian motion (B
Wasserstein barycentres A key motivation for the tree-structured setting that we will consider is the important Wasserstein barycentre problem (Agueh and Carlier, 2011), which provides a nat- ural notion of 'average' for probability distributions. Given l measures (Me) and weights (A1...., de) summing to I, the Wasserstein-2 barycentre is defined as
V = arg min
Σ 2: (1,,)),
(4)
where (3 v) denotes the optimal transport cost between l4 and v for the quadratic ground-cost.
The Wasserstein barycentre problem is notoriously difficult to solve, as involves several OT sub- problems for which one of the marginals can change in the optimisation. An elegant approach leverages the fixed-point property x = Σ, A,TT(z)) which holds at the solution (under mild as- sumptions), where each T; is the OT map from v to l4: (Agueh and Carlier, 20 Álvarez-Fsstban al., 2016). Intuitively, this states that 'each point in the support of the barycentre is the average of the corresponding points in the marginals'. This property has motivated an iterative fixed-point approach for barycentre computation, which involves iteratively updating a candidate barycentre by solving for the OT maps T; and constructing the next iterate as T using the pushforward map T(x) = Σ, T,T,(() (4). Such methods have been shown to be highly successful in practice (Álvarez- Esteban et al., 2016; Korotin al., 2022; Lindheim, 2023; Tanguy et al., 2024).
2.2 Iterative Markovian Fitting
We now outline recent flow-based generative modelling approaches for solving for the SB problem, which will form the basis for our approach. The SB solution psB can be characterised as the unique path measure that is both Markov, and mixture of bridges psB = P Jo,7. that has correct marginals PSB = No; P.B pr (Léonard, 2013). This property motivates the Iterative Markovian Fitting (IMF) procedure (Shi et al., 2023; Peluchetti, 2023), which solves for the SB solution by alternately projecting between Markovian processes and processes with the correct bridges. We refer to Shi et al. (2023) for full details of the IMF procedure, but recall here the basic presentation. We recall the following definitions.
Definition 2.1 (Reciprocal class, Reciprocal projection). A path measure II P(C) is in the reciprocal class (()) of Q if it is a mixture of bridges of Q conditional on their values at the endpoints, II Q 10,T. For path measure P E P(C). the reciprocal projection is defined to be the mixture of bridges according to its induced coupling, projR(q) (P) P.o.r
3

---

Definition 2.2 (Markovian class, Markovian projection). Let M denote the set of Markovian path measures associated to a diffusion of the form dX, v(t, X,) + o,dB, with v, locally Lipschitz. For reference process (a Bd Be)ttjo,TT the Markovian projection of measure IP E (R(Q) is defined to be (when well-defined) the path measure associated the SDE
Ep
[xr]x.] Xt.
dX,-
dt + adBr,
Xo ~ 40-
(5)
T-
can be shown that, under mild conditions, these definitions coincide with the following minimisa- tion problems over path measures (Shi et al., 2023),
projeco) (TP) arg min (TI(P)
proj arg min (T.
(T))
MEM
The IMF iterations are defined below, and converge to a unique fixed point which is the SB solution. p2++1 projm (T²"), p2n+2 projR(e) (122++)).
Training via bridge-matching The IMF procedure requires learning the vector field correspond- ing the Markovian projections (I)). This is done by training a neural network up with bridge-matching loss objective (Peluchetti, 2022), for which the drift in (5) is the optimum,
XT X1
[(
E (Xo.Xr)~ X(x.((X
(Xo,
- 1
(8)
Comparison to Iterative Proportional Fitting Traditionally, the standard way to solve the SB problem via the Iterative Proportional Fitting (IPF) procedure (Fortet, 1940; Kullback, 1968; Rüschendorf, 1995), which minimises the KL divergence between subsequent iterations while alter- nating the endpoint measure that is fixed.
p2a+1 arg min (LII)22") P.Pr=rr
p2x+2 = arg min (L
P:Bo==o
This method was implemented using diffusion model-based approaches in De Bortoli et al. (2021), giving the Diffusion Schrödinger Bridge (DSB) method (see also Vargas et al. (2021). Chen et al. (2022)). While allowing for arbitrary reference measures Q, this approach suffers from several shortfalls, such as only preserving both marginals at convergence (in comparison to IMF, which preserves both marginals each iteration), as well as expensive trajectory caching, and 'forgetting' of the original reference measure (Fernandes et al., 2022). As result, IPF approaches have largely been superseded by IMF approaches which avoid these issues.
2.3 Tree-structured Schrödinger bridge
In this work, we consider the tree- tree-structure Schrodinger Bridge problem (Haasler et al., 2021). In particular, we extend the IMF procedure to the dynamic Tree SB setting of Noble et al. (2023).
Stochastic processes on the tree In order to explain the dynamic Tree SB problem, we first need to define stochastic processes on tree. We will consider trees T (V. E) with vertex set V and edge set E, where É : E + R>0 is an edge-length function. One can extend the tree to a uniquely arewise-connected metric space in the natural way; the are connecting the two endpoints of an edge is identified with a line segment of corresponding length T = E(e), and they are connected according to the graph structure. Via slight abuse of notation we will also denote this metric space as T. As such, we can define the space C(T, IRd) of continuous paths from T into Rd, and we let P(Cr)= == P(C ))) denote the space of probability measures over such paths. We will present the methodology according to a directed tree T, rooted choose an ordered edge set E, corresponding to a depth-first traversal (though we will often omit the dependence on the notation). While this presentation may appear somewhat complex, the construction is quite natural; for an illustration of processes on stresstruc see Figures I and 4.
vertex r, for which we can
We will consider running SDEs along each edge according to the ordering E, in the directed tree structure. We sample X, ~ Pr and then sequentially simulate SDEs along each edge we traverse the directed tree. For each edge e = (u, v) E ε,, we run an SDE dX; = (1, Xf)d( + df for time / E [0, 7 initialised at X = X.. and after simulation we let X, Xς... Such stochastic processes induce path measure E P(CT) over the whole tree. Note that when the tree is a branch with 2 vertices and edge, this recovers the standard case described earlier.

---

Tree-structured Schrödinger bridges Now that we have constructed stochastic processes on the tree, the tree-structured Schrödinger Bridge problem is defined analogously the standard case. However, now the marginals may be fixed only at subset of vertices S C V (see Figure 1). For reference measure Q P(CT), the dynamic and static TreeSB problems are defined respectively as (TreeSBayn)
psB = arg min (KL((K P. = b4 Vi E S),
PEP(CT)
arg min (KL(II Qv)
I. = p Vie
(TreeSBatal
N)P((
As in the standard case, using the chain rule for KL divergence gives (under mild conditions) equiv- alence of the two problems: the dynamic SB solution can be expressed as a mixture of bridges over the static SB coupling, psn 1 aI 2. In this work, we aim solve the TreeSBayn problem.
Connection to multi-marginal optimal transport In the remainder of the work, we will consider the reference measure Q to be associated with running Brownian motions ( (0, along each edge e. The induced reference coupling Qy over the vertices is characterised by having independent Gaussian edge increments Y, - Y, ~ No0,, o"T((,))) By evaluating the KL term, the static TreeSB problem can therefore be rewritten as
11 B
arg min
[Ex=m[
Σ (u.))E
Xr, x.113] 20 HII)) p4 vie (10)
Plu.oy
IEP(()))VV))
This is precisely an cetrrop--reullrre cost function c(zy) = (m. тоा. | - = and entropy-regulariiation E = 20 (see Ap- pendix A for an overview). Note that each weight the reciprocal of the corresponding edge length.
marti-larginal OT problem with tree- rree-structure quadratic
Connection to Wasserstein barycentres There are many ways to define entropy-regularissed Wasserstein barycentres (see Appendix A for an overview). The work of Chizat (2023) unifies several of these, combining s-regularised OT costs with an entropic penalty term the barycentre weighted by r. Noble et al. (2023) show that the following doubly- regul-ree barycentre problem
V² arg min
Σ 22 (Ma, v) + (E - 1)=H(v)}.
(1I)
recovered from the TreeSB setting by considering a star-shaped tree with fixed marginals on the leaves and cost function c(Faz) = Ei= >1laro -x,13.
3 Iterative Markovian Fitting for tree-structured costs
In this work, we generalise the IMF procedure to the tree-based setting. Recall that standard IMF proceeds by iteratively Markovianising mixture of bridges, where each bridge is a stochastic pro- cess conditioned on endpoints zg and F1, and the mixing coupling Io.11 1 is over the endpoint values To. Similarly, we consider Markovianising mixtures of bridges, but instead the mixing coupling a joint distribution IIs over values at the observed vertices S, and the bridging processes are con- ditioned on these values. In the following, we justify that the IMF procedure extends to this setting. We then explain how the bridging and Markovianisation can performed in practice.
3.1 Markovian and reciprocal processes on the tree
Recall that the IMF procedure is motivated by the characterisation of the SB the unique Markov measure that is a mixture of its bridges. Our first contribution is to extend this characterisation to the tree-based setting. We begin with the following definitions, akin to those in the standard case.
Definition 3.1 (Markov class). Let MT denote the set of Markov path measures on the tree T, defined via a diffusion along each edge of the form dx; = (t, Lipschitz.
dBd with v², of locally
Definition 3.2 (Reciprocal class). II E P(CT) is in the reciprocal class Rs ( of Q for observed vertices 8 ifll Q1s That is, II a mixture of bridges of Q conditioned on the values S.
Observe that compared to the corresponding Definition 2.1 in standard IMF, the bridges are now conditioned on the values at vertices S, rather than on only the endpoints single edge. We now state the following result characterising the TreeSB solution.
5

---

Theorem 3.1 (TreeSB characterisationn. Under mild assumptions (in the Brownian case, namely that ddpx()) < oc and H(14) 0o for each S). there exists unique solution to the dynamic TreeSB problem (TreeSBayu). The solution is the unique process P that is both Markov and (sQ with correct marginals P, = fo E S.
3.2 Iterative Markovian Fitting the tree
Theorem 3.1 suggests that an analogous Iterative Markovian Fitting procedure could be used to solve the dynamic TreeSB problem; in this section, we show that this is indeed the case. The constructions and proofs proceed in much the same way as in Shi et al. (2023), but crucially rely on the following simple decomposition of the KL divergence according to the tree structure, which applies for measures in My and Rs(
Lemma 3.2 (KL, decomposition along tree). Take path measures P, TP that share a marginal the roo, P, Pr, and factorise along the tree. Under mild assumptions, we have the KL decomposition
KT(PPPP)
Σ (u,D)E,,
Ex,=P, [L((pu))) (11x,) (x))))
(12)
To extend the IMF procedure to the TreeSB problem, we require to extend the definitions and propo- sitions from the standard case the tree-based setting. The definition of the reciprocal class pro- jection remains much the same as in the standard case, the difference being that now the bridging processes are conditioned on their values vertices in S.
Definition 3.3 (Reciprocal projection). For process P E P(CT). the reciprocal projection is defined to be the induced mixture of bridges II = projes (P) PsQ.s. s-
Proposition 3.3. For I. P(CT), the reciprocal projection II* proj (P) solves the minimi- sation problem IIe arg (())) KL(PIIII.
The Markovian projection requires alterations account for the tree structure.
Definition 3.4 (Markovian projection). For IT E RR((Q), the Markovian projection onto MI, de-
noted M.* = (TT), is defined as
M II proj (II.).
(13)
Ie denotes the restriction of II to edge e, pr (II) denotes Markovian projection (Defini- tion 2.2), and H1 denotes the composition of the resulting measures according to the tree structure.
Proposition 3.4. Under mild assumptions, the Markovian projection M prrii solves the minimisation problem MI" arg minggemy MLIII.M...
The TreeIMF procedure then defined via the usual iterations, started from initialisation pd E (R)(Q) such that ppo = μ4 for E S: 2 proj (IP2"), p2=+2 projRs(() (1)2x))).
(14)
Note that this procedure recovers standard IMF when the tree has bridge structure (2 vertices, edge), just as TreeDSB recovers DSB in this setting.
We state the following result, which shows convergence of the IMF iterates to the TreeSB solution. This resembles Theorem in Shi et al. (2023), and the proof proceeds in largely the same way using the appropriate modifications and in particular leveraging the KL decomposition in Lemma 3.2. We defer the details to Appendix B.
Theorem 3.5 (Convergence of TreeIMF). Under mild conditions, the TreeSBayn solution Pa the unique fixed point of the TreeIMF iterates 1a, and have lim KL(P"" O.

---

3.3 Implementation
We have established the convergence of the TreeIMF procedure to the TreeSB solution. We now explain how the procedure be implemented in practice.
Constructing the bridge processes Recall that the reciprocal projection relies constructing bridges of the reference process Y associated to Q, conditioned on the values at vertices in S. Due the tree structure and the Brownian dynamics of Q, such bridges can be constructed by first sampling the values of YsUI ~ Qs+is at the unseen vertices in S = VVS, and then sampling Brownian bridges along each edge (u,v) E between Y, and Yu-
Note in particular that the conditional coupling Qs+is is tractable because the static coupling Qy is multivariate Gaussian. Specifically, by characterised via independent Gaussian edge increments Y,, Y, (N ofT"), (Y) & exp(- (00(--(noo) (u.))E 20 YY Yull²) x exp(- LY)
where L is the precision matrix given by
Eeee(H) क o-T(uoot
if v, where E(u) is set of edges incident to if (u, v) E otherwise.
Lu,0
Using the formula for conditional multivariate Gaussians in terms of the precision matrix, we have Ys ys (i, E), using the block matrix expressions á
and E (1s-se)))-- Based on the structure of the tree, we can therefore obtain a sample of the bridging process by: (1) calculating this conditional joint distribution Q& over the unseen vertices, (2) sampling Is according this distribution, and then (3) drawing Brownian bridges between Y., and
(
Y, along each edge in the tree.
Performing the Markovian projection The above construction allows us to construct samples from a reciprocal process II. Recall from Definition 3.4 that the Markovian projection then proceeds by performing individual Markovian projections along each edge. We therefore maintain neurally- parameterised drift function vg for each edge = (u, v), each of which is trained according to the bridge-matching loss
[(0e) = Jo (e M Xx
X, X, - 00, (Xe,
(15)
Bidirectional training After training the vector fields, let M denote the path measure of the re- sulting Markov process on T. Constructing the next reciprocal process requires samples from M. which can be obtained by running the learned diffusions from the root node r along the tree. Note that while in theory we have M, μr for each i S, errors will accumulate in practice. We therefore follow Shi et (2023) and learn both forward and backward diffusion processes along each edge, which give equivalent representations of the path measures via the corresponding time- reversals (see Appendix C for details). This enables simulation along the tree from any vertex in S to obtain samples from the coupling, helping to mitigate errors accumulating in the marginals. Note also that all edges train independently, so can be learned in parallel for faster computation.
In keeping with previous naming conven- tions, we call the proposed algorithm Tree Diffusion Schrödinger Bridge Matching (TreeDSBM). The method is summarised in Algorithm 1 with full implementation details in Appendix C, and we provide an illustration of the two-step procedure in Figure/ 1. The algorithm can be initialised any coupling H& over with correct marginals; a standard choice would be independent coupling n1 = Øres (i. We provide a simplified and more explicit version of the algorithm when used for barycentre computation in Appendix C.
Algorithm TreeDSBM
Input: Initial coupling 11g (e.g. independent), number of iterations N.
Let ITo = =I Q for ......... N 1] -
do
Learn 2| vector fields using (15) with II = II". to obtain Markovian process M Simulate MI*+1 from chosen root S to obtain samples from M
Let 11 = M Q using obtained samples from M
end
7

---

Text
TreDSIM
Trwntn
NTV QATTR
Table 2: Sinkhorn divergence to the "ground-truth barycentre, for barycen- tre samples generated from each Icaf vertex k. The Sinkhorn divergence is computed with entropy regularisation 0.01, using 5000 generated samples and 1500 points approximating the ground- truth (mean over runs).
D
Ground Truth
km2
Method
k == 0
k:
k k= 2
Figure 2: Comparison of the learned barycentre for TreeDSBM (6 IMF iterations) against TreeDSB (50 IPF iterations) and WIN. TreeDSB and TreeDSBM samples are generated from each leaf vertex k and for
TrecDSBM
1.14
1.05
1.08
(6 IMF)
0.07
+ 0.07
+ 0.11
TreeDSB
2.35
4.04
2.35
WIN we plot samples using the weighted- weighted-pushhorwaan pushforward expression for the barycentre. Also displayed is close approximation to the ground-truth, using the in-sample method of Cuturi and Doucet (2014).
(50 IPF)
WIN
E, AT.
1.17
3.4 Connection fixed-point Wasserstein barycentre algorithms
Recall that the unregularised Wasserstein-2 barycentre can be computed using the iterative fixed- point method, where each iteration proceeds by pushing-forward the current iterate through the map Σ, A,T, In our setting, the barycentre problem corresponds to a star-shaped tree with fixed marginals on the leaves. In this case, the bridges of Q are conditioned on each of the leaf vertices and the only unknown point is Yo at the centre vertex, so the conditional distribution Q&e simplifies to
Yo Vo~N( (.
Our method can therefore be viewed as a natural counterpart of fixed-point methods for barycentre computation, adapted to the case of flow-based entropic OT solvers. While a naive approach might consider using IMF to solve each OT sub-problent in the fixed-point scheme, resulting in nested iter- ations, our approach shows that the IMF and fixed-poing iterations can in fact be elegantly combined into a single iterative procedure, along with a well-understood theoretical grounding. In particular, the expensive OT map computations required for the fixed-point procedure can instead be switched out for inexpensive tridge-matching procedures. Each iteration is therefore cheap, and empirically found the TreeDSBM algorithm to retain the fast convergence property of the fixed-point proce- dure in terms of the number of iterations required.
4 Experiments
Synthetic 2d barycentre We first examine the performance of TreeDSBM in a low-dimensional synthetic example, using the experimental setup of Noble et al. (2023). We compute the (3 3) barycentre of a moon, spiral, and circle dataset with = 0.1 (recall TreeDSBM ran for 6 IMF iterations against TreeDSB ran for 50 IPF iterations (using checkpoints provided by Noble et al. (2023)). In Figure 2, we plot the obtained samples from both methods, and for comparison also display the barycentre obtained using the in-sample method from Cuturi and Doucet (2014) to give a close approximation to the ground-truth. To quantitatively assess per- formance, we report the Sinkhorn divergence (Genevay et al., 2018) relative to this 'ground-truth" in Table 2. The results show that TreeDSBM significantly improves over TreeDSB in this setting, approximating the barycentre to high degree of accuracy and at much lower computational cost (with good convergence after only a few IMF iterations). It is able to successfully capture the complex nature of the barycentre in this challenging example (previously suggested as potential limitation of dynamic solvers in Noble et al. (2023)), and the barycentres generated from different vertices k exhibit improved consistency. For details of the experimental setup, see Appendix D.I.
= We compare
We additionally report results for the iterative WIN method (Korotin et al., 2022). We see that TreeDSBM performs competitively with this strong baseline for continuous Wasserstein-2 solvers. The results reported for WIN are for the combined map (E, A,T;) as in our experiments the barycentre generator G#p was unable to fit the true barycentre accurately, were the maps T,- #14.. Additionally, we applied the W2CB (Korotin et al., 2021) and NOTWB (Kolesov et al.,
8

---

2024a) algorithms to this example, but were unable find hyperparameters to make the algorithms to converge to the correct solution. We hypothesise that the neural maps may struggle to model the discontinuous transports well, and also that this challenging example may result in difficult loss landscapes. We found TreeDSBM to exhibit stable training despite this challenging problem setting, and to also be the fastest of the algorithms to converge to the solution. We provide runtime analysis, along with further results and discussion, in Appendix D.I.
MNIST 2,4,6 barycentre We also compare performance of TreeDSBM with TreeDSB on a higher dimensional image
2
6 4
0
dataset, computing the (3 3
1 }) }-barycentre between MNIST
4 4
U
4
4
&
digits 2, 4, and 6 (LeCun et al., 2010). In this setting, TreeDSB reported in Noble et al. (2023) to exhibit training instability for low entropy regularisation E, causing it to struggle to match the marginal measures at the vertices. In contrast, the IMF approach of TreMDSBM ensures matching of these marginals, thus allowing the use of much smaller regularisation values. In fact. using a large regularisation for TreeDSBM would limit sample quality, because of the noise added when sampling from Q&ais in the reciprocal process construction. We there- fore use e = 0.02 for TreeDSBM with IMF iterations, and display a visual comparison with results reported in Noble et al. (2023) in Section 4. While it is difficult to validate the accuracy of such solutions, the obtained samples from TreDDSBM appear to display a greater resemblance to state-of-the-art barycentre methods for similar problems (see e.g. Kolesov et al. (2024a)), while again requiring significantly less computational cost than TreeDSB. We remark that the TreeDSBM samples in Figure 3b contain a small amount of noise, which is be expected as we solving for an entropy regulay-ee problem. In Appendix D.2 we also display samples with the noise reduced or removed, along with full experimental details and additional results.
4 9
0
4
4
4
4
6
0
4 4
3 0
A 4
4
4
6 4
3 6
2
4
0
2
(a) Tre
(b) TreeDSBM
(s = 0.5)
(s = 0.02)
Figure 3: Samples from the 2,4,6
MNIST barycentre.
Subset posterior aggregation The previous experiments have shown TreeDSBM to improve over its IPF counterpart TreeDSB. We now provide more detailed comparison with existing strongly- performing methods for continuous Wasserstein-2 barycentre estimation, namely WIN (Korotin et al., 2022), W2CB (Korotin et al., 2021), and NOTWB (Kolesov et al., 2024a). As a real-worid applications experiment, we compare the algorithms in the subset posterior aggregation setting, a standard experiment in the barycentre literature (Staib et al., 2017; Li et al., 2020; Fan et al., 2021; Korotin et al., 2021). The barycentre of subset posteriors is known to be close to the true posterior (Srivastava et al., 2018), and as such it can be efficient to compute posteriors on only subsets of datasets before aggregating them using barycentre algorithm. We consider the experimental setup and dataset used in Korotin et al. (2021) (the same dataset was also used previously in Li et al. (2020) and Fan et al. (2021)), which uses Poisson and negative-binomial regressions rental dataset (Fanaee-T) T, 2013), and report results for the BW2-UVF metric. We ran TreeDSBM for 4 IMF iterations with 8 = 0.001; for full experimental details, see Appendix D.3.
bike-
From Table 3, we see that all methods perform strongly. Given that we do not have perfect access to the ground truth barycentre, it is difficult to conclusively say which performs best, but we see that TreeDSBM certainly performs competitively with these state-of-the approaches. Moreover, we found Tre:DSBM to display fast training- both TrecDSBM and NOTWB obtained good results after only around 3 minutes of training, while W2CB took approximately 10 minutes and WIN took around 45 minutes. For more details regarding runtimes, see Appendix D.3.
ligees--dimnsiioal Gaussian experiments We also report results for computing the barycentre of Gaussian distributions, in increasingly high dimensions. This is a standard experiment in the literature, because in this setting the ground-truth barycentre is also Gaussian and the parameters can be calculated accurately using a fixed point method (Álvarez-Esteban et al., 2016).
We follow the experimental setup previously used in (Korotin et al., 2021, 2022; Kolesov et al., 2024a), in which 3 Gaussian distributions and its ground-truth (+. 3 bbbarycentre are randomly generated, for each dimension in (64, 96, 128]. We report results for the BW3 UVP and -VV metric,
1
metries. We see that all the methods again perform well in this setting. For the BW2- W2CB and NOTWB appear to have slight edge, but TreeDSBM is only slightly higher and is comparable with results for WIN. For the 7,2 UVP metric, results for TreeDSBM are higher than for W2CB and NOTWB, though the results are still low and are again comparable with WIN.

---

Table 3: 2-UV,, % for different algorithms on the subset posterior aggregation experiment in Korotin al. (2021), evaluated using 100,000 samples.
WIN
W2CB
NOTWB
TreeDSBM (Ours)
Poisson
0.014
0.026
0.023
0.008
Negative Binomial
0.009
0.024
0.018
0.012
Table 4: -V- % and - UVP, % for different algorithms on the high-dimensional Gaus- sian experiment, evaluated using 100,000 samples.
WIN
W2CB
NOTWB
TreeDSBM (Ours)
= 64
0.20
0.04
0,08
0.14
4 -VP,, % = 96
0.30
0.07
0.10
0.15
128
0.38
0.12
0.14
0.27
64
0.96
0.17
0.10
1.18
-UV-
96
1.20
0.20
0.10
1.13
128
1.46
0.25
0.13
1.23
Discussion
Related work Beyond the IPF and IMF approaches, other approaches for the SB problem include adversarial solvers (Kim et al., 2024; Gushchin et al., 2023). parameterisation via potentials for Gaussian mixtures (Korotin et al., 2024; Gushchin et al., 2024), and variational approaches (Deng et al., 2024). For Wasserstein barycentre computation, standard approaches ensure tractability by representing the solution with a finite set of points (either updating only weightings (Benamou et al., 2015; Solomon et al., 2015; Cuturi and Peyré, 2016; Staib et al., 2017: Dvurechenskii et al., 2018), or also the positions of the points in the support (Rabin et al., 2012; Cuturi and Doucet, 2014; effective in lower dimensions but scale
Claici et al., 2018; Luise et al., 2019)). Such approaches poorly, cannot be used to generate new samples, and do not capture the true continuous nature of the barycentre. To address these limitations, recent work has focused learning continuous approxi- mations. Li et al. (2020) optimise for the regularised dual potentials for general costs, while others parameterise Wasserstein-2 potentials with convex neural architectures using adversarial losses (Fan et al., 2021) or additional cycle-consistency regularisation (Korotin et al., 2021). Korotin et al. (2022) leverage the fixed-point property in Alvarez-FEseehi et al. (2016). We also highlight a recent line of works of Kolesov al. (2024b,a) which consider bi-level adversarial approaches for general cost functions. Our approach utilises non-adversarial bridge-matching loss objectives, which pro- vide stable training but requires multiple sequential IMF iterations. Finally, we emphasise that our approach tackles the tree-structured SB problem (Haasler et al., 2021; Noble et al., 2023), and thus applicable beyond only barycentre problems (we consider toy example in Appendix E.3).
re- stricted to quadratic cost functions for OT, and introduces entropic bias. Inference is expensive in comparison to methods that use a single function evaluation, and our method is not sampling-free requiring simulations of the current learned processes each iteration. Future advances in the flow and SB literatures can aid in addressing these limitations. We also provide an additional experi- ment in Appendix E that highlights possible limitation of our method in scenarios where there is simple shared structure between the known marginals and the barycentre.
Limitations Our method shares the same limitations as other flow- SB solvers.
Conclusion and future work We have extended the IMF procedure to the tree-structure SB prob- lem, providing scalable flow-based approach that particular can be used for entropic Wasser- stein barycentre computation. Our TreeDSBM algorithm displays improved performance over its IPF counterpart TreeDSB, and demonstrates that flow-based approaches for barycentre estimation can offer a compelling alternative to established continuous Wasserstein-2 barycentre algorithms. Future directions can investigate improved architectures and implementation techniques inspired by progress in the flow-men Row-matching literature, as well as extensions to other data modalities.
10

---

Acknowledgements
SDEs Theory". In: International Conference on SH supported by the EPSRC CDT in Modern Statistics and Statistical Machine Learning [grant number EP/S023151/11 PP is supported by the EPSRC CDT in Modern Statistics and Statistical Machine Learning EEP/S02351//]] a Google PhD Fellowship, and an NSERC Postgraduate Schol- arship (PGS D). GD was supported by the Engineering and Physical Sciences Research Council [grant number EP/Y018273/1J. The authors would like to thank James Thornto for helpful discus- sions.
References
Martial Agueh and Guillaume Carlier (2011). "Barycenters in the Wasserstein Space". In: SIAM Journal on Mathematical Analysis 43.2, pp. 904-924.
Michael Samuel Albergo and Eric Vanden-Eijnden (2023). "Building Normalizing Flows with Stochastic Interpolants". In: International Conference on Learning Representations.
Jason M. Altschuler and Enric Boix-Adsert (2022). "Wasserstein Barycenters Are NP-Hard to Com- pute". In: SIAM Journal on Mathematics of Data Science 4.1, pp. 179-203.
Pedro Alvarez-Esteban, E. del Barrio, J.A. Cuesta-Albertos, and C. Matrán (2016). "A fixed- point approach to barycenters in Wasserstein space". In: Journal of Mathematical Analysis and Applications 441 .2, pp. 744-762.
Brandon Amos, Lei Xu, and J. Zico Kolter (2017). "Input Convex Neural Networks". In: Interna- tional Conference on Machine Learning.
Jean-Dace Dean-David Benamou, Guillaume Carlier, Marco Cuturi, Luca Nenna, and Gabriel Peyré (2015).
"Iterative Bregman Projections for Regularized Transportation Problems". In: SIAM Journal on Scientific Computing 37.2, A1111-A1138
V. I. Bogachev, T. I. Krasovitskii, and S. V. Shaposhnikov (2021). "On uniqueness of probability solutions of the Fokker-Planck- Kolmogrro equation". In: Shornik: Mathematics 212.6, p. 745.
David Bolin, Alexandre B. Simas, and Jonas Wallin (2024). Markov properties of Gaussian random fields on compact metric graphs. arXiv: 2304. 03190 [math. PR].
James Bradbury, Roy Frostig, Peter Hawkins. Matthew James Johnson, Chris Leary. Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman and Qiao Zhang (2018). JAX: composable transformations of Python+NumPP programs. Version 0.3.13.
Charlotte Bunne, Stefan G. Stark, Gabriele Gut, Jacobo Sarabia del Castillo, Mitch Levesque, Kjong-Van Van Lehmann, Lucas Pelkmans, Andreas Krause, and Gunnar Rätsch (2023). "Learning Single- Aingle-Cell Perturbation Responses using Neural Optimal Transport". In: vol. 20. Nature Meth- ods, pp. 1759-1768.
Tianrong Chen, Guan-Horn Liu, Molei Tao, and Evangelos Theodorou (2023). "Deep Momentum Multi-Marginal Schrödinger Bridge". In: Advances Neural Information Processing Systems.
Tianrong Chen, Guan-Horng Liu, and Evangelos Theodorou (2022). "Likelihood Training of Schrödinger Bridge using Forward-Backe Learning Representations.
Yongxin Chen, Giovanni Conforti, Tryphon T. Georgiou, and Luigia Ripani (2019). "Multi Multi-marginal Schrödinger Bridges". In: Geometric Science of Information. Springer International Publishing, pp. 725-732.
Lénaïc Chizat (2023). Doubly Regularized Entropic Wasserstein Barycenters. arXiv: 2303 11844 [math.00] .CC..
Sebastian Claici, Edward Chien, and Justin Solomon (2018). "Stochastic Wasserstein Barycenters". In: International Conference on Machine Learning.
Adrien Corenflos, James Thornton, George Deligiannidis, and Arnaud Doucet (2021). "Differen- tiable partiele filtering via entropy-regulatiiedd optimal transport". In: International Conference on Machine Learning.

---

Marco Cuturi (2013). "Sinkhorn Distances: Lightspeed Computation of Optimal Transport". In: Advances in Neural Information Processing Systems.
Marco Cuturi and Arnaud Doucet (2014). "Fast Computation of Wasserstein Barycenters", In: In- ternational Conference on Machine Learning.
Marco Cuturi, Laetitia Paang--apaxanthos Yingtao Tian, Charlotte Bunne, Geoff Davis, and Olivier Teboul (2022). "Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein", In: arXiv preprint arXiv:2201. 12324.
Marco Cuturi and Gabriel Peyré (2016). "A Smoothed Dual Approach for Variational Wasserstein Probiems". In: SIAM Journal on Imaging Sciences 9.1, pp. 320-343.
Marco Cuturi, Olivier Teboul, and Jean-Philippe Vert (2019). "Differentiable ranking and sorting using optimal transport". In: Advances in Neural Information Processing Systems.
Paolo Dai Pra (1991). "A stochastic control approach to reciprocal diffusion processes". In: Applied Mathematics and Optimization 23.1, pp. 313-329.
Valentin De Bortoli, Iryna Korshunova, Andriy Mnih, and Arnaud Doucet (2024). "Schrodinger Bridge Flow for Unpaired Data Translation". In: Conference on Neural Information Processing Systems.
Valentin De Bortoli, James Thortoon, Jeremy Heng, and Arnaud Doucet (2021). "Diffusion Schrödinger Bridge with Applications Score-Based Generative Modeling". In: Advances in Neural Information Processing Systems.
Wei Deng, Weijian Luo, Yixin Tan, Marin Bilos, Yu Chen, Yuriy Nevmyvaka, and Ricky T. Q. Chen (2024). "Variational Schrödinger Diffusion Models". In: International Conference on Machine Learning.
Prafulla Dhariwal and Alexander Nichol (2021). "Diffusion Models Beat GANs on Image Synthe- sis". In: Advances in Neural Information Processing Systems.
Pavel Dvurechenskii, Darina Dvinskikh, Alexander Gasnikov, Cesar Uribe, and Angelia Nedich (2018). "Decentralize and Randomize: Faster Algorithm for Wasserstein Barycenters". In: Ad- vances in Neural Information Processing Systems.
Jiaojiao Fan, Amirhossein Taghvaei, and Yongxin Chen (2021). "Scalable Computations of Wasser- stein Barycenter via Input Convex Neural Networks". In: International Conference on Machine Learning.
Hadi T T (2013). Bike Sharing.
UCI
Machine
Learning
Repository,
DOI:
http://dbboorg//022443/////
Kilian Fatras, Younes Zine, Szymon Majewski, Rémi Flamary, Rémi Gribonval, and Nicolas Courty (2021). Minibatch optimal transport distances; analysis and applications. arXiv: 2101. 01792 [stat.ML]
David Lopes Fernandes, Francisco Vargas, Carl Henrik Ek, and Neill D. F. Campbell (2022). "Shoot- ing Schrödinger's 's Cat". In: Fourth Symposium on Advances in Approximate Bayesian Inference.
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar 7. Alaya, Aurā@lie Boisbunon, Stanislas Chambon, Lactitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, LÅÃo Gau- theron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy. levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer (2021). "POT: Python Optimal Transport". In: Journal of Machine Learning Research 22.78. pp. 1-8.
Robert Fortet (1940). "Résolution d'un système d' d'équations de M. Schrödinger". In: Journal de Mathématiques Pures et Appliquées 9e série, 19.1-4, pp. 83-105.
Aude Genevay, Marco Cuturi, Gabriel Peyré, and Francis Bach (2016). "Stochastic Optimization for Large-scale Optimal Transport". In: Advances in Neural Information Processing Systems.
Aude Genevay, Gabriel Peyré, and Marco Cuturi (2018). "Learning generative models with Sinkhorn divergences". In: International Conference on Artificial Intelligence and Statistics.
12

---

Nikita Gushchin, Sergei Kholkin, Evgeny Burnaev, and Alexander Korotin (2024). "Light and Op- timal Schrödinger Bridge Matching". In: International Conference on Machine Learning.
Nikita Gushchin, Alexander Kolesov, Alexander Korotin, Dmitry P Vetrov, and Evgeny Burnaev (2023). "Entropic Neural Optimal Transport via Diffusion Processes". In: Advances in Neural Information Processing Systems.
Isabel Haasler, Axel Ringh, Yongxin Chen, and Johan Karlsson (2021). "Multimarginal Optimal Transport with Tree-Structured Cost and the Schrödinger Bridge Problem". In: SIAM Journal on Control and Optimization 59.4, pp. 2428-2453.
Ben Hambly and Terry Lyons (2008). Some notes on trees and paths, arXiv: 0809 1365 [math. CA].
L.V. Kantorovich (1942). On translation of mass. Proceedings of the USSR Academy of Sciences, 37.
Tero Karras, Mika Aittala, Timo Aila, and Samuli Laine (2022). "Elucidating the Design Space of Diffusion-Bard Based Generative Models". In: Advances in Neural Information Processing Systems.
Beomsu Kim, Yu-Guan Hsich, Michal Klein, marco cuturi, Jong Chul Ye, Bahjat Kawar, and James Thomton (2025). "Simple ReFlow: Improved Techniques for Fast Flow Models". In: Interna- fional Conference on Learning Representations.
Beomsu Kim, Gihyun Kwon, Kwanyoung Kim, and Jong Chul Ye (2024). "Unpaired Image-to- Image Translation via Neural Schrödinger Bridge". In: International Conference on Learning Representations.
Diederick P Kingma and Jimmy Ba (2015). "Adam: A method for stochastic optimization". In: International Conference on Learning Representations.
Alexander Kolesov, Petr Mokrov, Igor Udovichenko, Milena Gazdieva, Gudmund Pammer, Evgeny Burnaev, and Alexander Korotin (2024a). "Estimating Barycenters of Distributions with Neural Optimal Transport". In: International Conference on Machine Learning.
Alexander Kolesov, Petr Mokrov, Igor Udovichenko, Milena Gazdieva, Gudmund Pammer, Anasta- sis Kratsios, Evgeny Burnaev, and Alexander Korotin (2024b). FEnergy-Guide Continuous En- tropic Barycenter Estimation for General Costs". In: Advances in Neural Information Processing Systems.
Alexander Korotin, Vage Egiazarian, Lingxiao Li, and Evgeny Burnaev (2022). "Wasserstein It- erative Networks for Barycenter Estimation". In: Conference on Neural Information Processing Systems.
Alexander Korotin, Nikita Gushchin, and Evgeny Burnaev (2024). "Light Schrödinger Bridge". In:
International Conference on Learning Representations.
Alexander Korotin, Lingxiao Li, Justin Solomon, and Evgeny Burnaev (2021). "Continuous Wasserstein-2 Barycenter Estimation without Minimax Optimization". In: International Confer- ence on Learning Representations.
Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev (2023). "Neural Optimal Trans- port". In: International Conference on Learning Representations.
S. Kullback (1968). "Probability Densities with Given Marginals". In: The Annals of Mathematical Statistics 39.4, pp. 1236-1243.
Yann LeCun, Corinna Cortes, and CJ Burges (2010). "MNIST handwritten digit database". In: ATT Labs [Online]. Available:
2.
Christian Léonard (2012). "Girsanov Theory Under Finite Entropy Condition". In: Séminaire de Probabilités XL/V. Berlin, Heidelberg: Springer Berlin Heidelberg, pp. 429-465.
Lingxiao Li, Aude Genevay, Mikhail Yurochkin, and Justin M Solomon (2020). "Continuous Reg- ularized Wasserstein Barycenters". In: Advances in Neural Information Processing Systems.
Johannes von Lindheim (2023). "Simple approximative algorithms for free-support Wasserstein barycenters". In: Computational Optimization and Applications 85.1, pp. 213-246.
13

---

Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le (2023). "Flow Matching for Generative Modeling". In: International Conference on Learning Represen- fations.
Xingchao Liu, Chengyue Gong, and Qiang Liu (2023). "Flow Straight and Fast: Learning to Gener- ate and Transfer Data with Rectified Flow". In: International Conference on Learning Represen- tations.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang (2015). "Deep Learning Face Attributes in the Wild." In: ICCV. IEEE Computer Society, pp. 3730-3738.
Giulia Luise, Saverio Salzo, Massimiliano Pontil, and Carlo Ciliberto (2019). "Sinkhorn Barycenters with Free Support via Frank-Wolfe Algorithm". In: Advances in Neural Information Processing Systems.
Christian Léonard (2013). survey of the Schrödinger problem and some of its connections with optimal transport. arXiv: 1308. 0215 [math.PR]]
Christian Léonard, Sylvie Roelly, and Jean Zambrini (Jan. 2014). "Reciprocal processes, measure- theoretical point of view". In: Probability Surveys 11.
Ashok Makkuva, Amirhossein Taghvaci, Sewoong Oh, and Jason Lee (2020). "Optimal transport mapping via input convex neural networks". In: International Conference on Machine Learning.
Gaspard Monge (1781). "Mémoire sur la théorie des déblais des remblais", In: Histoire de l'Académie Royale des Sciences, pp. 666-704.
Maxence Noble, Valentin De Bortoli, Arnad Doucet, and Alain Durmus (2023). "Tree-Based Diffu- sion Schrödinger Bridge with Applications to Wasserstein Barycenters". In: Conference on Neural Information Processing Systems.
Marcel Nutz (2021). Introduction to Entropic Optimal Transport. Lecture Notes.
F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Pretten- hofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay (2011), 1). "Scikit-lear: Machine Learning Python". In: Journal of Machine Learning Research 12, pp. 2825-2830.
Stefano Peluchetti (2022). Non-Denoising Forward-Time Diffusionss
Stefano Peluchetti (2023). "Diffusion Bridge Mixture Transports. Schrödinger Bridge Problems and Generative Modeling". In: Journal Machine Learning Research 24.374, pp. 1-51.
Gabriel Peyré and Marco Cuturi (2019). "Computational Optimal Transport". In: Foundations and Trends in Machine Learning 11 (5-6), pp. pp- 355-602.
Aram-Alexandre Pooladian, Heli Ben- Bem-Hamu, Carles Domingo-Enrich, Brandon Amos, Yaron Lip- man, and Ricky T. Q Chen (2023). "Multisample Flow Matching: Straightening Flows with Mini- batch Couplings". In: International Conference on Machine Learning.
Julien Rabin, Gabriel Peyré, Julie Delon, and Mare Bernot (2012). "Wasserstein Barycenter and Its Application to Texture Mixing". In: Scale Space and Variational Methods in Computer Vision. Berlin, Heidelherg: Springer Berlin Heidelberg, pp. 435-446.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox (2015). "U-Net: Convolutional Networks for Biomedical Image Segmentation". In: Medical Image Computing and Computer-Assisted Inter- vention - MICCAI 2015. Cham: Springer International Publishing, pp. 234-241.
Ludger Rüschendorf (1995). "Convergence of the Iterative Proportional Fitting Procedure". In: The Annals of Statistics 23.4, pp. 1160-1174.
Filippo Santambrogio (2015). Optimal transport for applied mathematicions. Birkhauser.
Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, Lia Lee, Jenny Chen, Justin Brumbaugh, Philippe Rigollet, Konrad Hochedlinger, Rudolf Jaenisch, Aviv Regev, and Eric S. Lander (2019). "Optimal-Transport Analysis of Single-Cell Gene Expression Identifies Developmental Trajecto- ries in Reprogramming". In: Cell 176.4.
14

---

Yunyi Shen, Renato Berlinghieri, and Tamara Broderick (2025). marginl Schrödinger Bridges with Iterative Reference Refinement". In: International Conference on Artificial Intel- ligence and Statistics.
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet (2023). "Diffusion Schrödinger Bridge Matching". In: Conference on Neural Information Processing Systems.
Sidak Pal Singh, Andreas Hug. Aymeric Dieuleveut, and Martin Jaggi (2020). "Context Mover's Distance & Barycenters: Optimal Transport of Contexts for Building Representations". In: Inter- national Conference on Artificial Intelligence and Statistics.
Justin Solomon, Fernando de Goes, Gabriel Peyré, Marco Cuturi, Adrian Butscher, Andy Nguyen,
Tao Du, and Leonidas Guibas (July 2015). "Convolutional wasserstein distances: efficient optimal transportation on geometric domains". In: ACM Trans. Graph. 34.4.
Justin Solomon, Raif Rustamov, Leonidas Guibas, and Adrian Butscher (2014). "Wasserstein Prop- agation for Semi-Supervised Learning". In: International Conference on Machine Learning.
Max Sommerfeld, Jörn Schrieber, Yoav Zemel, and Axel Munk (2019). "Optimal Transport: Fast Probabilistic Approximation with Exact Solvers". In: Journal of Machine Learning Research 20.105, pp. 1-23.
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever (2023). "Consistency Models". In: International Conference on Machine Learning.
Yang Song. Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole (2021). "Score-Basse Generative Modeling through Stochastic Differential Equations". In: International Conference on Learning Representations.
Sanvesh Srivastava, Cheng Li, and David B. Dunson (2018). "Scalable Bayes via Barycenter in Wasserstein Space". In: Journal of Machine Learning Research 19.8, pp- 1-35.
Matthew Staib, Sebastian Claici, Justin M Solomon, and Stefanie Jegelka (2017). In: Advances in Neural Information Processing Systems.
Eloi Tanguy, Julie Delon, and Nathaël Gozlan (2024). Computing Barycentres of Measures for Generic Transport Costs. arXiv: 2501 04016 [math.NA].
Alexander Tong, Kilian FATRAS, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Brooks, Guy Wolf, and Yoshua Bengio (2024a). "Improving and generalizing flow-based genera- tive models with minibatch optimal transport". In: Transactions on Machine Learning Research.
Alexander Y. Tong. Nikolay Malkin, Kilian Fatras, Lazar Atanackovic, Yanlei Zhang, Guillaume Huguet, Guy Wolf, and Yoshua Bengio (2024b). "Simulation-Free Schrödinger Bridges via Score and Flow Matching", In: International Conference on Artificial Intelligence and Statistics.
Francisco Vargas, Pierre Thodoroff, Austen Lamacraft, and Neil Lawrence (2021). "Solving Schrödinger Bridges via Maximum Likelihood". In: Entropy 23.9.
Jianbo Ye, Panruo Wu, James Z. Wang, and Jia Li (2017). "Fast Discrete Distribution Clustering Using Wasserstein Barycenter With Sparse Support". In: IEEE Transactions on Signal Processing 65.9, pp. 2317-2332.
15

---

Appendix
The Appendix is structured in the following way. In Appendix A, we define the notation used and provide further background information regarding the problems discussed in the paper. In Appendix B, we provide the proofs of the results stated in the main body. In Appendix C, we describe in more detail how to implement the TreeDSBM algorithm, and also discuss extensions of the methodology for improving convergence speed. In Appendix D, give a full overview of the experimental setups the main body, and include further discussions of the findings. In Appendix E we provide additional experiments that could not be included in the main body due to space constraints. Finally, Appendix F includes the licenses for assets used in this work.
A Background
A.1 Notation
We denote by T tree with an edge length function, T= EV,E, (), or its associated metric space via associating edges e with the line segment of length T = E(e) (see Section 2.3). We let P(C) = P(C((0, 77. )) and P(C+) == (TC ) denote the space of path measures on the interval T) and T respectively. For the interval [0. T. .M denotes the set of Markov measures, and R(Q) denotes the reciprocal class for reference measure Q E M (see Definitions 2.1 and 2.2). For path measure Q P(C), Qr denotes the marginal distribution at time t, and io denotes the diffusion bridge conditional on values at 0, T. We use II Io |0, j,, E P(C) to denote mixtures of bridges
IT.) - x q 1o, (- (-120,
drr). Similarly on the tree T, My denotes the set of Markov measures, and R (Q) denotes the reciprocal class for reference measure Q .My and observed vertices S с V (see Definitions 3.1 and 3.2). For a path measure Q E P(CT), Qu denotes the marginal vertex E V. Qy denotes the induced joint distribution
set of vertices V. and i denotes the bridge process conditional on values at vertices S. We denote mixtures of bridges as = 1s E (CC)), defined similarly to above. Finally, [L denotes the Kullback-Leible divergence between measures ग, म. and H(x)) denotes the differential entropy.
A.2 Optimal transport
We here provide brief overview of optimal transport; for more details we refer to Santambrogio (2015), Peyré and Cuturi (2019).
Monge For two measures No and Pl. the original OT formulation given by Monge (1781) aims to find a map T R + Rd pushing oo onto μı, that is μr = T # #Ho while minimising the total cost of transportation according to a given cost function e: RRd x RRd R.
I clx, (po))d
min #yo m
(16)
Kantorovich While intuitive, the Monge formulation of OT is restrictive in that it does not permit splitting of mass in the transportation. Instead, it is common to consider the more general Kan- torovich formulation (Kantorovich, 1942), which searches over joint distributions я E P(Rd Rd) with the correct marginals,
Jf dx,, (x.y)) y).
min K: Po p
(17)
Entropic regularisation
The regropy-regularisee OT problem adds an entropic penalty term to the Kantorovich objective, which smooths the resulting transport plan.
(l (l..))d(( y) (0-(HH(m))
min
The resulting problem has many desirable properties; it enables differentiability with respect to the inputs, relaxes constraints on the corresponding potentials (Genevay et al., 2016), and for discrete measures enables efficient computation using Sinkhorn's 's algorithm (Cuturi, 2013).
16

---

A.3 Multi-marginal optimal transport
One can generalise the above optimal transport problem to the case of multiple marginal distribu- tions. For cost function : (Rd)a=1 I + R and an 'observed' set S C {0..... (0......j, multi marginal op- timal transport (mmOT) searches over joint distributions (P((Rd)) matching the prescribed marginals on S to minimise the objective,
min I::s, 14 ViES (Ray++
c(Fo.))dT(m..))
(18)
As in the standard case, one can define an retropy-regulariiedd version of the mmOT problem.
(nreyee ((20=))d((=0)) - =l((T)
min TH ViES
(19)
Multi-marginal Schrödinger bridges We emphasise here the distinction between the multi- marginal OT and corresponding SB problem that we consider in this work, compared to the multi- marginal SB problem that has been considered recently in works such as Chen et al. (2019), Chen et al. (2023), and Shen et al. (2025) to name a few. These lines works aim to find processes that evolve in time while fitting to known marginals at several different timepoints, and are motivated by inferring population dynamics from snapshot data. We rather consider costs defined according to a tree structure, in which the marginals some of the vertices may be known (as considered in Haasler et al. (2021) and Noble et al. (2023)), with a particular focus on applications to Wasserstein barycentre problems.
A.4
Wasserstein barycentres
Given l measures (Hi..... pe) and weights (A1,..,, de) summing to I, the Wasserstein-2 barycentre (Agueh and Carlier, 2011) is defined as
V = arg min Σ (,,
(20)
(1t, 2)) denotes the minimum attained by the OT solution in (17) for quadratic cost func- tion c(x, y) = 2 v1l.. Observe that this can be cast as an mmOT problem by using star-shaped cost function c(xo:e)
where (2,,
Ei= >e - 2=1lž where the marginals are prescribed on the leaf nodes S S= (1...., l) and the centre vertex O is an unobserved measure (which at the solution is the barycen- tre). The Wasserstein barycentre is widely studied due to its importance in applications, including in Bayesian learning (Srivastava et al., 2018), clustering (Ye et al., 2017), and representation learning (Singh et al., 2020) to name a few.
Types of regularised barycentre There are many different ways to define an entropy-regularisee formulation of the Wasserstein barycentre problem. Some formulations add inner regularisation, which replaces the Wasserstein distances with an entropy-regularrseed version, while others also consider outer regularisation in which an entropic penalty on the barycentre is added. The (e, T)) rogula-regularised barycentre of Chizat (2023) unifies many of these problems, and aims to min- imise the objective
min Σ A. ())
(21)
where W (P;. v/ = m (J || y1l/d y) + cKL(T]| @ v)}. We refer to Chizat (2023) for more details regarding the different types of entropy regulay-regulariied Wasserstein barycentres.
Tree-structured costs
The case of tree-structured costs for multi-marginal optimal transport is often specifically studied in the literature (Haasler et al., 2021; Noble et al., 2023), as it recovers Wasserstein barycentres as special case as described above, and also arises the Wasserstein propagation problem (Solomon et al., 2014, 2015). In the discrete setting, Haasler et al. (2021) show that the tree structure can be leveraged to design an efficient Sinkhorn-based algorithm.
The metric space T Recall from Section 2.3 that we identify tree T (with vertices, edge set, and length function (V,E, f)) with a metric space, by associating each edge e with the interval [o, f(e)), which are connected according to the tree structure. The same construction is used in Noble et al. (2023) for defining the dynamic TreeSB problem that we study in this work. For a rigorous description of such constructions, see for example Hambly and Lyons (2008), Bolin et al. (2024).
17

---

B Proofs
In this section, we give proofs of the results stated in the main paper. We first provide the proof of the TreeSB characterisation Theorem 3.1. We will follow the presentation and proof techniques of Léonard (2013), making the appropriate changes to extend to the tree-structured case. We then prove the convergence of the TreeIMF procedure stated in Theorem 3.1, following the presentation and proof techniques of Shi et al. (2023).
B.1 Existence and uniqueness of TreeSB solution
We first provide results pertaining to the existence and uniqueness of the TreeSB solution. Note that while we follow the presentation of Léonard (2013), one could instead rely on the existence results in Noble al. (2023).
Let us define another static minimisation problem only over observed vertices in S as
II'S R
arg min (II QQs) I1, = #4 Vi S}.
(TreeSB stat)
(EP((()))1))
We note here that the reference measures that we consider in the various SB problems are un- bounded, we consider the reference Brownian motions to be in stationarity. We will make of properties of the KI. divergence defined relative to these measures; see the Appendix Léonard (2013) for a justification of why these properties still hold when the reference measures are un- bounded.
We first provide a result detailing the relationships between the dynamic and static tree-structured SB problems.
Proposition B.1 (Compare to Proposition 2.3 in Léonard (2013)). The tree-structured Schrödinger Bridge problems (TreeSBaynk, (TreeSBetat) and (TreeSB admit at most one solution pSB P(CT), B E (((RR)))) and 1 (((Rd))51) respectively.
PSB solves (TroeSBayn). then P solves (TreeSBata). Conversely, i solves (Tree then (TrecSBaya) solved by mixing Brownian bridges along each edge as P P
Moreover, PSB solves (TreeSBayu). then PSB solves (teesS Conversely, if iT solves (TreeSB stat), then (TreeSBayal is solved by the corresponding mixture of Q-bridges conditioned on values at S, pSB = 11 B. Q
Proof. The first statement follows from the strict convexity of the (TreeSBatat) (TreeSBayn). and (TreeSB stat) problems.
The second follows by using the chain rule for KL divergence, conditioning on the values at the vertices V. We obtain
KL(P|Q) KL + f / KLLP((
(22)
We see that (L > (Tv()v)Q with equality if and only if P(.IXy)) Qv) for Py- a.e. Xy (assuming KL(PIQ) < oo). Therefore, IP solves the dynamic problem if and only if decomposes as a mixture over bridges Q.py according to the coupling Pp solving the static problem, ie. psB 1I y (if this were true, then P would be valid solution with lower KL divergence relative to Q. contradicting optimality of pŠB)) Note that such bridges just consist of Brownian bridges along the individual edges.
The third part follows similarly the second. Consider the KL decomposition but instead condi- tioning only on the values on the observed vertices S,
(T((P/)Q) = (L.(Ps|Q.... + f KL.P((IXS)IQ(((X))))....
(23)
(R-))SS)
Now we have (L(PIIQ > (L( with equality if and only if P(. Xs) = Q(1Xss) for
Ps-a.e Xss (assuming (T.(PIQ) x)). So in particular, we have that pSB must be a mixture of Q-bridges according to its own coupling over S, i.e. psB -PSBQ...
18

---

In the following, we will utilise this equivalence between the (TreeSBayu) and (TreeSBatal prob- lems. We now present an auxiliary result giving a criterion for the tree-structured SB problems to have a solution, in the vein of Lemma 2.4 in Léonard (2013). This is a technical result required to deal with the fact that we are considering an unbounded reference measure Q
Lemma B.2 (Compare Lemma 2.4 in Léonard (2013)). Let B IRd - o0) be measurable function such that
ZB
exp
B-E(())
Qss Qs (dxss) <0;
(24)
(RO)IS)
iES
and each i 8 take Pai E P((R) such that Bddu < oo.
(-oc, oo) (from the previous result). The static and dynamic rree-structured SB problems (TreeSB stat) and (yyeeSBay)) the μa have a (unique) solution if and only (finf[TreeSB: stat) = inf inf (dyy < oo (that is, if and only if the marginals Ny are such that there exists some IT e ((((R))151) satisfying ITŞ = pi for each i E S, and KT.(TTO) ([ss) 00).
Note that inf (TreeSBaalt
(TreSSSaaay
Proof. In light of the equivalence in Proposition B.I, we can consider just the static prob- lem (Taee) Since the marginals pr are tight on Rd it easily follows that the closed constraint T(fu)iss) == {II P((Rd) Sl II, μe Vi S} uniformly tight and thus compact in P((Rd)isi. From the characterisation of KI. divergence with respect to the unbounded measure Qs (see Léonard (2013), Appendix A), we have KT(TT] (T]]] = ( S(Re)IS) Eies ((I.)dI((z)) zB, where Q is the normalised measure Q 2n exp( Ds B)Q.. For II ((sw)ss), we have
[L(UI = KL - E/ Bdp; - ZB-
iES Rd
The lower-semicontiniity of II + QL and the assumption I Bdu, oo for each i E S. together imply that [L(II[[s)) is lower bounded and lower semi-continuous on the compact set ((fi))sss. Thus the static problem (TreeSBes) stat) admits a solution if and only if inf (TreeSBa stat) < xc.
In light of Lemma B.2, we can state the following result which provides conditions for the existence of the (TreeSBaya) solution. Recall we will consider an unnormalised Brownian reference measure, which satisfies Q = Leb for each i E S, so the following result applies with m = Leb.
Proposition B.3 (Compare to Proposition 2.5 in Léonard (2013)). Suppose that Q = m for each i S. for positive measure m. We have the following results.
(a) For (TreeSBss) stat) and (TreeSBaya) to have solution, it is necessary have KL(H]llm) < xo for each i E S.
(b) For sufficient conditions: Suppose there exists measurable functions A, B Rd R200 satisfying
(i) Qs(das) 2 exp
Eies Alx;))dm@s(c
(i) JIRAYIAN exp (ap---En -Eies )
Qs(drs) < co,
(ii) [ya(A+ + B)dda < co for each S.
(iv) KL(Hi](m) < so for each E S.
Then there exists a unique solution to the SB problems (Tr and (TreeSBdyn)
Proof. The first statement (a) follows by applying the result in Lemma B.2 and using the fact that KL Leb) < [I(II)][ oc.
19

---

For the second, we take the independent coupling ITg iiEs Hi, and observe that (under the above assumptions) it satisfies the conditions Lemma B.2. Clearly, has the correct marginals T = 14 by construction, so it remains only to check that KT. Qs) < oc. By expanding the KI. divergence and using the inequalities in (i) and (iii), one sees that the inequalities in (iv) are sufficient to ensure KL [Qs) < xc.
L
B.2 Characterisation of TreeSB solution
We now show the characterisation of the solution to Equation (TreeSBayn) stated in Theorem 3.1, upon which the IMF procedure depends. We first recall that the solution PSB is mixture of bridges of Q, conditional on the values at vertices in S.
Proposition B.4 (Reciprocal process). The solution PSB to (TroeSBayu) (if it exists) in the reciprocal class Rs (Q), i.e. pSB =P PPSQQ.s Q
Proof. This is a consequence of the third part of Proposition B.1.
We now show that for a Markov reference process Q, the (TreeSBaya) solution is also Markov. Fol- lowing Léonard (2013), to present the results we will use the following characterisation of Marko- vianity for path measure P on the tree structure. Consider an edge E È and a time E [0, T*.. Consider the time te as splitting the continuous tree into two distinct sections (note that such a split- ting may not be unique, the chosen point may correspond to vertex; in such cases consider any such split into two distinct parts). Denote the restrictions of a process X to these two parts as X and X= Then we say P is Markov mean that P(X= E ,, X= E (X;) = P(X E (Xx>>> .. ..) for any such split (together with the technical assumption that some time-marginal of IP is a-finite; see discussion in Léonard (2013)). Note that the Brownian reference process considered in
the main paper is Markov.
Proposition B.5 (Markov process). For Markov reference measure Q the solution TPSB to (TreeSBaya) it exists) is Markov.
Proof. The proof follows that of Proposition 2.10 in Léonard (2013). We outline the following changes to the notation for our setting, then the argument follows the same way.
Consider an edge e E E and a time te E J. along with a corresponding split of the tree at time le as described above. We define the following notation: Let CĀ = {w w E Cr} and C be the spaces of continuous paths on the two sections of the tree respectively.
O fw= For path measure P E P,CT), let P = P- = 2) E P(CT) be the measure conditioned on the process X taking value I at the time te, and moreover define its restrictions to the two sections as Pla.* and St. respectively.
Similarly to Léonard (2013) we now make the following claim, from which the result follows.
Claim B.6. Fix a time on the tree le as described above. Fix a = E Rd, a measure μ E P(Rd), and path measures the 'before' and 'after' sections PP E P(C n (Xt. = z}) and 1p e P(CZ n (X,, = =z]). Consider minimising (I over path measures TP E P(CT) constrained to satisfy Pt. 14, pto. Polii aand Phe. ffaii. Then the objective KI(-IIQ) allains its unique
minimum at (P) = Sge plo 8 Sto.z u((z).
Given the claim, the result follows according to the following argument: Suppose for a contradiction that the SB solution, here denoted I,, was not Markov. Then, there exists some time te and cor- respond split of the tree such that P( # Ps (-IX,) & P2 (-IX))) Applying the above claim with μ PP.., we see that 1s* and IP have the same marginals at all time-point on the tree, but P attains strictly lower KL divergence (L |[Q < KL( ()), contradicting the optimality of P.
The proof of the claim uses Jensen's inequality and proceeds exactly as the proof of Claim 2.11 in Léonard (2013), with the appropriate notation changes.
20

---

We now provide the characterisation of the TreeSB solution in Theorem 3.1. While the previous re- sults have been for a general reference measure Q, we present the following results for Q associated to running Brownian motions (a Bi [E[[, along each edge, considered in the main paper.
Theorem 3.1 (TreeSB characterisationn, Under mild assumptions (in the Brownian case, namely that d)p=() < and H(ns) < oo for each i E S), there exists unique solution to the dynamic TreeSB problem (TroeSBayu) The solution is the unique process that is both Markov and in (s( with correct marginals P. = pa i E S.
Proof. As we are considering a Brownian reference process, the assumptions in question are that [(l=:[/dp;(zz) < oo, and H(14) < oo for each i S. One can then verify (as in De Bortoli et al. (2024), Lemma D.2) that the criteria in Proposition B.3 hold by taking functions A and B to be quadratic, from which uniqueness and existence of the solution follow. From Proposition B.4 we have that the solution is in (Q)), and from Proposition B.5 we have that it is Markov.
We now need to show that if a measure pPD is Markov and in R&((1) and has the correct marginals Pi = μi for E S, then it is the TreeSB solution. Note first that as 1d Markov and reciprocal, its restriction to each edge rem 2.14 in Léonard et al. (2014) (noting that the required criterion holds for the Brownian reference ar measure) we have that ~., (x.))) Qe-a.e. for some non-negative measurable functions fe. Recall that the path measures are composition of the path measures along each edge ILvv fi()),, Q-a.e. for some non-negative measurable functions f. (via relabelling of the functions). Note too that po is in ( so we can also express the Radon-Nikody. derivative as aQ h((X,)es) for some non non-negative measur- able function h. Equating the two expressions, we see that we in fact must have a decomposition only over vertices in S, 47 anr
(u, v) is also Markov and reciprocal along that edge. Thus, by Theo-
cording to the tree structure, so this means that dF ay
d=d
= Ilies (i(X,), Q-a.e. for some non-negative measurable functions (where IT denotes the static coupling of TPO over the vertices v\).
aq aQv
=
The remainder follows the standard argument characterising the SB solution using the decomposi- tion according to potentials (see e.g. Nutz. (2021)). Consider static couplings in the constraint set II E T(f4hes) == {II E P((R))M : II, = f4 Vi E S} such that [L[II[[v) < co. By the above, we have that En (log( doy choice of II (for precise statement taking care regarding the integrability of the potentials, follow the argument Proposition 2.17 in Nutz (2021)). Therefore, for any such we have
Exes f log fidus, which is in particular independent of the
KL(TIQQV) 2 (L( KL(IIIIIT)
(25)
dIT") En [log( aev
(26)
am ETT, [log( बेदy
(27)
= |I.
(28)
and we thus see that 11 is the minimiser of the (TreeSB.sta) problem.
B.3 Properties of the tree-structured projections
We now move on to proving properties of the TreeIMF procedure. We follow the presentation of Shi et al. (2023). We begin by proving the properties the reciprocal and Markovian projections defined in Definitions 3.3 and 3.4. For a full set of required assumptions, see the assumptions A. A.2 and A.3 in Shi et al. (2023) Appendix C.2, which are standard in the literature and we assume to hold along each edge.
The following result follows similarly to the standard IMF case.
Proposition 3.3. For P E P(CT), the reciprocal projection II" = projoss() (P) solves the minimi- sation problem I1" arg (IIP)III..
21

---

Proof. This follows from the KL decomposition used in Proposition 1.1, conditioning on the values on the observed vertices S:
[L (L(Ps
(RR9)1S)
KL XI
(29)
Given we are optimising over the reciprocal class (Q). we have that the bridges (T- Xs) QQ- (s) fixed. Thus, the minimiser is achieved by taking ITs Ps, that is IT projRs(()) (P))
To prove subsequent results, require the following decomposition of the KL divergence according to the tree structure. Note that we consider Markov and reciprocal processes on the tree, both of which factorise according to the tree structure the following decomposition be applied.
Lemma 3.2 (KL decomposition along tree). Take path measures P.
I that share a marginal the root, P, Py, and factorise along the tree. Under mild assumptions, we have the KL decomposition
KL(PIP)-
E (u,))E,
Ex.=P (L(P(u.)) (-Xx.) ([X (1x.)))
(12)
Proof. This a consequence of the iterative application of the chain rule for KL. divergence applied according to the tree structure, and the conditional independence caused by the tree structure. By first applying the chain rule for the KL divergence conditional on the value at the root vertex r, we have
KLL(PIPP) = (T (P,)) + EX, KT (I
(30)
Ex Ex [I((XX)))))(((((X,)]]
(31)
Recall that the edge set is depth-wise ordered. We can again apply the KL chain rule to the term inside the expectation, now conditioned on the process along the first edge 61.
(L (-) (-1x,)
+ Ex -P (1x.)) (L (X)I, (32)
We can iteratively apply similar decompositions as we traverse the edges according to the ordered edge set.
KL (P( Xei, X, X.) X ..., X x,)))
KL (-|X*t X (x) Pek- (-|Xet "..,,,.."""",,, X X,))
EExeSsweetH x.)) ( X X,) (-) Xei Xx))]
(33)
= -KL( 11 (-1
Excett (1
KL (P(-- X X-) X X,))]
(34)
where in the second line the factorisation property along the tree structure (here, s(e) denotes the starting vertex of an edge e). Applying such decompositions for each edge in the ordered edge set &, one obtains
KL.(PIIP)
x e pp 161x* x [Pow (Pe (-1x
(35)
Xey ((P(()) x
Σ Ex KT((P(u.)) (-1x,, ((c.)) ((..)))))))))) (u.))E,
(36)
as required.
22

---

We now consider the Markov projection defined in Definition 3.4. Note that the restriction of the TreeSB solution each edge itself an SB (because is Markov and reciprocal), and thus can be associated with an SDE (Dai Pra, 1991). In the definition of the Markov class in Definition 3.1 we define the Markov class My via considering SDEs with locally Lipschitz, drifts. The restriction to locally Lipschitz drifts is technical requirement for applying the entropic version of Girsanov's theorem; this is standard in the literature and does not affect our methodology. Following Shi et al. (2023), we now provide result showing that the Markov projection also solves minimisation problem.
Proposition 3.4. Under mild assumptions, the Markovian projection M projM(1l) solves the minimisation problem MI = arg miDMeMy MI).
Proof. Applying the KL decomposition in Lemma 3.2, we have
KL =
Σ (u.0)EE
Ex,
(L (M
(+
(37)
We now analyse the individual KL expressions KI (L (I along each edge (u, v), using the proof techniques of Proposition 2 in Shi et al. (2023).
In particular, applying the argument in the proof of Proposition 2 in Shi et al. (2023), one sees that each conditional process Io is Markov and can be associated with (Xf) (xE)T( given
X; = = /
/ Eug V log Qo (XG (X..
Xilds + f dB;
(38)
Therefore, letting the restriction of the Markov process to edge (denoted above as M be associ- ated with process dY; (t, ) + odB; such that (I (-IX (M (-IX,) )) < co, with v. locally Lipschitz, then (using e.g. Léonard (2012), Theorem 2.3) one obtains
1 20 /
(I.
Eng [IIo En Eng V log (x (x;x [x;) X
(t. x;,1 (39)
Thus, substituting back into (37) we have
Σ Ex.-H.
(40)
c=(E))E
Tc Engo
Σ
Ex, Ex 20
10 o,c
V log (XF - (Xf (xf 1]
dt
339
(41)
Te Ex,.X
Σ 339
|102 IIo' po.
log (xF1x))))) X5, ve 1) dt.
20
(42)
This expression is minimised by taking (t, o [V log (x+11)) = x] along each edge E. This corresponds to performing Markovian projection along each edge ac- cording to the coupling II", which is exactly the definition of the tree-based Markovian projection in Definition 3.4.
We also note that, as an instance of bridge matching, along each edge the process II; and its cor- responding Markovian projection MI satisfy the same Fokker-Planck equation (Peluchetti (2022), Theorem 2). Thus by the uniqueness of the solutions of the Fokker- Planck equations under A. and A.3 in Shi et al. (2023) (see e.g. Bogachev et al. (2021)), they share the same marginals M. 11f
23

---

B.4 TreeIMF convergence
We follow the presentation convergence Shi al. (2023), but with the appropriate modifications to the proofs for the tree-based setting.
Lemma B.7 (Pythagorean property compare to Shi et al. (2023), Lemma 6). Take Markovian
process MI E MT and reciprocal process (Q)... Under mild assumptions, if (II oo then we have
KL(M]|M) KL proj (I)) + (I)))))
(43)
If (MI oo then have
ML( = KL(II] (rojRs(() (L)) (Q)
(44)
Proof. Proof of (43); The first identity follows from algebraic manipulations of expressions for the relevant KL divergences. From the proof of Proposition 3.4, we have
To
MI
Σ 3EE
Emg,
I
|
V log (x))(X Xo] ue ue t) dt
0.a
(45)
(where we suppress the superscript on the X, for notational convenience). Likewise, it can be shown that
(T) = Σ ("
To
1 20
Ex,-
|
V log Qr) - t)12 de.
339
(46)
From another application of the expression in Proposition 3.4, we have
To
KL(II] (I)) = Σ
Io llg (X, Xo]
Ex
2
3EE
o"EnS Eng V Vlog (1)
df
(47)
Σ
Exn.Xinii,,
[I log (Xr Xo
330
Ex
|Bng
[V
log (Xx])((x)) dt,
(48)
where to obtain the second line, we have expanded out the square and taken expectations over Xo in the cross-term.
Using these expressions, by applying the same algebraic manipulations as the proof of Lemma in Shi et al. (2023) to each term in the summations, one obtains
KI[II]] proj (II) + KL (II) MM) = (I
(49)
required.
Proof of (44): The second part also follows similarly to Shi et al. (2023), but instead conditioning on the values at S. Let II projR (z) (IP) PsQ Using the change of measure formula for KL divergence and the fact that 11 and II have the same bridges conditional on S,
dII
KT.(PP|ITT> (T.(PP +
log
(w) (d)())
(50)
CT
dll's (xs))
KL(PP| +
log
(51)
(29)181
all
dil's
-T<(P| +
log(
dIs))
(52)
(20)18)
am
KL(PP KL,
(53)
as required. Note that the second line we have used the fact that IT and IT* have the same bridges conditional on S. and in the third line we have used that Ps 11 by construction.
24

---

Algorithm 2: TreeDSBM for Wasserstein barycentre computation
Input:
Initial coupling 1g over measures fAs (c.g. independent). Number of IMF iterations N,
Entropic regularisation parameter 0.
Construct initial reciprocal process II
IT 2.s. That is, Sample Ys from the initial coupling ITg over the marginals jug X,Y,,
Sample from the unknown central marginal as Yo
Training samples are obtained from Brownian bridges along each edge between Y; and Yo.
for n E ,O...., 1} do
Learn 2 El vector fields using bridge mathhen loss (15) along each edge, using samples from current reciprocal process IЛ", to obtain Markovian process M
Construct next reciprocal process I = M' 2 using samples from MG That is,
Simulate M|*+1 starting from a chosen root r E S (one of the known measures jr) to obtain
samples Ys from MSS
Sample from the unknown central marginal as Y, ~ (E Id);;
Training samples are obtained from Brownian bridges along each edge between Y; and Yo-
end
We can finally state the result regarding convergence of the TreeIMF iterates to the TreeSB solution.
Theorem 3.5 (Convergence of TreeIMF). Under mild conditions, the TreeSBayn solution 1' the unique fixed point of the TreeIMF ilerates P", and we have limn- (L(PP" [IP) 0.
Proof. In light of the Pythagorean property in Lemma B.7, this follows from the same compactness argument as in Proposition and Theorem 8 in Shi et (2023).
C Implementation details and extensions
Algorithm for barycentre setting In Algorithm we provide a simplified and more explicit version of Algorithm for the case of computing Wasserstein barycentres.
Illustration of TreeDSBM The TreeDSBM procedure for Wasserstein barycentre computation (that is, for a star-shaped tree) is illustrated Figure 1. We also provide a diagram for a more general tree structure in Figure 4.
C.1 Implementation details and design choices
We implement the TreeDSBM procedure using the JAX framework (Bradbury et al., 2018). Below, outline some of the design considerations for implementing the method.
Vector field parameterisation In our implementation, we use separate neural networks param- eterise the vector fields for each direction along each edge, totalling 2/ networks in total. Note that one could alternatively use a shared network along each edge with an additional binary input indicating the direction; this parameterisation was used in De Bortoli et al. (2024).
Loss function We incorporate both the forwards and backwards losses in Equation (15) into single loss function, and thus optimise the forward and backwards directions simultaneously. One could incorporate all edges into single loss function and train simultaneously, alternatively could parallelise the edge optimisations across devices because the edges are optimised independently. This is strength of approach and can lead large speed-ups in training, training time does not need to increase in proportion to the number of edges.
Simulation As we train both directions along each edge, simulation of the SDEs along the tree structure can initialised from any of the observed nodes S. Following Shi et al. (2023), one can rotate the starting vertex between IMF iterations. This helps to mitigate any drift that accumu- lates the marginals, as the coupling samples will only use true samples from the current starting marginal; for analysis of this, see De Bortoli et al. (2024). In our experiments, often did
25

---

M4
(a) Reciprocal process II: The Ys (x) are sampled from the current coupling IIs over S. Conditional on the Ys. points (x) at marginals V are sampled as Ys-Is Qs-is- Brownian bridges are drawn along the edges between the samples.
(b) Markovianised process P: Vector fields trained by bridge-matching along each edge. Sam- ples (x) from the next coupling IIs are obtained by simulating the resulting SDEs along the tree struc- ture, started at one of the known marginals.
Figure 4: The two stages of the TreeIMF procedure, for a non-star- star-shaped tree structure. On this tree, the marginals at the leaf vertices S {1, 3, 5, 6, 73 (blue) are fixed. The marginals at vertices V (red) are not fixed, and change during the procedure.
not observe noticeable drift in the marginals, and instead generated samples in the coupling by ini- tialising equally across the observed vertices in S. In our sampling, we used an Malur-mma discretisation scheme with 50 uniformly-spaced steps.
Initial coupling Unless otherwise stated, used the independent coupling I1 JES M4 as the initial coupling. We note that any coupling over V with correct marginals on could be used, and we discuss some possible alternatives in subsequent sections.
Architectures Other than image experiments, we use basic MLP-based vector field model. consists of an MLP spatial embedding with hidden layers [128, 256] to embed into dimension 32, a time embedding consisting of a sine positional-encoding and an MLP with hidden layers [128, 256] also embedding into dimension 32, before concatenating the embeddings and passing through an MLP with hidden layers [512, 256, 128]. This is the same architecture used in De Bortoli et al. (2021) and Noble et al. (2023). In image experiments, we use the UNet architecture (Ronneberger et al., 2015) with the improvements from Dhariwal and Nichol (2021), using the JAX implementation from Song al. (2023). For all experiments, with default parameters 0.9 and 0.999.
the Adam optimiser (Kingma and Ba, 2015)
We note that for pointcloud experiments, capacity of the neural networks is not a limiting fac- tor-any sufficienll--eppressii
network will work similarly well and fairly small MLP networks suffice, so using several networks for the different edges does not pose issues. One could also use a single network across the edges and additionally condition on the edge. This makes particular sense for problems with shared structure between edges, such as those for image data (and indeed these are settings for which the networks would be larger, and maintaining multiple networks could become a computational bottleneck).
Memory requirements The parallel nature of the TreeDSBM algorithm during training provides practitioners with a trade-off between memory consumption and wall-clock time. Namely, one can train the edges simultancously if compute allows (either on a single GPU if enough memory, or parallelised across GPUs). If this cannot be done, one can train sequentially instead (in which case memory requirements for training would be comparable with standard bridge-matching). We report GPU consumption for different experiments in Table 5 (for sequential and joint training). for the hyperparameters used in the paper. These can of course be changed significantly by changing hyperparameters such as batch size. We note that in some of these experiments (e.g. the Gaussian experiments), peak GPU usage is due to the simulation and storage of the training samples for subsequent IMF steps, rather than during the network training. This can be reduced significantly by simulating in smaller batches, or by updating the cache during training rather than simulating all beforehand.
26

---

Table 5: Comparing memory usage (in MB) for sequential and joint training, for the hyperparame- ters used.
Sequential Training (MB)
Joint Training (MB)
2d
447
687
Data Aggregation (Poisson)
705
1219
Gaussian (d = 64)
1239
1239
MNIST 2,4,6
4683
6407
We remark that TreeDSBM has improved memory requirements compared to TreeDSB. In TreeDSB, training time-reversal along edge requires saving entire trajectories simulated along the reverse direction, whereas TreeDSBM only requires storing endpoints (this is one of the most significant benefits of IMF over IPF).
Alternative reference measures We have presented the methodology according to using Brow- nian reference measures along each edge. However, the methodology extends to other reference process, as long as the respective bridging processes conditioned on the values at S are tractable (for example, this is true for Prnstein-Uhlenbee transport is however less simple beyond the Brownian case. See Shi et al. (2023) for more detailed treatment of this general case.
processes). The connection to quadratic-cost optimal
C.2 Methodology extensions for improving convergence speed
We now discuss possible extensions of the TreeDSBM algorithm improve convergence speed, in the vein of existing extensions of Schrödinger bridge methodology for the standard two-marginal setting.
Warmstarting with minibatch mmOT couplings The TreeIMF procedure can be initialised with any coupling ITg over S with correct marginals. For the standard SB problem, Tong er al. (2024b) note that the SB solution is a mixture of bridges mixed by static e-OT solution, and thus pro- pose a single iteration of bridge matching on samples generated by a static e-OT solver applied on minibatches (sce also Pooladian et al. (2023), Tong et al. (2024a), and Fatras et al. (2021)).
Similarly, one can initialise TreeIMF using samples obtained from static mmOT solvers applied to minibatches. Such procedure can speed up convergence to the TreeSB solution by initialising closer to the true solution. We remark, however, that such minibatching approaches can incur large enrors (particularly in higher dimensions or for small minibatches; for example the Wasserstein-1 error grows O(B-1/(24) (Sommerfeld et al., 2019)), and so the advantages of such methods are lessened dimensionality increases.
Flow-based IMF on the tree Iterative Markovian Fitting presents mathematically elegant ap- proach for solving the SB problem, with significant practical improvements over the IPF procedure. However, the iterative nature of the algorithm remains a downside each iteration of the two-step procedure involves first simulating the current Markovian process, and then retraining a neural net- work with the wridge-matching loss. In practice, one might wonder if it is possible to perform the simulations and bridge-matching procedure simultaneously to avoid the expensive iterative nature of the algorithm. The recent work of De Bortoli et al. (2024) answered this in the affirmative and propose the a-IMF procedure, which instead corresponds to discretisation of a continuous flow of processes that converge the SB. We anticipate one could design an analogous methodology for the tree-based setting that we consider; we leave such extensions for future work.
D Experimental details
Here, we provide experimental details and additional results and discussion regarding the experi- ments included in the main body.
D.1 Synthetic 2d barycentre
Datasets We follow the experimental setup of Noble et al. (2023). The marginals consist of moon, circle, and spiral datasets from scikit-learn (Pedregosa et al., 201 I), centred and scaled by a factor of 7.0. We aim to learn the ($. 3.)--bayceentrre of the dataset. This is a challenging problem- the
27

---

IMF O
C-011
10
Grownd Truth
IMF 1
15
10
IMF 2
15
Figure 5: TreMDSBM samples in the 2d experiment, com- paring different regularisation values E.
Figure
Progression of
TreeDSBM (e 0.1) barycen-
tre approximations through the IMF iterations.
Table 6: Progression of the Sinkhorn divergence to the ground truth, and the average Sinkhorn divergence to the marginals, during the IMF iterations (tea over 5 runs).
IMF
IMF 2
IMF 3
IMF 4
IMF 5
IMF 6
Fit to solution
17.7 + 0.2
1.28 + 0.04 1.21 + 0.07 1.13 = 0.05 1.12 + 0.05
1.09 + 0.06 0.20 0.02
Fit to marginals
0.16 + 0.03
0.17 + 0.03
0.18 0.04 0.19 = 0.03 0.19 + 0.03
lower-dimensional and discontinuous structures in the marginals complex and fragmented support, and the transport maps to the barycentre are highly discontinuous.
that the barycentre has
Hyperparameters For TreeDSBM, we use 0.1 and run for IMF iterations. For training the vector fields, we use 10,000 training steps and a batch size of 4096. We use the Adam optimiser (with default parameters 0.9, 0.999) with learning rate le-3 and exponential moving average parameter of 0.99. At inference we use the Fuler-Maruyana. scheme with 50 steps. We generate a batch of 10,000 training couplings for subsequent TreelMF iterations (a third simulated from each marginal). For the other algorithms, we use their default parameters provided in their respective codebases.
Comparison with alternative algorithms For the ground-truth, we use the in-sample free-support barycentre algorithm of Cuturi and Doucet (2014) implemented in Python Optimal Transport (Fla- mary et al., 2021), with 1500 datapoints. Note that the aim of this experiment not to outperform in-sample methods; it is known that such approaches perform well in low dimensions, but do not scale well as dimension increases. The in-sample method is used here to provide a close approxima- tion the ground-truth, allowing us to judge the success of the continuous Wasserstein-2 barycentre approaches that we compare.
For TreeDSB, we use the checkpoints provided by Noble al. (2023) which were trained for 50 IPF iterations.
We also report results for the WIN algorithm from Korotin al. (2022). This an iterative algorithm inspired by Álvarez-Esteban al. (2016); it pushes forward source latent distribution, P through a function G to give generative model y = G#p for the barycentre, and also learns maps T: and T, transporting from the generated barycentre to and from the marginals respectively. In our experiments, the barycentre generator were the maps T #4: We hypothesise that the neural maps struggle to model the discontinuous transports well. However, the combined map (2, X,T,) #v was able to give good approximation of the true barycentre, which are the results we report in Figure 2 and Table 2.
G# p was unable to fit the true barycentre accurately,
Additionally, we applied the W2CB (Korotin et al., 2021) and NOTWB (Kolesov et al., 2024a) algorithms to this example, but were unable find hyperparameters to make the algorithms to converge to the correct solution. Again, this may be due to the difficulty modelling the discontinuous transports with neural networks. We also anticipate that the loss landscapes caused by this complex
28

---

4
4
4
2
44
4
4 9
4 6 4
4
2
8
0
7
4 4
(a) Barycentre samples,
(b)
Samples
as
(c) Barycentre samples, 0.001.
0.02.
Σ, A.Y. = 0.02.
Figure 7: (a) The true TreeSB solution will have noise present due to reguliarrgullarisationn If one wishes to reduce this, one can instead (b) construct samples as Σ, A;Y, or (c) use a smaller e value.
example may have caused these methods to get stuck in local minima. appears that the iterative schemes of TreeDSBM, TreeDSB, and WIN aid in overcoming such issues.
Convergence speed In Figure 6, we demonstrate the progression of the TreeDSBM barycentre approximation as we run the IMF iterations. We run with 0.1, and plot samples generated from leaf O. In Table 6, also show how the Sinkhorn divergence to the ground truth evolves as the IMF iterations progress. After only two IMF iterations, TrecDSBM already gives a good approximation to the solution. Such behaviour reflects similar results reported in Lindheim (2023), which observes that iterative fixed-point approaches (Álvarez-Esteban et al., 2016) exhibit very fast convergence to the solution. TreeDSBM performs only a single bridge-matching iteration along each edge before updating the barycentre, rather than computing full OT maps (which would be expensive). As such, it strikes a good balance between the efficiency of iterative sixed-point-bee based approaches, without requiring full OT map computations before updating the barycentre approximation.
Runtime analysis We report approximate runtimes for the three methods that converged, All experiments were ran on a single Nvidia GeForce RTX 2080Ti GPU.
Our TreeDSBM implementation took approximately 1 minute for each IMF iteration when training the edges jointly, and took around minutes to run the 6 IMF iterations. Note that one could also obtain good results using fewer IMF iterations fewer training steps.
In contrast, the alternative methods were significantly slower. WIN required around 8000 training steps to obtain good barycentre approximation, which took approximately hour 20 minutes. The provided checkpoint for TreeDSB is for 50 IPF iterations, each of which would require 6 time- reversal training procedures, and would thus take significantly longer train than TreeDSBM.
Additional results In Figure 5, we plot TreeDSBM samples obtained for different values of rngropy-regullaiisation E. As expected, increasing leads to slight blurring bias in the solution. In Table we also report the quality of fit to the marginals, calculated by simulating from the 'moon' marginal to the centre and then out to the other leaf nodes, and averaging the resulting Sinkhorn divergences to ground-truth samples from these marginals. We see that in this experiment there is negligible drift accumulation in the marginals.
We overall found TreeDSBM to perform strongest in this experiment. Its bridge-matching objectives and iterative nature provide fast and stable training in what complex and challenging problem setting, and its dynamic-transport approach means that it is able to model the discontinuous transport maps accurately.
D.2 MNIST 2,4,6 barycentre
We also follow the experimental setup of Noble et al. (2023) and compute the (1 3: barycentre of the 2,4 and 6 digits in the MNIST dataset (LeCun et al., 2010).
Hyperparameters We use the UNet architecture of Song et al. (2023), with 64 channels, channel multiples of (1,2,2), attention at layers (16,8), and residual blocks at each layer. We train each bridge-matching procedure with 10,000 steps with batch size 64, at learning rate of 1e-4 and with
29

---

exponential moving average weight of 0.999. For subsequent IMF iterations, we simulate 8,192 coupling samples. We run for 4 IMF iterations, beyond which we did not see much change between iterations. The TreeDSB samples plotted in Section 4 are those displayed in Noble et al. (2023), as we did not have the computational resources to run TreeDSB to convergence in this setting.
Role of = In Section 4, the TreeDSBM barycentre has some noise present in the samples. We remark here that this noise should be present in an accurately computed solution, as we are solving for the e-TreeSB solution which adds entropic regularisation. If one wishes for less noise in the samples, one can run with a smaller antropy-regularinn samples by sampling Y, according to the learned coupling, and weighting them Σ; X;Y; (though note that would not be clear what kind of barycentre such samples would be from). We plot results for a smaller 0.001, and for the weighted coupling samples Σ; ArY; for 0.02 in Figure 7, to demonstrate that TreeDSBM can generate samples with minimal noise if so desired.
value e. Alternatively, one could generate
Convergence speed We provide an analysis of the convergence speed of TreeDSBM in terms of the number of IMF iterations required. Unfortunately there is no ground-truti barycentre to compare to in this example; this makes quantitative evaluation of the obtained barycentre difficult. We therefore instead assess the fit to the marginals (from which we have true samples) as a proxy for the success of the algorithm, along with the transport cost - the transport cost provides an indication of the optimality of the maps, while if the marginals are not fitted accurately, then the resulting barycentre will be unreliable. While not an ideal measure of the 'quality" of the barycentre itself, this does provide a quantitative and, importantly tractable proxy for the 'success' of the algorithm.
To this end, we report in Table 7 the transport cost and FID values for samples from the marginals, for the 4 IMF iterations (note that we train classifier and use the obtained features for the FID calculation, so these values should not be compared with those in other works). We initialise the sampling from 1000 unseen test samples of the digit 6, and report the FID values of the obtained 2s and 4s (averaged). We see that, as expected, the transport cost decreases as the IMF iterations proceed, indicating that the barycentre approximation is improving. We also observe that the FID scores increase slightly (though there is little visible difference) this is a consequence of the drift that can accrue in the marginals, and is consistent with the expected behaviour of DSBM from which this effect is inherited (this can be reduced by training for longer or using the techniques discussed in Appendix C.I).
Table 7: Progression of the total transport cost and fit the marginals (as measured by FID), during the IMF iterations. Note that the FID values are obtained using a trained classifier, so should not be compared to values in other works.
IMF 1
IMF 2
IMF 3
IMF 4
Transport cost
431
402
389
378
Ave. FID (2s and 4s)
61
82
89
93
D.3 Subset posterior aggregation
The previous experiments have shown TreeDSBM to improve over its IPF counterpart TreeDSB. In the following experiments, we provide a more detailed comparison with current groongly-peffomen methods for continuous Wasserstein-2 barycentre estimation, by reporting results for standard ex- periments in the literature. We include comparisons against the methods WIN (Korotin et al.. 2022), W2CB (Korotin et al., 2021), and the recent method NOTWB (Kolesov et al., 2024a). These meth- ods have demonstrated strong empirical performance in their respective works, and are chosen here to be representative of different approaches the literature- WIN is an iterative method inspired by Divarez-Esteban et al. (2016), W2CB an Input Convex Neural Network-basee approach (Amos et al., 2017; Makkuva et al., 2020), and NOTWB is based on recent Neural OT methodology (Korotin et al., 2023). We use the implementations links in Section F. We note that, as ever in barycentre studies, it is somewhat challenging to assess the performance of solvers due to the lack of the ground-truth solution (other than in certain specific examples). Here, we report results on standard experiments used in the literature.
their publicly available code, to which we provide the
30

---

Experimental setup We work with the experimental setup and dataset used in Korotin et al. (2021) (the same dataset was also used previously Li et al. (2020) and Fan et al. (2021)), which uses Poisson and negative-binomia regressions on a bike-rental dataset (Fanaee-TT 2013). The aim is to predict the hourly number of bike rentals using features including day of the week, weather conditions, and more. The dataset is 8-dimensional and is split into 5 distinct subsets each of size 100,000. The "ground-truth' barycentre consists 100,000 samples from the full dataset posterior.
Following the literature. we report the BW2 -VP metric between the "ground-truth and the ob- tained samples in Table 3. For methods that generate from each marginal, we report the average over generations from each marginal, and for WIN we report results for the barycentre generator. The BW3-UVe metric is defined
BW 2) %,
BW3 -UP(( 100
(54)
2 Var(D)
where the Bures- Wasserstein metric is defined as BW D) = (W (m). Ex)) for the respective means and covariances of the distributions.
Hyperparameters For TreeDSBM, we use = 0.001 and run for 4 IMF iterations. For training the vector fields, we use 2000 training steps and batch size of 4096. We use the Adam optimiser (with default parameters 0.9, 0.999) with learning rate le-3 and exponential moving average pa- rameter of 0.99. At inference we use the Fuler-Maruyama scheme with 50 steps. We generate a batch of 50,000 training couplings for subsequent TreeIMF iterations (10,000 simulated from each marginal).
For W2CB, we use learning rate of le-4 in the negative binomial setting. For WIN, in the negative binomial case we rescale the source z-mmer by a factor of 10.0 to match the scale of the data better; without this, it did not appear to converge. We run W2CB and WIN for 10000 training iterations, and NOTWB for 2500 iterations. Other than those mentioned, we use the default parameters provided in the respective codebases.
Runtime analysis: We report approximate runtimes for the different approaches; all experiments were ran single Nvidia GeForce RTX 2080Ti GPU. To compare approximate time taken, report time taken to for the methods to converge close to their final output cutput-chosen by monitoring the -W3-UVP metric and choosing the time beyond which it no longer decreases significantly (note these are not the amount time used in Table 3, which trained using the hyperparameters described above). This is somewhat subjective, but is a fairer comparison than just reporting times for running with default parameters. We remark that it may be possible to improve these runtimes with further hyperparameter tuning and by optimising the algorithm implementations, but investigating such optimisations is beyond the scope of this work.
The W2CB algorithm appeared to give good results after approximately 1000 training steps in both cases, which took around 10 minutes in our experiments. WIN converged after around 2500 itera- tions, which took around 45 minutes. NOTWB converged quickly after only around 200 iterations, which took around 2 minutes.
In both experiments, each TreelMF iteration for our TreeDSBM implementation took approximately 20 seconds when training the edges jointly. TreeDSBM converged well using 4 IMF iterations, and training took around 2 minutes 30 seconds (including time for simulating training samples for the next iteration). This is comparable with NOTWB, the fastest of the alternative methods.
Convergence speed We provide the values of the BW2 BW2-UVP metric as the IMF iterations progress in Table 8, and again we observe very fast convergence.
Table 8: Progression of the BW UVP metric during the IMF iterations, for the subset posterior aggregation experiment (sed over 5 runs).
IMF
IMF 2
IMF
IMF 4
Poisson
31.1 - 0.03
0.0085 + 0.0003
0.0075 + 0.0006
0.0076 + 0,0005
Negative Binomial
31.0 = 0.01
0.0123 = 0.0003
0.0118 + 0.0007
0.0121 0.0004
31

---

D.4 Higher-dimensional Gaussian experiments
Experimental setup We follow the experimental setup previously used (Korotin et al., 2021, 2022: Kolesov al., 2024a), in which Gaussian distributions and its ground-truth (3 1 3 })- barycentre are randomly generated, for each dimension in {64, 96, 128). We report results for 2 (which measures the quality of the overall generated barycentre) as described above, and additionally report the I,2 UVP metric, which measures the quality of the individual maps to the barycentre and is defined for each marginal as
1 T
L L' = 100
%
(55)
Var
where 7 denotes the learned map from the marginal to the barycentre, and T is the known ground truth mapping. Again, we provide the averages over the marginals.
Hyperparameterss For TreeDSBM, we use E = =le-4 and run for IMF iterations. For training the vector fields, we use 10,000 training steps and a batch size of 4096. We use learning rate 1c-3 and exponential moving average parameter 0.99. We generate batch of 50,000 training couplings for subsequent TreeIMF iterations (simulated equally from each marginal). Results reported for TreeDSBM are the average over runs.
For the alternative methods, we run W2CB with learning rate le-4, and otherwise use the default parameters provided in their respective codebases.
Convergence speed We provide the values of the -VP progress in Table 8, and again we observe very fast convergence.
UVP metric as the IMF iterations
Table 9: Progression of the UVP and -VP metries during the IMF iterations, for the Gaussian d = 64 experiment (st, over runs).
IMF 1
IMF 2
IMF 3
IMF 4
RW? UVP
16.0 0.01
0.12 + 0.01
0.13 + 0.02
0.14 0.03
-VP
16.7 + 0.01
1.19 + 0.02
1.18 + 0.02
1.18 + 0.03
Discussion of continuous Wasserstein-2 barycentre solver comparisons
Our experiments show that our TreeDSBM algorithm exhibits strong performance, and is competi- tive against -ft--o-tthhhrrrr methods for continuous Wasserstein-2 barycentre estimation in range of settings. In particular, TreeDSBM offers fast and stable training- even in complex settings due to its bridge-matn matching loss objectives, and comes with a well-understood theoretical analysis.
The best choice of barycentre algorithm may depend on the specific problem setting. For example, the transport maps exhibit complex behaviour (possibly due to lower-dimensional manifold- structures in the datasets) then the flow-based approach of TreeDSBM will likely perform strongly (such as in the 2d barycentre experiment in Section 4). Also, when fast training is required then our experiments suggest TreeDSBM is strong option. On the other hand, if fast inference is important then a one-step-generation solver such as NOTWB might be preferable. Note that one could incorporate distillation techniques from the flow-matching literature for improving the speed of TreeDSBM inference after training.
Overall, DSM offers compelling new addition the taxonomy of continuous Wasserstein-2 barycentre solvers, with distinctly different characteristies to alternative approaches due to its flow- based nature.
E Additional experiments
E.1
Further comments regarding computational considerations
Choice of entropy regularisation Choosing the entropy regularisation parameter is a perennial question in entropic OT. Standard methods to choose this value in commonly used OT libraries (for example, choosing in proportion to the costs) provide good guidance for choosing suitable
32

---

Table 10: Effect of entropy regularstggul parameter E in the 2d and data aggregation experiments.
1.0
0.3
0.1
2d, Sinkhorn-divegencce
1.24
0.99
1.02
1e-3
=3e-4
=1e4
Data Aggregation (Poisson), BW2 UVP
0.012
0.008
0.008
Table 1I: Effect of batch size in the 2d and data aggregation experiments.
Batch size
64
246
1024
4096
2d. Sinkhorn-divergence
1.57
1.24
1.04
1.04
Data Aggregation (Poisson), BW? UVP
0.032
0.017
0.013
0.012
values. Typically, one may want to choose E as small as possible the reduce the entropic bias. One advantage of TreeDSBM over TreeDSB is that it allows for much smaller epsilon (TreeDSB does not converge for too-small e, as simulated trajectories struggle reach the other marginals). We provide a visualisation of the role of in the 2d example in Figure 5, and also for two values of for the MNIST experiment in Figure 7. In Table 10, we also add some further quantitative results for different values, in the 2d and subset posterior aggregation settings.
Fitting to the marginals One of the limitations of our approach is that errors can accumulate in the marginals as IMF iterations proceed; this is limitation inherited from standard IMF and similar reflow methods. Standard techniques from the literature can used to mitigate this (such as rotating the starting marginal as in Shi et al. (2023), or using the projection methods in Kim et al. (2025)). It is therefore important that the bridge-matching steps fit the marginals accurately, and hyperparameters should be chosen accordingly. To provide indication of how the learned bridge- matching quality affects the overall solution, we provide results for varying batch size on the 2d and data aggregation experiments in Table 11, for the hyperparameters used in the paper. We have also provided results assessing how the fit to the marginals changes as the IMF iterations progress in the 2d and MNIST experiments in Tables 6 and respectively.
E.2 Ave! Celeba benchmark
In this section, we provide an example that illustrates a potential limitation of our approach. As previously discussed, it is difficult to evaluate performance of barycentre algorithms in high dimen- sions due to the lack a ground-truth. To combat this, Korotin et al. (2022) proposed the Ave, celeba! barycenter benchmark, which consists of 3 distributions of transformed CelebA faces (Liu et al., 2015), for which the ( 2. a Wassersteen-2 barycentre recovers the true CelebA dataset. The resulting dataset consists of around 67k samples in each marginal, and each image is shape 64 x 64 x 3.
We consider applying the TreeDSBM in this example. is known that performing between complex datasets such as images can be challenging, so for the first step we instead pre- train models using single bridge-matching iteration from standard Gaussian to each marginal. To obtain the next coupling for training, we run the process from one of of the marginals to the latent representation in the Gaussian, and then out to the other marginals. This aids in learning, as there is often good structure preserved between the obtained samples from each marginal. For subsequent iterations we also warmstart the parameters from these pretrained models. The experiments were conducted on Nvidia A100 GPUs on Google Colaboratory.
Hyperparameters We use the UNet architecture of Song et al. (2023), with 128 channels, channel multiples of (1,2,2,2), attention at layers (32,16,8), and 4 residual blocks at each layer.
We use a 0.01 and train each bridge-matching procedure with batch size 32, at a learning rate of le-4 and with exponential moving average weight of 0.999. For pretraining, we run for 20,000 training steps (which takes approximately 5 hours), and for subsequent IMF iterations we run for 10,000 steps (which each take around 2.5 hours).
33

---

<=1
k
Leaf
Barycenter
Figure 8: Samples from TreeDSBM applied to the Ave! Celeba benchmark. Many samples are transported well, but some pick up unwanted artifacts. We discuss the findings from this experiment below.
For training subsequent IMF iterations along each edge, we generate samples from the coupling from each datapoint in the corresponding marginal. This mitigates drift in the sample quality the marginals, as we always use true datapoints from the marginal during training. We run for 2 IMF iterations, and did not see much change for subsequent IMF iterations beyond this.
Discussion of results TreeDSBM is able scale to the high-ddmensional setting, but the obtained samples do not match the visual quality of state-of-the-arr results such as those reported in Kolesov al. (2024a). Observe that some of the generated samples are good, but some contain additional artifacts that should not be present. We anticipate that this is due to the initial pretraining cou- pling. When generating samples Y, from the pretraining coupling, we simulate from one marginal to the Gaussian latent, and then out to the other marginals. This often results in strong structural similarities between the obtained Y; which yields good barycentre samples for training the next IMF iteration. However, sometimes the coupling samples Y; do not resemble each other, and the resulting barycentre sample consists of separate overlaid images. This appears to be difficult for the algorithm to recover from, resulting in the artifacts visible in some of the generated images.
Overall this suggests limitation of TreeDSBM for this particular benchmark the true transport maps are in fact very simple in this specific example (primarily just colour changes), but it is diffi- cult for TreM to learn this because communication between the edges is infrequent and only occurs after each IMF iteration. In contrast, state-of-the-ar methods for this benchmark optimise all the maps together and with much more interaction between them, which we anticipate is better inductive bias for the shared structure present in this benehmark. Note that the fact that TreeDSBM optimises the edges separately is in fact a strength of the approach in many settings; it results in sta- ble training without needing adversarial objectives, and allows for speed-ups by training the edges simultaneously. However, this experiment suggests that this may be a limitation of our method in scenarios where the true maps exhibit lot of shared structure (as in the case in this example), as communication between edges occurs too infrequently to recognise this shared structure.
We remark that observed improved performance in this benchmark by using a shared architecture over the edges (conditioning on the edge and the direction), compared to using a different network along each edge. Such architecture makes sense in this example, there are shared features along the edges that the network can learn and this reduces the computational and memory cost. When using shared architecture, the network also appeared to create more consistent samples from the initial coupling. However, there are still unwanted artifacts present in many of the generated samples, and so alternative methods such as Kolesov et al. (2024a) are likely more suitable for settings such as these, as discussed above.
We anticipate that the performance of TreeDSBM in this setting could be improved through architec- tural changes and other implementation tricks from the flow-matching literature (for example, using preconditioned flow parameterisations (Karras et al., 2022), and techniques to mitigate marginal drift in reflow methods (Kim et al., 2025)). Such investigations offer promising directions for future work.
E.3 Beyond star-shaped trees
So far, we have demonstrated the empirical performance of TreeDSBM only on star-shaped trees, as we have focused on computing Wasserstein barycentres. Finally, we demonstrate that our TreeDSBM also works for non-star shaped trees, and thus has potential applications beyond only barycentre computation. We consider simple 2-dimensional example with the same tree structure
34

---

S
Figure 9: TreeDSBM applied to non-star-shaped tree.
shown in the TreeIMF diagram in Figure 4, with standard scikit-learn distributions the ob- served leaves and each edge having length I. As in the previous 2-dimensional barycentre problem, this is challenging task due to the discontinuous transport maps and fragmented supports the solution. We run TreeDSBM for 4 IMF iterations with = 0.1, and train each edge for 10,000 iterations with learning rate 1e-3 and exponential moving average parameter of 0.99. We plot the obtained measures in Figure 9, and see that TrecDSBM is again able to learn the complex mappings required for this setting. While direct applications of mon-star-shaped trees are less clear than the barycentre case, examples have been studied in Haasler et al. (2021) and Solomon et al. (2015), and they could have potential applications for modelling temporal behaviour of population dynamics, for example if populations were known to split according to a known structure. We leave investigating possible applications general tree-structured costs for future work.
Licenses
The following assets were used in this work.
TreeDSB (Noble et al., 2023), MIT License https://github.
WIN, Ave! Celeba dataset (Korotin et al., 2022), MIT License
W2CB (Korotin al., 2021), MIT License
NOTWB (Kolesov 2024a), MIT License
JAX Consistency (Song et al., 2023), Apache-2.0 License
Bike Sharing, UCI Machine Learning Repository (Fanaee-T, 2013), CC BY 4.0 License
https://archive ics. uci. www../.......daaaeaeeeeeeeeeeeee
MNIST digits classification dataset (LeCun et al., 2010), CC BY-SA 3.0 License
OTT-JXX (Cuturi al., 2022), Apache-2.0 License
Python Optimal Transport (Flamary et al., 2021), MIT License
35

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: In abstract and introduction, we state that we extend the IMF procedure to the tree-based SB setting. In Section 3 we demonstrate the theoretical soundness of our proposed approach, and we explain how to implement it in Section 3.3. In Section 4 we demonstrate the empirical performance of our approach, showing inherits the benefits IMF over IPF in the tree-based setting.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discuss the limitations of our approach in Section 5. We also provide an experiment illustrating potential limitation of our approach in Appendix E.
3. Theory assumptions and proof'
Question: For each theoretical result, does the paper provide the full of assumptions and a complete (and correct) proof?
Answer: [Yes]
We present simple statements of the theoretical results in the main text, with full assumptions and proof's included in Appendix B.
Justification:
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data provided not)?
Question:
Answer: [Yes]
Justification: We report full implementation details of our reported experiments in Ap- pendix C for reproducibility, including neural architecture and hyperparameter choices and data generation procedures.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
We provide code to run our experiments in the supplementary material. For alternative algorithms, we use the open-source code provided by authors to which we pro- vide links in Appendix
Justification:
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We include these details in Appendix D.
7. Experiment statistical significance
36

---

Question: Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Answer: [No]
Justification: It is not computationally feasible to include error bars for all results, due to long training times for some of the algorithms. For the TreeDSBM algorithm, we report results for mean=std over 5 runs in Appendix D.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: These details are included in Appendix D.
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: htaps://nnene http///eurrp/./e/ee//eeee
Answer: [Yes]
We have reviewed the Code of Ethics and confirm that our work conforms these guidelines. We see no potential harmful consequences of our work, and include extensive reproducibility details in Appendix D.
Justification:
Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Our work mostly theory and methodology focused for the tree-structured Schödinger Bridge problem. We do not see any immediate societal impacts, though certain applications may share similar societal consequences as flow-based generative models upon which our approach is based.
Justification:
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
Guidelines:
Our work poses no such risks.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We cite the works that we use, and include the licenses of code and datasets in Appendix F.
13. New assets
Are new assets introduced in the paper well documented and is the documenta- tion provided alongside the assets?
Question:
Answer: [NA]
37

---

Justification: The paper does not release new assets.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the pa- per include the full text of instructions given participants and screenshots, applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development in this research does not involve LLMs any important, original, or ntn-standard components.
38


Reviews:
[{"review_id": "zr2fyjC3nu", "reviewer": "Reviewer_3rwT", "summary": "This paper adapts an existing technique to estimate Schrödinger bridges between two distributions---the so-called Iterative Markovian Fitting (IMF)---to the case of more-than-two distribution (multi-marginal OT) with tree structured cost (of which a typical example is star-shaped cost, encoding the celebrated Wasserstein barycenter problem).\n\nThe work proposes a set of precise definition to adapt the IMF (so far only introduced in the context of two-marginals) to this more general setting. \n\nMarkovian projection (one of the two steps to set up the IMF) _on a tree_ can be decomposed into performing Markovian projection on each edge, each of them being implemented as training (in a Monte-Carlo way) a NN to minimize a quadratic loss. This yield a (reasonably) tractable implementation of that scheme in practice. \n\nNumerical experiments are conducted to showcase the feasibility of the method.", "questions": "As said in the Weakness section, I may be misunderstanding the numerical results showcased in Figure 3. May you elaborate a bit on this?", "limitations": "NA", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "# Strengths\n\n**Clarity:** The paper is clear and well written. \n\n**Significance:** I believe that this approach is promising, it offers new perspective on dynamic entropic multi-marginal OT and reduce the gap between implementation and the underneath theoretical objects.\n\n**Quality:** The paper is theoretically sound as far as I can tell. \n\n# Weaknesses\n\n**Originality:** One could perhaps argue that this approach is only mixing already established techniques. I nonetheless believe that this \"had to be done\" and would not consider this as being a strong limitation for publication. \n\n**Quality:** On the numerical side, I am somewhat surprised with the results of Figure 3., which looks of much lower quality than those of https://arxiv.org/pdf/2402.03828 (cited in the work) for instance. Am I misunderstanding something?", "was_revised": true, "final_justification": "After discussion with the authors, I stick with my positive rating. I believe that this is a competent paper, and the authors' rebuttal clarified my potential concerns. \n\nBecause of its somewhat \"incremental\" (in that it mixes known techniques, even though effort is required to do so), I would not rank it as a top-tier paper, but I still believe that it may belong to NeurIPS and keep my positive rating (If I could, I would probably have rated it 4.5).", "rebuttal": "We thank the reviewer for their positive and constructive comments, and we address their comments and questions below. We first provide a **brief rebuttal summary**:\n- **'Originality'**: We propose to add a paragraph to the introduction highlighting the fixed-point barycentre interpretation more prominently, to highlight that our approach occupies an important position in the literature.\n- **'Figure 3'**: The noise in the samples in Figure 3 is expected, due to the entropy regularisation. There are ways to reduce/avoid this - see Figure 7 in the Appendix.\n\nWe expand on these points in more detail below, and highlight any proposed changes for revisions in **bold**.\n\n## \"One could perhaps argue that this approach is only mixing already established techniques\"\n\nWe appreciate the reviewer for raising this point (and for also noting this is not a strong limitation). Currently the narrative of the paper focuses on TreeDSBM being an IMF counterpart to TreeDSB. We discuss below a slightly alternative perspective (which in fact initially motivated our approach), which we hope highlights more clearly why we believe our method occupies an important and prominent position in the literature.\n\nOne of the most successful and elegant approaches for Wasserstein-2 barycentre computation is the iterative fixed-point approach, which involves iteratively updating a candidate barycentre $\\nu_j$, by solving for the OT map $T_i$ from $\\mu_i$ to $\\nu_j$, and updating $\\nu_{j+1}$ via the construction $\\sum_i \\lambda_i Y_i$ (where the $Y_i$ come from the coupling induced by the maps $T_i$). This approach was popularised by the seminal paper [1], and has formed the basis of many algorithmic developments, including in machine learning in [2] (one of the current SOTA approaches). Indeed, it has strong performance, and has been observed to converge quickly in only a few iterations [3]. However, a downside is that it requires solving $K \\cdot N$ complete OT problems, where $K$ is the number of marginals and $N$ is the number of outer iterations, which is expensive as solving even a single OT problem can be challenging.\n\nThe IMF procedure [4] has recently gained popularity for solving the entropy-regularised OT problem, and proceeds by iteratively performing bridge matching on its induced couplings. It is therefore a prominent question whether IMF can be used in combination with a fixed-point approach for solving for Wasserstein-2 barycentres. A naive approach would be to insert IMF to solve for each OT subproblem, requiring $K \\cdot N$ complete IMF problems, which would in turn be $K \\cdot N \\cdot M$ bridge matching problems (moreover, making this approach grounded in theory would be non-trivial, as the fixed-point method holds only for the unregularised OT problem).\n\nThe key aspect of our approach is that, rather than nesting the fixed-point iterations and the IMF iterations, they can in fact be elegantly *combined into a single iterative procedure*, along with a well-understood theoretical grounding.  In particular, our approach shows that the expensive OT computations required for the fixed-point procedure (the main limitation of such approaches), can instead be switched out for inexpensive bridge-matching procedures. Each fixed-point iteration is therefore significantly cheaper, and empirically we found the TreeDSBM algorithm to retain the fast convergence of the fixed-point procedure in terms of the number of iterations required.\n\nIt is for this reason that we believe TreeDSBM is an *extremely natural algorithm for Wasserstein-2 barycentre computation*, and thus occupies a prominent position in the literature. We consider the fact it is a simple and elegant combination of the tree-structured framework with IMF to be one of its main strengths.\n\nWe believe that the paper would benefit from advertising this fixed-point perspective more prominently, as currently the emphasis focuses on being an IMF counterpart to TreeDSB. We thus **propose to add a paragraph in the introduction highlighting this motivation and interpretation.**\n\n## results of Figure 3\n\nWe currently provide a discussion regarding this point in Section D.2 of the Appendix (lines 978-985) and in Figure 7. We **propose adding a reference to this discussion to this in the main paper**. We recall the main points from this discussion below.\n\nThe results plotted in Figure 3 use a fairly large value of entropy-regularisation parameter $\\varepsilon=0.02$ (which was selected to be more comparable with TreeDSB, which requires a large $\\varepsilon$ to converge). This means that there *should* be noise present in an accurately computed solution, which is why the results look as they are. Indeed, the results shown in Figure 3 resemble those for other entropy-regularised barycentre approaches (see e.g. results for EGBARY in Figure 5 of the paper [5] linked by the reviewer).\n\nClearly this may not be desirable; fortunately this can be reduced/avoided in two ways (which we do in the Appendix):\n- Using a smaller $\\varepsilon$ value (e.g. $\\varepsilon=0.001$)\n- Generating samples as $\\sum_i \\lambda_i Y_i$, where $Y_i$ are samples from the coupling over the marginals induced by the learned transport maps.\n\nSamples using both these approaches are included in Figure 7 in the Appendix (note minor typo, caption of (a) should read $\\varepsilon=0.02$ rather than $0.2$). These samples are much more comparable with those in the work [5] referenced by the reviewer.\n\n## Other changes\n\nFinally, we highlight here some additional changes that we propose, following feedback from other reviewers.\n- Move the comparisons with other continuous barycentre solvers [2].[5],[6] to the main body (they are currently in Appendix E.1 and E.2 due to space constraints)\n- Additional ablations on the role of epsilon, batch size (see response to reviewer QVwU).\n\nWe hope that these comments have addressed your questions and concerns. If you have any further questions, please let us know!\n\n\n## References\n\n[1] Alvarez-Esteban et al, 2016, A fixed-point approach to barycenters in Wasserstein space\n\n[2] Korotin et al, 2022, Wasserstein Iterative Networks for Barycenter Estimation\n\n[3] von Lindheim, 2022, Simple Approximative Algorithms for Free-Support Wasserstein Barycenters\n\n[4] Shi et al, 2023, Diffusion Schrödinger Bridge Matching\n\n[5] Kolesov et al, 2024, Estimating Barycenters of Distributions with Neural Optimal Transport\n\n[6] Korotin et al, 2021, Continuous Wasserstein-2 Barycenter Estimation without Minimax Optimization"}, {"review_id": "6hVIiRgKUP", "reviewer": "Reviewer_QVwU", "summary": "This paper extends the standard iterative Markovian fitting approach to tree structured problems. One important feature of this is that it covers Wasserstein barycentres, gives rise to interesting applications.", "questions": "1) I think introduction overly focuses on technical aspects and relevant methods, but leaves a question mark for the nonspecialist reader why it is even of interest to look at these problems. Could you please provide a paragraph, aiming at an average NeurIPS reader, why your paper is of interest to this community?\n\n2) Including above - why is the tree setting of special interest beyond barycenters?\n\n3) I think, again, for someone who is not familiar with the very specialist earlier work, it is hard to understand all abstract definitions provided in this paper. Can authors provide a simple example, e.g., for a Gaussian setting (or whatever is appropriate) and make the problem fully explicit? This will make everything concrete for a reader who is not already working on SDEs on trees. Otherwise, the paper risks having an audience that is too narrow. It would be great if TreeDSBM (Algorithm 1) can be made explicit in a simple setting as an example.\n\n4) Can authors elaborate a bit the challenges and observations about the training process? How does the cost scale with dimension or number of data points? Is there any instability issues? These would be good to mention in limitations if any. It would be better to clearly demonstrate the stability claims for your method by comparing it to relevant adversarial methods - and demonstrate this point specifically.\n\n5) Is there any guideline to choose $\\varepsilon$ for your method - which seems to impact the results?\n\n6) Many hyper-parameters (training steps, batch size, network width, $\\varepsilon$) vary across experiments; no ablation on drift-network architecture, number of edges trained, or tolerance to poor initial couplings. Demonstrating failure cases better would give a better picture.\n\n7) How do introduced neural network architectures play with the assumptions required for theory? Please discuss.\n\n8) Is it at all possible to generalize your framework beyond quadratic costs and non-Brownian cases?", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths: There are multiple strengths of this work - notably\n\n1) The paper is very well written from a technical perspective. Mathematical structure is made very clear throughout, clear literature review with connections to IPF and other ideas.\n2) The paper goes beyond simplistic experiments, with MNIST and image experiments.\n\nWeaknesses: While mathematically clean and well-written, the paper risks being obscure to most of the NeurIPS readers, see some suggestions to improve clarity.", "was_revised": true, "final_justification": "My concerns are addressed but in general I think the paper's impact is limited and still too technical (the changes requested was a bit hard during a rebuttal period anyway). Thus I'd like to keep my score.", "rebuttal": "We thank the reviewer for their positive and constructive comments, and we address their comments and questions below. We first provide a **brief rebuttal summary**:\n- 'Overly technical presentation': We propose to:\n    - add a paragraph to the introduction regarding the importance of the Wasserstein barycentre problem.\n    - add a paragraph to the introduction emphasising the fixed-point interpretation, which we anticipate may be more intuitive for a reader less familiar with related works.\n    - add an explicit description of the algorithm in the simplest $\\frac{1}{K}$-barycentre case.\n- 'Challenges in training/hyperparameters': We address each of these in turn below, with reference to relevant experiments in the Appendix, and with new ablations on the role of batch size and $\\varepsilon$.\n\nWe expand on these points and the reviewer's questions in more detail below, and highlight any proposed changes for revisions in **bold**.\n\n## \"Overly technical presentation\"\n\nWe thank the reader for highlighting this, as (behind the mathematical notation) we believe our algorithm is in fact quite intuitive, so are keen to convey this to the reader.\n\n### \"why it is even of interest to look at these problems\"\n\nWe propose to **add a paragraph in the introduction explaining the importance of the Wasserstein-2 barycentre problem**. In particular, we will highlight its interpretation as the natural notion of ‘mean’ for probability distributions, and its applications in Bayesian learning, clustering, and representation learning to name a few.\n\n### \"overly focuses on technical aspects and relevant methods\", \"Can authors provide a simple example\"\n\nCurrently the narrative of the paper focuses on TreeDSBM being an IMF counterpart to TreeDSB, which we agree may be somewhat opaque for a reader unfamiliar with these works. Our approach was in fact initially motivated by trying to combine IMF with fixed-point approaches for barycentre computation (please see our response to Reviewer 3rwT). We believe that the fixed-point perspective — namely that each point in the 'average' measure, is the 'average' of its corresponding points in the marginals — may be more intuitive for a reader less familiar with related work.\n\nWhile discussed in the paper, we believe a less familiar reader would benefit from having this interpretation highlighed more prominently. We thus **propose to add a paragraph in the introduction highlighting this motivation and interpretation**, and to **add a fully explicit description of Algorithm 1 in the simplified case of barycentres with weights 1/K**. This is the simplest setting (as the barycentre is constructed via $Y_0$ centred at $\\sum_i \\frac{1}{K} Y_i$, i.e. the ‘average’ of its corresponding points in the marginals), and should be intuitive for the reader.\n\n### \"why is the tree setting of special interest beyond barycenters\"\n\nGiven our initial fixed-point motivation described above, the fact our approach extends to the tree setting was something of an added bonus that arose out of using the tree-based mathematical framework. Nevertheless, such problems have been utilised in works such as [1]. We conjecture that it could have applications in trajectory inference problems where populations might be known to 'split' according to a given structure (for example, maybe for single-cell dynamics). We unfortunately don't have the biological expertise to be confident in devising meaningful experiments for such settings, so instead provided a toy experiment in Appendix E.3 in which TreeDSBM appeared to work well, demonstrating fast and stable convergence. We have a brief paragraph discussing these wider problem settings in Appendix E.3, and **propose extending this to include the above discussion**, and **referencing this paragraph at the end of the Related Work section**. We hope future practitioners can build on our work for use in such settings.\n\n## Challenges in training/hyperparameters\n\nScaling with number of datapoints/dim:\n- As it uses bridge-matching (trained via gradient descent on minibatches), our approach inherits the same scalability with datapoints as other flow-matching style algorithms.\n- In terms of dimension, we included experiments for high-dimensional Gaussians in Appendix E.2 for $d=64, 96, 128$, and observed no significant performance drop as $d$ increased.\n\nInstability:\n- We did not encounter instability issues, which we attribute to using stable bridge-matching losses, and an iterative procedure rather than adversarial losses.\n- \"demonstrate the stability claims\": In the Appendix, we have also compared to a range of state-of-the-art solvers [2],[3],[4]. Regarding stability, we found the solvers [4] (adversarial loss) and [3] (single regularised loss) to fail on the 2d-experiment. We discuss such issues in lines 935-940 in the Appendix.\n- We propose **adding a sentence referencing this in the main body**.\n\nError in marginals:\n- We believe the main algorithmic limitation is that errors can accumulate in the marginals as IMF iterations proceed; this is a limitation inherited from standard IMF. We discuss this in lines 249-251 in the main paper, and in lines 851-858 in the Appendix. Rotating the starting marginal, as in [5], helps mitigate this.\n- It is therefore important that the bridge-matching steps fit the marginals accurately, which is affected by the hyperparameters training steps, batch size, network width. The effect of these parameters on bridge-matching is well-understood and they act in the expected ways. To provide an indication of how the learned bridge-matching quality affects the overall solution, we provide an ablation for batch size below on the 2d and data aggregation experiments for the hyperparameters used in the paper (similar effects occur for increasing training steps).\n\n| Batch size                                      | 64       | 256     | 1024   | 4096   |\n| - | - | - | - | - |\n| 2d, sinkhorn divergence $\\downarrow$            | 1.57     | 1.24    | 1.04   | 1.04   |\n| Data Aggregation (Poisson), BW2-UVP $\\downarrow$| 0.032    | 0.017   | 0.013  | 0.012  |\n\n- Network width: any sufficiently-expressive network will work similarly well. Fairly small MLP networks are sufficient for the pointcloud data experiments (i.e. 2$d$, subset aggregation, high-dim Gaussians).\n- We propose to **extend the discussion in Appendix lines 851-858**, by incorporating the above discussion and ablations.\n\nRole of $\\varepsilon$\n- This is a perennial question in entropic OT. Standard methods to choose $\\varepsilon$ in commonly used OT libraries (e.g. POT, OTT-JAX—done in proportion to the costs) provide good guidance for choosing suitable values. We add a few additional comments below.\n- Typically, one may want to use as small an epsilon as possible (for less entropic bias). One advantage of TreeDSBM over TreeDSB is that it allows for much smaller epsilon (TreeDSB does not converge for too-small epsilon as simulated trajectories do not reach the other marginals).\n- There is a visualisation of the role of epsilon in Figure 5 in the Appendix, and also for two values of epsilon for the MNIST experiment in Figure 7. We propose to **add quantitative results for varying epsilon to the Appendix**. Results for the 2d and data aggregation experiments are shown below.\n\n| $\\varepsilon$                                   | 1.0      | 0.3     | 0.1   |\n| - | -| - | - |\n| 2d, sinkhorn divergence $\\downarrow$            | 1.24     | 0.99    | 1.02   |\n|                                                 | 1e-3     | 3e-4    | 1e-4   |\n| Data Aggregation (Poisson), BW2-UVP $\\downarrow$| 0.012    | 0.008   | 0.008  |\n\n## Neural Architectures and Assumptions\n\nThe theoretical analysis requires the SDE drifts to be locally Lipschitz, which is inherited from DSBM (for which one reason is to make use of known results regarding the well-posedness of the Doob’s $h$-transform). This is a weak regularity requirement, and is satisfied by standard neural architectures.\n\n\n## Beyond Quadratic Costs, non-Brownian Case\n\nIn its current form, the algorithm works only for the quadratic cost due to the Brownian reference measure (as with DSBM). We provide some thoughts regarding this below.\n\n- While DSBM works only for the quadratic cost and $\\sigma>0$, [6] shows that ReFlow (i.e. $\\sigma=0$) converges to OT for a convex cost $c(x,y) = h(x-y)$ if one restricts the form of the drift, and uses a Bregman divergence loss. One could potentially replace the bridge-matching steps with this procedure. The barycentre update would have to be modified accordingly (likely from $\\sum_i \\lambda_i y_i$ to $argmin_z \\sum_i c(z, y_i)$ [7], though this may not be available in closed form).\n- Extensions of IMF to more general settings exist (e.g. state-dependent costs [8]). Incorporating such ideas into our framework is a promising direction.\n- Non-Brownian reference measures: as with DSBM, our method supports any reference process for which conditional bridges can be sampled (e.g. linear drift processes, see Appendix B of [5]). We focus on the Brownian case in the main paper due to its connection to OT.\n\n\nWe hope that these comments have addressed your questions and concerns. If you have any further questions, please let us know!\n\n## References\n\n[1] Haasler et al, 2021, Multi-marginal Optimal Transport with a Tree-structured cost and the Schrödinger Bridge Problem\n\n[2] Korotin et al, 2022, Wasserstein Iterative Networks for Barycenter Estimation\n\n[3] Korotin et al, 2021, Continuous Wasserstein-2 Barycenter Estimation without Minimax Optimization\n\n[4] Kolesov et al, 2024, Estimating Barycenters of Distributions with Neural Optimal Transport\n\n[5] Shi et al, 2023, Diffusion Schrödinger Bridge Matching\n\n[6] Liu, 2022, Rectified Flow: A Marginal Preserving Approach to Optimal Transport\n\n[7] Tanguy et al, 2025, Computing Barycentres of Measures for Generic Transport Costs\n\n[8] Liu et al, 2023, Generalized Schrödinger Bridge Matching"}, {"review_id": "7yxSXPSy4z", "reviewer": "Reviewer_HGM9", "summary": "This paper extends the Iterative Markovian Fitting (IMF) procedure to solve tree-structured Schrödinger Bridge (TreeSB) problems, proposing the TreeDSBM algorithm. The work addresses a gap in the literature where only IPF-based methods (TreeDSB) existed for this setting, while IMF approaches have shown superior properties in the standard two-marginal case. The TreeSB problem generalizes optimal transport to multiple marginals connected by a tree structure, with Wasserstein barycenters as a special case when considering star-shaped trees. The authors demonstrate theoretical convergence guarantees and show empirical improvements over TreeDSB on synthetic 2D data, MNIST digits, and high-dimensional face datasets.", "questions": "Scalability concerns: How does the computational cost scale with tree size compared to TreeDSB? You mention parallel training of edges, but what about memory requirements for storing $2\\times∣E∣$ networks and the overall wall-clock time comparison?\n\n⁠Failure mode analysis: The CelebA experiment suggests the method may not be suitable for certain types of optimal transport problems. Can you provide more principled criteria for when TreeDSBM is expected to work well versus poorly? This would help practitioners understand when to apply your method.\n\n⁠Convergence rate: While you show convergence, what can you say about the convergence rate compared to TreeDSB? The 6 vs 50 iterations comparison is compelling, but I'd like to see this validated across different problem settings and regularization parameters.", "limitations": "Yes.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\nIn my opinion the paper’s main strength is a very clean theoretical generalization of IMF to multi‑marginal settings, together with a convergence proof of the TreeIMF. The resulting algorithm is embarrassingly parallel on each edges, and at least on the reported tests, run much faster than IPF‑style solvers.\n\nEmpirically, the results convincingly demonstrate the advantages of IMF over IPF in the tree setting. The 2D synthetic experiment shows TreeDSBM achieving comparable quality to TreeDSB with 6 IMF iterations versus 50 IPF iterations, which is a substantial computational improvement. The connection to fixed-point barycentre methods (Section 3.4) provides nice intuition for the star-shaped case.\n\n\nWeaknesses:\n\nMy main concern is with the experimental validation, particularly the lack of rigorous quantitative evaluation on higher-dimensional problems. While the 2D case uses Sinkhorn divergence against a ground truth, the MNIST and CelebA experiments rely primarily on visual assessment. For a method claiming to advance the state-of-the-art, I would expect more comprehensive benchmarking.\nThe CelebA experiment reveals a significant limitation that I think deserves more critical examination. The authors acknowledge that their samples \"do not match the visual quality of state-of-the-art results\" and hypothesize this is due to the iterative nature being unsuitable for color alteration tasks. However, this raises questions about the general applicability of the approach - if the method struggles with an important class of transport problems, this significantly limits its practical utility.\nI'm also skeptical about some implementation details. The method requires learning 2∣E∣2|E|\n2∣E∣ vector fields (forward and backward for each edge), which could become computationally expensive for large trees. The authors mention parallelization but don't provide concrete scalability analysis or comparison of training times with TreeDSB.\n\nThe theoretical analysis, while correct, relies heavily on \"mild assumptions\" that are relegated to the appendix. For a theory-focused paper, I would prefer seeing the key assumptions stated explicitly in the main text, particularly since the tree setting introduces additional complexity compared to standard IMF.", "was_revised": true, "final_justification": "My concerns about Experimental validation have been addressed.  While questions remain about scalability, I believe the methodological contribution outweighs the practical limitations. Therefore I would like to increase the score.", "rebuttal": "We thank the reviewer for their feedback, and we address their comments and questions below. We first provide a **brief rebuttal summary**:\n- **Experimental validation**: There are several further experiments in the Appendix, many of which address concerns of the reviewer. In particular, there are comparisons with alternative barycentre algorithms, and discussion of runtimes. If accepted, we intend to use the additional space to move some of these to the main paper.\n- **Scalability**: We address the reviewer's concerns in turn below. In particular, in the majority of use-cases, network sizes will be fairly small and thus not pose a computational bottleneck; in cases where it does, we outline ways it can be improved.\n- **CelebA**: We have subsequently improved performance on this benchmark by using a shared architecture across the edges.\n\nWe expand on these points in more detail below, and highlight any proposed changes for revisions in **bold**.\n\n## Experimental validation\n\nWe note that there are several further experiments in the Appendix (Sections D and E), which could not be included in the main paper submission due to space constraints. We anticipate the reviewer may not have seen these, based on their comments. We hope these results allay the reviewer’s main concern regarding experimental validation.\n\nIn these experiments, we include comparisons with state-of-the-art methods WIN [1], W2CB [2] and NOTWB [3]. These methods are chosen to be representative of different approaches in the literature— WIN is an iterative fixed-point approach, W2CB is an Input Convex Neural Network-based method, and NOTWB is based on recent Neural OT methodology.\n\nIn particular, we include in Appendix E comparisons on a subset posterior aggregation experiment (a standard experiment in the continuous Wasserstein-2 barycentre literature, e.g. in [2],[4]), and also on high-dimensional Gaussian examples for $d=64,96,128$ (also a standard experiment, e.g. in [1],[2],[3]). In lines 1130-1145 we provide a discussion regarding the findings, highlighting the pros and cons of different approaches in the literature.\n\nImportantly, **we intend to move the subset posterior aggregation and high-dimensional Gaussian experiments to the main paper** if accepted.\n\nIn Appendix D there are also more detailed examinations of the experiments in the main paper, including runtime analysis, barycentre progression during IMF iterations, and the effect of $\\varepsilon$.\n- Run-time analysis: We discuss run-times of the various algorithms in the Appendix in lines 941-952, and lines 1079-1095, and TreeDSBM appeared favourable compared to existing barycentre solvers.\n- Ablations: We will add additional results concerning the role of batch size, and regularisation value $\\varepsilon$ to the Appendix—please see response to reviewer QVwU.\n\n## Implementation details and scalability\n\nNumber of edges:\n- The reviewer notes a concern that learning $2 |E|$ separate networks may cause scalability problems. We note that this is similar in many continuous Wasserstein barycentre algorithms, including TreeDSB and e.g. [1],[2],[3]. Fortunately, there are ways to improve this if necessary, and we also do not envisage this being a limiting factor for the vast majority of use cases.\n- As discussed in Appendix C, one can use a single network along each edge and additionally condition on the direction. As the reversed drifts should be approximately the negative of each other, this can be efficiently parameterised as $(-1)^{dir} v_{\\theta}(x,t,dir)$ (as done in [6], who show it has minimal impact on performance).\n- More generally, one could use a single shared network by additionally conditioning on the edge $e$. This makes particular sense for problems with shared structure between edges.\n- In our experiments, memory bottlenecks only arose in the CelebA experiment, which is exactly a scenario in which there is shared structure. We have subsequently used this single-network parameterisation and found improved performance (see below).\n- We propose to **extend the discussion in Appendix C** to include conditioning on edges too.\n\nMemory requirements, wall-clock time\n- Our algorithm offers a trade-off between memory consumption and run time due to its parallel nature. Namely, one can train edges simultaneously if compute allows (on a single GPU if enough memory, or parallelised across GPUs). If this cannot be done, one can train sequentially instead (in which case memory requirements for training would be comparable with standard bridge-matching).\n- In our experiments, we were (other than for CelebA) able to train simultaneously on a single GPU as the networks are not that large. Please see the response to reviewer ZjRa for further information regarding GPU usage. We anticipate that the majority of use-cases in practice would be for pointcloud data (e.g. the data aggregation experiment), which would use small MLP-like architectures with low memory requirements.\n- Comparison to TreeDSB: We note that TreeDSBM has two major benefits over TreeDSB regarding run-time and memory requirements. Firstly, TreeDSB cannot be parallelised in the same way, as each IPF iteration rotates the training around the edges sequentially. Secondly, TreeDSB requires saving entire trajectories to train the time-reversals, while TreeDSBM only requires storing the endpoints (this is one of the most significant benefits of IMF over IPF).\n\n## CelebA\n\nWe currently provide an extended discussion regarding the CelebA experiment in lines 1006-1029 in the Appendix.  In particular, we identify that the initial coupling (using bridge-matching models along each edge) does not always provided images that have structural similarity, so the resulting initial barycentre samples sometimes consist of distinct overlaid images. The subsequent IMF steps struggle to recover from this, resulting in the artifacts visible in some of the generations.\n\nIn this benchmark, the true maps are relatively simple (primarily just colour changes), but TreeDSBM struggles to recognise this simpler structure due to infrequent communication between edges (which occurs only after each IMF iteration). In contrast, methods which perform well on this benchmark optimise all the maps together and with much more interaction between them, providing more scope for identifying the shared structure present in this specific problem. We hope this helps clarify for practitioners when TreeDSBM is likely to perform well or poorly.\n\nWe hypothesised in the Appendix that different architectures could make TreeDSBM more competitive in this setting. We have since used a shared architecture as outlined in the previous section, which has improved the results (though they remain short of the best-performing methods for this problem). By learning features jointly across edges, the bridge-matching models behave more similarly along each edge, yielding coupling samples with greater structural similarity.\n\nWe thus propose to **include in the Appendix a comparison of the results for the two parameterisations, to illustrate the role that architecture can play** in such highly-structured settings. We emphasise that much of the behaviour discussed here is likely specific to this particular benchmark; we anticipate the majority of practical use-cases would be for pointcloud data, for which similar behaviour would not arise.\n\n\n## Convergence Rate\n\nWe are glad that you find the fast convergence of TreeDSBM compelling. As discussed in the response to reviewer 3rwT, we believe that the TreeDSBM algorithm inherits good convergence properties from the similar fixed-point barycentre algorithms, which are known to converge quickly (see e.g. [7], who show good results with just a single iteration).\n\nWe can include quantitative results regarding convergence speed in terms of IMF iterations in the Appendix. We provide such results for the 2d and data aggregation experiments below.\n\n| IMF step |1|2|3|4|5|6\n| - | - | - | - | - | - | - |\n| 2d, sinkhorn divergence $\\downarrow$ |17.2 |1.27 |1.19  |1.14 | 1.07 | 1.04 |\n| Data Agg. (Poisson), BW2-UVP $\\downarrow$| 31.1 | 0.010 | 0.011  | 0.012\n\n\nWe can also provide similar results for TreeDSB. Note this would be expected to converge much slower, particularly for small epsilon, as fitting of the marginals only arises through iteratively learning time-reversals— if the forward processes rarely reach the other marginals (as is the case for smaller epsilon) this is very slow.\n\n## Theoretical assumptions\n\nThe key assumptions for the Brownian case are that $\\int \\lVert x \\rVert^2 d \\mu_i < \\infty$ and $H(\\mu_i)< \\infty$ for each marginal $i$ (see proof of Theorem 3.1). We will **mention these in the main paper**. More generally, the required assumptions are stated in Proposition B.3 (for a general reference measure $\\mathbb{Q}$), as well as standard technical assumptions inherited from the DSBM paper [8, Appendix C.1].\n\nWe hope that the presence of experiments in the Appendix that are relevant to many of your questions, along with our comments and results above, and our comments in other rebuttals, have addressed your concerns. If you have any further questions, please let us know!\n\n## References\n\n[1] Korotin et al, 2022, Wasserstein Iterative Networks for Barycenter Estimation\n\n[2] Korotin et al, 2021, Continuous Wasserstein-2 Barycenter Estimation without Minimax Optimization\n\n[3] Kolesov et al, 2024, Estimating Barycenters of Distributions with Neural Optimal Transport\n\n[4] Fan et al, 2021, Scalable Computations of Wasserstein Barycenter via Input Convex Neural Networks\n\n[5] Haasler et al, 2021, Multi-marginal Optimal Transport with a Tree-structured cost and the Schrödinger Bridge Problem\n\n[6] De Bortoli et al, 2024, Schrödinger Bridge Flow for Unpaired Data Translation\n\n[7] von Lindheim, 2022, Simple Approximative Algorithms for Free-Support Wasserstein Barycenters\n\n[8] Shi et al, 2023, Diffusion Schrödinger Bridge Matching"}, {"review_id": "upJtTLclBK", "reviewer": "Reviewer_ZjRa", "summary": "This paper proposes TreeDSBM, a novel algorithm that extends the Iterative Markovian Fitting (IMF) procedure to solve Schrödinger Bridge (SB) problems defined over arbitrary tree-structured cost graphs. This generalization enables scalable solutions to multi-marginal entropy-regularized optimal transport problems with structured cost functions, such as those arising in Wasserstein barycentre computation. The authors provide a rigorous theoretical foundation for TreeDSBM. They also demonstrate empirical benefits of TreeDSBM over its IPF-based counterpart, TreeDSB, in selected barycentre estimation tasks.", "questions": "1. The proposed method seems applicable to any multi-marginal OT problem with tree-structured costs. Can the authors provide one or two additional use cases beyond barycentre computation to demonstrate this generality?\n\n2. It would be very helpful to compare against more recent continuous barycentre methods. Can the authors either expand the baseline set or justify their omission?\n\n3. Could the authors report training time, inference time, and GPU memory consumption for TreeDSBM versus TreeDSB or other methods? This would help in understanding the method’s practical viability.\n\n4. Was it inadvertently omitted in the submission?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 4, "strengths_and_weaknesses": "Strengths:\n\n1. The extension of IMF to general tree-structured costs is a clear and non-trivial innovation. Prior flow-based SB methods have largely been restricted to two-marginal settings or relied on IPF-like iterations. This work addresses an important gap and is a natural yet challenging generalization.\n\n2. The paper is rigorous and mathematically well-grounded. Key theoretical constructs such as KL decompositions on trees, convergence guarantees, and generalized projections are carefully introduced and justified.\n\nWeaknesses:\n\n1. Despite the general applicability of tree-structured cost models, the paper only evaluates TreeDSBM in the context of Wasserstein barycentre estimation. While this is a canonical and well-motivated use case, it raises questions about the method's general utility in broader machine learning applications such as structured generative modeling, domain adaptation, or trajectory inference.\n\n2. In the barycentre setting, the paper compares TreeDSBM only against TreeDSB and WIN. The lack of broader comparisons makes it harder to contextualize the method’s empirical competitiveness.\n\n3. A practical appeal of IMF over IPF is improved sample efficiency and training stability. However, the paper does not report computational time, memory usage, or wall-clock comparisons, which are critical for assessing scalability in high-dimensional or large-tree settings.\n\n4. Many technical discussions, proofs, and implementation details are deferred to the appendix, but the appendix seems to be missing. This makes it difficult to fully verify the claims or reproduce the results.", "was_revised": true, "final_justification": "Thank authors for their rebuttal and discussion. In particular, they have made a substantial effort to address the reviewers’ concerns regarding the quantitative evaluation. I recommend acceptance.", "rebuttal": "We thank the reviewer for their positive and constructive comments, and we address their comments and questions below. We first provide a **brief rebuttal summary**:\n- Appendix: The appendix is located in the zip folder under ‘Supplementary Material’: Then go `neurips_treeIMF_supplementary` -> `supplementary.pdf`. Please let us know if you have issues accessing it.\n- The Appendix contains many additional experiments and discussions relevant to the reviewer's questions, in particular comparisons with alternative continuous barycentre algorithms, and discussion of runtimes.\n- We provide additional information and discussion regarding computational resources in this rebuttal.\n\nWe expand on these points in more detail below, and highlight any proposed changes for revisions in **bold**.\n\n## Comparison with other continuous barycentre methods\n\nWe do in fact provide comparisons with other continuous Wasserstein-2 barycentre solvers in the Appendix. Namely, we include comparisons with state-of-the-art methods WIN [1], W2CB [2] and NOTWB [3]. These methods are chosen to be representative of different approaches in the literature— WIN is an iterative fixed-point approach, W2CB is an Input Convex Neural Network-based method, and NOTWB is based on recent Neural OT methodology.\n\nIn particular, we include in Appendix E comparisons on a subset posterior aggregation experiment (used as a standard experiment in continuous Wasserstein-2 barycentre approaches, e.g. in [2],[4]), and also on high-dimensional Gaussian examples for $d=64,96,128$ (again a standard experiment in the literature, e.g. in [1],[2],[3]). In lines 1130-1145 we provide a discussion regarding the findings, highlighting the pros and cons of different approaches in the literature.\n\nImportantly, if accepted **we intend to move the data aggregation and high-dimensional Gaussian experiments to the main paper**, using the extra page afforded to us.\n\n\n## Training time, inference time, GPU memory consumption\n\nTraining time:\n- We discuss run-times of the various algorithms that we compare in the Appendix in lines 941-952, and in lines 1079-1095, and TreeDSBM appeared favourable compared to existing barycentre solvers.\n\nInference time:\n- Inference time for TreeDSBM is slower than for methods that use a single function evaluation, as it involves solving an SDE (discussed in line 361 in main paper, and line 1139 in the Appendix). However, it can still be fast if one JIT-compiles the sampling procedure: transporting 1,000 samples along one of the maps in the data aggregation experiment takes 7.8 ms for the single-evaluation NOTWB method, and 15 ms for TreeDSBM when using 50 discretisation steps.\n\nGPU memory consumption\n- We comment here that the parallel nature of the algorithm during training means that a practitioner can have a trade-off between memory consumption and wall-clock time. Namely, one can train the edges simultaneously if compute allows (either on a single GPU if enough memory, or parallelised across GPUs). If this cannot be done, one can train sequentially instead (in which case memory requirements for training would be comparable with standard bridge-matching).\n- We report GPU consumption for different experiments below (for sequential and joint training), for the hyperparameters used in the paper. These can be changed significantly by changing hyperparameters such as batch size.\n            \n| Experiment                 | Sequential Training | Joint Training  |\n| -------------------------- | ---------- | ------- |\n| 2d                         | 447 MB     | 687 MB  |\n| Data Aggregation (Poisson) | 705 MB     | 1219 MB |\n| Gaussian (d=64)            | 1239 MB    | 1239 MB |\n| MNIST (2,4,6)              | 4683 MB    | 6407 MB |\n\n- We note that in some of these experiments (e.g. Gaussian), peak GPU usage is due to the simulation and storage of the training samples for subsequent IMF steps, rather than during the network training. This can be reduced significantly by simulating in smaller batches, or by updating the cache during training rather than simulating all beforehand.\n- We also note that TreeDSBM has two major benefits over TreeDSB regarding run-time and memory requirements. Firstly, TreeDSB cannot be parallelised in the same way, as the training of the networks in each IPF iteration rotates around the edges sequentially. Secondly, the memory required to store training data for TreeDSB is far higher, as training a time-reversal along an edge requires saving entire trajectories simulated along the reverse direction, whereas TreeDSBM only requires storing endpoints (this is one of the most significant benefits of IMF over IPF).\n\n## Application beyond barycentres\n\nOur initial motivation for our study was to utilise the IMF procedure in combination with the highly-performant fixed-point approaches for barycentre approximation. We include a detailed exposition of this perspective in the response to Reviewer 3rwT.\n\nWe found that the tree-based framework provided the correct mathematical framework to combine these approaches in a simple and elegant way. Given our initial fixed-point motivation, the fact that our approach extends to more general tree settings was something of an added bonus that arose out of utilising the tree-based framework. Nevertheless, such problems have been utilised in works such as [5]. We conjecture that it could have applications in trajectory inference problems where populations might be known to 'split' according to a given structure (for example, maybe for single-cell dynamics). We unfortunately don't have the biological expertise to be confident in devising meaningful experiments for such settings, so instead provided a toy experiment in Appendix E.3 in which TreeDSBM appeared to work well, demonstrating fast and stable convergence. We have a brief paragraph discussing these wider problem settings in Appendix E.3, and **propose extending this to include the above discussion**, and **referencing this paragraph at the end of the Related Work section**. We hope future practitioners can build on our work for use in such settings.\n\n\nWe thank the reviewer for their already kind feedback, and we hope that the contents of the Appendix and this rebuttal have addressed their comments. If you have any further questions, please let us know!\n\n## References\n\n[1] Korotin et al, 2022, Wasserstein Iterative Networks for Barycenter Estimation\n\n[2] Korotin et al, 2021, Continuous Wasserstein-2 Barycenter Estimation without Minimax Optimization\n\n[3] Kolesov et al, 2024, Estimating Barycenters of Distributions with Neural Optimal Transport\n\n[4] Fan et al, 2021, Scalable Computations of Wasserstein Barycenter via Input Convex Neural Networks\n\n[5] Haasler et al, 2021, Multi-marginal Optimal Transport with a Tree-structured cost and the Schrödinger Bridge Problem"}]

Author rebuttal:

Meta-review: The authors consider a Schrodinger bridge problem. They solve its generalized variant, namely, multi-marginal case with tree-structured cost. Wasserstein barycenter problem is a particular case of such tree-structured cost. They generalized a well-known Iterative Markovian Fitting approach to solve this problem.
Experimental verification confirmed ability of the approach to solve the Schrodinger bridge problem for tree-structured cost.


As a strengths of the paper I would highlight 
- generalization of IMF to solve a schrodinger bridge problem with tree-structured cost
- the text is well-written

As a weaknesses I would highlight
- somewhat weak empirical evaluation on multidimensional problems

The main reasons for the accept decision are
- generalization of a solution of the important problem to a more general and important case
- clear text and theoretical explanations
- interesting insights about why the approach is more efficient for classical Wasserstein barycenter problem (see discussion with reviewer 3rwT)


During the discussion with the reviewers several issues were raised that should be addressed in the final version of the paper. Here is the list of the main issues
- the topic of the paper is rather techincal, so some more intuitive discussions are needed (reviewer QVwU)
- add discussion of why the proposed approach is so important for classical Wasserstein barycenter problem (reviewer 3rwT)
- comments to questions of reviewers HGM9 and ZjRa should be extended and included in the appendix, especially those concerning experimental evaluation
Decision: Accept (poster)