Title: Tree-Based Diffusion Schrödinger Bridge with Applications to Wasserstein Barycenters
Venue: NeurIPS 2023
Accepted: yes
Keywords: Schrödinger bridge

optimal transport

diffusion model

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

Tree-Based Diffusion Schrödinger Bridge with Applications to Wasserstein Barycenters

Abstract

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

1 Introduction

In the last decade, computational Optimal Transport (OT) has shown great success with applications in various fields such as biology (Schiebinger et al., 2019; Bunne et al., 2022), shape correspondence (Su et al., 2015; Feydy et al., 2017; Eisenberger et al., 2020), control theory (Bayraktar et al., 2018; Acciaio et al., 2019) and computer vision (Schmitz et al., 2018; Carion et al., 2020). While OT commonly seeks at computing the transport plan that minimizes the cost of moving between two distributions, it can naturally be extended to the multi-marginal setting (mOT) when considering several distributions. This extension of OT has notably been studied in quantum chemistry (Cotar et al., 2013), clustering (Cuturi & Doucet, 2014) and statistical inference (Srivastava et al., 2018). In particular, a popular application in unsupervised learning of mOT with Euclidean cost consists in computing the Wasserstein barycenter of a set of probability distributions (Agueh & Carlier, 2011; Benamou et al., 2015; Álvarez-Esteban et al., 2016; Peyré et al., 2019).

Interior point methods can be used to solve OT and mOT problems but they come with computational challenges (Pele & Werman, 2009). In order to mitigate these limitations, one often considers an entropic regularization of OT, known as Entropic OT (EOT). This regularized formulation can be efficiently solved in discrete state-spaces using the celebrated Sinkhorn algorithm (Cuturi, 2013; Knight, 2008; Sinkhorn & Knopp, 1967), which admits a continuous state-space counterpart referred

to as the Iterative Proportional Fitting (IPF) procedure (Fortet, 1940; Kullback, 1968; Ruschendorf, 1995). In the case of a quadratic cost, EOT is equivalent to the static formulation of the Schrödinger Bridge (SB) problem (Schrödinger, 1932). Given a reference diffusion with finite time horizon T and two probability measures, solving SB amounts to finding the closest diffusion to the reference (in terms of Kullback-Leibler divergence on path spaces) with the given marginals at times t = 0 and t = T . This framework naturally arises in stochastic control (Dai Pra, 1991) where one aims at controlling the marginal distribution of a stochastic process at a fixed time. Recently, De Bortoli et al. (2021) introduced Diffusion Schrödinger Bridge (DSB), an approximation of a dynamic version of the IPF scheme on path spaces, see also Vargas et al. (2021); Chen et al. (2022). This methodology leverages advances in the field of denoising diffusion models (Song et al., 2021; Ho et al., 2020) in order to derive a scalable and efficient scheme to solve SB, and thus EOT.

Similarly to OT, mOT admits an entropic regularization (EmOT), which can be solved via a multimarginal generalization of Sinkhorn/IPF algorithm (Benamou et al., 2015; Marino & Gerolin, 2020). Recently, Haasler et al. (2021) proposed an extension of the static SB problem in discrete statespace to any multi-marginal tree-based setting. They notably made the correspondence between this formulation and EmOT, when the cost function writes as the sum of interaction energies onto the given tree structure, and introduced an efficient version of Sinkhorn algorithm to solve it.

Motivations and contributions. In this work, we investigate the continuous and dynamic counterpart of the tree-based framework from Haasler et al. (2021). To be more specific, we present an extension of the static SB formulation in continuous state-space to any multi-marginal tree-based setting, referred to as TreeSB. Then, we establish the equivalence between TreeSB and a formulation of EmOT relying on a (quadratic) tree-structured cost function, analogously to Haasler et al. (2021). Inspired by DSB, we develop TreeDSB, a dynamic counterpart of the multi-marginal IPF (mIPF) to solve it, by operating on path spaces and using score-based diffusion techniques. To bridge gaps in literature, we prove the convergence of mIPF iterations in a non-compact setting under mild assumptions, by extending results on IPF convergence (Ruschendorf, 1995). Finally, we illustrate our approach on examples of Wasserstein barycenters from statistical inference and image processing.

Although our approach can be applied to any tree, we focus on star-shaped trees . In this setting, we show that TreeSB reduces to a regularized Wasserstein barycenter problem. Our method comes with several benefits compared to existing works. First, it is out-of-sample, i.e. , it does not require re-running the full procedure when given a new data point. Second, our formulation of the Wasserstein barycenter problem obtained from TreeSB allows us to avoid numerical issues of having to choose the regularization too small, see Section 5. Finally, to the best of our knowledge, this is the first methodology to extend ideas from diffusion-based models to the computation of Wasserstein barycenters. In particular, we believe that the idea of iterative refinement, i.e. , solving the dynamic counterpart of a static problem, plays a key role in the efficiency and scalability of the method.

Notation. For any measurable space ( X , X ) , we denote by P ( X ) the space of probability measures defined on ( X , X ) . Unless specified, X is defined as the Borel sets on X . For any ℓ ∈ N , let P ( ℓ ) = P (( R d ) ℓ ) ; we denote P (1) by P . Assume that X = ( R d ) ℓ for some ℓ ∈ N . For any x ∈ X and any m,n ∈ { 0 , . . . , ℓ } such that m ≤ n , let x m : n = ( x m , x m +1 , . . . , x n ) . Let Leb be the Lebesgue measure. For any non-negative function f : X → R + , such that ∫ X f dLeb < + ∞ , define H( f ) = -∫ X f log f dLeb ∈ ( -∞ , + ∞ ] . For any distribution µ ∈ P ( X ) , we define the entropy of µ as H( µ ) = H(d µ/ dLeb) if µ ≪ Leb and H( µ ) = + ∞ otherwise. For any two arbitrary measures µ and ν defined on ( X , X ) , define the Kullback-Leibler divergence between µ and ν as KL( µ | ν ) = ∫ X log(d µ/ d ν )d µ -∫ X d µ + ∫ X d ν if µ ≪ ν and KL( µ | ν ) = + ∞ otherwise. For any T > 0 , we denote by C([0 , T ] , R d ) the space of continuous functions from [0 , T ] to R d . For any path measure P ∈ P (C([0 , T ] , R d )) , we denote by Ext( P ) ∈ P (2) the coupling between the extremal distributions of P , i.e. , Ext( P ) = P 0 ,T . Note that, for a given coupling π 0 ,T ∈ P (2) , there may exist several path measures P verifying Ext( P ) = π 0 ,T . For any undirected tree T = ( V , E ) with vertices V and edges E , we denote by { v, v ′ } (or { v ′ , v } ) the undirected edge between v ∈ V and v ′ ∈ V , if it exists. Given r ∈ V , we denote by T r = ( V , E r ) the directed version of T rooted in r , where the directed edges E r are uniquely defined from the edges E , see Appendix B for further details. In this case, the edge linking v ∈ V to v ′ ∈ V in T r is denoted by ( v, v ′ ) . Finally, for any integers ( n, K ) ∈ N × N ∗ , we define n mod( K ) as the the remainder of the Euclidean division of n by K .

2 Background and setting

Multi-marginal optimal transport. Let ℓ ∈ N ∗ . Given a cost function c : ( R d ) ℓ +1 → R , a subset S ⊂ { 0 , . . . , ℓ } and a family of probability measures { µ i } i ∈ S ∈ P | S | , mOT consists in solving

where π i is the i -th marginal of π , i.e. , π i ( A ) = π ( proj -1 i ( A )) for any A ∈ B ( R d ) , with proj i : x 0: ℓ ↦→ x i . Given some weights ( w i ) i ∈{ 1 ,...,ℓ } ∈ ( R + ) ℓ , the Wasserstein barycenter between the measures { µ i } i ∈ S is given by π ⋆ 0 in (mOT), in the case where S = { 1 , . . . , ℓ } and c ( x 0: ℓ ) = ∑ ℓ i =1 w i ∥ x 0 -x i ∥ 2 (Peyré et al., 2019). In particular, when w i = 1 /ℓ , the distribution π ⋆ 0 can be regarded as the Fréchet mean (Karcher, 2014) of the measures { µ i } i ∈ S for the Wasserstein distance of order 2. Similarly to OT, (mOT) can be relaxed using the following entropic regularization

where ε > 0 is a hyperparameter and ν is an arbitrary measure defined on (( R d ) ℓ +1 , B (( R d ) ℓ +1 )) .

Link with Schrödinger Bridge. We first recall the relationship between Schrödinger Bridge and EOT. Given T > 0 , Q a (reference) path measure, i.e. , Q ∈ P (C([0 , T ] , R d )) and two measures µ 0 , µ 1 ∈ P ( R d ) , solving the SB problem amounts to finding the path measure P ⋆ defined by

If Q is associated with a Stochastic Differential Equation (SDE) 2 , of the form d X t = -a X t d t +d B t , with a ≥ 0 , then it can be shown, see (Léonard, 2014, Proposition 1) that P ⋆ 0 ,T verifies

This is called the static formulation of SB. It can be shown that solving (static-SB) is equivalent to solving EOT with quadratic cost and regularization ε = 2sinh( aT ) /a if a > 0 , ε = 2 T if a = 0 . Moreover, since P ⋆ = P ⋆ 0 ,T ⊗ Q | 0 ,T , where Q | 0 ,T is the measure Q conditioned on initial and terminal conditions, solving the dynamic problem (SB) is equivalent to solving (static-SB).

Similarly, (EmOT) can be easily rewritten in a static multi-marginal SB fashion

with (d π 0 / dLeb)( x 0: ℓ ) ∝ exp[ -c ( x 0: ℓ ) /ε ](d ν/ dLeb)( x 0: ℓ ) , where π 0 is the reference measure.

Diffusion Schrödinger Bridge. Recently, De Bortoli et al. (2021) introduced Diffusion Schrödinger Bridge (DSB), a numerical scheme to solve (SB). It approximates the iterates of a dynamic version of the Iterative Proportional Fitting (IPF) scheme (Sinkhorn & Knopp, 1967; Knight, 2008; Peyré et al., 2019; Cuturi & Doucet, 2014), which can be described as follows: consider a sequence of path measures ( P n ) n ∈ N such that P 0 = Q and for any n ∈ N

This procedure alternatively projects between the measures with fixed initial distribution and the ones with fixed terminal distribution. For the first iteration, we get that P 1 = µ 1 ⊗ Q | T . Assuming that Q is given by d X t = f t ( X t )d t +d B t , with f : [0 , T ] × R d → R d , then P 1 is associated with the time-reversal of this SDE initialized at µ 1 . The time-reversal of an SDE has been derived under mild assumptions on the drift and diffusion coefficients (Haussmann & Pardoux, 1986; Cattiaux et al., 2021). In this case, we have ( Y T -t ) t ∈ [0 ,T ] ∼ P 1 , with Y 0 ∼ µ 1 and

where p t is the density of P 0 t w.r.t. the Lebesgue measure. The score ∇ log p t is estimated using score matching techniques (Hyvärinen, 2005; Vincent, 2011). The first iterate of DSB, P 1 , corresponds to a denoising diffusion model (Ho et al., 2020; Song et al., 2021). DSB iterates further and not only parameterizes the backward process but also the forward process. It can therefore be seen as a refinement of diffusion models drawing a bridge between generative modeling and optimal transport.

Tree-based framework. Consider an undirected tree T = ( V , E ) , with vertices V and edges E , such that V is identified with { 0 , . . . , ℓ } . Inspired by Haasler et al. (2021), we restrict our study of (EmOT), to the case where the cost function c is the tree-structured quadratic cost derived from T

where w v,v ′ is a weight on the edge { v, v ′ } , which links v to v ′ (and v ′ to v ). Furthermore, as in Haasler et al. (2021), we choose S , i.e. , the set of vertices of T with constrained marginals, to coincide with the leaves of T . This framework recovers important applications, from Wasserstein barycenters to Wasserstein propagation, see Solomon et al. (2014, 2015). We emphasize that it differs from an OT problem defined on the space of graphs (Chen et al., 2016). Here, each node represents a probability measure (observed or to be inferred) and each edge represents a coupling between two distributions.

We consider an arbitrary vertex r ∈ V and choose ν in (EmOT) such that (d ν/ dLeb)( x 0: ℓ ) = φ r ( x r ) , where φ r is a density defined on R d . Due to the form of ν and c , the reference measure π 0 in (mSB-like) is therefore a probability distribution which factorizes along T r = ( V , E r ) , the directed version of T rooted in r . We refer to Appendix B for more details on the notion of directed trees. In this setting, (EmOT) is equivalent to the tree-based problem

where π 0 v ′ | v ( · | x v ) = N( x v , ε/ (2 w v,v ′ )I d ) and π 0 r ≪ Leb with density φ r . In a manner akin to Haasler et al. (2021), we thus establish, in continuous state-space, the correspondence between (TreeSB), a static tree-based version of SB, and a version of EmOT with tree-structured cost (1). In our work, we make the following assumption on the constrained marginals { µ i } i ∈ S .

In what follows, we define K as the number of leaves of T , denoting S = { i 0 , . . . , i K -1 } , and define the horizon times T v,v ′ = ε/ (2 w v,v ′ ) for any { v, v ′ } ∈ E . For any i k ∈ S , we will denote by T k = ( V , E k ) the directed version of T rooted in the leaf i k . In the next section, we present our dynamic method to solve (TreeSB), called Tree-based Diffusion Schrödinger Bridge .

3 Tree-based Diffusion Schrödinger Bridge

In this section, we present a method to solve (TreeSB) in the case where r ∈ S , i.e. , r is a leaf of T . We refer to Appendix E for the extension to the case where r ∈ V \ S . Without loss of generality, see Appendix E, we assume that r = i K -1 and choose φ r = d µ i K -1 / dLeb , such that π 0 i K -1 = µ i K -1 .

Dynamic approach to mIPF. In order to approximate solutions of (TreeSB), we consider the multi-marginal extension of the IPF algorithm, denoted by mIPF. Namely, we define a sequence of probability distributions ( π n ) n ∈ N such that for any n ∈ N

where k n = ( n -1) mod( K ) and ( k n +1) is identified with n mod( K ) . We define a mIPF cycle as a sequence of K consecutive mIPF updates. In particular, each marginal constraint is considered exactly once during one mIPF cycle. In a practical setting, our main aim is to sample from the (mIPF) iterates at the lowest cost. Although these updates can be made explicit, see Marino & Gerolin (2020) for instance, direct sampling is unfeasible in practice when d is large. To overcome this limitation, we suggest to compute these iterates in a dynamic fashion with equivalent path measures.

Since π 0 factorizes along T , see (2), one can show that the iterates of (mIPF) also factorize along T , see Section 4. Since these iterates all have a constrained marginal, we obtain the following decomposition for any n ∈ N : π n = µ i kn ⊗ ( v,v ′ ) ∈ E kn π n v ′ | v where E k n denotes the set of edges of the directed tree T k n . Then, our approach consists in computing dynamic iterates, i.e. , path measures, along the edges of T that coincide on their extremal times with the static iterates ( π n ) n ∈ N . Namely, for any n ∈ N , for any edge ( v, v ′ ) ∈ E k n , we define a path measure P n ( v,v ′ ) ∈ P (C([0 , T v,v ′ ] , R d )) such that Ext( P n ( v,v ′ ) ) = π n v,v ′ , where Ext( P n ( v,v ′ ) ) stands for the joint distribution of P n ( v,v ′ ) at times 0 and T v,v ′ . In particular, it comes that π n v ′ | v = P n ( v,v ′ ) ,T v,v ′ | 0 . Using the tree-based form of the (mIPF) iterates, we can thus sample from π n by (i) following the directed edges of T k n , (ii) diffusing along them the corresponding path measures ( P n ( v,v ′ ) ) ( v,v ′ ) ∈ E kn and (iii) picking the samples on the vertices. When T is a bridge-shaped tree (2 vertices, 1 edge), it simply reduces to the dynamic reformulation of the IPF scheme. In what follows, we explain how to obtain our dynamic sequence.

Definition of the dynamic iterates. We first compute the iterate P 0 , corresponding to the dynamic version of π 0 defined (2), in Proposition 1. Then, we build the following iterates by recursion on n ∈ N and prove their well-posedness in Proposition 2.

Proposition 1. Let T K -1 = ( V , E K -1 ) , the directed tree associated with T = ( V , E ) and root i K -1 . Then, for any ( v, v ′ ) ∈ E K -1 , there exists P 0 ( v,v ′ ) ∈ P (C([0 , T v,v ′ ] , R d )) with Ext( P 0 ( v,v ′ ) ) = π 0 ( v,v ′ ) and such that P 0 ( v,v ′ ) | 0 is the distribution of ( B t ) t ∈ [0 ,T v,v ′ ] , recalling that T v,v ′ = ε/ (2 w v,v ′ ) .

Before deriving the dynamic counterpart of the (mIPF) iterates, we introduce several definitions. For any path measure P , we denote by P R the time-reversal of P . For any directed tree and any vertex v of this tree, p ( v ) refers to the (unique) parent of v , and c ( v ) to the unique child of v when it exists, see Appendix B for more details.

Let n ∈ N . Assume that we have defined the sequence of our dynamic iterates ( P m ( v,v ′ ) ) ( v,v ′ ) ∈ E km ,m ≤ n up to stage n .

Consider the path P n = { ( v j , v j +1 ) } J j =1 in the directed tree T k n such that v 1 = i k n and v J +1 = i k n +1 . In particular, for any ( v, v ′ ) ∈ E k n +1 , either ( v ′ , v ) ∈ P n or ( v, v ′ ) ∈ E k n \ P n . This is illustrated in Figure 1 when V = { 0 , 1 , 2 , 3 , 4 } , S = { 2 , 3 , 4 } , i k = 3 and i k +1 = 4 : in this case, P = { (3 , 1) , (1 , 0) , (0 , 4) } and (1 , 2) is the only edge common to E k and E k +1 .

Consider now the directed tree T k n +1 . We define the ( n +1) -th iterate of our dynamic sequence by recursion on the edges of this tree, following the breadth-first order. In this order, ( i k n +1 , c ( i k n +1 )) = ( v J +1 , v J ) is the first edge considered.

First, we define P n +1 ( v J +1 ,v J ) = µ i kn +1 ⊗ ( P n ( v J ,v J +1 ) ) R | 0 . In the case of a bridge-shaped tree, this is exactly the ( n +1) -th update described in DSB. Then, for any ( v, v ′ ) ∈ E k n +1 \{ ( v J +1 , v J ) } ,

Proposition 2. Consider the sequence of dynamic iterates defined by (a) and (b). Then, for any n ∈ N and any ( v, v ′ ) ∈ E k n , P n ( v,v ′ ) ∈ P (C([0 , T v,v ′ ] , R d )) and we have Ext( P n ( v,v ′ ) ) = π n ( v,v ′ ) .

Proposition 2 highlights the equivalence between the (mIPF) iterates and our dynamic iterates. These path measures are defined iteratively, by following the updates (a) and (b) along the edges of T . The key observation here is that the computation of each dynamic iterate reduces to a sequence of updates (b) on a path linking two leaves of T . We emphasize that our iterates could be similarly obtained by directly considering a dynamic formulation of (TreeSB) and introducing the formalism of deterministic time branching processes. We leave the study of this problem for future work. We now get into the details of our practical implementation, which relies on score-based methods.

