Schrödinger Bridge Matching for Tree-Structured Costs and Entropic Wasserstein Barycentres
This work addresses a problem in optimal transport and generative modeling for researchers and practitioners, offering an incremental extension to existing methods.
The authors tackled the problem of computing the Schrödinger Bridge for tree-structured costs and entropic Wasserstein barycentres by extending the Iterative Markovian Fitting procedure, resulting in an algorithm that inherits advantages over traditional methods and requires only simple bridge-matching steps.
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.