SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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  $\mu_{1}, \mu_{2}, \mu_{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\mathrm{mIPF}$  cycles  $\varepsilon = 0.1$  and we compare ourselves with Li et al. (2020) ( $\mathrm{L}_{2}$ -reg. equal to  $10^{-4}$ ) and Cuturi & Doucet (2014). We evaluate the methods using the  $\mathrm{BW}_2^2$ -UVP metric, where  $\mu^{\star}$  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  $\varepsilon \rightarrow 0$ . In our setting, we observe that if  $\varepsilon$ , 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.
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).
<sup>4</sup>https://archive.ics.uci.edu/ml/datasets/wine
We thank James Thornton for the DSB codebase<sup>5</sup> 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).
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 Peyre, 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 forward-backward 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.
<sup>5</sup>https://github.com/JTT94/diffusion_schrodinger_bridge
Schmitz, M. A., Heitz, M., Bonneel, N., Ngole, F., Coeurjolly, D., Cuturi, M., Peyre, 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.
Schrodinger, E. Sur la théorie relativiste de l'électron et l'interprétable 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., Peyre, 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, L., 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  $\mathsf{E}$ , we equivalently refer to the cardinal of  $\mathsf{E}$  as  $\operatorname{card}(\mathsf{E})$  or  $|\mathsf{E}|$ . Let  $(\mathsf{X},\mathcal{X})$  be a measurable space. For any  $x \in (\mathbb{R}^d)^{\ell +1}$  and any  $m \in \{0,\dots ,\ell \}$ , let  $x_{-m} = (x_0,\ldots ,x_{m - 1},x_{m + 1},\ldots ,x_\ell)$ . For any family of measures  $\{\nu_j\}_{j\in \{0,\dots ,\ell \}}$  defined on  $(\mathsf{X},\mathcal{X})$  and any  $i \in \{0,\dots ,\ell \}$ , let  $\nu_{-i} = \bigotimes_{j\in \{0,\dots ,\ell \} \setminus \{i\}}\nu_j$ . Let  $I = \{i_1,\dots ,i_q\} \subset \{1,\dots ,\ell \}$  and  $\mu \in \mathcal{P}^{(\ell)}$  such that  $\mu \ll \mathrm{Leb}$ . We define  $I^c = \{1,\dots ,\ell \} \backslash I$  and denote it by  $\{i_1^c,\dots ,i_{\bar{q}}^c\}$  where  $\bar{q} = \ell -q$ . We denote the marginal of  $\mu$  along  $I$  by  $\mu_I$ , i.e.,  $\mu_I \in \mathcal{P}^{(q)}$  and we have for any  $A \in B((\mathbb{R}^d)^q)$ ,  $\mu_I(\mathsf{A}) = \int_{\mathsf{X}}\mu (x)\prod_{j = 1}^{q}\delta_{x_{i_j}}(\mathsf{A}_j)\mathrm{d}x$ . In addition, note that  $\mu_I \ll \mathrm{Leb}$ . We denote the conditional distribution of  $\mu$  given  $I$  by  $\mu_{|I}(\cdot |\cdot)$ , i.e.,  $\mu_{|I}(\cdot |\cdot) \in \mathcal{P}^{(\bar{q})}\times (\mathbb{R}^d)^q$  and we have for any  $y\in (\mathbb{R}^d)^q$  and any  $A\in B((\mathbb{R}^d)^{\bar{q}})$ ,  $\mu_{|I}(\mathsf{A}|y) = \int_{\mathsf{X}}\mu (x) / \mu_I(y)\prod_{j = 1}^{q}\delta (x_{i_j} - y_j)\prod_{j' = 1}^{\bar{q}}\delta_{x_{i_{j'}}'}(A_{j'})\mathrm{d}x.$  Remark that for any  $y\in (\mathbb{R}^d)^q, \mu_{|I}(\cdot |y)\ll \mathrm{Leb}$ . For any subset  $J\subset I^c$  with card  $(\mathsf{J}) = q_{\mathsf{J}}$ , we also define  $\mu_{|\mathsf{J}|}(\cdot |\cdot)\in \mathcal{P}^{(q_1)}\times (\mathbb{R}^d)^q$  such that for any  $y\in (\mathbb{R}^d)^q, \mu_{|\mathsf{J}|}(\cdot |y) = \{\mu_{|\mathsf{J}|}(\cdot |y)\}_{\mathsf{J}}.$  For a collection of functions  $\{f_i\}_{i\in I}$ , with  $|C| <   |1,\dots ,n|$  and  $n\in \mathbb{N}$  such that  $f_{i}:  {\mathbb{R}}^{d}\to {\mathbb{R}}.$  we define  $\oplus_{i\in I}f_i:(\mathbb{R}^d)^n\to {\mathbb{R}}$  such that for any  $x = (x_1,\dots ,x_n)\in (\mathbb{R}^d)^n,  \oplus_{i\in I}f(x) =  \sum_{i\in I}f_i(x_i).$
B Introduction to trees
Undirected tree. An undirected graph  $\mathsf{T} = (\mathsf{V},\mathsf{E})$ , with vertices  $\mathsf{V}$  and edges  $\mathsf{E}$ , is said to be an undirected tree if it is acyclic and connected (Valiente, 2002, Definition 1.19.). In particular, we have  $\operatorname{card}(\mathsf{E}) = \operatorname{card}(\mathsf{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  $\mathsf{T}' = (\mathsf{V}',\mathsf{E}')$  is a sub-tree of  $\mathsf{T}$  if  $\mathsf{T}'$  is an undirected tree with vertices  $\mathsf{V}' \subset \mathsf{V}$  and edges  $\mathsf{E}' \subset \mathsf{E}$ . For any vertex  $v \in \mathsf{V}$ , we define the set of its neighbours  $\mathsf{N}_v$  as the set of vertices  $v' \in \mathsf{V}$  such that  $\{v,v'\} \in \mathsf{E}$ . The integer  $\operatorname{card}(\mathsf{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  $\mathsf{V}_{\mathsf{L}} \subset \mathsf{V}$ . The (unique) path in  $\mathsf{T}$  between two vertices  $v$  and  $v'$  is the sequence of two-by-two distinct edges  $\{\{v_i,v_{i+1}\}\}_{i=1}^n$  (with  $n \geq 1$ ) such that  $v_k = v_{k+1}$  for any  $k \in \{1,\dots,n\}$  such that  $k = 0 \bmod(2)$ ,  $v_1 = v$  and  $v_{n+1} = v'$ . This path can be seen as a linear sub-tree of  $\mathsf{T}$ , and we define  $n$  as the length of this path. We say that  $\mathsf{T}$  is weighted if there exists a map  $w: \mathsf{E} \mapsto \mathbb{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  $\mathsf{T}$  is said to be rooted in  $r \in \mathsf{V}$  if  $r$  defines a partial ordering  $\leq_{\mathsf{T},r} \subset \mathsf{V} \times \mathsf{V}$  such that for any  $v_1,v_2 \in \mathsf{V}, v_1 \leq_{\mathsf{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  $\mathsf{T}_r = (\mathsf{V},\mathsf{E}_r)$  rooted in  $r\in \mathsf{V}$ . Any directed edge  $e\in \mathsf{E}_r$  from  $v_{1}\in \mathsf{V}$  to  $v_{2}\in \mathsf{V}$  is denoted by  $(v_{1},v_{2})$ .  $\mathsf{T}_r$  is said to be a directed tree rooted in  $r$  if (i) the underlying undirected graph  $\mathsf{T} = (\mathsf{V},\mathsf{E})$  is an undirected tree rooted in  $r$  and (ii) any  $(v_{1},v_{2})\in \mathsf{E}_{r}$  is directed according to the partial ordering  $\leq_{\mathsf{T},r}$ , i.e.,  $\{v_1,v_2\} \in \mathsf{E}$  and  $v_{1}\leq_{\mathsf{T},r}v_{2}$ . For any vertices  $(v,v^{\prime})\in \mathsf{V}\times \mathsf{V}$  such that  $v\leq_{\mathsf{T},r}v^{\prime}$ , the (unique) path in  $\mathsf{T}_r$  from  $v$  to  $v^{\prime}$ , denoted by  $\mathrm{path}_{\mathsf{T}_r}(v,v')$ , is defined as the directed version of the path in  $\mathsf{T}$  between  $v$  and  $v^{\prime}$  (viewed as a sub-tree of  $\mathsf{T}$ ), which is rooted in  $v$ . We say that  $\mathsf{T}_r$  is weighted, if  $\mathsf{T}$  is weighted and the edges of  $\mathsf{T}_r$  have the same weights as the corresponding undirected edges of  $\mathsf{T}$ . For any  $(v_{1},v_{2})\in \mathsf{E}_{r}$ , we denote this weight by  $w_{v_1,v_2}$ . We say that  $\mathsf{T}_r$  is the (unique) directed version of  $\mathsf{T}$  rooted in  $r$ . It is endowed with a canonical vertex numbering  $\zeta :\mathsf{V}\to \{0,\dots ,\mathrm{card}(\mathsf{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  $\mathsf{T}$ , i.e., if  $v_{1}\leq_{\mathsf{T},r}v_{2}$ ,  $\zeta (v_{1})\leq \zeta (v_{2})$ , and satisfies  $\zeta (r) = 0$ . In the rest of the paper, we will write in an equivalent manner  $v$  or  $\zeta (v)$ .
For any vertices  $(v_{1}, v_{2}) \in \mathsf{E} \times \mathsf{E}$  such that  $v_{1} \leq_{\mathsf{T}, r} v_{2}$ ,  $\mathrm{path}_{\mathsf{T}_r}(v_1, v_2)$  corresponds to the ordered set of edges in  $\mathsf{E}_r$  which define the ordered path between two vertices  $v_{1}$  and  $v_{2}$ . For any vertex  $v \in \mathsf{V}$ , we define:
(a) the set of its children  $\mathsf{C}_v$  as the set of vertices  $v' \in \mathsf{V}$  such that  $(v, v') \in \mathsf{E}_r$ . In particular, for any  $v \in \mathsf{V}_L$ , the set of leaves, one has  $\mathsf{C}_v = \emptyset$ .
(b) its parent as the unique vertex  $p(v)$  such that  $(p(v), v) \in \mathsf{E}_r$ , if  $v \neq r$  (the parent of the root is not defined).
Note that  $\mathsf{N}_r = \mathsf{C}_r$  and, for any vertex  $v\in \mathsf{V}\backslash \{r\}$ $\mathsf{N}_v = \{p(v)\} \cup \mathsf{C}_v$
Definition 8 (Tree-structured directed Probabilistic Graphical Model (PGM)). Consider a directed tree  $\mathsf{T}_r = (\mathsf{V},\mathsf{E}_r)$ . The directed PGM induced by  $\mathsf{T}_r$  (Koller & Friedman, 2009, Definition 3.4.), denoted by  $\mathcal{P}_{\mathsf{T}_r}$ , is the family of distributions  $\pi \in \mathcal{P}^{(|\mathsf{V}|)}$  which have a Markovian factorization along  $\mathsf{T}_r$ , i.e.,
$$
\mathcal {P} _ {\mathrm {T} _ {r}} = \left\{\pi \in \mathcal {P} ^ {(| \mathrm {V} |)}: \pi = \pi_ {r} \bigotimes_ {(v, v ^ {\prime}) \in \mathrm {E} _ {r}} \pi_ {v ^ {\prime} | v} \right\}.
$$
Lemma 9. Consider an undirected tree  $\mathsf{T} = (\mathsf{V},\mathsf{E})$ . Let  $(r,r^{\prime})\in \mathsf{V}\times \mathsf{V}$ . Let  $\mathsf{T}'$  be a sub-tree of  $\mathsf{T}$  with vertices  $\mathsf{V}'$  such that  $r^{\prime}\in \mathsf{V}^{\prime}$ . Denote by  $\mathsf{T}_{r'}'$  the directed version of  $\mathsf{T}'$  rooted in  $r^\prime$ . Then, for any  $\pi \in \mathcal{P}_{\mathsf{T}_r}$ , we have  $\pi_{\mathsf{V}'}\in \mathcal{P}_{\mathsf{T}_{r'}}$ .
Proof. Let  $(r,r^{\prime})\in \mathsf{V}\times \mathsf{V}$ . We denote by  $\mathsf{T}_r = (\mathsf{V},\mathsf{E}_r)$ , respectively  $\mathsf{T}_{r'} = (\mathsf{V},\mathsf{E}_{r'})$ , the directed version of  $\mathsf{T}$  rooted in  $r$ , respectively  $r^\prime$ . We define the paths  $\mathsf{P}_{r,r^{\prime}} = \mathrm{path}_{\mathsf{T}_r}(r,r')\subset \mathsf{E}_r$  and  $\mathsf{P}_{r',r} = \mathrm{path}_{\mathsf{T}_{r'}}(r',r)\subset \mathsf{E}_{r'}$ . It is easy to see that
(a)  $\mathsf{E}_r\backslash \mathsf{P}_{r,r'} = \mathsf{E}_{r'}\backslash \mathsf{P}_{r',r}$
(b)  $\mathsf{P}_{r,r'} = \{(v_2,v_1):(v_1,v_2)\in \mathsf{P}_{r',r}\} ,$
(c)  $\mathsf{P}_{r',r} = \{(v_2,v_1):(v_1,v_2)\in \mathsf{P}_{r,r'}\}$
Let  $\pi \in \mathcal{P}_{\mathsf{T}_r}$ . First note that for any  $(v_{1}, v_{2}) \in \mathsf{E}_{r}$ , we have by Bayes decomposition  $\pi_{v_1} \pi_{v_2|v_1} = \pi_{v_2} \pi_{v_1|v_2} = \pi_{v_1, v_2}$ . Then it comes
$$
\begin{array}{l} \pi = \pi_ {r} \bigotimes_ {(v _ {1}, v _ {2}) \in \mathsf {E} _ {r}} \pi_ {v _ {2} | v _ {1}} \\ = \pi_ {r} \bigotimes_ {(v _ {1}, v _ {2}) \in \mathsf {P} _ {r, r ^ {\prime}}} \pi_ {v _ {2} | v _ {1}} \bigotimes_ {(v _ {1}, v _ {2}) \in \mathsf {E} _ {r} \backslash \mathsf {P} _ {r, r ^ {\prime}}} \pi_ {v _ {2} | v _ {1}} \\ = \pi_ {r} \bigotimes_ {\left(v _ {2}, v _ {1}\right) \in \mathsf {P} _ {r ^ {\prime}, r}} \pi_ {v _ {2} | v _ {1}} \bigotimes_ {\left(v _ {1}, v _ {2}\right) \in \mathsf {E} _ {r ^ {\prime}} \backslash \mathsf {P} _ {r ^ {\prime}, r}} \pi_ {v _ {2} | v _ {1}} \\ = \pi_ {r ^ {\prime}} \bigotimes_ {(v _ {1}, v _ {2}) \in \mathsf {P} _ {r ^ {\prime}, r}} \pi_ {v _ {2} | v _ {1}} \bigotimes_ {(v _ {1}, v _ {2}) \in \mathsf {E} _ {r ^ {\prime}} \backslash \mathsf {P} _ {r ^ {\prime}, r}} \pi_ {v _ {2} | v _ {1}} \\ = \pi_ {r ^ {\prime}} \bigotimes_ {(v _ {1}, v _ {2}) \in \mathsf {E} _ {r ^ {\prime}}} \pi_ {v _ {2} | v _ {1}}, \\ \end{array}
$$
and therefore, we have  $\pi \in \mathcal{P}_{\mathrm{T}_{r^{\prime}}}$
Let  $\mathsf{T}'$  be a sub-tree of  $\mathsf{T}$  with vertices  $\mathsf{V}'$  such that  $r' \in \mathsf{V}'$ . First note that  $\mathsf{E}_{r'}' \subset \mathsf{E}_{r'}$ . Using the previous computation, we have for any  $\mathsf{A} \in \mathcal{B}((\mathbb{R}^d)^{|V'|})$ ,
$$
\begin{array}{l} \pi_ {\mathsf {V} ^ {\prime}} (\mathsf {A}) = \int_ {(\mathbb {R} ^ {d}) ^ {| \mathsf {V} |}} \pi_ {r ^ {\prime}} (x _ {r ^ {\prime}}) \bigotimes_ {(v _ {1}, v _ {2}) \in \mathsf {E} _ {r ^ {\prime}}} \pi_ {v _ {2} | v _ {1}} (x _ {v _ {2}} | x _ {v _ {1}}) \prod_ {v ^ {\prime} \in \mathsf {V} ^ {\prime}} \delta_ {x _ {v ^ {\prime}}} (\mathsf {A} _ {v ^ {\prime}}) \mathrm {d} x \\ = \int_ {(\mathbb {R} ^ {d}) ^ {| \mathrm {V} | - | \mathrm {V} ^ {\prime} |} \left\{\pi_ {r ^ {\prime}} \left(\mathsf {A} _ {r ^ {\prime}}\right) \bigotimes_ {(v _ {1}, v _ {2}) \in \mathsf {E} _ {r ^ {\prime}} ^ {\prime}} \pi_ {v _ {2} | v _ {1}} \left(\mathsf {A} _ {v _ {2}} | x _ {v _ {1}}\right) \right\} \bigotimes_ {(v _ {1}, v _ {2}) \in \mathsf {E} _ {r ^ {\prime}} \backslash \mathsf {E} _ {r ^ {\prime}} ^ {\prime}} \pi_ {v _ {2} | v _ {1}} \left(x _ {v _ {2}} | x _ {v _ {1}}\right) \mathrm {d} x _ {\mathrm {V} \backslash \mathrm {V} ^ {\prime}} \\ = \left\{\pi_ {r ^ {\prime}} \bigotimes_ {(v _ {1}, v _ {2}) \in \mathsf {E} _ {r ^ {\prime}} ^ {\prime}} \pi_ {v _ {2} | v _ {1}} \right\} (\mathsf {A}), \\ \end{array}
$$
which proves that  $\pi_{\mathsf{V}'} \in \mathcal{P}_{\mathsf{T}_{r'}'}$ .
Discretized undirected tree. Let  $N \geq 1$ . Consider an undirected tree  $\mathsf{T} = (\mathsf{V}, \mathsf{E})$  with weights  $w$ . We say that  $\mathsf{T}^{(N)} = (\mathsf{V}^{(N)}, \mathsf{E}^{(N)})$  is a  $N$ -discretized version of  $\mathsf{T}$  if it is an undirected tree with weights  $w^{(N)}$  such that
(a)  $\mathsf{V}^{(N)} = \mathsf{V}\bigsqcup_{\substack{e\in \mathsf{E},\\ k\in \{1,\ldots ,N - 1\}}}\{v_e^k\} ,$
(b)  $\mathsf{E}^{(N)} = \cup_{e\in \mathsf{E}}\cup_{k = 0,\ldots ,N - 1}\left\{\{v_e^k,v_e^{k + 1}\} \right\}$  with the convention that the vertices  $v_{e}^{N}$  and  $v_{e}^{N}$  are defined such that  $\{v_e^0,v_e^N\} = e$
(c)  $\sum_{e\in \mathrm{path}_{\mathsf{T}}(v,v^{\prime})}1 / w_{e}^{(N)} = 1 / w_{v,v^{\prime}}$  , if  $\{v,v^{\prime}\} \in \mathsf{E}$
Remark that the leaves of  $\mathsf{T}^{(N)}$  are exactly the original leaves of  $\mathsf{T}$  and that  $\mathsf{T}^{(1)} = \mathsf{T}$ . The non-uniqueness of  $\mathsf{T}^{(N)}$  comes from the freedom of choice on the weights of its edges.
Discretized directed tree. Let  $N \geq 1$ . Consider a directed tree  $\mathsf{T}_r = (\mathsf{V}, \mathsf{E}_r)$  rooted in  $r \in \mathsf{V}$  with weights  $w$ . We say that  $\mathsf{T}_r^{(N)} = (\mathsf{V}^{(N)}, \mathsf{E}_r^{(N)})$  is a  $N$ -discretized version of  $\mathsf{T}_r$  if it is the directed version of  $\mathsf{T}^{(N)}$  rooted in  $r$ , where  $\mathsf{T}^{(N)}$  is a  $N$ -discretized version of the underlying undirected tree of  $\mathsf{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  $\mathrm{C}_0^\infty (\mathbb{R}^d)$  is the space of infinitely differentiable real-valued functions which vanish at infinity. In addition, we have that  $\mathcal{S}_{+}^{d}$  is the space of  $d\times d$ , symmetric, non-negative matrices.
Definition 10. Let  $T > 0$  or  $T = +\infty$ ,  $\sigma : [0, T) \times \mathbb{R}^d \to \mathcal{S}_+^d$  and  $b : [0, T) \times \mathbb{R}^d \to \mathbb{R}^d$ , locally bounded measurable functions. We define the infinitesimal generator,  $\mathcal{A}$ , given for any  $f \in C_0^\infty(\mathbb{R}^d)$ ,  $t \in [0, T)$  and  $x \in \mathbb{R}^d$  by
$$
\mathcal {A} _ {t} (f) (x) = \left\langle b _ {t} (x), \nabla f (x) \right\rangle + \frac {1}{2} \left\langle \sigma_ {t} (x) \sigma_ {t} (x) ^ {\top}, \nabla^ {2} f (x) \right\rangle . \tag {4}
$$
We say that a probability measure  $\mathbb{P}$  satisfies the martingale problem for  $\mathcal{A}$  if for any  $t\in [0,T)$  and  $f\in \mathrm{C}_0^\infty (\mathbb{R}^d)$ , we have that  $(f(\mathbf{X}_t) - \int_0^t\mathcal{A}_s(f)(\mathbf{X}_s)\mathrm{d}s)_{s\in [0,t]}$  is a  $\mathbb{P}$ -martingale.