SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $p_{v,v^{\prime},t}$  is the density of  $\mathbb{P}_{(v,v^{\prime}),t}^{n}$  w.r.t. the Lebesgue measure.
Following the construction of our dynamic iterates, we now explain how the sequence  $\{\pi_{(v,v')}^n\}_{n\in \mathbb{N}^*, (v,v')\in \mathsf{E}_{k_n}}$  is recursively defined. Let  $n \in \mathbb{N}$ ,  $k_n = (n - 1) \bmod(K)$ . Define the path  $\mathsf{P}_n = \mathrm{path}_{\mathsf{T}_{i_{k_n}}} (i_{k_n}, i_{k_n + 1})$ . Then, for any  $(v, v') \in \mathsf{E}_{k_n + 1}$ ,
(a) if  $(v,v^{\prime})\in \mathsf{E}_{k_n}\backslash \mathsf{P}_n$  , then  $\pi_{(v,v^{\prime})}^{n + 1,N} = \pi_v^{n + 1}\bigotimes_{m = 0}^{N - 1}\pi_{(v,v^{\prime}),m + 1|m}^{n,N},$
(b) if  $(v', v) \in \mathsf{P}_n$ , then  $\pi_{(v, v')}^{n+1, N} = \pi_v^{n+1} \bigotimes_{m=0}^{N-1} \pi_{(v', v), N-m-1|N-m}^{n, N}$ .
These computations may be obtained by considering the sequence given by (mIPF) to solve the multi-marginal Tree-SB problem associated to  $\mathsf{T}^{(N)} = (\mathsf{V}^{(N)},\mathsf{E}^{(N)})$ , the  $N$ -discretized version of  $\mathsf{T}$  (see Appendix B) with weights  $w_{e_m}^{(N)} = 2\gamma_m / \varepsilon$ , which is given by
$$
\pi^ {\star} = \operatorname {a r g m i n} \left\{\mathrm {K L} \left(\pi \mid \pi^ {0, N}\right): \pi \in \mathcal {P} ^ {\left(\mathrm {V} ^ {(N)}\right)}, \pi_ {i} = \mu_ {i}, \forall i \in \mathrm {S} \right\},
$$
with  $\pi^{0,N} = \pi_r^0\bigotimes_{(v,v')\in \mathsf{E}_r}\pi_{(v,v'),1:N|0}^{0,N}$ .
To approximate the IPF recursion given by (a) and (b), we use on each edge of  $\mathsf{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 \in \mathbb{N}$ , any  $(v, v') \in \mathsf{E}_{k_n}$  with  $k_n = (n - 1) \bmod(K)$ , we have
$$
\pi_ {(v, v ^ {\prime}), m + 1 | m} ^ {n, N} (\cdot | x _ {m}) = \mathrm {N} \big (F _ {m, v, v ^ {\prime}} ^ {n} (x _ {m}), \gamma_ {m} \mathrm {I} _ {d} \big).
$$
Let  $n \in \mathbb{N}$ . Consider the path  $\mathsf{P}_n = \mathrm{path}_{\mathsf{T}_{i_{k_n}}}(i_{k_n}, i_{k_n + 1})$ . Let  $(v, v') \in \mathsf{E}_{k_n + 1}$ . Define  $p^n = \pi_{(v, v')}^{n,N}$  and  $m_N = N - m - 1$ . Then, if  $(v', v) \in \mathsf{P}_n$ , we have
$$
F _ {m, v, v ^ {\prime}} ^ {n + 1} = \operatorname {a r g m i n} _ {\mathrm {F} \in \mathrm {L} ^ {2} \left(\mathbb {R} ^ {d}, \mathbb {R} ^ {d}\right)} \tag {31}
$$
$$
\mathbb {E} _ {p _ {m _ {N}, m _ {N} + 1} ^ {n}} [ \| \mathrm {F} (X _ {m _ {N} + 1}) - (X _ {m _ {N} + 1} + F _ {m _ {N}, v ^ {\prime}, v} ^ {n} (X _ {m _ {N}}) - F _ {m _ {N}, v ^ {\prime}, v} ^ {n} (X _ {m _ {N} + 1})) \| ^ {2} ],
$$
otherwise, we have  $F_{m,v,v'}^{n+1} = F_{m,v,v'}^n$ .
In practice, we use two neural networks per edge  $\{v,v^{\prime}\} \in \mathsf{E}$ , one for each possible direction of the edge, such that  $F_{v,v^{\prime}}(\theta_{v,v^{\prime}}^{n},m,x)\approx F_{m,v,v^{\prime}}^{n}(x)$  and  $F_{v^{\prime},v}(\theta_{v^{\prime},v}^{n},m,x)\approx F_{m,v^{\prime},v}^{n}(x)$ . For any  $\{v,v^{\prime}\} \in \mathsf{E}$ , the parameter  $\theta_{v,v^{\prime}}^{n}$  is updated at iteration  $n$  via the score matching loss defined by (31) in Proposition 26 if  $(v,v^{\prime})\in \mathrm{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 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,\dots ,i_{K - 1}\}$  is randomly shuffled at each new mIPF cycle.
<sup>6</sup>https://github.com/JTT94/diffusion_schrodinger_bridge
G.1 General experimental setup
Implementation of Algorithm 1 in practice. Let  $n \in \mathbb{N}$ , with  $k_{n} = (n - 1) \bmod(K)$ ,  $k_{n} + 1 = n \bmod(K)$ . Consider the path  $\mathsf{P}_n = \mathrm{path}_{\mathsf{T}_{i_{k_n}}} (i_{k_n}, i_{k_n + 1})$ . Assume that we are provided with a dataset  $\mathsf{D}_{i_{k_n}}$ , which contains  $M$  samples from  $\pi_{i_{k_n}}^n$ . Following Lines 7-9 in Algorithm 1, we apply processes (a) and (b) recursively on the edges  $(v, v') \in \mathsf{P}_n$ .
(a) Sampling step (Line 7). For any  $x_0 \in \mathsf{D}_v$ , we sample from the diffusion trajectory (30) given by the Euler Maruyama discretization of  $\mathbb{P}_{v,v'}^n$  starting from  $x_0$ . This gives us  $M \times N$  trajectory samples. We then store the last iterate of each trajectory in a new dataset  $\mathsf{D}_{v'}$ , which thus approximates  $\pi_{v'}^n$ .
(b) Training step (Lines 8-9). In order to avoid heavy computation, we approximate the mean-matching 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  $\theta_{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'\} \in \mathsf{E}$ . We now give details on the design of the time schedule  $\{\gamma_k\}_{k=1}^N$  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  $\gamma_k = \gamma_0 + (2k/N)(\bar{\gamma} - \gamma_0)$  for any  $k \in \{0, \dots, N/2\}$  (the rest of the sequence being obtained by symmetry) where  $\gamma_0$  is a free parameter and  $\bar{\gamma}$  is determined by  $\sum_{k=1}^{N} \gamma_k = T_{v,v'}$ . In our experiments, we set  $N = 50$  and  $\gamma_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, MLP-Block (1a) has shape  $(d, 128, \max(256, 2d))$ , MLPBlock (1b) has shape  $(32, 128, \max(256, 2d))$ , and MLPBlock (2) has shape  $(2 \times \max(256, 2d), \max(256, 2d), \max(128, d), d)$ , where  $d$  denotes the dimension of input data. We optimize the networks with ADAM (Kingma & Ba, 2014) with learning rate  $10^{-4}$  and momentum 0.9. For each of the networks, we set the batch size to 4,096 and the number of iterations to 10,000 for the synthetic datasets and 15,000 for the subset posterior aggregation task. Our experiments ran on 1 Intel Xeon CPU Gold 6230 20 cores @ 2.1 Ghz CPU.
Figure 6: Architecture of the "Basic Model".
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  $\varepsilon$  to its lowest value such that the algorithm is stable. Finally, for sake of fairness with our method, we initialise the barycenter measure with  $\pi_r^0$  when solving the problem ( $\mu_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, 2d)$ ,  $\max(128, 2d)$ ,  $\max(128, 2d)$ ). The activation functions are ReLu. We optimize the networks with ADAM (Kingma & Ba, 2014) with learning rate  $10^{-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,\varepsilon)\in \{(2,0.1),(16,0.2),(64,0.5),(128,1.0),(256,2.0)\}$ . We run TreeDSB for  $10\mathrm{mIPF}$  cycles with regularization parameter  $\varepsilon = 0.1$ , starting from the central node initialized to a Gaussian distribution  $\mu_0$  chosen as detailed in Appendix F with  $\alpha = 1$ . Thus, we solve the regularized Wasserstein barycenter problem ( $\mu_0$ -regWB), which contains an additional regularization with respect to  $\mu_0$ . This choice is justified, since the non-regularized barycenter is known to be a Gaussian distribution, and  $\mu_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\mathrm{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 \cdot 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 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  $\varepsilon = 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\mathrm{mIPF}$  cycles with regularization parameter  $\varepsilon = 0.1$ . We start from the central node with a Gaussian distribution  $\mu_0$  chosen as detailed in Appendix F with  $\alpha = 1$ , and thus solve the barycenter formulation ( $\mu_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.
<sup>7</sup>https://archive.ics.uci.edu/ml/datasets/wine
Figure 7: Evaluation of the sampling process for wine-homogeneous (left) and wine-heterogeneous (right). We display the Autocorrelation function on 500 lags (above) and the evolution over the iterations of ULA of the negative log-likelihood (NLL) evaluated on each training dataset (below). In particular, the samples are decorrelated and the NLL has a satisfying profile.
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  $\varepsilon$ .
To do so, we consider three different choices of regularization in TreeDSB: (i)  $\varepsilon = 0.2$  (50 mIPF cycles), see Figure 8, (ii)  $\varepsilon = 0.1$  (50 mIPF cycles), see Figure 9 and (iii)  $\varepsilon = 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  $\varepsilon$ , 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  $\varepsilon$  decreases.
Figure 8: Estimated 2D barycenter obtained by TreeDSB with  $\varepsilon = 0.2$  (50 mIPF cycles). First row: starting from Swiss-roll. Second row: starting from Circle. Third row: starting from Moons.
Figure 9: Estimated 2D barycenter obtained by TreeDSB with  $\varepsilon = 0.1$  (50 mIPF cycles). First row: starting from Swiss-roll. Second row: starting from Circle. Third row: starting from Moons.
Figure 10: Estimated 2D barycenter obtained by TreeDSB with  $\varepsilon = 0.05$  (60 mIPF cycles). First row: starting from Swiss-roll. Second row: starting from Circle. Third row: starting from Moons.
For purpose of illustration, we provide in Figure 11 the barycenter obtained by state-of-the-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  $\varepsilon = 0.5$  (Cuturi & Doucet, 2014) and (iii) entropic-regularized convolutional Wasserstein barycenter with  $\varepsilon = 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.
Figure 11: Estimated 2D barycenter obtained by in-sample algorithms. From left to right: Cuturi & Doucet (2014) (non-regularized), Cuturi & Doucet (2014) (regularized), Solomon et al. (2015).
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\mathrm{mIPF}$  cycles with  $\varepsilon = 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  $\varepsilon$  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).
Figure 12: Tree DSB results for MNIST digits 0 and 1, after  $15\mathrm{mIPF}$  cycles with  $\varepsilon = 0.5$
Figure 13: From left to right: Fan et al. (2020), Korotin et al. (2021) and Li et al. (2020).
Digits 2,4 and 6. In Figure 14, we report the results obtained by running TreeDSB on MNIST digits 2,4 and 6, for  $10\mathrm{mIPF}$  cycles with  $\varepsilon = 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\mathrm{mIPF}$  cycles with  $\varepsilon = 0.2$ , starting from MNIST-6. We display 30 samples from the reconstructed margins (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\mathrm{mIPF}$  cycles with  $\varepsilon = 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\mathrm{mIPF}$  cycles with  $\varepsilon = 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  $\varepsilon$  decreases. Note that the reconstructed MNIST digits are less truthful to the original datasets when  $\varepsilon$  is low.
Figure 14: Tree DSB results for MNIST digits 2,4 and 6, after  $10\mathrm{mIPF}$  cycles with  $\varepsilon = 0.5$ . First row: starting from MNIST-2. Second row: starting from MNIST-4. Third row: starting from MNIST-6.
Figure 15: Tree DSB results for MNIST digits 2,4 and 6, after  $10\mathrm{mIPF}$  cycles with  $\varepsilon = 0.2$ , starting from MNIST-6. First row: samples from the reconstructed marginals. Second row: samples from the estimated barycenter.
Figure 16: Tree DSB results for MNIST digits 0,1 and 4, after  $10\mathrm{mIPF}$  cycles with  $\varepsilon = 0.5$ . First row: samples from the reconstructed marginals, starting from MNIST-1. Second row: samples from the estimated barycenter, starting from MNIST-0. Third row: samples from the estimated barycenter, starting from MNIST-1.
Figure 17: Tree DSB results for MNIST digits 0,1 and 4, after  $10\mathrm{mIPF}$  cycles with  $\varepsilon = 0.2$ . First row: samples from the reconstructed marginals, starting from MNIST-0. Second row: samples from the estimated barycenter, starting from MNIST-0. Third row: samples from the estimated barycenter, starting from MNIST-1.