Deep Momentum Multi-Marginal Schrödinger Bridge
This work addresses the problem of trajectory inference in high-dimensional systems for researchers in fields like computational biology, though it appears incremental as it builds on existing Schrödinger Bridge models.
The paper tackles the challenge of reconstructing population dynamics from unlabeled samples at coarse time intervals by introducing DMSB, a framework that learns smooth measure-valued splines for stochastic systems, outperforming baselines on synthetic and real-world single-cell RNA sequence datasets.
It is a crucial challenge to reconstruct population dynamics using unlabeled samples from distributions at coarse time intervals. Recent approaches such as flow-based models or Schrödinger Bridge (SB) models have demonstrated appealing performance, yet the inferred sample trajectories either fail to account for the underlying stochasticity or are $\underline{D}$eep $\underline{M}$omentum Multi-Marginal $\underline{S}$chrödinger $\underline{B}$ridge(DMSB), a novel computational framework that learns the smooth measure-valued spline for stochastic systems that satisfy position marginal constraints across time. By tailoring the celebrated Bregman Iteration and extending the Iteration Proportional Fitting to phase space, we manage to handle high-dimensional multi-marginal trajectory inference tasks efficiently. Our algorithm outperforms baselines significantly, as evidenced by experiments for synthetic datasets and a real-world single-cell RNA sequence dataset. Additionally, the proposed approach can reasonably reconstruct the evolution of velocity distribution, from position snapshots only, when there is a ground truth velocity that is nevertheless inaccessible.