Approximation of the dynamic iterates. The time-reversal operated in the update (b) can be computed explicitly, see Haussmann & Pardoux (1986) for instance. Indeed, assuming that P n ( v ′ ,v ) is associated with d X t = f t,v ′ ,v ( X t )d t + d B t with X 0 ∼ π n v ′ , then, under mild conditions, its time-reversal ( P n ( v ′ ,v ) ) R is associated with d Y t = {-f T -t,v ′ ,v + ∇ log p v ′ ,v,T -t } ( Y t )d t + d B t with Y 0 ∼ π n +1 v , where p v ′ ,v,t is the density of P n ( v ′ ,v ) ,t w.r.t. the Lebesgue measure. The score ∇ log p v ′ ,v,T -t can then be approximated using score-matching techniques (Hyvärinen, 2005; Vincent, 2011) which are now ubiquitous in diffusion models (Song et al., 2021) and used in DSB De Bortoli et al. (2021). Therefore, at iteration ( n +1) , the update (b) is similar to the one of DSB for each edge on the path joining i k n and i k n +1 . In practice, we parameterize the drifts f t,v,v ′ for any { v, v ′ } ∈ E with neural networks f t,θ v,v ′ and use the mean-matching loss introduced by De Bortoli et al. (2021). Note that doing so, we obtain 2 | E | neural networks. The whole procedure consisting in computing our dynamic iterates using the DSB framework is called Tree-based Diffusion Schrödinger Bridge (TreeDSB) and is summarized in Algorithm 1.

Algorithm 1 TreeDSB (Training)

The algorithm is initialized with f t,θ v,v ′ = 0 for all { v, v ′ } ∈ E . This corresponds to Brownian motion dynamics when sampling at the first iteration of TreeDSB, see Proposition 1. Note that in Algorithm 1, when we sample from P n ( v j ,v j +1 ) , we update f t,θ v j +1 ,v j which will be used to sample from P n +1 ( v j +1 ,v j ) in the next iterations. In order to sample from the dynamics P n ( v j ,v j +1 ) , we consider its Euler-Maruyama discretization, see Appendix F for more details. We describe the different steps of the algorithm in the case of a toy example below, see Figure 2 for an illustration.

TreeDSB on a toy tree. We consider a star-shaped tree with three leaves denoted { 1 , 2 , 3 } and its central node { 0 } . Following (2), we define π 0 with r = 3 and φ r = (d µ 3 / dLeb) . During the first iteration of TreeDSB, T is rooted at vertex 3 and we compute samples from the forward path P 0 = { (3 , 0) , (0 , 1) } with Brownian motions, see Proposition 1, in order to learn the backward path { (1 , 0) , (0 , 3) } . In the next iteration, we re-root the tree T at vertex 1 and consider the forward path P 1 = { (1 , 0) , (0 , 2) } , where the edges (1 , 0) and (0 , 2) are respectively given by the first iteration and the initialisation. This highlights that TreeDSB does not require to update the whole tree . The following iterations are done similarly. At each iteration n ∈ N , we sample from π n by first sampling from µ k n at leaf i k n and then following the parameterized SDEs on the directed edges of T k n .

4 Theoretical properties of mIPF

In this section, we study some of the theoretical properties of the static iterates ( π n ) n ∈ N , that are equivalent to our dynamic iterates according to Proposition 2. In the case where the cost function c is bounded in (EmOT), results of convergence of (mIPF) exist (Marino & Gerolin, 2020; Carlier, 2022). However, our setting does not satisfy their assumptions, since our transport cost is quadratic and the measures are defined on R d . In what follows, we provide the first non-quantitative convergence results for (mIPF) in a non-compact setting.

For the rest of the section, we consider a static formulation of the multi-marginal Schrödinger bridge problem which is more general than (TreeSB), defined as

where S ⊂ { 0 , . . . , ℓ } , π 0 ∈ P , { µ i } i ∈ S ∈ P | S | . We consider the following set of assumptions.

A1. There exists a family of measures { ν i } i ∈{ 0 ,...,ℓ } defined on ( R d , B ( R d )) such that π 0 ≪ ⊗ ℓ i =0 ν i with density h = d π 0 / (d ⊗ ℓ i =0 ν i ) and µ i ≪ ν i with density r i = d µ i / d ν i for any i ∈ S .

A3. There exists a family of probability measures { ˜ µ j } j ∈{ 0 ,...,ℓ }\ S such that π 0 ∼ ˜ π 0 , where ˜ π 0 = ⊗ i ∈ S µ i ⊗ j ∈{ 0 ,...,ℓ }\ S ˜ µ j .

In particular, (static-mSB) recovers (TreeSB) by considering ν i = Leb for any i ∈ { 0 , . . . , ℓ } and h ( x 0: ℓ ) = φ r ( x r ) exp[ -c ( x 0: ℓ ) /ε ] in A 1. We detail in Appendix D how A 2 and A 3 can be met in (TreeSB). Under these assumptions, the multi-marginal Schrödinger Bridge exists.

Proposition 3. Assume A 1 and A 2 . Then, there exists a unique solution π ⋆ to (static-mSB) . In addition, assume A 3 . Then, there exists a family { ψ ⋆ i } i ∈ S of measurable functions ψ ⋆ i : R d → R such that

In order to establish the existence and uniqueness result of Proposition 3, we extend results from Nutz (2021) to the multi-marginal setting. A consequence of Proposition 3 is that the iterates of (mIPF) can be described using potentials.

Corollary 4. Assume A 1 , A 2 and A 3 . Let ( π n ) n ∈ N be the sequence given by (mIPF) . Then, for any n ∈ N ∗ with k n = ( n -1) mod( K ) and q n ∈ N such that n = q n K + k n +1 , there exists a family of measurable functions { ψ q n +1 i 0 , . . . , ψ q n +1 i kn , ψ q n i kn +1 , . . . , ψ q n i K -1 } such that

In the tree-based setting, Corollary 4 explains why the (mIPF) iterations preserve the tree-based Markovian nature of π 0 . We now prove that the marginal π n i converges to µ i for any i ∈ S , as n goes to infinity, i.e. , we have marginal convergence on the leaves of T .

Proposition 5. Assume A 1 and A 2 . Let ( π n ) n ∈ N be the sequence given by (mIPF) . Then, we have lim n →∞ ∥ π n i -µ i ∥ TV = 0 for any i ∈ S .

The previous result does not ensure the convergence of ( π n ) n ∈ N to the solution to (static-mSB). In particular, Proposition 5 does not provide the convergence of the marginals on the nodes v ∈ V \ S , which is key to compute regularized Wasserstein barycenters with TreeDSB. Relying on additional assumptions, we now derive the convergence of (mIPF).

A5. There exist ¯ c ∈ (0 , ∞ ) such that exp( ψ n i k -ψ n +1 i k ) ≤ ¯ c , for any n ∈ N , any k ∈ { 0 , . . . , K -2 } .

These assumptions can be seen as multi-marginal extensions of the ones of Ruschendorf (1995), see Appendix D for a discussion and examples.

Proposition 6. Assume A 1 , A 2 , A 3 , A 4 and A 5 . Let ( π n ) n ∈ N be the sequence given by (mIPF) . Then, we have lim n →∞ ∥ π n -π ⋆ ∥ TV = 0 , where π ⋆ is given in Proposition 3.

To the best of our knowledge, Proposition 6 is the first convergence result of (mIPF) without assuming that the space is compact or that the cost is bounded. We highlight that traditional techniques to prove the convergence of IPF cannot be easily extended to the multi-marginal setting as pointed by Carlier (2022). In the case of bounded cost, quantitative results exist (Marino & Gerolin, 2020; Carlier, 2022). We leave the study of such results in the unbounded cost setting for future work.

5 Application to Wasserstein barycenters

Although Algorithm 1 can be applied to trees T with fixed marginals on the leaves, one case of particular interest is star-shaped trees, i.e. , trees with a central node, denoted by index 0 , and such that S = { 1 , . . . , ℓ } (see Figure 2 for an illustration with ℓ = 3 ). In this section, we draw a link between (TreeSB) and regularized Wasserstein barycenters. We recall the definition of the Wasserstein distance of order 2 with ε -entropic regularization between µ and ν (Peyré et al., 2019, Chapter 4)

In this work, we consider the ( ℓε, ( ℓ -1) ε ) -doubly-regularized Wasserstein-2 barycenter problem (Chizat, 2023) defined as follows

where ( w i ) i ∈{ 1 ,...,ℓ } ∈ (0 , + ∞ ) ℓ . The following proposition shows the equivalence between the barycenter problem (regWB) and the multi-marginal Schrödinger bridge problem (TreeSB) over T . In particular, it allows us to use TreeDSB to estimate the solution µ ⋆ ε of (regWB).

Proposition 7. Let ε > 0 . Assume A 0. Also assume that T is a star-shaped tree with central node indexed by 0 , and that the reference measure of (TreeSB) defined in (2) verifies r = i K -1 and φ r = d µ i K -1 / dLeb > 0 . Under A 2 , (regWB) has a unique solution π ⋆ 0 , where π ⋆ solves (TreeSB) .

The proof of this result is postponed to Appendix D. More generally, we show in Appendix D that, for any tree T , (TreeSB) is equivalent to a regularized version of the Wasserstein propagation problem (Solomon et al., 2014, 2015). Moreover, we present in Appendix E an extension of Proposition 7 in the case where the chosen root r is not a leaf of T . We finally emphasize that the formulation of (regWB) leads to a minimization of the entropy of the barycenter. In particular, this allows us to choose ε reasonably large in TreeDSB, which is a stability advantage compared to other regularized methods which do not consider this further regularization.

6 Related work

Diffusion Schrödinger Bridge. Schrödinger Bridges (Schrödinger, 1932) have been extensively studied using tools from stochastic control and probability theory (Léonard, 2014; Dai Pra, 1991; Chen et al., 2021). More recently, algorithms were proposed to efficiently approximate such bridges in the context of machine learning. In particular, De Bortoli et al. (2021) proposed DSB while Vargas et al. (2021); Chen et al. (2022) developed related algorithms. In Chen et al. (2023), the authors study a multi-marginal version of DSB in a linear tree-based setting, where the set of observed nodes is the whole set of vertices. However, contrary to our setting, Chen et al. (2023) introduced a momentum variable. This allows for smoother trajectories which are desirable for single-cell trajectories applications and correspond to some spline interpolation in the space of probability measures (Chen et al., 2018). A general framework for tree-based static Schrödinger Bridges on discrete state-spaces was given in Haasler et al. (2021). In this work, we extend their formulation to a dynamic and continuous setting, see Appendix D for more a thorough comparison.

Wasserstein barycenters. The notion of Wasserstein barycenter was first introduced in Rabin et al. (2012) and then later studied in Agueh & Carlier (2011). The algorithms to solve this problem can be split into two families: the in-sample based approaches and the parametric ones. In-sample approaches require access to all the measures µ i which are assumed to be empirical measures (Cuturi &Doucet, 2014; Benamou et al., 2015; Solomon et al., 2015). Related to this class of algorithms is the semi-discrete approach, which aims at computing a Wasserstein barycenter between continuous distribution but rely on a discretization of the barycenter (Claici et al., 2018; Staib et al., 2017; Mi et al., 2020). Most recent approaches do not rely on a discrete representation of the barycenter, but instead parameterize it using neural networks. These approaches can be further split into two categories. First, measure-based optimization approaches parameterize the measures using a neural network. This is the case of Cohen et al. (2020), where the barycenter is given by a generative model, which is then optimized . Fan et al. (2020) introduce an optimization procedure which relies on a min-max-min problem using the framework of Makkuva et al. (2020). More recently, Korotin et al. (2022) considered a fixed point-based algorithm introduced in Álvarez-Esteban et al. (2016) to update a generative model parametrizing the barycenter. On the one hand, potential-based methods rely on a dual formulation of the barycenter. Korotin et al. (2021) parameterized the dual potentials using Input Convex Neural Network and considered regularizing losses imposing conjugacy and congruency. On the other hand, Li et al. (2020) consider a dual version of the regularized Wasserstein barycenter problem contrary to other works. Our approach applied to start-shaped trees also approximates a regularized Wasserstein barycenter. However, contrary to Li et al. (2020), we do not consider a parameterization of the potentials in the static setting but instead, parameterize the drift of an associated dynamic formulation using Schrödinger bridges. To the best of our knowledge TreeDSB is the first approach leveraging DSB-like algorithms to compute Wasserstein barycenters.

7 Experiments

In our experiments 3 , we illustrate the performance of TreeDSB to compute entropic regularized Wasserstein barycenters for various tasks . We choose to compare our method with state-of-the-art regularized algorithms: fast free-support Wasserstein barycenter (fsWB) (Cuturi & Doucet, 2014) , and continuous regularized Wasserstein barycenter (crWB) (Li et al., 2020). In all of our settings, we consider a star-shaped tree with K leaves and edge weights that are equal to 1 /K , resulting in a

sequential training procedure over 2 K neural networks. The initial diffusion is always a Brownian motion parameterized as explained in Proposition 1. Hence, the time horizon on each edge is defined by T = Kε/ 2 . The order of the leaves is randomly shuffled between the mIPF cycles. We consider 50 steps for the time discretization on [0 , T ] . We refer to Appendix G for details on the choice of the schedule, the architecture of the neural networks and the settings of our experiments.

Synthetic two dimensional datasets. We first illustrate TreeDSB in a synthetic two dimensional setting. We consider three different datasets Swissroll (vertex 0, starting node r ), Circle (vertex 2) and Moons (vertex 3) and compute their Wasserstein barycenter (vertex 1) by running TreeDSB for 50 mIPF cycles with ε = 0 . 1 . In Figure 3, we show the estimated densities of the datasets on the leaves of the tree (we emphasize that the distributions plotted on each leaf are generated from the central barycenter measure). In Figure 4, we observe the consistency between the barycenters generated from the different leaves. In Appendix G, we present additional results for this setting.

15

-10

mIPF cycle:50,fromvertex1 to 0

15

10

5

0

-5

10

15

15 10

-5

0

5

10

15

mlPF cycle:50,from vertex1 to 2

mlPF cycle:50,from vertex1 to 3

15

15

10

10

5

5

-5

10

10

15

15

0

5

10

15

-15 10

5

0

5

10

15

mlPFcycle:50,fromvertex0to1

mlPFcycle:50,fromvertex2to1

mlPFcycle:50,fromvertex3to1

15

15

15

10

10

10

5

5

5

0

0

0

-5

-5

-5

-10

-10

-10

-15

-15

15

15

-10

-5

0

5

10

15

15

-10

-5

0

5

10

15

15

-10

-5

0

5

10

15

Synthetic Gaussian datasets. Next, we consider three independent Gaussian distributions with zero mean and random non-diagonal covariance matrices whose conditional number is less than 10, following Fan et al. (2020). In this case, the non-regularized barycenter can be exactly computed. To evaluate the performance of the algorithms, we use the Bures-Wasserstein Unexplained Variance Percentage (UVP), following (Korotin et al., 2021, Section 5). Given a target distribution µ ⋆ ∈ P and some approximation µ ∈ P , we define

