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

2025-06-20

AI Summary

This work addresses the Schrödinger Bridge (SB) problem under tree-structured costs and the entropy-regularized Wasserstein barycenter computation. We propose the first generalization of the Iterative Markov Fitting (IMF) framework to multi-marginal SB problems with arbitrary tree topologies. Methodologically, we unify the modeling of entropy-regularized optimal transport across arbitrary tree-structured marginals and design a streaming, dynamic fixed-point iteration scheme—overcoming scalability and adaptivity limitations of classical IPF and static barycenter algorithms. Key contributions include: (1) establishing the first IMF-based solution paradigm for tree-structured SB; (2) enabling online barycenter updates with rapid convergence; and (3) preserving IMF’s computational efficiency and numerical stability while supporting arbitrary tree structures. Experiments demonstrate significant improvements in convergence speed and numerical robustness, providing a novel tool for distribution alignment in streaming generative modeling.

Technology Category

• Machine Learning: Imitation Learning & Inverse Reinforcement Learning
• Search and Optimization: Sampling/Simulation-based Search
• Constraint Satisfaction and Optimization: Distributed CSP/Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Systems and Infrastructure for Web, Mobile and WoT: Experiences and lessons learnt from Web-based algorithms and system deployments

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 specific case of Wasserstein barycentres, our approach can be viewed as extending fixed-point approaches for barycentre computation to the case of flow-based entropic OT solvers.

Problem

• Extends IMF to solve tree-structured Schru00f6dinger Bridge problems
• Generalizes optimal transport for multi-marginal tree-based costs
• Enables flow-based entropic OT solvers for Wasserstein barycentres

Innovation

• Extends IMF to tree-structured SB problem
• Solves SB via sequential bridge-matching steps
• Enables flow-based entropic OT barycentre computation