Momentum Multi-Marginal Schr\"odinger Bridge Matching

TL;DR

The paper introduces 3MSBM, a novel multi-marginal Schr"odinger bridge framework that learns smooth trajectories satisfying multiple constraints, improving the modeling of complex, long-range temporal dependencies in stochastic systems.

Contribution

It extends existing bridge matching methods to multi-marginal settings, enabling the capture of long-range dependencies and enhancing scalability and convergence.

Findings

Outperforms existing methods in real-world applications

Effectively captures complex temporal dynamics

Improves convergence and scalability of matching algorithms

Abstract

Understanding complex systems by inferring trajectories from sparse sample snapshots is a fundamental challenge in a wide range of domains, e.g., single-cell biology, meteorology, and economics. Despite advancements in Bridge and Flow matching frameworks, current methodologies rely on pairwise interpolation between adjacent snapshots. This hinders their ability to capture long-range temporal dependencies and potentially affects the coherence of the inferred trajectories. To address these issues, we introduce \textbf{Momentum Multi-Marginal Schr\"odinger Bridge Matching (3MSBM)}, a novel matching framework that learns smooth measure-valued splines for stochastic systems that satisfy multiple positional constraints. This is achieved by lifting the dynamics to phase space and generalizing stochastic bridges to be conditioned on several points, forming a multi-marginal conditional stochastic optimal control problem. The underlying dynamics are then learned by minimizing a variational objective, having fixed the path induced by the multi-marginal conditional bridge. As a matching approach, 3MSBM learns transport maps that preserve intermediate marginals throughout training, significantly improving convergence and scalability. Extensive experimentation in a series of real-world applications validates the superior performance of 3MSBM compared to existing methods in capturing complex dynamics with temporal dependencies, opening new avenues for training matching frameworks in multi-marginal settings.