Schrödinger Bridge Matching for Tree-Structured Costs and Entropic Wasserstein Barycentres

TL;DR

This work tackles scalable computation of Schrödinger Bridges for tree-structured, multi-marginal costs and their application to entropic Wasserstein barycentres. It extends Iterative Markovian Fitting (IMF) to the TreeSB setting, producing TreeDSBM, a flow-based method that alternates Markovian projections and bridge-conditioned recursions on a tree. Theoretical results establish existence, uniqueness, and convergence to the TreeSB solution, while experiments demonstrate faster convergence and competitive accuracy compared with IPF-based and other continuous barycentre methods across synthetic 2D, MNIST, subset-posterior, and high-dimensional Gaussian tasks. The approach enables scalable, sample-based barycentre computation that leverages simple bridge-matching steps rather than expensive OT map computations, with practical implications for multi-marginal OT and Bayesian aggregation in high dimensions.

Abstract

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

Figures (10)

• Figure 1: The two stages of TreeIMF. On this tree, the marginals at the leaf vertices $\mathcal{S}$ (blue) are fixed. The marginal at the vertex in $\mathcal{V} \backslash \mathcal{S}$ (red) is not fixed, and can change during the procedure.
• Figure 2: Comparison of the learned barycentre for TreeDSBM (6 IMF iterations) against TreeDSB (50 IPF iterations) and WIN. TreeDSB and TreeDSBM samples are generated from each leaf vertex $k$, and for WIN we plot samples using the weighted-pushforward expression for the barycentre. Also displayed is a close approximation to the ground-truth, using the in-sample method of cuturidoucet14.
• Figure 3: Samples from the 2,4,6 MNIST barycentre.
• Figure 4: The two stages of the TreeIMF procedure, for a non-star-shaped tree structure. On this tree, the marginals at the leaf vertices $\mathcal{S} = \{1,3,5,6,7\}$ (blue) are fixed. The marginals at vertices $\mathcal{V} \backslash \mathcal{S}$ (red) are not fixed, and change during the procedure.
• Figure 5: TreeDSBM samples in the 2$d$ experiment, comparing different regularisation values $\varepsilon$.

Theorems & Definitions (34)

• Definition 2.1: Reciprocal class, Reciprocal projection
• Definition 2.2: Markovian class, Markovian projection
• Definition 3.1: Markov class
• Definition 3.2: Reciprocal class
• Theorem 3.1: TreeSB characterisation
• Lemma 3.1: KL decomposition along tree
• Definition 3.3: Reciprocal projection
• Proposition 3.1
• Definition 3.4: Markovian projection
• Proposition 3.1