Multi-marginal temporal Schrödinger Bridge Matching from unpaired data

TL;DR

MMtSBM introduces a scalable, multi-marginal temporal Schrödinger Bridge framework that extends diffusion Schrödinger Bridge Matching to multiple time marginals via a factorized Iterative Markovian Fitting procedure. By exploiting a factorized reciprocal class and Markov projections, it provides convergence guarantees and practical algorithms for high-dimensional trajectory reconstruction from unpaired data. The approach yields state-of-the-art performance on high-dimensional transcriptomic benchmarks and can recover continuous video dynamics in image spaces, including MNIST-scale morphing and complex cellular imagery. This work offers a principled, efficient pathway to infer hidden temporal evolution from static samples in domains ranging from single-cell biology to video synthesis.

Abstract

Many natural dynamic processes -- such as in vivo cellular differentiation or disease progression -- can only be observed through the lens of static sample snapshots. While challenging, reconstructing their temporal evolution to decipher underlying dynamic properties is of major interest to scientific research. Existing approaches enable data transport along a temporal axis but are poorly scalable in high dimension and require restrictive assumptions to be met. To address these issues, we propose Multi-Marginal temporal Schrödinger Bridge Matching (MMtSBM) from unpaired data, extending the theoretical guarantees and empirical efficiency of Diffusion Schrödinger Bridge Matching (arXiv:2303.16852) by deriving the Iterative Markovian Fitting algorithm to multiple marginals in a novel factorized fashion. Experiments show that MMtSBM retains theoretical properties on toy examples, achieves state-of-the-art performance on real-world datasets such as transcriptomic trajectory inference in 100 dimensions, and, for the first time, recovers couplings and dynamics in very high-dimensional image settings. Our work establishes multi-marginal Schrödinger bridges as a practical and principled approach for recovering hidden dynamics from static data.

Figures (7)

• Figure 1: Top row: epoch $0$ (only noisy flow matching). Bottom row: epoch $5$ (after MMtSBM training). From left to right: snapshots at times $(0,0.5,1,1.3)$. True marginal times $(t_0, t_1, t_2)=(0, 1, 2)$. The order of the $3$ true marginals is: $t_0=$ dark blue; $t_1=$ red; $t_2=$ light blue. Generated samples are in green. In the background is the quiver plot of the learned score network.
• Figure 2: Evolution of mean, variance, and covariance in the multi-marginal $50d$ Gaussian transport. Dash lines are the theoretical true values.
• Figure 3: Video generated by MMtSBM on MNIST, backward direction. Starting image is from the test set. From left to right: generation at time $t=4, 3.5, 3, 2.5, 2, 1.5, 1, 0.5, 0$. Integer times are marginal times.
• Figure 4: Ground truth biotine examples at training marginal times $t=0,1,2,3,4,5,6$, from left to right.
• Figure 5: Video generated by MMtSBM on biotine, forward direction. To be read in reading order: top left $\rightarrow$ top right, then bottom left $\rightarrow$ bottom right. Generations at times $t=0, 0.5, 1, 1.5, ..., 5.5, 6$. Top-left starting image is from the test set.

Theorems & Definitions (28)

• Definition 3.1: Static formulation
• Proposition 3.2: Dynamic--static equivalence
• Proposition 3.3: Markovianity
• Proposition 3.4: Form of the solution
• Definition 3.2: Factorized reciprocal class and projection
• Proposition 3.5: Local reciprocal structure of the factorized class
• Definition 3.3: Markovian projection in the factorized setting
• Proposition 3.6: Variational characterization of the factorized Markovian projection
• Proposition 3.7
• Conjecture 3.1: Analogue of leonard2014survey Theorem 2.12