Schr\"odinger Bridge Matching for Tree-Structured Costs and Entropic Wasserstein Barycentres

TL;DR

This paper extends flow-based Schr"odinger Bridge methods to handle tree-structured costs and Wasserstein barycentres, enabling scalable and efficient solutions for complex multi-marginal optimal transport problems.

Contribution

It introduces an IMF-based algorithm for tree-structured Schr"odinger Bridge problems, generalizing existing methods to multi-marginal and barycentre cases with improved properties.

Findings

Algorithm inherits advantages of IMF over IPF.

Effective for multi-marginal problems with tree-structured costs.

Enables scalable computation of Wasserstein barycentres.

Abstract

Recent advances in flow-based generative modelling have provided scalable methods for computing the Schr\"odinger 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.