Research map / Learning theory and algorithms / Optimal transport
Schrödinger Bridge Matching for Tree-Structured Costs and Entropic Wasserstein Barycentres
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Extends Iterative Markovian Fitting (IMF) to tree-structured entropic multi-marginal Schrödinger bridges (TreeDSBM), yielding a flow-based fixed-point-style solver for entropic Wasserstein barycentres via bridge-matching.
Problem
Scalably solving the tree-structured dynamic Schrödinger bridge (entropic multi-marginal OT with tree costs) and entropic Wasserstein-2 barycentres; the existing tree solver TreeDSB is IPF-based and struggles with stability at low entropy.
Approach
Generalises IMF from a single edge to a tree: alternate per-edge Markovian projection (bridge-matching drifts) and reciprocal projection (mixture of bridges conditioned on observed-vertex values), with a KL decomposition along the tree proving convergence. For barycentres this merges the fixed-point update with flow-based OT, replacing OT map solves with bridge-matching.
Key results
- On the 2D moon/spiral/circle barycentre, TreeDSBM with 6 IMF iterations significantly improves over TreeDSB with 50 IPF iterations at much lower computational cost.
- On subset posterior aggregation (Poisson), TreeDSBM achieves BW2^2-UVP 0.008% vs WIN 0.014%, W2CB 0.026%, NOTWB 0.023%.
- On high-dimensional Gaussian barycentres, TreeDSBM is competitive with WIN, W2CB and NOTWB, with W2CB and NOTWB showing a slight edge on BW2^2-UVP.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use TreeDSBM instead of WIN, W2CB, or NOTWB | Continuous barycentre estimation · Subset posterior aggregation, Poisson | BW_2^2-UVP (%): 0.014 -> 0.008; 0.026 -> 0.008; 0.023 -> 0.008 | Helped |
| Use TreeDSBM instead of WIN | Continuous barycentre estimation · Subset posterior aggregation, Negative Binomial | BW_2^2-UVP (%): 0.009 -> 0.012 | Hurt |
| Use TreeDSBM instead of W2CB or NOTWB | Continuous barycentre estimation · Subset posterior aggregation, Negative Binomial | BW_2^2-UVP (%): 0.024 -> 0.012; 0.018 -> 0.012 | Helped |
| Use TreeDSBM instead of W2CB or NOTWB | Continuous barycentre estimation · High-dimensional Gaussian barycentre, BW_2^2-UVP (%) across d=64, 96, 128 | BW_2^2-UVP (%): 0.04, 0.07, 0.12 -> 0.14, 0.15, 0.27; 0.08, 0.10, 0.14 -> 0.14, 0.15, 0.27 | Hurt |
| Use TreeDSBM instead of WIN | Continuous barycentre estimation · High-dimensional Gaussian barycentre, BW_2^2-UVP (%) across d=64, 96, 128 | BW_2^2-UVP (%): 0.20, 0.30, 0.38 -> 0.14, 0.15, 0.27 | Helped |
| Use TreeDSBM instead of W2CB or NOTWB | Continuous barycentre estimation · High-dimensional Gaussian barycentre, L^2-UVP (%) across d=64, 96, 128 | L^2-UVP (%): 0.17, 0.20, 0.25 -> 1.18, 1.13, 1.23; 0.10, 0.10, 0.13 -> 1.18, 1.13, 1.23 | Hurt |
| Use TreeDSBM instead of WIN | Continuous barycentre estimation · High-dimensional Gaussian barycentre, L^2-UVP (%) across d=64, 96, 128 | L^2-UVP (%): 0.96, 1.20, 1.46 -> 1.18, 1.13, 1.23 | Mixed |
| Replace TreeDSB with TreeDSBM | Tree-structured barycentre solver · Synthetic 2d moon, spiral, and circle barycentre | Sinkhorn divergence | Helped |
| Replace TreeDSB with TreeDSBM | MNIST 2,4,6 barycentre · MNIST 2,4,6 | Sample quality | Helped |
| Use lower entropy regularisation with TreeDSBM than TreeDSB | MNIST 2,4,6 barycentre · MNIST 2,4,6 | Marginal matching | Helped |
| Replace TreeDSB with TreeDSBM | Tree-structured barycentre solver · MNIST 2,4,6 | Computational cost | Helped |
Surprising findings
- On the Gaussian experiments, TreeDSBM's L2-UVP is higher than W2CB and NOTWB, so it is not uniformly best against specialised barycentre solvers.
- W2CB and NOTWB could not be made to converge on the 2D example without hyperparameter tuning, while TreeDSBM trained stably.
Limitations
- Restricted to quadratic cost functions and introduces an entropic bias.
- Inference is expensive and not sampling-free; each IMF iteration requires simulating learned SDEs.
- Sampling noise in the reciprocal projection means large regularisation degrades sample quality, so small epsilon is needed.
Experimental setup
Datasets: 2D moon, spiral, circle synthetic data MNIST digits 2, 4, 6 bike-rental dataset (subset posterior aggregation) randomly generated Gaussians in dimensions 64, 96, 128
Benchmarks: 2D moon/spiral/circle barycentre subset posterior aggregation high-dimensional Gaussian barycentres
Metrics: Sinkhorn divergence BW2^2-UVP L2-UVP
Baselines: TreeDSB WIN W2CB NOTWB
Builds on: Iterative Markovian Fitting TreeDSB Diffusion Schrödinger Bridge bridge matching
Closest papers by meaning
- Exponential Convergence Guarantees for Iterative Markovian Fitting (NeurIPS 2025 · similarity 0.91)
- Light and Optimal Schrödinger Bridge Matching (ICML 2024 · similarity 0.91)
- Discrete Diffusion Schrödinger Bridge Matching for Graph Transformation (ICLR 2025 · similarity 0.89)
- Feedback Schrödinger Bridge Matching (ICLR 2025 · similarity 0.89)
- Generalized Schrödinger Bridge Matching (ICLR 2024 · similarity 0.89)
- Categorical Schrödinger Bridge Matching (ICML 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.