Diffusion Schrödinger Bridge Matching

Abstract

Solving transport problems, i.e. finding a map transporting one given distribution to another, has numerous applications in machine learning. Novel mass transport methods motivated by generative modeling have recently been proposed, e.g. Denoising Diffusion Models (DDMs) and Flow Matching Models (FMMs) implement such a transport through a Stochastic Differential Equation (SDE) or an Ordinary Differential Equation (ODE). However, while it is desirable in many applications to approximate the deterministic dynamic Optimal Transport (OT) map which admits attractive properties, DDMs and FMMs are not guaranteed to provide transports close to the OT map. In contrast, Schrödinger bridges (SBs) compute stochastic dynamic mappings which recover entropy-regularized versions of OT. In this work, we introduce Iterative Markovian Fitting (IMF), a new methodology for solving SB problems, and Diffusion Schrödinger Bridge Matching (DSBM), a novel numerical algorithm for computing IMF iterates.

Summary

This paper introduces Iterative Markovian Fitting (IMF), a new theoretical framework for computing Schrödinger bridges that alternates between projecting onto the space of Markov processes (Markovian projection) and projecting onto the reciprocal class of the reference measure (reciprocal projection). Unlike Iterative Proportional Fitting which alternates between marginal constraints (preserving only one marginal at a time), IMF preserves both marginal constraints at every iteration. Convergence to the unique SB is proved via a Pythagorean theorem for KL divergence.

The practical algorithm DSBM implements IMF iterates using simple regression losses in the spirit of bridge matching and flow matching. At each iteration, DSBM learns forward and backward drifts by regressing onto conditional scores of the reference bridge, avoiding time-discretization issues of prior DSB methods. The authors show DSBM-IPF recovers IPF iterates of DSB, while DSBM-IMF generalizes rectified flow (recovered as the deterministic σ→0 limit). Flow Matching and Bridge Matching are recovered as first-iteration cases, establishing DSBM as a unifying framework.

Experiments demonstrate superiority over DSB on 2D tasks, high-dimensional Gaussians (d=50), MNIST/EMNIST transfer, CelebA 64×64, AFHQ 512×512, and unpaired fluid flow downscaling.

Key Contributions

• Iterative Markovian Fitting: alternating Markovian and reciprocal projections, preserving both marginals throughout
• DSBM algorithm implementing IMF via regression losses without time-discretization artifacts
• Theoretical convergence guarantees via Pythagorean theorem and monotone KL decrease
• Unifying framework: recovers DDMs, Bridge Matching, Flow Matching, Rectified Flow, OT-CFM, and DSB as special cases
• Forward-backward alternation preventing bias accumulation
• Probability flow ODE for the learned SB enabling likelihood computation

Methodology

IMF alternates: (1) Markovian projection — find Markov process M* minimizing reverse-KL to current path measure, preserving all time marginals; (2) Reciprocal projection — sample joint (X₀, X_T) from M and reconstruct reference bridge. DSBM learns drifts via regression on bridge conditional scores. Initialized with either independent coupling (DSBM-IMF) or reference process coupling (DSBM-IPF). Forward-backward alternation exploits time-symmetry of the Markovian projection.

Key Findings

• DSBM outperforms DSB on all 2D datasets in Wasserstein distance and path energy
• No variance/covariance estimation errors in d=50 (unlike DSB and Rectified Flow)
• Much lower KL to true SB than DSB and SB-CFM in d=20 and d=50
• ~30% more computationally efficient than DSB on MNIST transfer
• Diffusion σ controls quality/alignment tradeoff; optimal σ scales with image resolution
• Stochasticity (σ > 0) essential for general transport tasks; Rectified Flow fails on non-trivial problems

Important References

1. Diffusion Schrödinger Bridge with Applications to Score-Based Generative Modeling — Prior DSB algorithm that DSBM improves upon
2. Flow Straight and Fast — Rectified Flow recovered as σ→0 limit of DSBM-IMF
3. Score-Based Generative Modeling through Stochastic Differential Equations — DDMs as first DSBM-IPF iteration

Atomic Notes

• Schrödinger bridge
• Iterative Markovian Fitting
• Markovian projection
• reciprocal class
• Iterative Proportional Fitting
• bridge matching

---

paper