Oh SnapMMD! Forecasting Stochastic Dynamics Beyond the Schrödinger Bridge's End
Renato Berlinghieri, Yunyi Shen, Jialong Jiang, Tamara Broderick
TL;DR
The paper tackles forecasting stochastic dynamics from snapshot data, where trajectories are unobserved, by introducing SnapMMD, a framework that learns SDEs through direct joint state-time distribution matching using Maximum Mean Discrepancy. This approach enables inference of unknown, state-dependent volatility, handles incomplete observations, and yields an interpretable velocity field plus an RKHS-based $R^2$ diagnostic for model fit. Across synthetic and real datasets—including Lotka–Volterra, repressilator variants, Gulf of Mexico currents, and PBMC immune activation—the method delivers superior forecasting and competitive interpolation against Schrödinger-bridge baselines while providing robust diagnostics and interpretable dynamics. The work has practical impact for analyzing time-course data in biology and related dynamic systems, with code available for reproducibility and further application.
Abstract
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ödinger-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 $R^2$-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.