# Oh SnapMMD! Forecasting Stochastic Dynamics Beyond the Schr"odinger Bridge's End

arXiv: https://arxiv.org/abs/2505.16082 · HTML: https://mlsanity.com/paper/2505.16082/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 38 references that are in the map.
- First posted on arXiv 21 May 2025.

## Builds on (most cited first, 24 of 24)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [Diffusion Schr"odinger Bridge with Applications to Score-Based Generative Modeling](https://mlsanity.com/paper/2106.01357/) (2021)
- [FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models](https://mlsanity.com/paper/1810.01367/) (2018)
- [Scalable Gradients for Stochastic Differential Equations](https://mlsanity.com/paper/2001.01328/) (2020)
- [TrajectoryNet: A Dynamic Optimal Transport Network for Modeling Cellular Dynamics](https://mlsanity.com/paper/2002.04461/) (2020)
- [Solving Schr"odinger Bridges via Maximum Likelihood](https://mlsanity.com/paper/2106.02081/) (2021)
- [Deep Generative Learning via Schr"odinger Bridge](https://mlsanity.com/paper/2106.10410/) (2021)
- [Proximal Optimal Transport Modeling of Population Dynamics](https://mlsanity.com/paper/2106.06345/) (2021)
- [Simulation-free Schr"odinger bridges via score and flow matching](https://mlsanity.com/paper/2307.03672/) (2023)
- [Manifold Interpolating Optimal-Transport Flows for Trajectory Inference](https://mlsanity.com/paper/2206.14928/) (2022)
- [Meta Flow Matching: Integrating Vector Fields on the Wasserstein Manifold](https://mlsanity.com/paper/2408.14608/) (2024)
- [Scalable Unbalanced Optimal Transport using Generative Adversarial Networks](https://mlsanity.com/paper/1810.11447/) (2018)
- [Deep Momentum Multi-Marginal Schr"odinger Bridge](https://mlsanity.com/paper/2303.01751/) (2023)
- [Multi-marginal Schr"odinger Bridges with Iterative Reference Refinement](https://mlsanity.com/paper/2408.06277/) (2024)
- [The data-driven Schroedinger bridge](https://mlsanity.com/paper/1806.01364/) (2018)
- [Estimating Barycenters of Measures in High Dimensions](https://mlsanity.com/paper/2007.07105/) (2020)
- [Neural Structure Learning with Stochastic Differential Equations](https://mlsanity.com/paper/2311.03309/) (2023)
- [Identifying Drift, Diffusion, and Causal Structure from Temporal Snapshots](https://mlsanity.com/paper/2410.22729/) (2024)
- [Trajectory Inference with Smooth Schr"odinger Bridges](https://mlsanity.com/paper/2503.00530/) (2025)

## Built on by (most cited first, 5 of 5)

- [A Call to Lagrangian Action: Learning Population Mechanics from Temporal Snapshots](https://mlsanity.com/paper/2605.08550/) (2026)
- [Gradient-flow SDEs have unique transient population dynamics](https://mlsanity.com/paper/2505.21770/) (2025)
- [Optimal Multimarginal Schr"odinger Bridge: Minimum Spanning Tree over Measure-valued Vertices](https://mlsanity.com/paper/2509.10626/) (2025)
- [CellBRIDGE: Learning Cellular Trajectories via Interaction-Aware Alignment](https://mlsanity.com/paper/2605.30635/) (2026)
- [Simulation-Free Learning of Population Dynamics with Wasserstein Lagrangian Residuals](https://mlsanity.com/paper/2610.03679/) (2026)

## Abstract (arXiv)

Scientists often want to make predictions beyond the observed time horizon of "snapshot" data following latent stochastic dynamics. For example, in time course single-cell mRNA profiling, scientists have access to cellular transcriptional state measurements (snapshots) from different biological replicates at different time points, but they cannot access the trajectory of any one cell because measurement destroys the cell. Researchers want to forecast (e.g.) differentiation outcomes from early state measurements of stem cells. Recent Schr"odinger-bridge (SB) methods are natural for interpolating between snapshots. But past SB papers have not addressed forecasting -- likely since existing methods either (1) reduce to following pre-set reference dynamics (chosen before seeing data) or (2) require the user to choose a fixed, state-independent volatility since they minimize a Kullback-Leibler divergence. Either case can lead to poor forecasting quality. In the present work, we propose a new framework, SnapMMD, that learns dynamics by directly fitting the joint distribution of both state measurements and observation time with a maximum mean discrepancy (MMD) loss. Unlike past work, our method allows us to infer unknown and state-dependent volatilities from the observed data. We show in a variety of real and synthetic experiments that our method delivers accurate forecasts. Moreover, our approach allows us to learn in the presence of incomplete state measurements and yields an R2-style statistic that diagnoses fit. We also find that our method's performance at interpolation (and general velocity-field reconstruction) is at least as good as (and often better than) state-of-the-art in almost all of our experiments.