SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Consider now the directed tree  $\mathsf{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  $\mathbb{P}_{(v_{J + 1},v_J)}^{n + 1} = \mu_{i_{k_n + 1}}\otimes (\mathbb{P}_{(v_J,v_{J + 1})}^n)_{|0}^R$ . In the case of a bridge-shaped tree, this is exactly the  $(n + 1)$ -th update described in DSB. Then, for any  $(v,v^{\prime})\in \mathsf{E}_{k_n + 1}\backslash \{(v_{J + 1},v_J)\}$ ,
(a) either  $(v,v^{\prime})\in \mathsf{E}_{k_n}\backslash \mathsf{P}_n$  , and we define  $\mathbb{P}_{(v,v^{\prime})}^{n + 1} = \mathbb{P}_{(p(v),v),T_{p(v),v}}^{n + 1}\otimes \mathbb{P}_{(v,v^{\prime})|0}^{n},$
(b) or  $(v', v) \in \mathsf{P}_n$ , and we define  $\mathbb{P}_{(v, v')}^{n+1} = \mathbb{P}_{(p(v), v), T_{p(v), v}}^{n+1} \otimes (\mathbb{P}_{(v', v)}^n)_{|0}^R$ .
Proposition 2. Consider the sequence of dynamic iterates defined by (a) and (b). Then, for any  $n \in \mathbb{N}$  and any  $(v,v') \in \mathsf{E}_{k_n}, \mathbb{P}_{(v,v')}^n \in \mathcal{P}(\mathrm{C}([0,T_{v,v'},\mathbb{R}^d))$  and we have  $\operatorname{Ext}(\mathbb{P}_{(v,v')}^n) = \pi_{(v,v')}^n$ .
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  $\mathsf{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  $\mathsf{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  $\mathbb{P}_{(v',v)}^n$  is associated with  $\mathrm{d}\mathbf{X}_t = f_{t,v',v}(\mathbf{X}_t)\mathrm{d}t + \mathrm{d}\mathbf{B}_t$  with  $\mathbf{X}_0\sim \pi_{v'}^n$ , then, under mild conditions, its time-reversal  $(\mathbb{P}_{(v',v)}^n)^R$  is associated with  $\mathrm{d}\mathbf{Y}_t = \{-f_{T - t,v',v} + \nabla \log p_{v',v,T - t}\} (\mathbf{Y}_t)\mathrm{d}t + \mathrm{d}\mathbf{B}_t$  with  $\mathbf{Y}_0\sim \pi_v^{n + 1}$ , where  $p_{v',v,t}$  is the density of  $\mathbb{P}_{(v',v),t}^n$  w.r.t. the Lebesgue measure. The score  $\nabla \log p_{v',v,T - t}$  can then be approximated using score-matching techniques (Hyvarinen, 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'\} \in E$  with neural networks  $f_{t,\theta_{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)
1: Input:  $\mathsf{T} = (\mathsf{V},\mathsf{E}),\{\mu_i\}_{i\in \mathsf{S}},\{\theta_{v,v'}\}_{\{v,v'\} \in \mathsf{E}},N\in \mathbb{N}$   
2: for  $n = 0,\ldots ,N$  do  
3: Let  $k_{n} = (n - 1)\mathrm{mod}(K)$   
4: Get path between  $i_{k_n}$  and  $i_{k_{n + 1}}$ $\mathsf{P}_n = \{v_j,v_{j + 1}\}_{j = 1}^J$   
5: while not converged do  
6: for  $j = 1,\dots ,J$  do  
7: Sample from  $\mathbb{P}_{v_j,v_{j + 1}}^n$  (Euler-Maruyama)  
8: Compute mean matching loss  $\ell (\theta_{v_{j + 1},v_j})$   
9:  $\theta_{v_{j + 1},v_j}\gets$  Gradient Step  $(\ell (\theta_{v_{j + 1},v_j}))$   
10: Update  $f_{t,\theta_{v_{j + 1},v_j}}$   
11: end for  
12: end while  
13: end for  
14: Output:  $\{\theta_{v,v'}\}_{\{v,v'\} \in \mathsf{E}}$
The algorithm is initialized with  $f_{t,\theta_{v,v'}} = 0$  for all  $\{v, v'\} \in 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  $\mathbb{P}_{(v_j, v_{j+1})}^n$ , we update  $f_{t,\theta_{v_{j+1}, v_j}}$  which will be used to sample from  $\mathbb{P}_{(v_{j+1}, v_j)}^{n+1}$  in the next iterations. In order to sample from the dynamics  $\mathbb{P}_{(v_j, v_{j+1})}^n$ , 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  $\pi^0$  with  $r = 3$  and  $\varphi_r = (\mathrm{d}\mu_3 / \mathrm{dLeb})$ . During the first iteration of TreeDSB,  $\mathsf{T}$  is rooted at vertex 3 and we compute samples from the forward path  $\mathsf{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  $\mathsf{T}$  at vertex 1 and consider the forward path  $\mathsf{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\in \mathbb{N}$ , we sample from  $\pi^n$  by first sampling from  $\mu_{k_n}$  at leaf  $i_{k_n}$  and then following the parameterized SDEs on the directed edges of  $\mathsf{T}_{k_n}$ .
Figure 2: Illustration of one mIPF cycle solved by TreeDSB for a toy star-shaped tree. At each iteration, our method learns the backward stochastic process (dotted arrows) that goes from the target leaf (green-circled), corresponding to the constrained marginal, to the current root of the tree (red-circled) by using samples from the forward stochastic process (solid arrows).
4 Theoretical properties of mIPF
In this section, we study some of the theoretical properties of the static iterates  $(\pi^n)_{n\in \mathbb{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  $\mathbb{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
$$
\pi^ {\star} = \operatorname {a r g m i n} \left\{\mathrm {K L} \left(\pi \mid \pi^ {0}\right): \pi \in \mathscr {P} ^ {(\ell + 1)}, \pi_ {i} = \mu_ {i}, \forall i \in S \right\}, \quad (\text {s t a t i c - m S B})
$$
where  $\mathsf{S} \subset \{0, \dots, \ell\}$ ,  $\pi^0 \in \mathcal{P}$ ,  $\{\mu_i\}_{i \in \mathbb{S}} \in \mathcal{P}^{|S|}$ . We consider the following set of assumptions.
A1. There exists a family of measures  $\{\nu_i\}_{i\in \{0,\dots ,\ell \}}$  defined on  $(\mathbb{R}^d,\mathcal{B}(\mathbb{R}^d))$  such that  $\pi^0\ll \bigotimes_{i = 0}^{\ell}\nu_{i}$  with density  $h = \mathrm{d}\pi^{0} / (\mathrm{d}\bigotimes_{i = 0}^{\ell}\nu_{i})$  and  $\mu_{i}\ll \nu_{i}$  with density  $r_i = \mathrm{d}\mu_i / \mathrm{d}\nu_i$  for any  $i\in \mathsf{S}$ .
A2.  $\{\pi \in \mathcal{P}^{(\ell +1)}:\mathrm{KL}(\pi \mid \pi^0) <   \infty ,\pi_i = \mu_i,\forall i\in \mathsf{S}\} \neq \emptyset .$
A3. There exists a family of probability measures  $\{\tilde{\mu}_j\}_{j\in \{0,\dots ,\ell \} \setminus S}$  such that  $\pi^0\sim \tilde{\pi}^0$  , where  $\tilde{\pi}^{0} = \bigotimes_{i\in S}\mu_{i}\bigotimes_{j\in \{0,\ldots ,\ell \} \setminus S}\tilde{\mu}_{j}$
In particular, (static-mSB) recovers (TreeSB) by considering  $\nu_{i} = \mathrm{Leb}$  for any  $i\in \{0,\dots ,\ell \}$  and  $h(x_{0:\ell}) = \varphi_r(x_r)\exp [-c(x_{0:\ell}) / \varepsilon ]$  in A1. We detail in Appendix D how A2 and A3 can be met in (TreeSB). Under these assumptions, the multi-marginal Schrodinger Bridge exists.
Proposition 3. Assume A1 and A2. Then, there exists a unique solution  $\pi^{\star}$  to (static-mSB). In addition, assume A3. Then, there exists a family  $\{\psi_i^\star\}_{i\in \mathbb{S}}$  of measurable functions  $\psi_{i}^{\star}:\mathbb{R}^{d}\to \mathbb{R}$  such that
$$
\left(\mathrm {d} \pi^ {\star} / \mathrm {d} \pi^ {0}\right) = \exp \left[ \bigoplus_ {i \in \mathsf {S}} \psi_ {i} ^ {\star} \right] \quad \pi^ {0} \text {- a . s .}
$$
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 A1, A2 and A3. Let  $(\pi^n)_{n\in \mathbb{N}}$  be the sequence given by (mIPF). Then, for any  $n\in \mathbb{N}^*$  with  $k_{n} = (n - 1)\bmod (K)$  and  $q_{n}\in \mathbb{N}$  such that  $n = q_nK + k_n + 1$ , there exists a family of measurable functions  $\{\psi_{i_0}^{q_n + 1},\ldots ,\psi_{i_{k_n}}^{q_n + 1},\psi_{i_{k_n + 1}}^{q_n},\ldots ,\psi_{i_{K - 1}}^{q_n}\}$  such that
$$
\left(\mathrm {d} \pi^ {n} / \mathrm {d} \pi^ {0}\right) \left(x _ {0: \ell}\right) = \exp \left[ \bigoplus_ {j = 0} ^ {k _ {n}} \psi_ {i _ {j}} ^ {q _ {n} + 1} \left(x _ {i _ {j}}\right) \bigoplus_ {j = k _ {n} + 1} ^ {K - 1} \psi_ {i _ {j}} ^ {q _ {n}} \left(x _ {i _ {j}}\right) \right] \quad \pi^ {0} - a. s.
$$
In the tree-based setting, Corollary 4 explains why the (mIPF) iterations preserve the tree-based Markovian nature of  $\pi^0$ . We now prove that the marginal  $\pi_i^n$  converges to  $\mu_{i}$  for any  $i\in S$ , as  $n$  goes to infinity, i.e., we have marginal convergence on the leaves of  $\mathsf{T}$ .
Proposition 5. Assume A1 and A2. Let  $(\pi^n)_{n\in \mathbb{N}}$  be the sequence given by (mIPF). Then, we have  $\lim_{n\to \infty}\| \pi_i^n -\mu_i\|_{\mathrm{TV}} = 0$  for any  $i\in S$ .
The previous result does not ensure the convergence of  $(\pi^n)_{n\in \mathbb{N}}$  to the solution to (static-mSB). In particular, Proposition 5 does not provide the convergence of the marginals on the nodes  $v\in \mathsf{V}\backslash \mathsf{S}$ , which is key to compute regularized Wasserstein barycenters with TreeDSB. Relying on additional assumptions, we now derive the convergence of (mIPF).
A4.  $\bigoplus_{i\in \mathbb{S}}\mathrm{L}^1 (\mu_i)\subset \mathrm{L}^1 (\pi^\star)$  is closed.
A5. There exist  $\bar{c}\in (0,\infty)$  such that  $\exp (\psi_{i_k}^n -\psi_{i_k}^{n + 1})\leq \bar{c},$  for any  $n\in \mathbb{N}$ , any  $k\in \{0,\dots ,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 A1, A2, A3, A4 and A5. Let  $(\pi^n)_{n\in \mathbb{N}}$  be the sequence given by (mIPF). Then, we have  $\lim_{n\to \infty}\| \pi^n -\pi^\star \|_{\mathrm{TV}} = 0$ , where  $\pi^{\star}$  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  $\mathsf{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,\dots ,\ell \}$  (see Figure 2 for an illustration with  $\ell = 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  $\varepsilon$ -entropic regularization between  $\mu$  and  $\nu$  (Peyre et al., 2019, Chapter 4)
$$
W _ {2, \varepsilon} ^ {2} (\mu , \nu) = \inf  \left\{\int \| x _ {1} - x _ {0} \| ^ {2} \mathrm {d} \pi \left(x _ {0}, x _ {1}\right) - \varepsilon \mathrm {H} (\pi): \pi \in \mathscr {P} ^ {(2)}, \pi_ {0} = \mu , \pi_ {1} = \nu \right\}. \tag {3}
$$
In this work, we consider the  $(\ell \varepsilon, (\ell - 1)\varepsilon)$ -doubly-regularized Wasserstein-2 barycenter problem (Chizat, 2023) defined as follows
$$
\mu_ {\varepsilon} ^ {\star} = \arg \min  \left\{\sum_ {i = 1} ^ {\ell} w _ {i} W _ {2, \varepsilon / w _ {i}} ^ {2} (\mu , \mu_ {i}) + (\ell - 1) \varepsilon \mathrm {H} (\mu): \mu \in \mathscr {P} \right\}, \quad \text {(r e g W B)}
$$
where  $(w_{i})_{i\in \{1,\dots ,\ell \}}\in (0, + \infty)^{\ell}$ . 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  $\mu_{\varepsilon}^{\star}$  of (regWB).
Proposition 7. Let  $\varepsilon >0$ . Assume A0. Also assume that  $\mathsf{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  $\varphi_r = \mathrm{d}\mu_{i_{K - 1}} / \mathrm{dLeb} > 0$ . Under A2, (regWB) has a unique solution  $\pi_0^\star$ , where  $\pi^{\star}$  solves (TreeSB).
The proof of this result is postponed to Appendix D. More generally, we show in Appendix D that, for any tree  $\mathsf{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  $\mathsf{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  $\varepsilon$  reasonably large in TreeDSB, which is a stability advantage compared to other regularized methods which do not consider this further regularization.
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  $\mu_{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 Alvarez-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.
In our experiments<sup>3</sup>, 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
$^{3}$ Code available at https://github.com/maxencenoble/tree-diffusion-schrodinger-bridge.
sequential training procedure over  $2K$  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\varepsilon /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~\mathrm{mIPF}$  cycles with  $\varepsilon = 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.
Figure 3: Estimated densities on the leaves.
Figure 4: From left to right: barycenter estimated from the leaves Swiss-roll, Circle and Moons.
Synthetic Gaussian datasets. Next, we consider three independent Gaussian distributions with zero mean and random non-diagonal covariance matrices whose conditional number is less than 10, following Fan et al. (2020). In this case, the non-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  $\mu^{\star} \in \mathcal{P}$  and some approximation  $\mu \in \mathcal{P}$ , we define
$$
\mathrm{BW}_{2}^{2}\text{-UVP} (\mu ,\mu^{\star}) = 100\cdot 2\mathrm{BW}_{2}^{2}(\mu ,\mu^{\star}) / \operatorname {Var}(\mu^{\star})\% ,
$$
where  $\mathrm{BW}_2^2 (\mu ,\mu^\star) = W_2^2 (\mathrm{N}(\mathbb{E}[\mu ],\mathrm{Cov}(\mu)),\mathrm{N}(\mathbb{E}[\mu^\star ],\mathrm{Cov}(\mu^\star))$
Method	d=2	d=16	d=64	d=128	d=256
fsWB (Cuturei & Doucet, 2014)	0.06±0.01	2.86±0.06	11.12±0.06	14.47±0.07	17.41±0.05
crWB (Li et al., 2020)	0.02±0.01	1.52±0.11	11.41±0.73	5.75±0.02	18.27±0.54
Tree DSB	0.63±0.26	1.07±0.58	1.39±0.07	1.92±0.02	2.62±0.07
Table 1: Gaussian setting: comparison with the regularized methods crWB and fsWB.
In this setting, we choose  $\mu^{\star}$  to be the non-regularized barycenter and assess the dependency w.r.t. the dimension of the algorithms using the  $\mathrm{BW}_2^2$ -UVP metric. In Table 1, we compare ourselves with the two regularized methods Li et al. (2020) ( $\mathrm{L}_2$ -reg. equal to  $10^{-4}$ ) and Cuturi & Doucet (2014). We run TreeDSB for  $10\mathrm{mIPF}$  cycles with  $\varepsilon = 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\mathrm{mIPF}$  cycles with  $r$  that corresponds to the digit 6 and  $\varepsilon = 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.
Figure 5: Samples from the estimated MNIST 2-4-6 marginals and from their Wasserstein barycenter.
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.
Method	Without het.	With het.
fsWB (Cuturi & Doucet, 2014)	12.95±0.35	14.43±0.51
crWB (Li et al., 2020)	20.66±0.71	23.06±0.12
Tree DSB	8.69±0.12	8.90±0.68
Table 2: Bayesian fusion setting: comparison with the regularized methods crWB and fsWB.