where BW 2 2 ( µ, µ ⋆ ) = W 2 2 (N( E [ µ ] , Cov( µ )) , N( E [ µ ⋆ ] , Cov( µ ⋆ )) .

In this setting, we choose µ ⋆ to be the non-regularized barycenter and assess the dependency w.r.t. the dimension of the algorithms using the BW 2 2 -UVP metric. In Table 1, we compare ourselves with the two regularized methods Li et al. (2020) ( L 2 -reg. equal to 10 -4 ) and Cuturi & Doucet (2014). We run TreeDSB for 10 mIPF cycles with ε = 0 . 1 . Bold numbers represent the best values up to statistical significance. While Li et al. (2020) and Cuturi & Doucet (2014) enjoy better performance in low dimensions ( d = 2 ), TreeDSB outperforms these methods as the dimension increases.

MNIST Wasserstein barycenter. We then turn to an image experiment using MNIST dataset (LeCun, 1998). Here, an image is not considered as a 2D-dimensional distribution as in Cuturi & Doucet (2014) and Li et al. (2020), but as a sample from a high-dimensional probability measure ( d = 784) . We aim at computing a Wasserstein barycenter between the digits 2 , 4 and 6 . To do so, we

run TreeDSB for 10 mIPF cycles with r that corresponds to the digit 6 and ε = 0 . 5 . In Figure 5, we display samples from the estimated marginals on the leaves, to assess the reconstruction of the digits 2 , 4 and 6 , and samples from the barycenter, obtained by diffusing from the leaf corresponding to the digit 6 . Our results prove the scalability of TreeDSB to the high-dimensional setting, compared to state-of-the-art regularized methods. Additional results on MNIST dataset are given in Appendix G.

D

6

6

6

6

6

22222

2a222

22222

2222

441

4

ht

4

4449

4

4444

hhthh

Subset posterior aggregation. Finally, we evaluate TreeDSB in the context of Bayesian fusion (Srivastava et al., 2018), also called posterior aggregation. Given a Bayesian model and a dataset partitioned into several shards, this task aims at recovering the full data posterior distribution from the posterior distributions computed on each shard.

In particular, it has been proved that the barycenter of the subdataset posteriors is close to the full data posterior under mild assumptions (Srivastava et al., 2018). Here, we consider a logistic regression model applied to the wine dataset 4 ( d = 42 ) and proceed as follows. We first split this dataset into 3 subsets, with or without heterogeneity, and estimate the posterior parameters on each shard. Then, we draw samples from the obtained logistic distributions to define µ 1 , µ 2 , µ 3 . Then, we compute the Wasserstein barycenter of these measures, and compare it to the posterior computed on the full dataset. As in the synthetic Gaussian experiment, we run TreeDSB for 10 mIPF cycles ε = 0 . 1 and we compare ourselves with Li et al. (2020) ( L 2 -reg. equal to 10 -4 ) and Cuturi & Doucet (2014). We evaluate the methods using the BW 2 2 -UVP metric, where µ ⋆ is the estimated full data posterior, and report the results in Table 2. In both settings, we observe that our method outperforms existing regularized methods to compute Wasserstein barycenters.

Limitations. One of the main limitation of entropic regularized OT approach is that their behavior is usually badly conditioned as ε → 0 . In our setting, we observe that if ε , or equivalently T , is too low then the algorithm becomes less stable as the training of the models slows down. In the future, we plan to mitigate this issue by incorporating fixed point techniques like the one used in Korotin et al. (2022). Finally, since our algorithm is based on DSB (De Bortoli et al., 2021), it suffers from the same limitations. In particular, training different neural networks iteratively incurs some bias in the SDE which is harmful for large number of mIPF iterations.

8 Discussion

In this paper, we introduced Tree-based Diffusion Schrödinger Bridge (TreeDSB) a scalable scheme to approximate solutions of entropic-regularized multi-marginal Optimal Transport (mOT) problems. Our methodology leverages tools from the diffusion model literature and extends Diffusion Schrödinger Bridge (De Bortoli et al., 2021). In particular, it approximates the iterates of the multimarginal Iterative Proportional Fitting (mIPF) algorithm, for which we prove its convergence under mild assumptions. We illustrate the efficiency of TreeDSB for image processing and Bayesian fusion, using the link between mOT and Wasserstein barycenters. In future work, we would like to study quantitative convergence bounds for mIPF in the unbounded cost setting. Another line of work would be to scale TreeDSB to higher dimensional problems building on recent developments in the diffusion model and flow matching community (Lipman et al., 2023; Peluchetti, 2023; Shi et al., 2023).

Acknowledgments

We thank James Thornton for the DSB codebase 5 and useful discussions. AD acknowledges support from the Lagrange Mathematics and Computing Research Center. AD and MN would like to thank the Isaac Newton Institute for Mathematical Sciences for support and hospitality during the programme The mathematical and statistical foundation of future data-driven engineering when work on this paper was undertaken. MN acknowledges funding from the grant SCAI (ANR-19-CHIA-0002).

References

Acciaio, B., Backhoff-Veraguas, J., and Carmona, R. Extended mean field control problems: stochastic maximum principle and transport perspective. SIAM journal on Control and Optimization , 57 (6):3666-3693, 2019.

Agueh, M. and Carlier, G. Barycenters in the Wasserstein space. SIAM Journal on Mathematical Analysis , 43(2):904-924, 2011.

Álvarez-Esteban, P. C., Del Barrio, E., Cuesta-Albertos, J., and Matrán, C. A fixed-point approach to barycenters in Wasserstein space. Journal of Mathematical Analysis and Applications , 441(2): 744-762, 2016.

Bayraktar, E., Cox, A. M., and Stoev, Y. Martingale optimal transport with stopping. SIAM Journal on Control and Optimization , 56(1):417-433, 2018.

Benamou, J.-D., Carlier, G., Cuturi, M., Nenna, L., and Peyré, G. Iterative Bregman projections for regularized transportation problems. SIAM Journal on Scientific Computing , 37(2):A1111-A1138, 2015.

Bunne, C., Papaxanthos, L., Krause, A., and Cuturi, M. Proximal optimal transport modeling of population dynamics. In International Conference on Artificial Intelligence and Statistics , pp. 6511-6528. PMLR, 2022.

Carion, N., Massa, F., Synnaeve, G., Usunier, N., Kirillov, A., and Zagoruyko, S. End-to-end object detection with transformers. In Computer Vision-ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part I 16 , pp. 213-229. Springer, 2020.

Carlier, G. On the linear convergence of the multimarginal Sinkhorn algorithm. SIAM Journal on Optimization , 32(2):786-794, 2022.

Cattiaux, P., Conforti, G., Gentil, I., and Léonard, C. Time reversal of diffusion processes under a finite entropy condition. arXiv preprint arXiv:2104.07708 , 2021.

Chen, T., Liu, G.-H., and Theodorou, E. A. Likelihood training of Schrödinger bridge using forwardbackward SDEs theory. International Conference on Learning Representations , 2022.

Chen, T., Liu, G.-H., Tao, M., and Theodorou, E. A. Deep momentum multi-marginal Schrödinger bridge. 2023. doi: 10 . 48550/ARXIV . 2303 . 01751.

Chen, Y., Georgiou, T., Pavon, M., and Tannenbaum, A. Robust transport over networks. IEEE Transactions on Automatic Control , 62(9):4675-4682, 2016.

Chen, Y., Conforti, G., and Georgiou, T. T. Measure-valued spline curves: An optimal transport viewpoint. SIAM Journal on Mathematical Analysis , 50(6):5947-5968, 2018.

Chen, Y., Georgiou, T. T., and Pavon, M. Optimal transport in systems and control. Annual Review of Control, Robotics, and Autonomous Systems , 4, 2021.

Chizat, L. Doubly regularized entropic Wasserstein barycenters. arXiv preprint arXiv:2303.11844 , 2023.

Claici, S., Chien, E., and Solomon, J. Stochastic Wasserstein barycenters. In International Conference on Machine Learning , pp. 999-1008. PMLR, 2018.

Cohen, S., Arbel, M., and Deisenroth, M. P. Estimating barycenters of measures in high dimensions. arXiv preprint arXiv:2007.07105 , 2020.

Cotar, C., Friesecke, G., and Klüppelberg, C. Density functional theory and optimal transportation with Coulomb cost. Communications on Pure and Applied Mathematics , 66(4):548-599, 2013.

Csiszár, I. I-divergence geometry of probability distributions and minimization problems. The Annals of Probability , pp. 146-158, 1975.

Cuturi, M. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in Neural Information Processing Systems , 26, 2013.

Cuturi, M. and Doucet, A. Fast computation of Wasserstein barycenters. In International Conference on Machine Learning , pp. 685-693. PMLR, 2014.

Dai Pra, P. A stochastic control approach to reciprocal diffusion processes. Applied Mathematics and Optimization , 23(1):313-329, 1991.

De Bortoli, V., Thornton, J., Heng, J., and Doucet, A. Diffusion Schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems , 34: 17695-17709, 2021.

Dupuis, P. and Ellis, R. S. A weak convergence approach to the theory of large deviations . John Wiley & Sons, 2011.

Eisenberger, M., Toker, A., Leal-Taixé, L., and Cremers, D. Deep shells: Unsupervised shape correspondence with optimal transport. Advances in Neural Information Processing Systems , 33: 10491-10502, 2020.

Fan, J., Taghvaei, A., and Chen, Y. Scalable computations of Wasserstein barycenter via input convex neural networks. arXiv preprint arXiv:2007.04462 , 2020.

Feydy, J., Charlier, B., Vialard, F.-X., and Peyré, G. Optimal transport for diffeomorphic registration. In Medical Image Computing and Computer Assisted Intervention- MICCAI 2017: 20th International Conference, Quebec City, QC, Canada, September 11-13, 2017, Proceedings, Part I 20 , pp. 291-299. Springer, 2017.

Flamary, R., Courty, N., Gramfort, A., Alaya, M. Z., Boisbunon, A., Chambon, S., Chapel, L., Corenflos, A., Fatras, K., Fournier, N., et al. POT: Python optimal transport. The Journal of Machine Learning Research , 22(1):3571-3578, 2021.

Fortet, R. Résolution d'un système d'équations de M. Schrödinger. Journal de Mathématiques Pures et Appliquées , 19(1-4):83-105, 1940.

Haasler, I., Ringh, A., Chen, Y., and Karlsson, J. Multimarginal Optimal Transport with a treestructured cost and the Schrödinger bridge problem. SIAM Journal on Control and Optimization , 59(4):2428-2453, 2021.

Haussmann, U. G. and Pardoux, E. Time reversal of diffusions. The Annals of Probability , pp. 1188-1205, 1986.

Ho, J., Jain, A., and Abbeel, P. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851, 2020.

Hyvärinen, A. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research , 6(4), 2005.

Karcher, H. Riemannian center of mass and so called Karcher mean. arXiv preprint arXiv:1407.2087 , 2014.

Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Knight, P. A. The Sinkhorn-Knopp algorithm: convergence and applications. SIAM Journal on Matrix Analysis and Applications , 30(1):261-275, 2008.

Kober, H. A theorem on Banach spaces. Compositio Mathematica , 7:135-140, 1940.

Koller, D. and Friedman, N. Probabilistic Graphical Models: Principles and Techniques . MIT press, 2009.

Korotin, A., Li, L., Solomon, J., and Burnaev, E. Continuous Wasserstein-2 barycenter estimation without minimax optimization. arXiv preprint arXiv:2102.01752 , 2021.

Korotin, A., Egiazarian, V., Li, L., and Burnaev, E. Wasserstein iterative networks for barycenter estimation. arXiv preprint arXiv:2201.12245 , 2022.

Kullback, S. Probability densities with given marginals. The Annals of Mathematical Statistics , 39 (4):1236-1243, 1968.

LeCun, Y. The MNIST database of handwritten digits. http://yann. lecun. com/exdb/mnist/ , 1998.

Léonard, C. A survey of the Schrödinger problem and some of its connections with optimal transport. Discrete & Continuous Dynamical Systems-A , 34(4):1533-1574, 2014.

Li, L., Genevay, A., Yurochkin, M., and Solomon, J. M. Continuous regularized Wasserstein barycenters. Advances in Neural Information Processing Systems , 33:17755-17765, 2020.

Lipman, Y., Chen, R. T., Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. International Conference on Learning Representations , 2023.

Luo, Z.-Q. and Tseng, P. On the convergence of the coordinate descent method for convex differentiable minimization. Journal of Optimization Theory and Applications , 72(1):7-35, 1992.

Makkuva, A., Taghvaei, A., Oh, S., and Lee, J. Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning , pp. 6672-6681. PMLR, 2020.

Marino, S. D. and Gerolin, A. An optimal transport approach for the Schrödinger bridge problem and convergence of Sinkhorn algorithm. Journal of Scientific Computing , 85(2):27, 2020.

Mi, L., Yu, T., Bento, J., Zhang, W., Li, B., and Wang, Y. Variational Wasserstein barycenters for geometric clustering. arXiv preprint arXiv:2002.10543 , 2020.

Minsker, S., Srivastava, S., Lin, L., and Dunson, D. Scalable and robust Bayesian inference via the median posterior. In International Conference on Machine Learning , pp. 1656-1664. PMLR, 2014.

Nichol, A. Q. and Dhariwal, P. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning , pp. 8162-8171, 2021.

Nutz, M. Introduction to entropic optimal transport, 2021.

Pele, O. and Werman, M. Fast and robust Earth mover's distances. In 2009 IEEE 12th International Conference on Computer Vision , pp. 460-467. IEEE, 2009.

Peluchetti, S. Diffusion bridge mixture transports, Schrödinger bridge problems and generative modeling. arXiv preprint arXiv:2304.00917 , 2023.

Peyré, G., Cuturi, M., et al. Computational optimal transport: With applications to data science. Foundations and Trends® in Machine Learning , 11(5-6):355-607, 2019.

Rabin, J., Peyré, G., Delon, J., and Bernot, M. Wasserstein barycenter and its application to texture mixing. In Scale Space and Variational Methods in Computer Vision: Third International Conference, SSVM 2011, Ein-Gedi, Israel, May 29-June 2, 2011, Revised Selected Papers 3 , pp. 435-446. Springer, 2012.

Ruschendorf, L. Convergence of the Iterative Proportional Fitting procedure. The Annals of Statistics , pp. 1160-1174, 1995.

Schiebinger, G., Shu, J., Tabaka, M., Cleary, B., Subramanian, V., Solomon, A., Gould, J., Liu, S., Lin, S., Berube, P., et al. Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell , 176(4):928-943, 2019.

Schmitz, M. A., Heitz, M., Bonneel, N., Ngole, F., Coeurjolly, D., Cuturi, M., Peyré, G., and Starck, J.-L. Wasserstein dictionary learning: Optimal transport-based unsupervised nonlinear dictionary learning. SIAM Journal on Imaging Sciences , 11(1):643-678, 2018.

Schrödinger, E. Sur la théorie relativiste de l'électron et l'interprétation de la mécanique quantique. Annales de l'Institut Henri Poincaré , 2(4):269-310, 1932.

Shi, Y., De Bortoli, V ., Campbell, A., and Doucet, A. Diffusion Schrödinger bridge matching. arXiv preprint arXiv:2303.16852 , 2023.

Sinkhorn, R. and Knopp, P. Concerning nonnegative matrices and doubly stochastic matrices. Pacific Journal of Mathematics , 21(2):343-348, 1967.

Solomon, J., Rustamov, R., Guibas, L., and Butscher, A. Wasserstein propagation for semi-supervised learning. In International Conference on Machine Learning , pp. 306-314. PMLR, 2014.

Solomon, J., De Goes, F., Peyré, G., Cuturi, M., Butscher, A., Nguyen, A., Du, T., and Guibas, L. Convolutional Wasserstein distances: Efficient optimal transportation on geometric domains. ACM Transactions on Graphics (ToG) , 34(4):1-11, 2015.

Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. International Conference on Learning Representations , 2021.

Srivastava, S., Li, C., and Dunson, D. B. Scalable Bayes via barycenter in Wasserstein space. The Journal of Machine Learning Research , 19(1):312-346, 2018.

Staib, M., Claici, S., Solomon, J. M., and Jegelka, S. Parallel streaming Wasserstein barycenters. Advances in Neural Information Processing Systems , 30, 2017.

Stroock, D. W. and Varadhan, S. S. Multidimensional diffusion processes , volume 233. Springer Science & Business Media, 1997.

Su, Z., Wang, Y., Shi, R., Zeng, W., Sun, J., Luo, F., and Gu, X. Optimal mass transport for shape matching and comparison. IEEE Transactions on Pattern Analysis and Machine Intelligence , 37 (11):2246-2259, 2015.

Valiente, G. Algorithms on Trees and Graphs , volume 112. Springer, 2002.

Vargas, F., Thodoroff, P., Lamacraft, A., and Lawrence, N. Solving Schrödinger bridges via maximum likelihood. Entropy , 23(9):1134, 2021.

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. Attention is all you need. Advances in Neural Information Processing Systems , 30, 2017.

Vincent, P. A connection between score matching and denoising autoencoders. Neural computation , 23(7):1661-1674, 2011.

Appendix organization

First, additional notation is introduced in Appendix A. Then, we briefly recall some notions on undirected and directed trees in Appendix B. Similarly, martingale problems are introduced in Appendix C. The proofs of the main manuscript and additional theoretical results on Tree Schrödinger Bridges are given in Appendix D. Additional details on our consideration of the tree-based static SB problem are described in Appendix E. Details on the implementation of TreeDSB are given in Appendix F and the experiments are investigated in Appendix G.

A Additional notation

For any finite set E , we equivalently refer to the cardinal of E as card( E ) or | E | . Let ( X , X ) be a measurable space. For any x ∈ ( R d ) ℓ +1 and any m ∈ { 0 , . . . , ℓ } , let x -m = ( x 0 , . . . , x m -1 , x m +1 , . . . , x ℓ ) . For any family of measures { ν j } j ∈{ 0 ,...,ℓ } defined on ( X , X ) and any i ∈ { 0 , . . . , ℓ } , let ν -i = ⊗ j ∈{ 0 ,...,ℓ }\{ i } ν j . Let I = { i 1 , . . . , i q } ⊂ { 1 , . . . , ℓ } and µ ∈ P ( ℓ ) such that µ ≪ Leb . We define I c = { 1 , . . . , ℓ }\ I and denote it by { i c 1 , . . . , i c ¯ q } where ¯ q = ℓ -q . We denote the marginal of µ along I by µ I , i.e. , µ I ∈ P ( q ) and we have for any A ∈ B (( R d ) q ) , µ I ( A ) = ∫ X µ ( x ) ∏ q j =1 δ x i j ( A j )d x . In addition, note that µ I ≪ Leb . We denote the conditional distribution of µ given I by µ | I ( ·|· ) , i.e. , µ | I ( ·|· ) ∈ P (¯ q ) × ( R d ) q and we have for any y ∈ ( R d ) q and any A ∈ B (( R d ) ¯ q ) , µ | I ( A | y ) = ∫ X µ ( x ) /µ I ( y ) ∏ q j =1 δ ( x i j -y j ) ∏ ¯ q j ′ =1 δ x i c j ′ ( A j ′ )d x . Remark that for any y ∈ ( R d ) q , µ | I ( ·| y ) ≪ Leb . For any subset J ⊂ I c with card( J ) = q J , we also define µ J | I ( ·|· ) ∈ P ( q J ) × ( R d ) q such that for any y ∈ ( R d ) q , µ J | I ( ·| y ) = { µ | I ( ·| y ) } J . For a collection of functions { f i } i ∈ I , with I ⊂ { 1 , . . . , n } and n ∈ N such that f i : R d → R , we define ⊕ i ∈ I f i : ( R d ) n → R such that for any x = ( x 1 , . . . , x n ) ∈ ( R d ) n , ⊕ i ∈ I f ( x ) = ∑ i ∈ I f i ( x i ) .

B Introduction to trees

Undirected tree. An undirected graph T = ( V , E ) , with vertices V and edges E , is said to be an undirected tree if it is acyclic and connected (Valiente, 2002, Definition 1.19.). In particular, we have card( E ) = card( V ) -1 . The undirected edge between two nodes v 1 and v 2 is similarly denoted by { v 1 , v 2 } or { v 2 , v 1 } . We say that T ′ = ( V ′ , E ′ ) is a sub-tree of T if T ′ is an undirected tree with vertices V ′ ⊂ V and edges E ′ ⊂ E . For any vertex v ∈ V , we define the set of its neighbours N v as the set of vertices v ′ ∈ V such that { v, v ′ } ∈ E . The integer card( N v ) is referred to as the degree of v . The vertices with degree 1 are called leaves , and we denote the set of leaves by V L ⊂ V . The (unique) path in T between two vertices v and v ′ is the sequence of two-by-two distinct edges {{ v i , v i +1 }} n i =1 (with n ≥ 1 ) such that v k = v k +1 for any k ∈ { 1 , . . . , n } such that k = 0 mod(2) , v 1 = v and v n +1 = v ′ . This path can be seen as a linear sub-tree of T , and we define n as the length of this path. We say that T is weighted if there exists a map w : E ↦→ R + ; in this case, w ( { v 1 , v 2 } ) , or equivalently w ( { v 2 , v 1 } ) (also denoted by w v 1 ,v 2 or w v 2 ,v 1 ) is called the weight of the edge { v 1 , v 2 } . The tree T is said to be rooted in r ∈ V if r defines a partial ordering ≤ T ,r ⊂ V × V such that for any v 1 , v 2 ∈ V , v 1 ≤ T ,r v 2 if the node v 1 lies on the unique path between r and v 2 .

Directed tree. Consider a directed graph T r = ( V , E r ) rooted in r ∈ V . Any directed edge e ∈ E r from v 1 ∈ V to v 2 ∈ V is denoted by ( v 1 , v 2 ) . T r is a said to be a directed tree rooted in r if (i) the underlying undirected graph T = ( V , E ) is an undirected tree rooted in r and (ii) any ( v 1 , v 2 ) ∈ E r is directed according to the partial ordering ≤ T ,r , i.e. , { v 1 , v 2 } ∈ E and v 1 ≤ T ,r v 2 . For any vertices ( v, v ′ ) ∈ V × V such that v ≤ T ,r v ′ , the (unique) path in T r from v to v ′ , denoted by path T r ( v, v ′ ) , is defined as the directed version of the path in T between v and v ′ (viewed as a sub-tree of T ), which is rooted in v . We say that T r is weighted , if T is weighted and the edges of T r have the same weights as the corresponding undirected edges of T . For any ( v 1 , v 2 ) ∈ E r , we denote this weight by w v 1 ,v 2 . We say that T r is the (unique) directed version of T rooted in r . It is endowed with a canonical vertex numbering ζ : V →{ 0 , . . . , card( V ) -1 } , corresponding to a depth-first traversal of its nodes, starting from the root r (Valiente, 2002, Definition 3.1.). This numbering is consistent with the partial ordering on T , i.e. , if v 1 ≤ T ,r v 2 , ζ ( v 1 ) ≤ ζ ( v 2 ) , and satisfies ζ ( r ) = 0 . In the rest of the paper, we will write in an equivalent manner v or ζ ( v ) .

For any vertices ( v 1 , v 2 ) ∈ E × E such that v 1 ≤ T ,r v 2 , path T r ( v 1 , v 2 ) corresponds to the ordered set of edges in E r which define the ordered path between two vertices v 1 and v 2 . For any vertex v ∈ V , we define:

(a) the set of its children C v as the set of vertices v ′ ∈ V such that ( v, v ′ ) ∈ E r . In particular, for any v ∈ V L , the set of leaves, one has C v = ∅ .

̸

(b) its parent as the unique vertex p ( v ) such that ( p ( v ) , v ) ∈ E r , if v = r (the parent of the root is not defined).

Note that N r = C r and, for any vertex v ∈ V \{ r } , N v = { p ( v ) } ∪ C v .

Definition 8 (Tree-structured directed Probabilistic Graphical Model (PGM)) . Consider a directed tree T r = ( V , E r ) . The directed PGM induced by T r (Koller & Friedman, 2009, Definition 3.4.), denoted by P T r , is the family of distributions π ∈ P ( | V | ) which have a Markovian factorization along T r , i.e.,

Lemma 9. Consider an undirected tree T = ( V , E ) . Let ( r, r ′ ) ∈ V × V . Let T ′ be a sub-tree of T with vertices V ′ such that r ′ ∈ V ′ . Denote by T ′ r ′ the directed version of T ′ rooted in r ′ . Then, for any π ∈ P T r , we have π V ′ ∈ P T ′ r ′ .

Proof. Let ( r, r ′ ) ∈ V × V . We denote by T r = ( V , E r ) , respectively T r ′ = ( V , E r ′ ) , the directed version of T rooted in r , respectively r ′ . We define the paths P r,r ′ = path T r ( r, r ′ ) ⊂ E r and P r ′ ,r = path T r ′ ( r ′ , r ) ⊂ E r ′ . It is easy to see that

Let π ∈ P T r . First note that for any ( v 1 , v 2 ) ∈ E r , we have by Bayes decomposition π v 1 π v 2 | v 1 = π v 2 π v 1 | v 2 = π v 1 ,v 2 . Then it comes

and therefore, we have π ∈ P T r ′ .

Let T ′ be a sub-tree of T with vertices V ′ such that r ′ ∈ V ′ . First note that E ′ r ′ ⊂ E r ′ . Using the previous computation, we have for any A ∈ B (( R d ) | V ′ | ) ,

which proves that π V ′ ∈ P T ′ r ′ .

Discretized undirected tree. Let N ≥ 1 . Consider an undirected tree T = ( V , E ) with weights w . We say that T ( N ) = ( V ( N ) , E ( N ) ) is a N -discretized version of T if it is an undirected tree with weights w ( N ) such that

(b) E ( N ) = ∪ e ∈ E ∪ k =0 ,...,N -1 { { v k e , v k +1 e } } with the convention that the vertices v N e and v N e are defined such that { v 0 e , v N e } = e ,

Remark that the leaves of T ( N ) are exactly the original leaves of T and that T (1) = T . The nonuniqueness of T ( N ) comes from the freedom of choice on the weights of its edges.

Discretized directed tree. Let N ≥ 1 . Consider a directed tree T r = ( V , E r ) rooted in r ∈ V with weights w . We say that T ( N ) r = ( V ( N ) , E ( N ) r ) is a N -discretized version of T r if it is the directed version of T ( N ) rooted in r , where T ( N ) is a N -discretized version of the underlying undirected tree of T r .

C Background on martingale problems

In this section, we introduce the background on Stochastic Differential Equations (SDEs) and weak solutions of SDEs following the framework of (Stroock & Varadhan, 1997, Section 10.1, page 249). We recall that C ∞ 0 ( R d ) is the space of infinitely differentiable real-valued functions which vanish at infinity. In addition, we have that S d + is the space of d × d , symmetric, non-negative matrices.

Definition 10. Let T > 0 or T = + ∞ , σ : [0 , T ) × R d →S d + and b : [0 , T ) × R d → R d , locally bounded measurable functions. We define the infinitesimal generator , A , given for any f ∈ C ∞ 0 ( R d ) , t ∈ [0 , T ) and x ∈ R d by

We say that a probability measure P satisfies the martingale problem for A if for any t ∈ [0 , T ) and f ∈ C ∞ 0 ( R d ) , we have that ( f ( X t ) -∫ t 0 A s ( f )( X s )d s ) s ∈ [0 ,t ] is a P -martingale.

In the main document, see Section 2, we say that 'a path measure P is associated with d X t = b ( t, X t )d t + σ ( t, X t )d B t with ( B t ) t ≥ 0 a d -dimensional Brownian motion' if P solves the martingale problem associated with A given by (4). Unless specified, we always assume that such a path measure exists and is unique. Below, we recall the following theorem, see (Stroock & Varadhan, 1997, Theorem 10.2.2), which gives sufficient conditions for the existence and uniqueness of solutions to the martingale problem.

Theorem 11. Assume that for any x ∈ R d we have

In addition, assume that there exists C > 0 such that for any x ∈ R d

Then, there exists a unique solution to the martingale problem with initialization x 0 ∈ R d .

D Theoretical results on Tree Schrödinger Bridges

We respectively provide in Appendix D.1, Appendix D.2 and Appendix D.3 the proofs of the results of the main manuscript presented in Section 3, Section 4 and Section 5. Finally, we make a detailed comparison between our setting and the framework of Haasler et al. (2021) in Appendix D.4. In the rest of this section, we consider an undirected tree T = ( V , E ) , where | V | = ℓ +1 , and some subset S ⊂ V which we denote by S = { i 0 , . . . , i K -1 } . We define S c = V \ S .

D.1 Proofs of Section 3

Proposition 1 is straightforward to obtain by combining the definition of the Brownian motion with the definition of π 0 given in (2). The following lemma details the recursion relation between the (mIPF) iterates, which is key to prove Proposition 2.

Lemma 12. Let ( π n ) n ∈ N be the sequence given by (mIPF) . Let n ∈ N , k n = ( n -1) mod( K ) , k n +1 = n mod( K ) . Denote by T k n , respectively T k n +1 with edges E k n +1 , the directed version of T rooted in i k n , respectively in i k n +1 . We have:

Proof. We show the result (i) by recursion on n ∈ N , and will deduce (ii) from the proof. Using (2), we first have π 0 ∈ P T r , where r is chosen as i K -1 , see Section 3. Thus, we obtain the result (i) at step n = 0 . Assume now that π n ∈ P T kn for some n ∈ N .

Consider the paths P n = path T kn ( i k n , i k n +1 ) and P n +1 = path T kn +1 ( i k n +1 , i k n ) . Note that these two paths have the same length, denoted by J , and contain the same vertices, denoted by V n . Let π ∈ P ( ℓ +1) such that KL( π | π n ) < + ∞ . We have the following decomposition

Hence, the ( n +1) -th iterate of (mIPF) is given by π n +1 = π n +1 V n ⊗ π n | V n , with

Since π n ∈ P T kn , we have (i) π n | V n = ⊗ ( v,v ′ ) ∈ E kn \ P n π n v ′ | v and (ii) π n V n ∈ P P n +1 by Lemma 9, where P n +1 is viewed as a directed tree rooted in i k n +1 . Defining V n +1 = V n \{ i k n +1 } , we thus have π n V n = π n i kn +1 ⊗ π n V n +1 | i kn +1 where π n V n +1 | i kn +1 = ⊗ ( v,v ′ ) ∈ P n +1 π n v ′ | v .

Let π ∈ P ( J +1) such that π i kn +1 = µ i kn +1 and KL( π | π n V n ) < + ∞ . Similarly to the previous computation, we have the following decomposition

Therefore, we obtain

Noting that E k n \ P n = E k n +1 \ P n +1 and recalling that π n +1 = π n +1 V n ⊗ π n | V n , it finally comes

Therefore, π n +1 ∈ P T kn +1 , which achieves the recursion for (i), and we obtain (ii) by (5).

Hence, Lemma 12 shows that the (mIPF) iterates admit a Markovian factorization on T , and can be defined recursively using the edges of T . We now provide the proof of Proposition 2.

Proof of Proposition 2. We will prove this result by recursion on n ∈ N . Observe that the initialisation is directly given by Proposition 1. Assume now that the result of Proposition 2 stands for some n ∈ N . Let k n = ( n -1) mod( K ) , k n + 1 = n mod( K ) . Denote by T k n with edges E k n , respectively T k n +1 with edges E k n +1 , the directed version of T rooted in i k n , respectively in i k n +1 . For any vertex v of T k n +1 , we define p ( v ) as the (unique) parent of v and c ( v ) as the unique child of v when it exists. Consider the ( n +1) -th dynamic iterate defined by (a) and (b), i.e. , ( P n +1 ( v,v ′ ) ) ( v,v ′ ) ∈ E kn +1 . To prove that this iterate has the properties stated in Proposition 2, we proceed by recursion on the edges of T k n +1 , following the bread-first order in T k n +1 . In this order, the edge ( i k n +1 , c ( i k n +1 )) is the first to be considered. Remark that c ( i k n +1 ) is well defined since i k n +1 is a leaf of T .

Here, we denote T c ( i kn +1 ) ,i kn +1 by T . By construction, we have P n +1 ( i kn +1 ,c ( i kn +1 )) = µ i kn +1 ⊗ ( P n ( c ( i kn +1 ) ,i kn +1 ) ) R | 0 . By recursion assumption, P n ( c ( i kn +1 ) ,i kn +1 ) ∈ P (C([0 , T ] , R d )) since

( c ( i k n +1 ) , i k n +1 ) ∈ E k n . Then, P n +1 ( i kn +1 ,c ( i kn +1 )) is a well defined path measure on [0 , T ] . By definition of the (mIPF) sequence, we have µ i kn +1 = π n +1 i kn +1 . By recursion assumption, we also have that Ext( P n ( c ( i kn +1 ) ,i kn +1 ) ) = π n c ( i kn +1 ) ,i kn +1 . Hence, it comes that ( P n ( c ( i kn +1 ) ,i kn +1 ) ) R T | 0 = π n c ( i kn +1 ) | i kn +1 = π n +1 c ( i kn +1 ) | i kn +1 , where the last equality comes from Lemma 12. Finally, we obtain that Ext( P n +1 ( i kn +1 ,c ( i kn +1 )) ) = π n +1 i kn +1 ,c ( i kn +1 ) , which proves the initialisation.

Assume now that P n +1 is well defined and has the right properties, up to some edge in T k n +1 . Consider the following edge, denoted by ( v, v ′ ) ∈ E k n +1 , in the breadth-first order. By edge recursion, we have that Ext( P n +1 ( p ( v ) ,v ) ) = π n +1 p ( v ) ,v , and thus P n +1 ( p ( v ) ,v ) ,T p ( v ) ,v = π n +1 v . Define the path P n = path T kn ( i k n , i k n +1 ) . Then, we face two cases.

(i) Either ( v, v ′ ) ∈ E k n \ P n . Then, we have by (a) that

In particular, P n +1 ( v,v ′ ) is a well defined path measure on [0 , T v,v ′ ] . Since ( v, v ′ ) ∈ E k n , Ext( P n ( v,v ′ ) ) = π n v,v ′ by recursion assumption. In particular, P n ( v,v ′ ) ,T v,v ′ | 0 = π n v ′ | v = π n +1 v ′ | v where the last equality comes from Lemma 12. We thus have Ext( P n +1 ( v,v ′ ) ) = π n +1 v,v ′ .

(ii) Or ( v ′ , v ) ∈ P n . Then, we have by (b) that

In particular, P n +1 ( v,v ′ ) is a well defined path measure on [0 , T v,v ′ ] . Here, ( v ′ , v ) ∈ E k n and thus, Ext( P n ( v ′ ,v ) ) = π n v ′ ,v by recursion assumption. In particular, ( P n ( v ′ ,v ) ) R T v ′ ,v | 0 = π n v ′ | v = π n +1 v ′ | v where the last equality comes from Lemma 12. We thus have Ext( P n +1 ( v,v ′ ) ) = π n +1 v,v ′ .

This achieves the recursion.

D.2 Proofs of Section 4

Remark on assumption A1. Although A 1 is not needed to establish the result of Proposition 3, Corollary 4 and Proposition 5 , it is however crucial in the proof of convergence of (mIPF) stated in Proposition 6. Nevertheless, we choose to keep A 1 as an assumption in the statement of every theoretical result presented in Section 4 for sake of clarity.

Additional definitions. We define the set P S = ∩ i ∈ S P i , where P i = { π ∈ P ( ℓ +1) : π i = µ i } , i.e. , P S is the set of all probability measures π ∈ P ( ℓ +1) which verify

for any family of bounded measurable functions { f i } i ∈ S ∈ C( R d , R ) K . Since R d is separable, there exists a dense family of functions { f k i } k ∈ N ∗ ,i ∈ S , with f k i ∈ L ∞ ( µ i ) for any k ∈ N ∗ and any i ∈ S , such that π ∈ P S if and only if

or equivalently, upon centering f k i ,

In the rest of the section, we consider such family { f k i } k ∈ N ∗ ,i ∈ S .

For any n ∈ N ∗ , we also define P n S = ∩ i ∈ S P n i , where P n i = { π ∈ P ( ℓ +1) : ∫ ( R d ) ℓ +1 f k i ( x i )d π ( x 0: ℓ ) = 0 , ∀ k ∈ { 1 , . . . , n }} . In particular, we have

Finally, (static-mSB) can be rewritten as

Proof of Proposition 3 and Corollary 4. In this part of the section, we present an extension of the theoretical results from Nutz (2021) to the multi-marginal setting. We first present two technical results, Lemma 13 and Lemma 14, which are respectively adapted from (Nutz, 2021, Lemma 2.10.) and (Nutz, 2021, Lemma 2.11.).

Lemma 13. Let { ˜ µ j } j ∈ S c be a family of probability measures defined on ( R d , B ( R d )) . We define ˜ π 0 = ⊗ i ∈ S µ i ⊗ j ∈ S c ˜ µ j . Let A ∈ ⊗ ℓ m =0 B ( R d ) such that ˜ π 0 ( A ) = 1 . Then, for ˜ π 0 -almost any x ⋆ ∈ A , there exists a family of sets { X 0 m } ℓ m =0 ⊂ ( R d ) ℓ +1 such that

Proof. Consider such set A . We define for any m ∈ { 0 , . . . , ℓ } the set

where A u m = { y ∈ ( R d ) ℓ : ( y 0 , . . . , y m -1 , u, y m , . . . , y ℓ -1 ) ∈ A } .

Take i ∈ S . Assume that µ i ( X i ) < 1 . We recall that ˜ π 0 = ˜ π 0 -i ⊗ µ i . Using Fubini's theorem and that ∫ A x i i d˜ π 0 -i ( x -i ) < 1 for any x i ̸∈ X i , we have

which is absurd. Therefore, we obtain µ i ( X i ) = 1 , and similarly, we have ˜ µ j ( X j ) = 1 for any j ∈ S c . For any y ∈ ( R d ) ℓ , any m ∈ { 0 , . . . , ℓ } , we define the set

Let i ∈ S . We have by Fubini's theorem

̸

where the last equality comes from the fact that µ i ( X i ) = 1 for any i ∈ S , ˜ µ j ( X j ) = 1 for any j ∈ S c and that ˜ π 0 = ⊗ i ∈ S µ i ⊗ j ∈ S c ˜ µ j . Consequently, there exists a measurable set A -i ⊂ ∏ ℓ m =0 m = i X i such that the following properties hold: (a) µ i ( ¯ A y i ) = 1 for any y ∈ A -i , (b) ˜ π 0 -i ( A -i ) = 1 . Similarly, this result holds for any j ∈ S c , i.e. , there exists a measurable set A -j ⊂ ∏ ℓ m =0 m = j X i such that the following properties hold: (a) ˜ µ j ( ¯ A y j ) = 1 for any y ∈ A -j , (b) ˜ π 0 -j ( A -j ) = 1 . We consider such sets { A -m } ℓ m =0 for the rest of the proof and finally define the set

where ˜ A m = A -m ×{ u ∈ ¯ A y m : y ∈ A -m } . By definition, we have ˜ A ⊂ A ∩ ∏ ℓ m =0 X m , using the fact that ˜ A m ⊂ A for any m ∈ { 0 , . . . , ℓ } . In addition, for any i ∈ S , we get by Fubini's theorem

and similarly, we get ˜ π 0 ( ˜ A j ) = 1 for any j ∈ S c . We can deduce that ˜ π 0 ( ˜ A ) = 1 since ˜ π 0 ( ˜ A c ) ≤ ∑ ℓ m =0 ˜ π 0 ( ˜ A c m ) = 0 .

̸

̸

Let x ⋆ ∈ ˜ A . In particular, x ⋆ ∈ A . Wedefine the set A 0 = A ∩ ( ∏ ℓ m =0 X 0 m ) , where X 0 m = X m ∩ ¯ A x ⋆ -m m for any m ∈ { 0 , . . . , ℓ } . We now establish the result of Lemma 13.

We first prove (a). Let i ∈ S . Since x ⋆ ∈ ˜ A , we have x ⋆ ∈ ˜ A i and therefore x ⋆ -i ∈ A -i . By definition of A -i , we obtain that µ i ( ¯ A x ⋆ -i i ) = 1 and thus,

which gives µ i ( X 0 i ) = 1 , and similarly, we have ˜ µ j ( X 0 j ) = 1 for any j ∈ S c .

We now prove (b). Let m ∈ { 0 , . . . , ℓ } . Since x ⋆ ∈ ˜ A ⊂ A , we get x ⋆ m ∈ ¯ A x ⋆ -m m . Using that ˜ A ⊂ A ∩ ℓ m =0 X m , we get x ⋆ ∈ A 0 . Let x ∈ A 0 . We denote x m = ( x ⋆ 0 , . . . , x ⋆ m -1 , x m , x ⋆ m +1 , . . . , x ⋆ ℓ ) . We need to show that x m ∈ A and x m ∈ ∏ ℓ j =1 X 0 j = ∏ ℓ j =1 ( X j ∩ ¯ A x ⋆ -m j ) . First, since x j m = x j or x ⋆ j for any j ∈ { 0 , . . . , ℓ } , and x ∈ A 0 and x ⋆ ∈ A 0 , we get that for any j ∈ { 0 , . . . , ℓ } , x j m ∈ X j . Similarly, for any j ∈ { 0 , . . . , ℓ -1 } , x j m ∈ ¯ A x ⋆ m j . Therefore, we get that x m ∈ ∏ ℓ j =1 ( X j ∩ ¯ A x ⋆ -m j ) . Since x m ∈ A x ⋆ -m m (because x ∈ ∏ ℓ j =1 ( X j ∩ ¯ A x ⋆ -m j ) ), we get that x ∈ A , which concludes the proof.

Lemma 14. Let A 0 ⊂ ( R d ) ℓ +1 . For any m ∈ { 0 , . . . , ℓ } , we denote X 0 m = proj m ( A 0 ) . We make the following assumptions.

(a) Assume there exists x ⋆ ∈ A 0 such that for any x ∈ A 0 , for any m ∈ { 0 , . . . , ℓ } , we have ( x ⋆ 0 , . . . , x ⋆ m -1 , x m , x ⋆ m +1 , . . . , x ⋆ ℓ ) ∈ A 0 .

(b) Assume there exists a family of functions { φ n i k } n ∈ N ∗ ,k ∈{ 0 ,...,K -1 } with φ n i k : X 0 i k → [ -∞ , + ∞ ) such that for any n ∈ N ∗ and any k ∈ { 0 , . . . , K -2 } , we have φ n i k ( x ⋆ i k ) = 0 .

(c) Denote F n ( x ) = ∑ K -1 k =0 φ n i k ( x i k ) for any x ∈ A 0 . Assume that for any x ∈ A 0 , F ( x ) = lim n →∞ F n ( x ) exists and is such that F ( x ) ∈ [ -∞ , + ∞ ) with F ( x ⋆ ) ∈ R .

Then, for any i ∈ S , for any x i ∈ X 0 i , φ i ( x i ) = lim n →∞ φ n i ( x i ) exists and is such that φ i ( x i ) ∈ [ -∞ , + ∞ ) .

Proof. Consider A 0 ⊂ ( R d ) ℓ +1 such that assumptions (a), (b) and (c) hold. Remark that we have F n ( x ⋆ ) = φ n i K -1 ( x ⋆ i K -1 ) .

Let x ∈ A 0 . We denote x m = ( x ⋆ 1 , . . . , x ⋆ m -1 , x m , x ⋆ m +1 , . . . , x ⋆ l ) for any m ∈ { 0 , . . . , ℓ } . In particular, we have x m ∈ A 0 by assumption (a). Let us define

Using assumption (c), we have φ i ( x i ) ∈ [ -∞ , + ∞ ) for any i ∈ S . Let k ∈ { 0 , . . . , K -2 } . We have by definition of F n ,

̸

where we used assumption (b) in the last equality. Since x i k ∈ A 0 and x ⋆ ∈ A 0 , we have by assumption (c),

Furthermore, by combining the definition of F n with assumption (b), we have

which concludes the proof.

In what follows, before proving Proposition 3, we respectively show in Proposition 15 and Proposition 16 how A 2 and A 3 can be satisfied in the case where π 0 ∈ P T r , as in (2), that is

Proposition 15. Let π 0 ∈ P T r . Assume that π 0 r = µ r if r ∈ S or π 0 r = N( m r , σ r Id) , with m r ∈ R d and σ r > 0 if r ∈ S c . In addition, assume that for any ( v, v ′ ) ∈ E r , π 0 v ′ | v ( ·| x v ) = N( x v , σ v,v ′ Id) with σ v,v ′ > 0 . Finally, assume that for any i ∈ S , ∫ R d ∥ x ∥ 2 d µ i ( x ) < + ∞ and H( µ i ) < + ∞ . Then A 2 is satisfied.

Proof. Let π = ⊗ i ∈ S µ i ⊗ i ∈ S c ν i with ν i any Gaussian measure with positive definite covariance matrix. First, we have that

For any ( v, v ′ ) ∈ E r , there exists C v,v ′ ≥ 0 such that

We conclude the proof upon remarking that KL( π r | π 0 r ) < + ∞ .

Proposition 16. Let π 0 ∈ P T r . Assume that π 0 r = µ r if r ∈ S or π 0 r = N( m r , σ r Id) , with m r ∈ R d and σ r > 0 if r ∈ S c . In addition, assume that for any ( v, v ′ ) ∈ E r , π 0 v ′ | v ( ·| x v ) = N( x v , σ v,v ′ Id) with σ v,v ′ > 0 . Finally, assume that for any i ∈ S , µ i admits a positive density w.r.t. the Lebesgue measure. Then A 3 is satisfied.

Proof. We have that π 0 admits a positive density w.r.t the Lebesgue measure. Letting ˜ π 0 = ⊗ i ∈ S µ i ⊗ j ∈ S c ˜ µ j where ˜ µ j which admits a positive density w.r.t. the Lebesgue measure for any j ∈ S c , we get that ˜ π 0 admits a positive density w.r.t. the Lebesgue measure and therefore π 0 ∼ ˜ π 0 , which concludes the proof.

Using the preliminary results presented above, we are now ready to prove Proposition 3.

Proof of Proposition 3. Assume A 1 and A 2. Since P S is convex and closed in total-variation norm, there exists a probability distribution π ⋆ solution to (7), or equivalently to (static-mSB), by using A 2 with (Csiszár, 1975, Theorem 2.1.). Moreover, this solution is unique by strict convexity of KL( · | π 0 ) .

We now turn to the proof of existence of potentials defining (d π ⋆ / d π 0 ) , by adapting the arguments of (Nutz, 2021, Section 2.3.). Define ν n = argmin { KL( π | π 0 ) : π ∈ P n S } for any n ∈ N ∗ . Since { P n S } n ∈ N ∗ ⊂ P ( ℓ +1) is a decreasing sequence of sets that are convex and closed in total-variation norm such that (6) holds, we get from (Nutz, 2021, Proposition 1.17.) with A 2 that

Following (Nutz, 2021, Example 1.18), there exists a family of bounded measurable functions { φ n i } n ∈ N ∗ ,i ∈ S with φ n i : R d → R such that for any n ∈ N ∗

We consider such family { φ n i } n ∈ N ∗ ,i ∈ S for the rest of the proof. By combining (8) and (9), we obtain, up to extraction,

or equivalently

We now define the following sets

Using (10), we have π 0 ( A ⋆ ) = 1 . Using A 3, it comes ˜ π 0 ( A ⋆ ) = 1 . Moreover, we also get that π ⋆ ( B ⋆ ) = 1 by (10). Thus, it comes π 0 ( B ⋆ ) > 0 , and ˜ π 0 ( B ⋆ ) > 0 using A 3.

We then apply Lemma 13 to ˜ π 0 and A = A ⋆ . Since ˜ π 0 ( B ⋆ ) > 0 , it implies that there exists x ⋆ ∈ B ⋆ and a measurable set A 0 ⊂ B ⋆ verifying the properties (a) and (b). Following (Nutz, 2021, Corollary 2.12), we may assume without loss of generality in the statement of Lemma 13 that the sets X 0 m are measurable with ∏ ℓ m =0 X 0 m ⊂ A . In this case, we obtain that µ i (proj i ( A 0 )) = 1 for any i ∈ S .

We now aim at applying Lemma 14 to the set A 0 . Remark that A 0 directly satisfies assumption (a). For any n ∈ N ∗ , consider the following transformation of the functions { φ n i } i ∈ S

For any i ∈ S , we restrict φ n i k to X 0 i k , so that the family { φ n i } n ∈ N ∗ ,i ∈ S now verifies assumption (b). Finally, since A 0 ⊂ A ⋆ and x ⋆ ∈ B ⋆ , we directly obtain assumption (c).

Therefore, Lemma 14 may be applied. It provides us with the family of functions { φ i } i ∈ S defined by φ i : X 0 i → [ -∞ , + ∞ ) with φ i = lim n →∞ φ n i µ i -a.s. for any i ∈ S . Since µ i (proj i ( A 0 )) = 1 for any i ∈ S , we may extend the functions φ i to R d . In particular, we can find a family of functions { ψ ⋆ i } i ∈ S with ψ ⋆ i : R d → [ -∞ , + ∞ ) such that ψ ⋆ i = φ i µ i -a.s. Note that these functions are measurable as limits of measurable functions.

Since π 0 ∼ ˜ π 0 by A 3, (10) turns into

Finally, we show that the functions ψ ⋆ i are µ i -a.s. finite. Let i ∈ S . Let us define A i = { x i ∈ R d : ψ ⋆ i ( x i ) = -∞} . Using (11), we obtain (d π ⋆ / d π 0 )( A i × ( R d ) ℓ ) = 0 . Since π ⋆ i = µ i , we have

which gives the result.

We now turn to the proof of Corollary 4, which states that the iterates of (mIPF) can be expressed via potentials, in the same manner as the solution π ⋆ to (static-mSB).

Proof of Corollary 4. Assume A 1, A 2 and A 3. We prove the result of this corollary by recursion on n ∈ N ∗ . First take n = 1 . In this case, the first iteration of (mIPF) is a multi-marginal SB problem of the form (static-mSB) where S = { i 0 } with reference measure π 0 . Therefore, using A 2 and A 3, we can apply Proposition 3 and obtain existence of ψ 1 i 0 : R d → R such that

By taking ψ 0 i k = 0 for k ∈ { 1 , . . . , K -1 } , we thus obtain the result at step n = 1 .

Now assume that the result is verified for some n ∈ N ∗ , with k n = ( n -1) mod( K ) . We define k n + 1 = n mod( K ) and q n ∈ N as the quotient of the Euclidean division of n by K . In this case, the ( n +1) -th iteration of (mIPF) is a multi-marginal SB problem of the form (static-mSB) where S = { i k n +1 } with reference measure π n . Using (13), we have that A 2 is satisfied for this new (static-mSB) problem. A 1 and A 3 are satisfied for this problem, given the form of π n . Therefore, we can apply Proposition 3 and obtain existence of ψ q n +1 i kn +1 : R d → R such that

By assumption, we have that π n ≪ π 0 . Hence, we obtain π n +1 ≪ π 0 and thus,

By combining (12) with the result of the recursion at step n , we directly obtain the result at step n +1 , which achieves the proof.

Proofs of Proposition 5 and Proposition 6. In this part of the section, we establish the proofs of results related to the convergence of (mIPF), respectively Proposition 5 and Proposition 6, which can be seen as a natural extension of (Ruschendorf, 1995, Proposition 2.1.) and (Ruschendorf, 1995, Theorem 3.1.).

Proof of Proposition 5. Under A 1 and A 2, we obtain by Proposition 3 existence and uniqueness of a solution to (static-mSB), which we denote by π ⋆ . Since π ⋆ ∈ P S , using recursively (Csiszár, 1975, Theorem 3.12.), the fact that { π i k = µ i k : π ∈ P ( | V | ) } is convex for any k ∈ { 0 , . . . , K -1 } and (mIPF), we obtain

Therefore, we have

∞

i

=1

∑

KL(

π

i

|

π

i

-

1

)

KL(

π

⋆

|

π

0

)

<

∞

and thus,

Let n ∈ N ∗ with n > 2 K , k ∈ { 0 , . . . , K -1 } and let q n ∈ N be the quotient of the Euclidean division of n -1 by K . We define n k = q n K + k + 1 with ( n k -1) = k mod( K ) if n k ≤ n . Otherwise, we set n k = ( q n -1) K + k +1 with ( n k -1) = k mod( K ) . Note that we always have | n -n k | ≤ 2 K . In particular, we have π n k i k = µ i k by definition of (mIPF). Therefore, we obtain

where each term goes to 0 as n → + ∞ in the last inequality by (14), which achieves the proof.

We now turn to the proof Proposition 6, which requires several preliminary technical results. For the rest of this section, we define, for any n ∈ N , q n as the quotient of the Euclidean division of n -1 by K (in particular, q 0 = -1 ).

Schrödinger equations. Under A 1, A 2 and A 3, we know from Proposition 3 that the unique solution π ⋆ to (static-mSB) can be π 0 -a.s. written as (d π ⋆ / d π 0 ) = exp[ ⊕ i ∈ S ψ ⋆ i ] , where { ψ ⋆ i } i ∈ S are measurable potentials, referred to as Schrödinger potentials . These functions are determined by the fixed-point Schrödinger equations

which are obtained by marginalising π ⋆ along its constrained marginals. This family of potentials is not unique. Indeed, for any family of real numbers { λ i k } k ∈{ 0 ,...,K -2 } , we have

where ˜ ψ i k = ψ ⋆ i k + ˜ λ i k for any k ∈ { 0 , . . . , K -1 } with ˜ λ i k = λ i k if k ∈ { 0 , . . . , K -2 } and ˜ λ i K -1 = -∑ K -2 i =0 λ i k .

Remark on the initialisation of (mIPF) . Consider a probability measure ¯ π 0 ∈ P ( ℓ +1) of the form

where { ψ 0 i } i ∈ S is a family of measurable potentials with ψ 0 i : R d → R such that ∣ ∣ ∫ R d ψ 0 i d µ i ∣ ∣ < ∞ for any i ∈ S . Then, for any π ∈ P S , we have

Hence, (static-mSB) is equivalent to the multi-marginal SB problem

We refer to (Peyré et al., 2019, Proposition 4.2) for the EOT counterpart of this result. This means that the solutions of the multi-marginal Schrödinger Bridge problem are invariant by multiplication of the reference measure by potentials on the fixed marginals. Consequently, the initialisation of the (mIPF) sequence may be chosen as ¯ π 0 instead of π 0 .

For sake of clarity, we now refer to the reference probability measure of (static-mSB) as ¯ π or π -1 and to the initialisation of the (mIPF) iterates as π 0 .

Solving (mIPF) with potentials. To prove the convergence of the (mIPF) iterates to the solution π ⋆ given by Proposition 3, we first rewrite these iterates with potentials, following the form of π ⋆ .

To do so, we recursively define the sequence of potentials { ψ n i } n ∈ N ,i ∈ S by

and for any n ∈ N ∗ and k ∈ { 0 , . . . , K -1 }

recalling that q n is the quotient of the Euclidean division of n -1 by K .

We now define the sequence of probability measures { π n } n ∈ N by

In particular, we have (d π 0 / d¯ π ) = exp[ ⊕ K -1 ℓ =0 ψ 0 i ℓ ] = exp[ ψ 0 i K -1 ] , and thus ∫ R d ψ 0 i K -1 d µ i K -1 = KL( µ i K -1 | ¯ π i K -1 ) . Consequently, π 0 can be chosen as the initialisation of (mIPF), following the previous remark, if we assume that KL( µ i K -1 | ¯ π i K -1 ) < ∞ . In (TreeSB) with r = i K -1 , the latter assumption is directly verified since we choose ¯ π i K -1 = µ i K -1 .

Let n ∈ N , with k n = ( n -1) mod( K ) , k n +1 = n mod( K ) . Using (16) and (17), we get that π n i kn = µ i kn . Moreover, we have

with the convention that ψ -1 i K -1 = 0 . In particular, we obtain that π n +1 | i kn +1 = π n | i kn +1 .

In conclusion, the sequence { π n } n ∈ N defined in (18) verifies π n +1 = µ i kn +1 π n | i kn +1 for any n ∈ N . By decomposition property of the Kullback-Leibler divergence, this sequence solves (mIPF) with initialisation π 0 . We consider such iterates in the following.

Since π n i kn = µ i kn , we have that

Before proving a multi-marginal counterpart to (Ruschendorf, 1995, Lemma 4.1), we state and prove the following result.

Proposition 17. Let π 0 , π 1 two probability measures on R d such that π 0 ≪ π 1 . Then, denoting f = d π 0 / d π 1 , the following assertions are equivalent:

(b) ∫ R d | log( f )( x ) | d π 0 ( x ) < + ∞

(c) ∫ R d log( f )( x ) ✶ f ( x ) > 1 d π 0 ( x ) < + ∞ If one of these conditions is satisfied then ∫ R d | log( f )( x ) | d π 0 ≤ KL( π 0 | π 1 ) + 2 / e .

Proof. First, note that

∫ R d | log( f )( x ) | ✶ f< 1 d π 0 ( x ) ≤ ∫ R d | log( f )( x ) f ( x ) | ✶ f< 1 d π 1 ( x ) ≤ 1 / e , (21) where we have used that for any u ∈ [0 , 1] , | u log( u ) | ≤ 1 / e . We have that (b) implies (c). Using the previous result we have that (c) implies (b). Hence (c) and (b) are equivalent. In addition, it is clear that (b) implies (a). Finally (this is more of a convention), we have that KL( π 0 | π 1 ) = ∫ R d log( f )( x ) ✶ f ( x ) > 1 d π 0 ( x ) + ∫ R d log( f )( x ) ✶ f ( x ) < 1 d π 0 ( x ) < + ∞ . Using (21) this implies (c). Finally, we have which concludes the proof.

We begin with the following lemma which controls the integral of the potentials uniformly w.r.t. n ∈ N . It can be seen as the multi-marginal counterpart of (Ruschendorf, 1995, Lemma 4.1).

Lemma 18. Assume A 4 . There exist { c i } i ∈ S ∈ (0 , + ∞ ) K such that for any function f : ( R d ) ℓ +1 → R of the form f = ⊕ i ∈ S f i , we have

For any n ∈ N ⋆ , we have

Proof. First, we have that (22) is a direct consequence of (Kober, 1940, Theorem 1) and A 4. Let us now prove (a). Using (20), we have

where the last inequality follows the proof of Proposition 5.

Since the first term in the inequality of (b) is equal to KL( π ⋆ | π nK ) , we obtain (b) using that KL( π ⋆ | π nK ) ≤ KL( π ⋆ | ¯ π ) following the proof of Proposition 5.

Let us now prove (c). Since KL( π ⋆ | ¯ π ) < ∞ , using Proposition 17, we have that ⊕ i ∈ S ψ ⋆ i ∈ L 1 ( π ⋆ ) . From (b) and Proposition 17, we also get that ⊕ i ∈ S ( ψ ⋆ i -ψ n i ) ∈ L 1 ( π ⋆ ) , and thus ∫ ( R d ) ℓ +1 ∣ ∣ ⊕ i ∈ S ( ψ ⋆ i -ψ n i ) ∣ ∣ d π ⋆ ≤ C 0 with C 0 > 0 . Therefore, we have

Using (22), we conclude with A 4 that for any i ∈ S , we have

which concludes the proof of (c).

The next lemma gives an explicit expression for KL( π n | ¯ π ) . It can be seen as the multi-marginal counterpart of (Ruschendorf, 1995, Lemma 4.2).

Lemma 19. For any n ∈ N , with k n = ( n -1) mod( K ) , we have

Proof. Let n ∈ N , with k n = ( n -1) mod( K ) . Using (18), we have

Consider m ∈ { k n + 1 , . . . , K -1 } . Let m n be the closest integer to n such that m n > n and m = ( m n -1) mod( K ) . By (19), we have

Using (19) recursively, we obtain

where we recall that π m n i m = µ i m .

Consider now ℓ ∈ { 0 , . . . , k n -1 } . Let ℓ n be the closest integer to n such that ℓ n > n and ℓ = ( ℓ n -1) mod( K ) . By (19), we have

and using (19) recursively, we obtain

where we recall that π ℓ n i ℓ = µ i ℓ . We conclude the proof upon combining (23), (24) and (25).

We are now ready to prove a uniform integrability result which is the multi-marginal counterpart of (Ruschendorf, 1995, Lemma 4.4). Before stating Lemma 21, we prove the following well-known lemma. We recall that a sequence (Ψ n ) n ∈ N such that for any n ∈ N , Ψ n ∈ L 1 ( µ ) , is uniformly integrable w.r.t. µ if (i) sup n ∈ N ∫ R d | Ψ n | d µ < + ∞ and (ii) for any ε > 0 , there exists K > 0 such that for any n ∈ N , ∫ B(0 ,K ) c | Ψ n | d µ ≤ ε .

Lemma 20. Let f : R → R , convex and non-decreasing on [ A, + ∞ ) with A > 0 and lim x → + ∞ f ( x ) /x = + ∞ . Assume that sup n ∈ N ∫ R d f ( | Ψ n | )d µ < + ∞ . Then, (Ψ n ) n ∈ N is uniformly integrable w.r.t. µ .

Proof. Since f is convex, using Jensen's inequality, we get that sup n ∈ N f ( ∫ R d | Ψ n | d µ ) < + ∞ and since lim x → + ∞ f ( x ) /x = + ∞ we have sup n ∈ N ∫ R d | Ψ n | d µ < + ∞ . Let ε > 0 , there exists K > 0 such that for any x > K , x ≤ εf ( x ) /B with B = sup n ∈ N ∫ R d f ( | Ψ n | )d µ < + ∞ . Therefore, we have for any n ∈ N

which concludes the proof.

Lemma 21. Assume A 4 and A 5 . Then, { exp[ ⊕ i ∈ S ψ n i ] } n ∈ N is uniformly integrable w.r.t. ¯ π .

Proof. It is enough to show that the sequence { f (exp[ ⊕ i ∈ S ψ n i )] } n ∈ N is bounded in L 1 (¯ π ) , where f : u ↦→ u log( u ) is continuous, convex and such that lim u →∞ f ( u ) /u = + ∞ , see Lemma 20. Let n ∈ N . We have

With the preliminary results stated above, we are now ready to prove Proposition 6.

Proof of Proposition 6. Using A 4 and A 5, we have, by Lemma 21, uniform integrability of { exp[ ⊕ i ∈ S ψ n i ] } n ∈ N in L 1 (¯ π ) . Therefore, the sequence { π nK } n ∈ N is relatively compact with respect to the weak topology of σ (L 1 (¯ π ) , L ∞ (¯ π )) , denoted as the τ -topology. We recall that lim n →∞ KL( π nK +1 | π nK ) = 0 . This implies that { π nK +1 } n ∈ N is also relatively τ -compact. By trivial recursion, we obtain that the sequences { π nK + k } n ∈ N , where k ∈ { 2 , . . . , K -1 } are also relatively τ -compact. Therefore, { π n } n ∈ N is relatively τ -compact and τ -sequentially compact.

Weconsider an increasing function Φ : N → N such that { π m } m ∈ Φ( N ) is a τ -convergent subsequence, and we denote by ˜ π its limit for this topology. In particular, ˜ π ∈ P S by Proposition 5. We assume without loss of generality that Φ( N ) ⊂ K N .

Using the lower semi-continuity of the Kullback-Leibler divergence (Dupuis & Ellis, 2011, Lemma 1.4.3), we get

Consider k ∈ { 0 , . . . , K -2 } . By (19), we have

and thus,

With Proposition 5, we obtain that { exp[ ψ n i k -ψ n +1 i k ] } n ∈ N converges to 1 in L 1 ( µ i k ) . In addition using the uniform integrability of { ψ n i k } n ∈ N and A 5, we get

Wedenote m = Kℓ . Since KL( π m | ¯ π ) = ∫ R d ψ ℓ i K -1 d µ i K -1 + ∑ K -2 k =0 ∫ R d ψ ℓ i k exp[ ψ ℓ i k -ψ ℓ +1 i k ]d µ i k by Lemma 19, we finally have

where the last inequality comes from Lemma 18.

Since ˜ π i = µ i for any i ∈ S , using Proposition 5, we have ˜ π = π ⋆ by uniqueness of π ⋆ . Hence, π ⋆ is the only limit point of { π n } n ∈ N in the τ -topology. In particular, KL( π n | ¯ π ) → KL( π ⋆ | ¯ π ) . Since P S is convex, this last result implies ∥ π ⋆ -π n ∥ TV → 0 , see the proof of Theorem 2.1 in Csiszár (1975).

We finish this section by highlighting that A 5 is stronger than (Ruschendorf, 1995, B1). A natural extension of the latter assumption would consist of having a guarantee on the ( K -1) first potentials given by (17), as presented below.

Under A 6, (Ruschendorf, 1995, Lemma 4.3) can be adapted as written below.

Lemma 22. Assume A 6. Then, for any n ∈ N ∗

(a) for any k ∈ { 0 , . . . , K -2 } , there exists α n,k ∈ N such that

(b) there exists α n,K -1 ∈ N such that

where { α n,k } n ∈ N ∗ ,k ∈{ 0 ,...,K -1 } is a strictly increasing sequence that can be explicitly defined.

Proof. We prove the result by recursion on n ∈ N ∗ .

Take n = 1 . Let k ∈ { 0 , . . . , K -2 } . We define α 1 ,k = 0 and directly obtain (a) by A 5 since ψ 0 i k = 0 . Let us prove (b). We have by (17)

Using (a) at rank n = 1 , we have

and therefore, we obtain (b) by taking α 1 ,K -1 = 0 . Let us assume that the result is verified for some n ∈ N ∗ . We have

Using (a) and (b) at rank n , we have

Therefore, we obtain

Now, we define α n +1 , 0 = ∑ K -1 k =1 α n,k to obtain (a) for k = 0 . Consider now k ∈ { 1 , . . . , K -2 } . Following the same steps as above, we recursively define

which gives (a) at rank n +1 . Let us now prove (b) at rank n +1 . We have

Using (a) at rank n +1 , we obtain

Therefore, by taking α n +1 ,K -1 = ∑ K -2 k =0 α n +1 ,k , we obtain (b), which concludes the proof.

Unfortunately, Lemma 22 only yields non-vacuous bounds in the case K = 2 . Indeed, when K > 2 , the sequence { α n,k } n ∈ N ∗ ,k ∈{ 0 ,...,K -1 } leads to increase the bounds on the quantities exp[ ψ n -1 i k -ψ n i k ] , which motivates the use of A 5.

D.3 Proof of Section 5

For the rest of this section, we consider the multi-marginal Schrödinger bridge problem given by (TreeSB) and establish in Proposition 24 the correspondence with the regularized Wasserstein propagation problem presented in Solomon et al. (2014, 2015). We first state a technical result.

Lemma 23. Let ε > 0 . Assume that π 0 is given by (2) , where r ∈ V is chosen arbitrarily. Then, for any π ∈ P T r , we have

where we recall that C v = { v ′ ∈ V : ( v, v ′ ) ∈ E r } .

Proof. Since π, π 0 ∈ P T r , we obtain the following decomposition

0

We finally obtain the result by using the definition of π 0 and noticing that ∫ R d H( π v ′ | v ( ·| x v ))d π v ( x v ) = H( π v,v ′ ) -H( π v ) for any ( v, v ′ ) ∈ E r .

Proposition 24. Let ε > 0 and µ 0 ∈ P such that µ 0 ≪ Leb . Assume that π 0 is given by (2) , where r ∈ V is chosen arbitrarily, and that φ r = d µ 0 / dLeb . Also assume A 2. Then, the set of marginals of the solution to (TreeSB) is exactly the solution to the entropic-regularized Wasserstein Propagation problem (Solomon et al., 2014, 2015) defined by

where we recall that C v = { v ′ ∈ V : ( v, v ′ ) ∈ E r } .

Proof. Assume that π 0 is given by (2), where r ∈ V is chosen arbitrarily, and that φ r = d µ 0 / dLeb . In particular, we have π 0 r = µ 0 . Moreover, it is clear that π 0 verifies A 1, and A 3 by Proposition 16.

By definition of the regularized Wasserstein distance given in (3), we have for any { ν v } v ∈ V ∈ P ℓ +1

In particular, we have F ( { π v } ) ≤ G ( π r , { π v,v ′ } ) for any π ∈ P ( ℓ +1) . We now prove the result of Proposition 24 in two steps denoted by Step 1 and Step 2 .

Step 1. Let us not assume A 2 for now. In this case, we prove in Step 1.a and Step 1.b that solving (WP) is equivalent to solving a modified version of (TreeSB) given by

Remark that any solution to ( T r -TreeSB) is a solution to (TreeSB), but the converse result may not be true.

Step 1.a: (WP) = ⇒ ( T r -TreeSB) . Consider a solution { ν ⋆ v } v ∈ V to (WP). For any ( v, v ′ ) ∈ E r , W 2 2 ,ε/w v,v ′ ( ν ⋆ v , ν ⋆ v ′ ) is well defined and thus, there exists ν ( v,v ′ ) ∈ Π( ν ⋆ v , ν ⋆ v ′ ) such that

Using the gluying lemma, we build the probability measure π ⋆ = ν ⋆ r ∏ ( v,v ′ ) ∈ E r ν ( v,v ′ ) v ′ | v such that (i) π ⋆ ∈ P T r , and (ii) π ⋆ v,v ′ and ν ( v,v ′ ) have the same distribution for any ( v, v ′ ) ∈ E r . In particular, we have π ⋆ i = µ i for any i ∈ S .

Let us show now that π ⋆ is a solution to ( T r -TreeSB). Let π ∈ P T r such that π i = µ i for any i ∈ S . We have

Therefore, π ⋆ is a solution to ( T r -TreeSB).

Step 1.b: ( T r -TreeSB) = ⇒ (WP) . Consider now a solution π ⋆ to ( T r -TreeSB). Since π ⋆ ∈ P T r , we have π ⋆ = π ⋆ r ∏ ( v,v ′ ) ∈ E r π ⋆ v ′ | v and π ⋆ i = µ i for any i ∈ S .

Let us show that { π ⋆ v } v ∈ V is a solution to (WP). Let { ν v } v ∈ V ∈ P ℓ +1 such that ν i = µ i for any i ∈ S .

Let { ν ( v,v ′ ) } ( v,v ′ ) ∈ E r be a family of probability measures such that ν ( v,v ′ ) ∈ P (2) , ν ( v,v ′ ) v = ν v , ν ( v,v ′ ) v ′ = ν v ′ for any ( v, v ′ ) ∈ E r .

Using the gluying lemma, we build the probability measure π = ν r ∏ ( v,v ′ ) ∈ E r ν ( v,v ′ ) v ′ | v , such that (i) π ∈ P T r and (ii) π v,v ′ and ν ( v,v ′ ) have the same distribution for any ( v, v ′ ) ∈ E r . We have

By taking the infimum in the previous inequality over the families { ν ( v,v ′ ) } ( v,v ′ ) ∈ E r , we obtain by (26) that

and therefore, { π ⋆ v } v ∈ V is a solution to (WP).

Step 2. We now assume A 2. By Proposition 3, there exists a unique solution π ⋆ ∈ P ( ℓ +1) to (TreeSB) such that we π 0 -a.s. have (d π ⋆ / d π 0 ) = exp[ ⊕ i ∈ S ψ ⋆ i ] , where { ψ ⋆ i } i ∈ S are measurable potentials with ψ ⋆ : R d → R . Since π 0 ∈ P T r , we also have π ⋆ ∈ P T r , i.e. , the potentials { ψ ⋆ i } i ∈ S do not modify the Markovian nature of π 0 . Therefore, π ⋆ is also the unique solution to ( T r -TreeSB). Using the equivalence between ( T r -TreeSB) and (WP) established in Step 1 , we finally obtain the result of Proposition 24.

In particular, Proposition 7 directly derives from Proposition 24 by taking r = i K -1 and µ 0 = µ i K -1 .

D.4 Comparison with Haasler et al. (2021)

In their work, Haasler et al. (2021) study the static and discrete-state counterpart of our approach. Given a state space X such that | X | = n +1 with n ∈ N , they establish a correspondence between multi-marginal EOT with a general tree-based cost and discrete-time multi-marginal static Schrödinger bridge, and provide an efficient method to solve these problems. In this section, we provide details on their framework and give a precise comparison between our theory and their results.

To be coherent with the setting of Haasler et al. (2021), we adapt here some of our notation. Let us define Z ( q ) = R ( n +1) q + . For any q ∈ N ∗ , the set of probability measures on X q is defined as P ( q ) = { M ∈ Z ( q ) : ⟨ M, 1 ⟩ = 1 } . We denote P = P (1) . For any tensors M,P ∈ Z ( q ) , the KullbackLeibler divergence between M and P is defined as KL( M | P ) = ⟨ M log( M/P ) -M + P, 1 ⟩ and the

entropy of M is defined as H( M ) = -KL( M | 1 ) , where the operations are meant componentwise. In the rest of the section, we consider an undirected tree T = ( V , E ) with | V | = ℓ +1 such that V may be identified with { 0 , . . . , ℓ } .

Details on the results of Haasler et al. (2021). In their paper, the authors consider a cost tensor C ∈ Z ( ℓ +1) that factorizes along T , i.e. , for any { j 0 , . . . , j ℓ } with for any i ∈ { 0 , . . . , ℓ } , j i ∈ { 0 , . . . , n } , we have

where C { v,v ′ } ∈ Z (2) is a cost matrix for transportation between the marginals at vertices v and v ′ , see (Haasler et al., 2021, Eq. (3.1)). In particular, this cost can be seen as the discrete counterpart of the tree-based cost introduced in (1) in the quadratic setting.

Given a subset S ⊂ V with | S | = K and a set of marginals { µ i } i ∈ S ∈ P K , Haasler et al. (2021) study the EOT problem associated to T , see (Haasler et al., 2021, Eq. (2.4)), which is given by

This problem may be solved with Sinkhorn algorithm (Cuturi, 2013; Knight, 2008; Sinkhorn & Knopp, 1967), for which the authors provide an efficient implementation adapted to the tree-based setting, see (Haasler et al., 2021, Algorithm 3.1). Moreover, they state the convergence of their method in (Haasler et al., 2021, Theorem 3.5), as a direct consequence of the results presented in Luo &Tseng (1992).

In (Haasler et al., 2021, Section 4.2), it is assumed that S corresponds to the set of the leaves of T , as we do, and it is shown an equivalence between (discrete-EmOT) and the discrete-state static SB problem stated in (Haasler et al., 2021, Eq 4.2), which is given by

where T r = ( V , E r ) is the directed version of T rooted in an arbitrary vertex r ∈ S , and A ( v,v ′ ) = exp( -C ( v,v ′ ) /ε ) ∈ Z (2) for any ( v, v ′ ) ∈ E r . Remark that A ( v,v ′ ) may not necessarily be a transition probability matrix.

Finally, Haasler et al. (2021) provide two main numerical experiments. In (Haasler et al., 2021, Section 5.2), they consider a tree with 15 vertices, 14 edges and 8 leaves, combined to the state-space X = { 0 , 1 } 50 × 50 , and solve the corresponding (discrete-EmOT) problem for the quadratic cost. In (Haasler et al., 2021, Section 6), they apply their methodology to estimate ensemble flows on a hidden Markov chain. Given τ ∈ N ∗ , they consider a tree T with τ internal vertices (modeling the distribution of N agents at time t ∈ { 1 , . . . , τ } ), that are linearly linked, and such that each of these vertices is independently linked to S leaves of T (modeling observations at time t ∈ { 1 , . . . , τ } ). In this setting, the state space is given by X = { 1 , . . . , 100 } N . They solve the formulation (discrete-TreeSB) where the reference measure is chosen as a random walk.

Comparison with our results. We now establish remarks on the main differences between our methodology and the work of Haasler et al. (2021).

First of all, the continuous state-space counterpart of (discrete-TreeSB) is given by

where π 0 is a reference measure which factorizes along T r . In this case, π v,v ′ , π v and π 0 v ′ | v in (28) respectively correspond to the continuous version of M ( v,v ′ ) , ν v and A ( v,v ′ ) in (discrete-TreeSB). In contrast, our formulation of the multi-marginal Tree Schrödinger Bridge problem given in (TreeSB) is a minimization problem over all probability measures π ∈ P ( ℓ +1) , and is not restricted to the distributions that admit a Markovian factorization along T as in (28). Hence, our framework may be considered more general. Remark that under A 1, A 2 and A 3 , Proposition 3 states that (TreeSB) admits a unique solution π ⋆ ≪ π 0 such that (d π ⋆ / d π 0 ) can be written with potentials. Then, π ⋆ ∈ P T r since π 0 ∈ P T r , and (TreeSB) is then equivalent to (28).

Furthermore, (EmOT) is more general than the continuous version of (discrete-EmOT), which we can recover by taking any measure ν of the form (d ν/ dLeb) = exp[ ⊕ i ∈ S φ i ] in (EmOT), where { φ i } i ∈ S is a family of potentials such that ∣ ∣ ∫ R d φ i d µ i ∣ ∣ < ∞ for any i ∈ S . As a consequence, our setting allows us to choose the root r ∈ V \ S for the SB problem, whereas Haasler et al. (2021) only consider the case where r ∈ S . In the latter case, we establish in Appendix E that r can be chosen arbitrarily, as stated by (Haasler et al., 2021, Corollary 4.3).

Finally, TreeDSB deeply differs from the framework of Haasler et al. (2021) due its dynamic nature. Although we solve the same tree-based static SB problem (up to continuous/discrete state-space consideration), our approach consists in computing dynamic iterates ( i.e. , path measures) using diffusion-based methods instead of static iterates ( i.e. , distributions) using Sinkhorn algorithm. This paradigm is at the core of the DSB (De Bortoli et al., 2021) methodology, and offers an efficient approach to tackle high-dimensional settings, where Sinkhorn algorithm would fail.

Here, we present some advantages of the method proposed by Haasler et al. (2021) compared to ours. First, Haasler et al. (2021) may choose any kind of tree-based cost in practice, while our methodology only holds for the quadratic cost. This limitation is shared with all approaches based on the DSB (De Bortoli et al., 2021) methodology. Indeed, since the cost is determined by the reference path measure, we often choose quadratic costs associated with Brownian motions or Ornstein-Uhlenbeck processes. Moreover, Haasler et al. (2021) may consider various inhomogeneous (discrete) state spaces for the vertices of T , as presented in their numerical experiments. In our case, this approach is not compatible with our diffusion-based method. Finally, unlike Haasler et al. (2021), our method is not scalable with the number of vertices or edges in T due to computational limits. This limitation is common to all multi-marginal approaches which rely on neural networks to parameterize the potential and/or the distributions of the multi-marginal OT method, see Li et al. (2020); Fan et al. (2020); Korotin et al. (2022, 2021) for instance.

E Further results on TreeSB

Choice of the root r in (TreeSB) . We recall that the reference measure π 0 considered in (TreeSB), which is defined in (2), verifies π 0 ∈ P T r for some fixed root r ∈ V and π 0 r ≪ Leb with density φ r . Moreover, we have π 0 v ′ | v ( · | x v ) = N( x v , ε/ (2 w v,v ′ )I d ) for any ( v, v ′ ) ∈ E r , and thus, π 0 is entirely determined by the choice of the root r and the density on the corresponding vertex φ r .

As presented in Appendix D.1, we recall that (TreeSB) is equivalent to any multi-marginal Tree-SB problem with a reference measure ¯ π 0 given by (15), i.e. , ¯ π 0 writes as (d¯ π 0 / d π 0 ) = exp[ ⊕ i ∈ S ψ 0 i ] , where { ψ 0 i } i ∈ S is a family of measurable potentials with ψ 0 i : R d → R such that ∣ ∣ ∫ R d ψ 0 i d µ i ∣ ∣ < ∞ for any i ∈ S . In the case where r is chosen as a leaf of T , this result implies that (TreeSB) is unchanged if

(a) φ r = d ν/ dLeb where ν ∈ P is such that KL( µ r | ν ) < ∞ ,

(b) r is replaced by r ′ ∈ S , as long as H( µ r ) < ∞ and H( µ r ′ ) < ∞ .

Therefore, under A 0, the setting chosen in Section 3 is equivalent to any other setting where r is arbitrarily chosen in S and φ r = d ν/ dLeb where KL( µ r | ν ) < ∞ .

Consider now the case where r ∈ S c , i.e. , r is not a leaf of T . Then, the choice of φ r can not be made arbitrarily anymore, since it determines a further regularization on the r -th marginal of the solution to (TreeSB). In this setting, the sequence defined by (mIPF) is unchanged. Hence, TreeDSB proceeds in the same manner as presented in Section 3, except for the first iteration, which we detail now.

Let us define P = path T i 0 ( i 0 , r ) , where T i 0 = ( V , E i 0 ) is the directed version of T rooted in i 0 . We recall that first iterate of (mIPF) is defined by

Following the proof of Lemma 12, it is clear that

where we emphasize that P = { ( v, v ′ ) ∈ E i 0 : ( v ′ , v ) ∈ E r } . Therefore, Proposition 2 still applies between r and i 0 , by considering r instead of i K -1 . In practice, this means that the first iteration of TreeDSB consists in computing the time reversal of the path measures P 0 ( v ′ ,v ) for any ( v, v ′ ) ∈ P .

Extension of the regularized Wasserstein barycenter problem (regWB) . Consider the regularized Wasserstein-2 barycenter problem defined as follows

where ( w i ) i ∈{ 1 ,...,ℓ } ∈ (0 , + ∞ ) ℓ and µ 0 ∈ P is a reference measure. This formulation admits a further regularization compared to (regWB), which tends to make µ ⋆ ε closer to µ 0 . In particular, given a Wasserstein barycenter problem onto a star-shaped tree, the formulation ( µ 0 -regWB) may be more adapted than (regWB) if we have an a priori on the form of the regularized barycenter. In the case where µ 0 = N(0 , σ 2 0 I d ) , letting σ 0 →∞ , we recover the ( ℓε, ( ℓ -1) ε ) doubly-regularized Wasserstein barycenter problem (regWB). In the same spirit as Proposition 7, we can derive the following result from Proposition 24, which proves that ( µ 0 -regWB) can be solved with TreeDSB.

Proposition 25. Let ε > 0 and µ 0 ∈ P such that µ 0 ≪ Leb . Assume A 0. Also assume that T is a star-shaped tree with central node indexed by 0 , and that the reference measure of (TreeSB) defined in (2) verifies r = 0 and φ r = d µ 0 / dLeb > 0 . Under A 2 , ( µ 0 -regWB) has a unique solution π ⋆ 0 , where π ⋆ is the solution to (TreeSB) .

Below, we provide practical guidelines to parameterize µ 0 when it is chosen as a Gaussian distribution.

Gaussian design of µ 0 in ( µ 0 -regWB) . Consider an undirected star-shaped tree T with K + 1 vertices and leaves { 1 , . . . , K } . In order to incorporate the marginal constraints in the penalization brought by µ 0 when it is a Gaussian distribution, we set its mean to ∑ K i =1 E [ µ i ] /K and its diagonal covariance matrix as α × ( ∑ K i =1 diag(Cov[ µ i ]) -1 /K ) -1 , where the inverse operation is componentwise and α is a positive hyperparameter. This choice of variance helps to correctly explore the state-space at the very first iteration of TreeDSB, which is key to ensure numerical stability. In this setting, (TreeSB) verifies A 2 and A 3, by Proposition 15 and Proposition 16. In particular, we use this approach for two of our experiments: synthetic Gaussian datasets and Bayesian fusion, see Appendix G.

F Algorithmic techniques

Time discretization in TreeDSB. Denote k n = ( n -1) mod( K ) for any n ∈ N . Let T = ( V , E ) be a weighted undirected tree and consider the multi-marginal Schrödinger bridge problem (TreeSB) associated to this tree. We recall that for any { v, v ′ } ∈ E , we define T v,v ′ = ε/ (2 w v,v ′ ) .

Consider the path measures { P n ( v,v ′ ) } n ∈ N , ( v,v ′ ) ∈ E kn recursively defined by (a) and (b). By combining Proposition 1, Proposition 2 and results on time reversal theory (Haussmann & Pardoux, 1986), we obtain by recursion that for any n ∈ N , any ( v, v ′ ) ∈ E k n , P n ( v,v ′ ) is associated with a Stochastic Differential Equation on [0 , T v,v ′ ] given by

Let N ∈ N ∗ . In order to sample from the dynamics (29) at iteration n ∈ N , we consider its Euler-Maruyama discretization on ( N +1) time steps,

where Z m ∼ N(0 , I d ) for any m ∈ { 1 , . . . , N } , t m = ∑ m i =1 γ i , and { γ m } N m =1 ∈ (0 , ∞ ) N is a time schedule such that ∑ N m =1 γ m = T v,v ′ . This results in approximating the path measure P n ( v,v ′ ) by the joint distribution π n,N ( v,v ′ ) ∈ P ( N +1) defined by

where π n,N ( v,v ′ ) ,m +1 | m ( ·| x m ) = N( x m + γ m +1 f n t m ,v,v ′ ( x m ) , γ m +1 I d ) for any m ∈ { 0 , . . . , N -1 } . If N is chosen large enough, then π n,N ( v,v ′ ) ,m and P n ( v,v ′ ) ,t m have approximately the same distribution

for any m ∈ { 0 , . . . , N } . Consequently, ( P n ( v,v ′ ) ) R is naturally approximated by the joint distribution ˜ π n,N ( v,v ′ ) ∈ P ( N +1) defined by

If N is chosen large enough, we obtain that

where p v,v ′ ,t is the density of P n ( v,v ′ ) ,t w.r.t. the Lebesgue measure.

Following the construction of our dynamic iterates, we now explain how the sequence { π n ( v,v ′ ) } n ∈ N ∗ , ( v,v ′ ) ∈ E kn is recursively defined. Let n ∈ N , k n = ( n -1) mod( K ) . Define the path P n = path T i k n ( i k n , i k n +1 ) . Then, for any ( v, v ′ ) ∈ E k n +1 ,

These computations may be obtained by considering the sequence given by (mIPF) to solve the multi-marginal Tree-SB problem associated to T ( N ) = ( V ( N ) , E ( N ) ) , the N -discretized version of T (see Appendix B) with weights w ( N ) e m = 2 γ m /ε , which is given by

To approximate the IPF recursion given by (a) and (b), we use on each edge of T the score-matching approach of De Bortoli et al. (2021), which avoids heavy computations of score approximations. The next proposition is direct adaptation of (De Bortoli et al., 2021, Proposition 3).

Proposition 26. Assume that for any n ∈ N , any ( v, v ′ ) ∈ E k n with k n = ( n -1) mod( K ) , we have

Let n ∈ N . Consider the path P n = path T i k n ( i k n , i k n +1 ) . Let ( v, v ′ ) ∈ E k n +1 . Define p n = π n,N ( v,v ′ ) and m N = N -m -1 . Then, if ( v ′ , v ) ∈ P n , we have

otherwise, we have F n +1 m,v,v ′ = F n m,v,v ′ .

In practice, we use two neural networks per edge { v, v ′ } ∈ E , one for each possible direction of the edge, such that F v,v ′ ( θ n v,v ′ , m, x ) ≈ F n m,v,v ′ ( x ) and F v ′ ,v ( θ n v ′ ,v , m, x ) ≈ F n m,v ′ ,v ( x ) . For any { v, v ′ } ∈ E , the parameter θ n v,v ′ is updated at iteration n via the score matching loss defined by (31) in Proposition 26 if ( v, v ′ ) ∈ path T i k n ( i k n , i k n +1 ) , see Algorithm 1.

G Additional experimental results and details

The numerical experiments presented in Section 7 are obtained by our own Pytorch implementation, which is inspired from the code 6 provided by De Bortoli et al. (2021). We first provide information on the general setting of our experiments in Appendix G.1, and then give details on each of them in Appendix G.2 along with additional results. We recall that a mIPF cycle is defined as a subset of K consecutive iterations of (mIPF) and that the order of the leaves given by { i 0 , . . . , i K -1 } is randomly shuffled at each new mIPF cycle.

G.1 General experimental setup

Implementation of Algorithm 1 in practice. Let n ∈ N , with k n = ( n -1) mod( K ) , k n +1 = n mod( K ) . Consider the path P n = path T i k n ( i k n , i k n +1 ) . Assume that we are provided with a dataset D i kn , which contains M samples from π n i kn . Following Lines 7-9 in Algorithm 1, we apply processes (a) and (b) recursively on the edges ( v, v ′ ) ∈ P n .

(a) Sampling step (Line 7). For any x 0 ∈ D v , we sample from the diffusion trajectory (30) given by the Euler Maruyama discretization of P n v,v ′ starting from x 0 . This gives us M × N trajectory samples. We then store the last iterate of each trajectory in a new dataset D v ′ , which thus approximates π n v ′ .

(b) Training step (Lines 8-9). In order to avoid heavy computation, we approximate the meanmatching loss (31) by an unbiased estimator obtained by subsampling b elements from the full trajectories computed in the sampling process, see (De Bortoli et al., 2021, Eq. (97)-(98)). Here, b refers to the batch-size parameter of the neural networks. Then, we perform gradient descent to optimize the parameter θ v ′ ,v , which parameterizes the backward drift on the edge ( v, v ′ ) .

To avoid any bias issue, the whole trajectories obtained at process (a) are refreshed at a certain frequency over the training iterations of the neural networks by once again simulating the diffusion (30). In our experiments, this refresh occurs each 500 iterations.

Setting of the time discretization. The number of time-steps N in the time discretization of the diffusions is chosen to be even and identical for each of the edges of the tree. Let { v, v ′ } ∈ E . We now give details on the design of the time schedule { γ k } N k =1 related to the edge { v, v ′ } , see Appendix F. Following De Bortoli et al. (2021), we choose this sequence to be invariant by time reversal and consider γ k = γ 0 +(2 k/N )(¯ γ -γ 0 ) for any k ∈ { 0 , . . . , N/ 2 } (the rest of the sequence being obtained by symmetry) where γ 0 is a free parameter and ¯ γ is determined by ∑ N k =1 γ k = T v,v ′ . In our experiments, we set N = 50 and γ 0 = 10 ¯5 .

Sampling improvement. In our code, we implemented the corrector scheme of Song et al. (2021) and the probability flow -based sampling approach detailed in (De Bortoli et al., 2021, Section H.3), but did not observe any significant improvement in our experiments using one of these techniques.

Choice of the architectures of the neural networks. In the case of the experiments related to synthetic datasets (two-dimensional toy datasets, Gaussian distributions) and to the subset posterior aggregation task, we implement the same architecture as presented in (De Bortoli et al., 2021, Figure 3). We refer to this model as 'Basic Model' and detail it in Figure 6. In the 'Basic Model', the PositionalEncoding block applies the sine transform described in Vaswani et al. (2017), with output dimension equal to 32, and each MLP Block represents a Multilayer Perceptron Network. In particular, MLPBlock (1a) has shape ( d, 128 , max(256 , 2 d )) , MLPBlock (1b) has shape (32 , 128 , max(256 , 2 d )) , and MLPBlock (2) has shape (2 × max(256 , 2 d ) , max(256 , 2 d ) , max(128 , d ) , d ) , where d denotes the dimension of input data. We optimize the networks with ADAM (Kingma & Ba, 2014) with learning rate 10 -4 and momentum 0 . 9 . For each of the networks, we set the batch size to 4,096 and the number of iterations to 10,000 for the synthetic datasets and 15,000 for the subset posterior aggregation task. Our experiments ran on 1 Intel Xeon CPU Gold 6230 20 cores @ 2.1 Ghz CPU.

MLPBlock (1a)

Concatenate

MLPBlock (2)

Output

k

PositionalEncoding

MLPBlock (1b)

In the case of the experiments related to MNIST dataset, we use a reduced UNET architecture based on Nichol & Dhariwal (2021), where we set the number of channels to 64 rather than 128. We implement an exponential moving average of network parameters across training iterations, with rate 0.999. We optimize the networks with ADAM (Kingma & Ba, 2014) with learning rate 10 -4 and momentum 0 . 9 . Finally, we set the batch size to 256 and the number of training iterations to 30,000. Our experiments ran using 1 Nvidia A100.

Details on regularized state-of the art methods. We run the fsWB algorithm (Cuturi & Doucet, 2014) with the implementation provided by Flamary et al. (2021). For each experiment, we run 100 Sinkhorn iterations with 1500 samples for each dataset ( i.e. , the maximum number of samples that it can generate) and set the regularization parameter ε to its lowest value such that the algorithm is stable. Finally, for sake of fairness with our method, we initialise the barycenter measure with π 0 r when solving the problem ( µ 0 -regWB) for synthetic Gaussian datasets and Bayesian fusion. To run the crWB algorithm (Li et al., 2020), we use the code provided by the authors. We consider the quadratic regularization, which is shown to be empirically more stable than entropic regularization. Following Fan et al. (2020), we choose the potential networks to be fully connected neural networks with 3 hidden layers of shape (max(128 , 2 d ) , max(128 , 2 d ) , max(128 , 2 d )) . The activation functions are ReLu. We optimize the networks with ADAM (Kingma & Ba, 2014) with learning rate 10 -4 for the subset posterior aggregation task and 10 -3 for the Gaussian experiment. Finally, we set the batch size to 4,096 and the number of training iterations to 50,000. We highlight that fsWB and crWB solve a regularized Wasserstein barycenter problem, which does not contain an additional penalization term on the entropy of the barycenter, contrary to TreeDSB.

G.2 Details on the experiments

Synthetic Gaussian datasets. For each dimension that we consider, we generate three different triplets of random non-diagonal covariance matrices whose condition number is less than 10. We then run the algorithms on each triplet and aggregate the obtained results. The Gaussian datasets contain 1,500 samples for fsWB, and 10,000 samples for crWB and TreeDSB. We run fsWB with the following settings ( d, ε ) ∈ { (2 , 0 . 1) , (16 , 0 . 2) , (64 , 0 . 5) , (128 , 1 . 0) , (256 , 2 . 0) } . We run TreeDSB for 10 mIPF cycles with regularization parameter ε = 0 . 1 , starting from the central node initialized to a Gaussian distribution µ 0 chosen as detailed in Appendix F with α = 1 . Thus, we solve the regularized Wasserstein barycenter problem ( µ 0 -regWB), which contains an additional regularization with respect to µ 0 . This choice is justified, since the non-regularized barycenter is known to be a Gaussian distribution, and µ 0 can be seen as an a priori for the regularized barycenter. For each of the three settings, we keep the best result among the 30 mIPF iterations. In this setting, TreeDSB and crWB have roughly the same training time.

Subset posterior aggregation. When considering a dataset splitted into several subdatasets, a common paradigm in bayesian inference consists in running Monte Carlo Markov Chain methods separately on these subdatasets, and then merge the obtained posteriors to recover the full posterior. The barycenter of these subdataset posteriors is proved to be close to the full data posterior under mild assumptions (Srivastava et al., 2018). In our setting, we consider the posterior aggregation problem for the logistic regression model associated to the wine dataset 7 (d = 42) with 3 subdatasets. We consider here two splitting methods: (i) either, data is uniformly splitted between 3 subdatasets with respect to the label distribution, denoted by wine-homogeneous , or (ii) data is splitted with some heterogeneity according to a Dirichlet distribution whose parameter is randomly chosen, denoted by wine-heterogeneous . Following Korotin et al. (2021), we use the stochastic approximation trick so that the subset posterior samples do not vary consistently from the full posterior in covariance (Minsker et al., 2014). We implement the Unadjusted Langevin Algorithm (ULA) to sample from each subdataset posterior and from the full posterior. In each case, we run ULA for 5 . 5 · 10 6 iterations with a well chosen step-size, and obtain 9,900 samples after applying a burn-in of order 10% and then a thinning of size 500. We provide in Figure 7 some metrics which assess the quality of this sampling process. We recall that the the full posterior samples serve as ground truth in this experiment.

The results presented in Table 2 were computed as follows. For fsWB, we first subsample 1,500 samples out of the 9,900 samples from each posterior, and then run the algorithm with ε = 0 . 5 . We repeat three times this procedure and then aggregate the results. In the case of crWB and TreeDSB, we run the algorithms three times with various seeds. Similarly to the Gaussian setting, we run TreeDSB for 10 mIPF cycles with regularization parameter ε = 0 . 1 . We start from the central node with a Gaussian distribution µ 0 chosen as detailed in Appendix F with α = 1 , and thus solve the barycenter formulation ( µ 0 -regWB). For each of the three settings, we keep the best result among the 30 IPF iterations. In this setting, TreeDSB and crWB have roughly the same training time.

ACF-dim=0(including burn-in&thinning)

1.0

data:full

data: 1

data:2

0.8

data:3

Autocorrelation

0.6

0.4

0.2

0.0

0

100

200

300

400

500

Lag

NLLevaluatedoneachtraindataset

103

data:full

data: 1

data:2

data:3

102

100

101

102

103

104

105

106

107

ACF-dim=0(includingburn-in &thinning)

1.0

data:full

data: 1

data:2

0.8

data:3

Autocorrelation

0.6

0.4

0.2

0.0

0

100

200

300

400

500

Lag

NLLevaluatedoneachtraindataset

103

data:full

data:1

data:2

data:3

102

101

100

101

102

103

104

105

106

107

MCiteration

MCiteration

Synthetic two-dimensional datasets. In this setting, we consider three different datasets ( Swiss-roll , Circle and Moons ) that each contain 10,000 samples. Since we do not have an a priori on the shape of the barycenter between these datasets, we consider the regularized Wasserstein barycenter problem (regWB), i.e. , r is chosen as a leaf and corresponds to one of the input datasets. We emphasize that this experiment is not intended to demonstrate the superiority of TreeDSB to compute 2D Wasserstein barycenters, but is rather meant to illustrate that (a) the marginals of the leaves are well recovered by the algorithm, see Figure 3, and that (b) the obtained barycenter is consistent when diffusing from the different leaves, see Figure 4. In all our experiments on 2D datasets, we observed that (a) was persistently verified without difficulty. In this section, we rather aim at illustrating (b) by providing additional results which assess the quality of the barycenter obtained by TreeDSB with respect to the choice of the starting leaf r and to the choice of the regularization parameter ε .

To do so, we consider three different choices of regularization in TreeDSB: (i) ε = 0 . 2 (50 mIPF cycles), see Figure 8, (ii) ε = 0 . 1 (50 mIPF cycles), see Figure 9 and (iii) ε = 0 . 05 (60 mIPF cycles), see Figure 10. For each of these settings, we run TreeDSB with the starting leaf r chosen as Swiss-roll (first row), Circle (second row) or Moons (third row), and display the final barycenter obtained by diffusing from Swiss-roll (first column), Circle (second column) and Moons (third column). Note that the vertex 0 always corresponds to the starting leaf, the vertex 1 to the barycenter node and that Figure 4 corresponds to the first row of Figure 9.

We can make the following observations. First, the estimated barycenter is always coherent within each row, which assesses the convergence of our method. Then, for each value of ε , the TreeDSB barycenter is rather consistent between the rows, i.e. , the choice of the starting leaf does not have a meaningful impact on our method. Finally, as expected, we observe that the support of the barycenter is less and less diffuse as long as ε decreases.

mIPFcycle:50,fromvertex0to1

mIPFcycle:50,fromvertex2to1

mIPFcycle:50,fromvertex3to1

15

15

15

10

10

10

5

5

5

0

0

-5

-5

-5

-10

-10

-10

15

-15

-15

-15

5-10-5

0

5

10

15

-15

-10

0

5

10

15

15

-10-5

0

5

10

mIPFcycle:50,fromvertex2to1

mIPF cycle:50,from vertex0 to 1

mIPFcycle:50,fromvertex3to 1

15

15

15

10

10

10

5

0

0

0

-5

-5

-5

-10

-10

-10

15

-15

-15

-15 -10

-5

0

10

15

-15-10

-5

0

5

10

15

15

-10

-5

0

5

10

mlPF cycle:50,from vertex2to1

mIPFcycle:50,fromvertex3to1

mIPFcycle:50,from vertex0to1

15

15

15

10

10

10

5

5

5

0

0

0

-5

-5

-5

-10

-10

-10

15

15

15

-15-10

0

5

-10

0

-10

5

-5

1015

15

-5

1015

15

-5

1015

mIPFcycle:50,fromvertex0to1

mIPFcycle:50,fromvertex2to1

mlPFcycle:50,fromvertex3to1

15

15

15

10

10

10

5

5

5

0

0

0

-5

-5

-5

-10

-10

-10

15

-15

15

-15-10-5

0

1015

-15-10-5

0

5

1015

15

-10

-5

0

10

mIPFcycle:50,fromvertex2to1

mIPFcycle:50,fromvertex0 to1

mIPFcycle:50,fromvertex3to1

15

15

15

10

10

10

5

0

0

0

-5

-5

-5

-10

-10

-10

15

15

-15

-15-10-5

0

5

10

15

-15-10

-5

0

5

10

15

-15-10

-5

0

5

10

15

mlPFcycle:50,fromvertex2to1

mlPFcycle:50,fromvertex3to1

mlPFcycle:50,fromvertex0to1

15

15

15

10

10

10

5

5

5

0

0

0

-5

-5

-5

-10

-10

-10

15

-15

-15

-15-10-5

-5

0

1015

-15 -10

1015

-15-10-50

1015

mIPFcycle:60,fromvertex0to1

mIPFcycle:60,from vertex2to1

mIPFcycle:60,from vertex3to1

15

15

15

10

10

10

5

5

0

0

-5

-5

-5

-10

-10

-10

-15

-15

-15

-15

5-10-5

0

5

10

-15-10-5

0

10

15

-10-5

0

5

10

mIPFcycle:60,fromvertex2to1

mIPFcycle:60,from vertex0 to1

mIPFcycle:60,fromvertex3to1

15

15

15

10

10

10

5

5

0

0

0

-5

-5

-5

-10

-10

-10

15

-15

-15

-15 -10 -5

0

1015

-15-10

-5

0

1015

15

-10

-5

0

5

10

15

mlPFcycle:60,from vertex2to1

mIPF cycle:60,from vertex3to1

mIPF cycle:60,from vertex 0to1

15

15

15

10

10

10

5

5

5

0

0

0

-5

-5

-5

-10

10

-10

15

15

15

0

0

-15-10

-5

5

1015

-15-10

-5

10

15

15

-10

-5

10

15

For purpose of illustration, we provide in Figure 11 the barycenter obtained by state-of-the-art two-dimensional in-sample methods that are available in POT library (Flamary et al., 2021): (i) nonregularized free-support Wasserstein barycenter (Cuturi & Doucet, 2014), (ii) entropic-regularized free-support Wasserstein barycenter (fsWB) with ε = 0 . 5 (Cuturi & Doucet, 2014) and (iii) entropicregularized convolutional Wasserstein barycenter with ε = 5 . 10 -4 (Solomon et al., 2015), which is specifically designed for images. We notably observe that TreeDSB cannot capture the full complexity of the 2D barycenter compared to these methods. We infer that this gap comes from the dynamic nature of TreeDSB, since increasing the number of training iterations per IPF iteration or improving the complexity of the neural networks did not bring any significant change in our results. Finally, we recall that the methods (i), (ii) and (iii) do not scale well with dimension, and have to be completely run again when new data samples are available, contrary to TreeDSB.

Free Support exact barycenter in 2D

Free Support Sinkhorn barycenter in 2D

Q

Convolutional Sinkhorn barycenterin 2D

15

15

10

10

20 -

5

40

0

0

60-

-5

-5

-10

10

80

15

15

-15

-10

-5

Q

5

10

15

15

-10

-5

0

5

10

15

Q

20

40

60

80

100

120

MNIST Wasserstein barycenter. This setting can be qualified as high-dimensional , since the data dimension is d = 784 . Here, each digit dataset contains 1,000 samples. As in the two-dimensional setting, we do not have an a priori on the shape of the barycenter between MNIST digits, and thus consider the formulation (regWB), where the root r is chosen as a leaf. We propose below several experiments to assess the scalability of TreeDSB to this setting.

Digits 0 and 1. In Figure 12, we report the results obtained by running TreeDSB on MNIST digits 0 and 1, for 15 mIPF cycles with ε = 0 . 5 , starting from the leaf MNIST-0. We display 25 samples from the reconstructed MNIST-0 marginal (first column), from the reconstructed MNIST-1 marginal (fourth column), from the estimated barycenter by diffusing from MNIST-0 (second column) and diffusing from MNIST-1 (third column). We notably observe that the digits are well recovered and that the barycenter samples are consistent. We draw the reader's attention to the fact that TreeDSB showed numerical unstability with a regularization value ε lower than 0 . 5 . For purpose of illustration, we display in Figure 13 the Wasserstein barycenter obtained by non-regularized methods from Fan et al. (2020) and Korotin et al. (2021), and by the regularized approach from Li et al. (2020).

G

0

D

0

0

0

0

0

1

1

1

厂

C

0

0

0000

Digits 2,4 and 6. In Figure 14, we report the results obtained by running TreeDSB on MNIST digits 2,4 and 6, for 10 mIPF cycles with ε = 0 . 5 . Here, we consider three settings which differ by the starting leaf r in the algorithm: MNIST-2 (first row), MNIST-4 (second row), or MNIST-6 (third row). For each of these settings, we display 30 samples from the estimated barycenter by diffusing from MNIST-2 (first column), diffusing from MNIST-4 (third column) and diffusing from MNIST-6 (third column). We notably observe a global consistency of the barycenter samples across the various settings. In Figure 15, we report the results obtained by running TreeDSB on MNIST digits 2,4 and 6, for 10 mIPF cycles with ε = 0 . 2 , starting from MNIST-6. We display 30 samples from the reconstructed marginals (first row), from the estimated barycenter (second row) by diffusing from MNIST-2 (first column), diffusing from MNIST-4 (second column) and diffusing from MNIST-6 (third column). As expected, we observe less noisy barycenter samples compared to Figure 14, while still well recovering MNIST digits.

Digits 0,1 and 4. In Figure 16, we report the results obtained by running TreeDSB on MNIST digits 0,1 and 4, for 10 mIPF cycles with ε = 0 . 5 . We consider two settings which differ by the starting leaf r in the algorithm: MNIST-0 (second row) and MNIST-1 (first/third rows), for which we display samples from the reconstructed measures (first row). In Figure 17, we report the results obtained by running TreeDSB on MNIST digits 0,1 and 4, for 10 mIPF cycles with ε = 0 . 2 . We consider two settings which differ by the starting leaf r in the algorithm: MNIST-0 (first/second row), for which display samples from the reconstructed measures (first row), and MNIST-1 (third row). For all of these settings, we display 30 samples from the estimated barycenter by diffusing from MNIST-0 (first column), diffusing from MNIST-1 (third column) and diffusing from MNIST-4 (third column). Similarly to the digits 2-4-6, we observe consistency within the barycenter samples, unconditionally to the starting leaf, and less noise as ε decreases. Note that the reconstructed MNIST digits are less truthful to the original datasets when ε is low.

H

4

4

力

4

3

口

4

4

f

y

4

4

4

h

4

十

4

4

7

4

4

+

4

4

4

4

+

牛

￥

+

h

4

力

+

h

4

4

f

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors develop a novel method for approximating Wasserstein barycenters, which is based on diffusion models. This method combines ideas from the recently introduced Diffusion Schrodinger bridge and a framework for multi-marginal optimal transport with tree-structured costs.\nThus, although the paper focuses on the barycenter problem (corresponding to a star shaped tree), the proposed approach can be used for any multi-marginal optimal transport problem, where the cost is composed of pairwise terms that are structured as a tree.\nA diffusion Schrodinger bridge type algorithm is presented that approximates the multimarginal iterative proportional fitting (mIPF) algorithm. The authors present novel convergence results of mIPF, which in contrast to previous works neither assume that the space is compact nor that the cost is bounded.\nFinally, the application to Wasserstein barycenters is evaluated for several experiments.\n\nSTRENGTHS:\nI see the strengths of this paper in its originality and significance. The paper extends results for Diffusion Schrodinger bridges to a class of multi-marginal problems, including Wasserstein barycenters. Moreover, the paper extends results for tree structured multi-marginal optimal transport problems by developing a diffusion based method for these, which makes the problem scalable to high-dimensional settings.\n\nWEAKNESSES:\nThe clarity of presentation should be improved in many parts of the paper. Please see comments below.\n\nQUESTIONS:\nMy main question concerns the different update schemes mentioned in Section 3. The (mIPF) algorithm requires sampling from the whole tree. The authors argue that each (mIPF) update can be done on a path of the tree. In lines 160-173 it is described that one first samples along all edges of the path, and then learns the reverse dynamics for the path. In contrast, Algorithm 1 samples and learns the reverse dynamics for each edge on the path individually, and this is repeated iteratively along the path. That is, three different update schemes are mentioned:\na) sample along whole tree and learn reverse dynamics;\nb) sample along whole path and learn reverse dynamics;\nc) iterate over the edges on the path, and for each edge sample along each and learn reverse dynamics.\nPlease clarify why these are equivalent, especially the equivalence with scheme c) is not obvious to me.\n\n\nThe following list is ordered mostly chronological, and not according to relevance.\n\n1. In Section 2 the authors should make sure that the background material and setting are introduced in a clear and concise manner. This includes my following comments, but I suggest to carefully work through the whole section.\n- Are all the details on dynamic and static versions of Schrodinger bridges really needed? E.g., it is not clear what $\\Phi$ and $a$ are, and I believe they are not used anywhere in the paper anyways, so this comment could be skipped.\n- The link between (EmOT) and Schrodinger bridges is not well explained. A tree is defined in the second sentence, but it is only explained how it is related to (TreeSB) in the last sentence. Also you probably need a reference to equation (1) somewhere?\n- The final paragraph \"our framework\" could be used to give a much more helpful summary of the main idea of this work. I suggest to state more clearly the motivation, contribution, and main result, and tone down the technical details (which are hard to understand at this point).\n\n2. Proposition 3 should be presented in a self-sufficient way. It is not completely obvious here what $\\mathbb{P}^n$ and $\\pi^n$ denote.\n\n3. What is the convergence criterion in Algorithm 1, line 5?\n\n4. I feel that Section 4 does not really follow the flow of the paper, and I wonder why the authors don't stress the convergence result more as a main contribution of the paper. Also, I'm wondering if the result relies on any additional assumptions compared to previous results? Moreover, I'm a bit concerned about the fact that the one page of Section 4 relies on 10 pages of proofs, which are deferred to the appendix (and which I was not able to check). I find it problematic to have to hide 95% of the derivation for a strong theoretical result in the supplementary material. It might be preferable to not include this result in the paper due to space limitations, and publish it separately. This would also guarantee a thorough peer review of the result. In my opinion the paper makes a strong enough contribution without this Section. \n\n5. For the synthetic two dimensional dataset it would be interesting to also see the true (possibly regularized) Wasserstein barycenter for reference.\n\n6. For the MNIST application, I suggest to clarify in the main paper that you consider the problem in 784 dimensional space. This may not be clear, since images could be seen as distributions of 2-dimensional data, similar to the first numerical experiment.\n\n7. What do you mean by one IPF cycle, and how do you chose the number of IPF cycles in the different experiments?\n\n8. Since the paper presents methods for optimal transport problems with general tree structure, it would be very interesting to also include some experiments that are not a Wasserstein barycenter problem.\n\n9. Typos: \n43: hyperparameter and;\n112: Section 5. We;\n220 applied to trees\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nAuthors develop a new method for computing multi-marginal optimal transport in the case where the cost function is tree structured. Their algorithm cleverly exploits the dynamic formulation of the problem. The correctness of the algorithm is thoroughly demonstrated, and benefits are shown over a number of toy examples as well as competitive baselines involving the computation of the Wasserstein Barycenter.\n\nSTRENGTHS:\nIt is a solid paper (theoretically and experimentally). Although the mathematical part is mostly technical and algorithm-related (but correct me if I am wrong), it is still quite sophisticated and praise deserving. Results section is compelling.\n\nWEAKNESSES:\nThere are some weaknesses: first, the paper is slightly hard to read and while there are some helpful guides (as the first figure) at some points the paper gets helplessly hard to follow. Second, the paper should better explain what the contribution over the already existing methods (i.e Hasslear et al). It is not clear from the reading what this adds on top of that. Third, the paper is not compelling on its motivation. Why if the object of interest is a tree-based cost, all experiments focus on the barycenter problem? it gives the impression that the paper addresses a problem that doesn't exist in practice. Fourth, while the experimental results are super interesting, there is a lack of commentary as to why we should expect this.\n\nQUESTIONS:\n1) What is the difference between your method and the one of Haessler et al?\n2)Could authors spell out applications where multimarginal OT inference in tree structured costs is important? Why is the focus only on the Barycenter?\n3)Authors are improving upon the competitive baselines such as Cuturi and Doucet (table 1). Why the method was expected to bring this benefit? is it related to the tree-ness or the fact that it implements a Schrodinger Bridge?. \n4)Likewise, why better results are obtained in data fussion? (Table 2). It seems to be implied that this is due to better performance in high dimensions. Why this should be the case? \n5)Line 175 seems to be an error in the Ext() statement.\n6)Type in line 305\n\nLIMITATIONS:\nn.a.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

