Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Momentum Multi-Marginal Schrödinger Bridge Matching

Panagiotis Theodoropoulos, Augustinos D. Saravanos, Evangelos A. Theodorou, Guan-Horng Liu

TL;DR

3MSBM addresses the challenge of reconstructing trajectories from sparse multi-snapshot data by lifting dynamics to phase space and solving a multi-marginal Schrödinger bridge. It introduces a two-stage, simulation-free training scheme: (i) compute a conditional multi-marginal bridge (3MBB) enforcing multiple positional constraints, and (ii) learn a drift a_t^\theta that matches this path while preserving all marginals, enabling scalable high-dimensional trajectory interpolation. The framework provides closed-form conditional accelerations with convergence guarantees to the unique mmSB, and demonstrates superior temporal coherence and marginal fidelity across LV, Gulf of Mexico, Beijing air quality, and single-cell datasets compared to state-of-the-art baselines. These results underscore 3MSBM's potential for robust dynamic interpolation in domains with sparse temporal observations, from biology to geoscience, while highlighting future work on dense image/video settings.

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ödinger 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.

Momentum Multi-Marginal Schrödinger Bridge Matching

TL;DR

3MSBM addresses the challenge of reconstructing trajectories from sparse multi-snapshot data by lifting dynamics to phase space and solving a multi-marginal Schrödinger bridge. It introduces a two-stage, simulation-free training scheme: (i) compute a conditional multi-marginal bridge (3MBB) enforcing multiple positional constraints, and (ii) learn a drift a_t^\theta that matches this path while preserving all marginals, enabling scalable high-dimensional trajectory interpolation. The framework provides closed-form conditional accelerations with convergence guarantees to the unique mmSB, and demonstrates superior temporal coherence and marginal fidelity across LV, Gulf of Mexico, Beijing air quality, and single-cell datasets compared to state-of-the-art baselines. These results underscore 3MSBM's potential for robust dynamic interpolation in domains with sparse temporal observations, from biology to geoscience, while highlighting future work on dense image/video settings.

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ödinger 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.

Paper Structure

This paper contains 46 sections, 9 theorems, 83 equations, 19 figures, 11 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Schrödinger Bridge
  4. Momentum Multi-Marginal Schrödinger Bridge
  5. Momentum Multi-Marginal Schrödinger Bridge Matching
  6. Intermediate Path Optimization
  7. Example - Bridge for 3 marginals
  8. Bridge Matching for Momentum Systems
  9. Convergence
  10. Training Scheme
  11. Experiments
  12. Lotka-Volterra
  13. Gulf of Mexico
  14. Beijing Air Quality
  15. Single cell sequencing
  16. ...and 31 more sections

Key Result

Theorem 3.1

Consider the following momentum system interpolating among multiple marginals We define the value function as $V_t(m_t) := \frac{1}{2} m_t^T P_t^{-1} m_t + m_t^T P_t^{-1} r_t$, where $P_t$, $r_t$ are the second- and first-order approximations, respectively. This formulation admits the following optimal control expression $u^\star_t(m_t) = -{\mathbf{g}} {\mathbf{g}}^T P_t^{-1} where $P_{n^+}:=\lim

Figures (19)

  • Figure 1: Trajectory comparison between by vanilla SB, and our 3MSBM.
  • Figure 2: Visualization of the Dynamic Principle. $P_t$, $r_t$ in Eq. (\ref{['eq:ODEs']}) are solved backward, propagating the influence of future pinned points in preceding segments, through the intermediate constraints.
  • Figure 3: Bridge materializations among 3 pinned points.
  • Figure 4: Iterative propagation of dynamics using the optimal conditional acceleration to approximate the velocity profile $\pi_n(v_n|x_n)$ at the intermediate marginals.
  • Figure 5: Trajectory comparison on LV among SBIRR, MMFM, Smooth SB, and our 3MSBM.
  • ...and 14 more figures

Theorems & Definitions (18)

  • Theorem 3.1: SOC representation of Multi-Marginal Momentum Brownian Bridge (3MBB)
  • Remark 3.3
  • Proposition 3.4
  • Theorem 3.5
  • Theorem B.1: SOC representation of Multi-Marginal Momentum Brownian Bridge (3MBB)
  • proof
  • Proposition B.2
  • proof
  • Remark B.3
  • Proposition B.4
  • ...and 8 more