comments: SUMMARY:
This paper generalizes the diffusion Schrödinger bridge framework to multidimensional problems with tree-structured quadratic cost function. The method builds on representing each egde in the graph as two SDEs (one for each direction) parameterized by neural networks. These are used to propagate the information (from the known marginals) in the graph. Under a set of assumptions the method is shown to converge in TV.

STRENGTHS:
It is a nice extension of the Diffusion Schrödinger Bridge framework, which allows for addressing a relatively large class of multimarginal optimal transport problems where each marginal is a probability measure. Also it is nice that convergence of the method is provided.

WEAKNESSES:
I think it would be nice with a little bit more information on the neural nets, about their structure and how they are used and updated. See the question below.

QUESTIONS:
I do not see how the assumptions A1-A5 uses the neural nets. The neural nets typically have finite prespecified parameters. Can you explain how one can represent any measure (satisfying the assumptions), which are infinite dimensional objects.

LIMITATIONS:
The limitations are adequately addressed.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors consider the multi-marginal optimal transport (mOT) with entropic regularization and a tree-structured quadratic cost. The authors propose to address it by developing the Tree-based Diffusion Schrodinger Bridge (TreeDSB), i.e., dynamic and continuous state-space of the multimarginal Sinkhorn algorithm. The authors apply the proposed approach to compute Wasserstein barycenter, i.e., equivalent to mOT with star-shaped tree.\n\nSTRENGTHS:\n+ The authors propose Tree-based Diffusion Schrodinger Bridge, a dynamic and continuous state-space counterpart of the multimarginal Sinkhorn algorithm.\n+ The authors illustrate some advantages of the proposed approach for Wasserstein barycenter.\n\nWEAKNESSES:\n+ The motivation of developing the dynamic and continuous state-space counterpart of the multimarginal Sinkhorn algorithm and experiments seem quite weak. For Wasserstein barycenters and problems illustrated in the experiments, it suffices to address the multimarginal OT (i.e., the static problem).\n+ Although there is some discussion in the appendix, it is essential that the submission is very closely related to (Hassler et al. 2021). It is an important part and should elaborate with details in the main manuscript to emphasize the contribution.\n+ The notations seem heavy and hard to follow. It is better to simplify the notations or give more explanation for the notations when they are introduced. At the present, it is hard to follow the meaning of the notation system.\n\nQUESTIONS:\n+ It is not clear why one needs to address the dynamic and continuous state-space of the multimarginal Sinkhorn. It suffices to consider the multimarginal OT directly for Wasserstein barycenters and problems in the experiments. Could the authors comment on this?\n+ Although the authors give some discussion about the relation with Hassler et al. (2021) in the appendix, it is essential to clarify their difference in the main manuscript (e.g., in related work section) to emphasize the contribution. Could the authors summary their main differences? (since they are very closely similar).\n+ In line 110, $S$ is the set of probability measures on the “leaves” of $T$. However, in line 118, why $r \\in S$ (it seems that $r$ is the root of tree T as in line 113)? So, what does the tree $T$ look like?\n+ Please explain the Figure 2, its description is too short and hard to follow.\n+ As in line 157-158, there are $2|E|$ neural networks, how are these networks related?\n+ Kitagawa and Pass (2015) showed the equivalent between multimarginal OT and the Wasserstein barycenter. Could the authors comment the result in Proposition 8 and the claim in line 271-272 with those in Kitagawa and Pass (2015)?\n\nRef: Kitagawa, J. and Pass, B., 2015, September. The multi-marginal optimal partial transport problem. In Forum of Mathematics, Sigma (Vol. 3, p. e17). Cambridge University Press.\n+ In line 284, how does the choice of $\\epsilon$ affect the result, and how to choose it in applications?\n+ For results in Tables 1, 2, Figures 6, 7, it is better to also report the corresponding time consumption?\n+ It seems unclear how to choose $\\epsilon$ (hyperparameters) for the barycenter problem? Does it affect the result? Some hyperparamters seem missing, e.g., how many free supports are used for fsWB?\n+ For results in Figures 6 and 7, how to interpret these results? (especially Figure 7)\n+ in line 330, “a scalable scheme”? it is unclear whether the proposed approach is scalable from the proposed algorithm and experiments? Please elaborate the scalability of the proposed method.\n\nSome minor points:\n+ In line 2, what is $C([O, T], R^d)$? There is no explanation for the new notation $C( , )$\n+ in line 78, it is better to explain the Stochastic Differential Equation form, e.g., there is no explanation about several new notations, $X_t, a, B_t$\n+ Please explain the Figures. The description for each Figure is very short, it is unclear the meaning of each Figure, e.g., how to understand the Figures 2, 3?\n+ In line 125, what is the role of $K$? \n+ In Proposition 1, what is the relation between $T_{K-1}$ and $T$ (although there is some explanation in line 71-72, it is hard to follow it). It is better to consider to simplify the notations or explain the notations with more details.\n\n---\n\nThank you for the rebuttal.\n\nLIMITATIONS:\nThe authors have discussed the limitation of the work. However, there is no discussion about the potential negative societal impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nSchröninger bridges (SB) can be understood as dynamical counterpart of the Entropic Optimal Transport problem. \n\nThis paper introduces the Tree Diffusion Schröninger Bridge (TDSB) model, which can be seen as a multi-marginal generalization of SB when the cost $c$ has a tree structure: given $\\ell$ marginals $\\pi_1,\\dots,\\pi_\\ell$, and a tree $T = (V,E)$ with $|V| = \\ell$ nodes, one has $c(x_1,\\dots,x_\\ell) = \\sum_{(v,v') \\in E} \\| x_v - x_{v'}\\|^2$. Let $S$ denote the leaves of $T$. \n\nIn that setting, the goal of the paper is to solve the TreeSB problem:\n$$ \\min_\\pi \\mathrm{KL}(\\pi | \\pi^0),\\quad \\pi_i = \\mu_i,\\ \\forall i \\in S. $$\nwhere the $\\mu_i$ are prescribed marginales, and $\\pi^0 = \\pi_r^0 \\bigotimes_{(v,v')} \\pi^0_{v,v'}$ is a reference measure that factorize along $T$, and $r \\in V$ is identified as a root of $T$, and the factors $\\pi^0_{v,v'}$ are chosen as normal distributions. \n\nIf $T$ has a star-structure, this problem is equivalent to the (doubly-regularized) Wasserstein barycenter problem (in sense that, at optimality, the marginal of $\\pi$ in the center of the star is the Wasserstein barycenter of the marginals on the leaves). \n\nThe contribution of the work is to solve this problem using an iterative algorithm along _paths_ of $T$ which can be summarized as a succession of $\\mathrm{KL}$ projections where only one of the marginals are constrained; if the tree is only made of two nodes (and one edge) we retrieve the Sinkhorn algorithm. In practice, the transition along a given edge (during one step of the algorithm) are performed by SDE of the form $\\dot X_t = f_\\theta(X_t) \\mathrm{d}t + \\mathrm{d} B_t$, where $B_t$ is the Brownian motion, and $f_\\theta$ is a drift that is learned (and parametrized by a neural network). This leads to an iterative scheme where one repetitively select path between two leaves and learns the drift to solve the KL projection along the path. \n\nFrom a theoretical standpoint, authors prove that (assuming the intermediate step converge, e.g. the drifts $f_\\theta$ end up being optimal) the proposed algorithm solves (asymptotically) the TDSB problem.  \n\nFrom a numerical perspective, they showcase the practicability of their approach on synthetic and simple benchmark datasets.\n\nSTRENGTHS:\n- I find the developed approach quite elegant. It mixes different notions in optimal transport literature (multi-maginal OT, entropic OT/Schrödinger bridge, doubly-regularized Wasserstein barycenters, parametrization of \"transport maps\" (the drifts) by neural networks, etc.). It is nice to see all these approaches implemented in a _practical_ setting (which can be implemented). \n- Comparison with previous work is well presented.  \n- The convergence results of the proposed approach are strong.\n\nWEAKNESSES:\n## Numerical experiments \n\nMy main remark would be about the numerical experiments. They are quite limited in their comparison with state-of-the-art methods and rather acts like proof-of-concept. In details: \n- They do not showcase how the proposed approach solves new problems. Though reasonable, none of the results showcase feel \"groundbreaking\". This may just be a matter of exposition, but right now, loosely speaking, I am not \"impressed\" by the experimental results. For instance, the choice of synthetic datasets (Figure 4) is questionable in that the actual barycenter is not very intuitive and thus it is hard to visually tell whether the displayed output (Figure 5) is good or not; in addition, the \"consistency\" of the proposed approach is not striking (I admit that there is a decent similarity between the three outputs, but we cannot faithfully consider that the model did converge). \n- The comparison with SotA methods is incomplete; e.g. there is no running time involved. Given the (maybe only superficial) complexity of the proposed approach---such as the need for a possibly huge amount of NN (which are not very small according to Section G, as they are MLPs), one may be feared by a possibly huge running time (compared to simpler approaches like, say, the one of Cuturi and Doucet). If there is a computational tradeoff (better result for longer training time), it must be mentioned and discussed in details. \n- I would expect a comparison with the _Iterative Bregman Projection_ of (Benamou et al., 2015), given the \"similarity\" (in that they are iterative KL projections) with the proposed approach (and I would consider IPB as being \"more SotA\" than Cuturi&Doucet, for instance; for instance the IPB paper claims in Section 1.3 that \"IPB converges (...) orders of magnitude faster than Cuturi&Doucet\"). \n- A bit minor, but the description of the _subset posterior aggregation_ experiment is very minimalist. It is hard for the reader that is not already familiar with such experiment to have an idea of what is done there.\n\nQUESTIONS:\n- My main question, which can be developed (see below) would be _why should one use this approach rather than existing ones?_ (e.g. IPB, fsWb, etc.). Of course, the paper gives some insight, but:\n\n1. As said above, there is no insight on the computational burden of the proposed approach. \n2. While I find very interesting that the approach works on any tree-structured coast, the work only focuses on Wasserstein barycenter, somewhat limiting its impact (that is, one strength of this work is that it solves a more general problem than previous Wasserstein-barycenter-dedicated methods). Giving an example + experiment where a (non-stared) tree cost is involved would be insightful (it would give a situation where the proposed approach as no challenger, as far as I know!).\n\nLIMITATIONS:\nauthors have addressed the limitations of their approach. \n\nI do not identify ethical concern specific to this work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-07"}