Pith. sign in

REVIEW 3 major objections 4 minor 12 references

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

T0 review · 3 major / 4 minor · reviewed 2026-08-07 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

Pith's one-line read This paper argues that forecasting latent stochastic dynamics from snapshot data can be done by fitting the joint state-time distribution with a maximum mean discrepancy loss, and demonstrates the approach in synthetic and real experiments.

desk verdict Genuinely new MMD-based snapshot forecasting method with strong experiments, but the volatility-inference claim outruns the paper's own identifiability analysis. read the letter →

arxiv 2505.16082 v1 pith:ZM4EHQLI submitted 2025-05-21 stat.ML cs.LGstat.ME

Renato Berlinghieri , Yunyi Shen , Jialong Jiang , Tamara Broderick This is my paper · ORCID
classification stat.MLcs.LGstat.ME
keywords maximummeandiscrepancysnapshotdatastochasticdifferentialequationsforecastingSchrödingerbridgetrajectoryinferencesingle-cellRNA-seqstate-dependentvolatility
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

Many scientific measurements, especially in single-cell biology, provide snapshots of a population at discrete times but never the trajectory of any individual, because measuring destroys the sample. This paper proposes SnapMMD, a framework that learns the underlying stochastic differential equation by matching the joint distribution of cell state and measurement time against the observed snapshots, using a maximum mean discrepancy (MMD) loss. The paper argues that this direct distributional fit makes forecasting beyond the final observed time point possible and lets the method infer unknown, state-dependent volatility—something previous Schrödinger-bridge methods cannot do, since their forecasts reduce to following a pre-set reference process or require a fixed volatility. Experiments on synthetic gene circuits, predator-prey systems, ocean currents, and real T-cell activation data show the method forecasting more accurately than the Schrödinger-bridge baselines, handling partially observed state variables, and producing an R²-style fit diagnostic.

What carries the argument

The load-bearing object is Proposition 3.1: for joint distributions $f(y,t)=f(y\mid t)h(t)$ and $g(y,t)=g(y\mid t)h(t)$ with kernel $K((y,t),(y',t'))=K_y(y,y')\delta(t-t')$, the squared MMD equals $\sum_{t\in\mathcal{T}} h(t)^2\,\operatorname{MMD}_{K_y}^2(f(\cdot\mid t),g(\cdot\mid t))$. This decomposition lets SnapMMD replace trajectory-level inference with per-time distribution matching, estimated by unbiased U-statistics and optimized by gradient-based methods through the SDE simulation. The R² diagnostic uses the weighted mixture of the empirical snapshot distributions as the baseline, so it measures how much better the fitted model explains the time-dependent state distributions than a single time-independent barycenter.

What would settle it

Generate data from an SDE with state-dependent volatility and from a second SDE that has the same time marginals but a different volatility (for example, the original volatility multiplied by an orthogonal matrix); if SnapMMD cannot tell the two systems apart from their snapshots, then the recovered volatility depends entirely on the user's family restriction. Alternatively, fit the repressilator data with a constant-volatility family and show that the one-step-ahead forecast MMD on held-out snapshots degrades measurably.

Watch

Extended reading notes

Core claim

The paper's central claim is that forecasting from snapshot data fails with standard Schrödinger bridges because their objective—minimizing Kullback-Leibler divergence to a reference process—forces either a pre-specified reference dynamic or a fixed, state-independent volatility; forecasting then just propagates the last observed marginal forward under that reference. SnapMMD instead fits a parametric SDE by minimizing MMD between the model's predicted joint distribution over state and time and the empirical joint distribution of the snapshots. With a kernel that factors into a state kernel times $\delta(t-t')$, the joint MMD decomposes into a weighted sum of per-time MMDs, turning the problem into a distributional least-squares regression across the observed times. The resulting framework can learn drift and volatility simultaneously, use only the observed dimensions when some state variables are missing, and diagnose fit with an R²-style metric that compares the model's MMD to that of a kernel barycenter of the empirical distributions.

Load-bearing premise

The method presupposes that the user can specify a parametric family of SDEs that contains the true drift and volatility functions, because Appendix E shows that infinitely many drift and volatility pairs produce the same marginal distributions over time.

Editorial extensions

If this is right

  • One-step-ahead forecasts beyond the observed time horizon become accurate when the chosen SDE family contains the true dynamics, outperforming the reported Schrödinger-bridge baselines in the experiments.
  • Unknown and state-dependent volatility can be inferred from snapshot data alone, as the repressilator and Lotka-Volterra experiments demonstrate.
  • Incomplete state observations can still support forecasting: matching the marginal of the observed dimensions is sufficient when the full multi-dimensional SDE is the model family.
  • The R²-style metric provides an interpretable fit diagnostic that can be used for early stopping and model selection.
  • Interpolation and velocity-field reconstruction are at least as good as, and often better than, the state-of-the-art methods in most experiments, even though the method was designed for forecasting.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • If SnapMMD holds up, trajectory inference in general might be reframed as distributional regression on state-time pairs, possibly extending to settings with non-i.i.d. measurement times or aggregated population measurements.
  • The identifiability limitation implies a practical guideline: users should prefer small, physically motivated model families and use the R² diagnostic to detect families that merely fit marginals without identifying the underlying dynamics.
  • A natural testable extension is to evaluate SnapMMD on multi-step and longer-horizon forecasts, since the paper evaluates one-step-ahead forecasts with a larger time increment, and the comparison to Schrödinger bridges may shift at longer horizons.
Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

3 major / 4 minor

Summary. The paper proposes SnapMMD, a framework for learning a parametric SDE from population-level snapshot data by minimizing a weighted sum of maximum mean discrepancy (MMD) terms between the model's simulated marginal distributions and empirical distributions at each observed time point. The objective is derived from a joint state–time matching perspective and is shown to reduce to a least-squares-like form (Proposition 3.1). The authors claim that SnapMMD provides accurate forecasts beyond the observed horizon, can infer unknown and state-dependent volatilities, handles incomplete state measurements, and yields an R^2-style diagnostic for model fit. Experiments are conducted on synthetic Lotka–Volterra and repressilator systems, a synthetic repressilator with missing protein observations, real ocean-current data from the Gulf of Mexico, and a real PBMC single-cell RNA-seq dataset. The method is compared against Schrödinger bridge baselines for forecasting and several interpolation baselines.

Significance. If the central claims hold, SnapMMD would be a practically valuable contribution to trajectory inference and forecasting from snapshot data, particularly for single-cell biology. The paper is clearly written, the core MMD decomposition (Proposition 3.1) is clean and correctly proven, and the authors provide code and detailed experimental descriptions. The R^2-style diagnostic and the handling of incomplete state measurements are useful additions. However, the load-bearing claims of consistently better forecasts and inference of unknown state-dependent volatility are weakened by (i) the absence of any identifiability analysis for the parametric families actually used, (ii) the fact that the objective only matches marginals at observed times, so forecasts beyond the horizon are not determined by the loss, and (iii) a direct counterexample in the Gulf of Mexico experiment where the baseline SBIRR-ref achieves lower forecast MMD. These issues are fixable within the manuscript's scope, so the paper warrants a major revision rather than rejection.

major comments (3)
  1. [§5, Appendix E, Eq. (4)] The advertised capability of inferring unknown, state-dependent volatilities is not supported. Appendix E explicitly concedes that the Fokker–Planck equation does not identify the pair (b0, g0) from the marginal distributions in general, and the paper provides no identifiability condition or verification for the parametric families actually used in the experiments (Eqs. A9, A10, A11, A13). In particular, the semiparametric model in Eq. (A11) has a manifest scale ambiguity: only the product M f_theta enters the drift, so scaling M and inversely scaling f_theta leaves the drift unchanged. Moreover, the loss in Eq. (4) matches marginals only at the I observed time points, so parameters that produce identical marginals at those times but different dynamics afterward are indistinguishable; forecasts beyond t_I are therefore not determined by the objective. The paper never reports fitted drift or volatility parameters against ground truth in the synthetic experiments, only downstream distributional metrics. To support the volatility-inference claim, the authors should either prove or empirically demonstrate identifiability for each model family used, and report parameter recovery in the synthetic settings.
  2. [Table 1, Gulf of Mexico forecast row, §4.4] The headline claim that SnapMMD 'consistently provides better forecasts than competitors' is directly contradicted by the Gulf of Mexico experiment: in the forecast row of Table 1, SnapMMD has MMD 0.66±0.03 while SBIRR-ref achieves 0.35±0.03. The authors attribute this to the RBF kernel preferring a diffuse cloud, but this explanation is problematic because the same RBF-MMD is exactly the training objective minimized by SnapMMD. At minimum, the paper should report forecast performance with additional metrics or kernels and temper the 'consistently better' claim, or provide a more principled justification for why MMD is the appropriate primary metric despite this failure case.
  3. [§4, 'Metrics of success'; Eq. (4)] The primary evaluation metric for forecasting and interpolation is the same MMD objective used for training (Eq. 4), with an RBF kernel and median-heuristic length scale. While the held-out time points prevent direct circularity, the comparison is still biased in SnapMMD's favor relative to baselines that do not optimize this metric. The Gulf of Mexico result shows that this bias is not always decisive, but the paper should include a sensitivity analysis (e.g., different kernel bandwidths, EMD as the primary metric, or a separate metric not aligned with the loss) to support the claim of superior forecasting.
minor comments (4)
  1. [Definition 3.1] Definition 3.1 contains a typo: 'Let f_theta(·|t) denote the model-predicted state distribution at time f_theta(·|t)' should read 'at time t'.
  2. [Table 1 and Table 2] The formatting of Tables 1 and 2 is inconsistent: some entries are shown as '0.01±0.01' without a leading zero in the standard deviation, and the color highlighting is not explained in the caption for Table 2. Please standardize the numeric formatting and caption conventions.
  3. [§D.7.5, Table 12] The caption of Table 12 says 'parametric model family' but the section is about the semiparametric family; this is presumably a copy-paste error and should be corrected.
  4. [§4.3] In the ReprProtein experiment, the baselines fit an mRNA-only model while SnapMMD fits the full mRNA–protein model. This is an unfair comparison that should be acknowledged more explicitly in the main text, even though the incomplete-state handling is a genuine advantage of the method.

Circularity Check

0 steps flagged · score 1.0 of 10

No significant circularity: the forecasting derivation is self-contained and the held-out predictions are not fitted values; identifiability is an acknowledged limitation, and the MMD evaluation overlap is a fairness concern, not a circular step.

full rationale

SnapMMD's core derivation is not circular. The training objective (Eq. 4) minimizes a weighted sum of per-time MMDs between model-simulated and empirical state distributions at the observed times t_1,...,t_I. Forecasting evaluates the fitted SDE at a held-out time t_{I+1} that does not appear in the loss; the forecast distribution is produced by simulating the SDE forward, so it is not a refitted or algebraically forced version of the training data. Proposition 3.1, which reduces the joint state-time MMD to a weighted sum of conditional MMDs, is a standard identity proved in Appendix B.2 and is not used to define the target quantities in terms of themselves. The R^2 statistic (Definition 3.1) is a training-time goodness-of-fit diagnostic used for early stopping and architecture selection; it is not presented as a forecast and does not enter the held-out evaluation. Self-citations to Shen et al. (2024, 2025) appear only as baselines (SB-forward, SBIRR) and as related work; they are not invoked to justify SnapMMD's identifiability, uniqueness, or predictive claims. Proposition B.1, which characterizes SB forecasting as reference propagation, relies on Lavenant et al. (2024), not on the authors' own prior work. The paper's own Section 5 and Appendix E explicitly concede that drift and volatility are not identified from marginal distributions in general (the Fokker-Planck invariance under h with div(h pi_t)=0 and under g0 A for orthogonal A), and that a complete identifiability characterization is out of scope; the parametric families used (Eqs. A9, A10, A11, A13) are not shown to be identifiable, and Eq. A11 has a manifest scale ambiguity between M and f_theta. This is a substantive correctness and validation gap, especially for the advertised volatility inference, but it is an identifiability limitation rather than a circular derivation. The use of RBF-MMD as both the training loss and the primary evaluation metric (with EMD as a secondary metric) could mildly favor the method being trained, but the forecast time point is held out and EMD and drift-MSE results corroborate the qualitative conclusions; this is a fairness concern, not evidence that predictions reduce to fitted inputs by construction. Overall, no load-bearing argument in the paper reduces to its own inputs or to an unverified self-citation chain.

Assumptions & free parameters 6 free parameters · 5 assumptions · 0 invented entities

The central claim rests on standard SDE regularity conditions, the i.i.d. state-time sampling assumption, and, crucially, the restriction to a user-specified parametric model family. The free parameters are the fitted SDE parameters and kernel/hyperparameter choices; no new entities are introduced. The identifiability analysis admits that volatility is not uniquely recoverable in general.

free parameters (6)
  • Lotka-Volterra drift and volatility params = alpha, beta, gamma, delta, sigma (estimated via MMD loss)
    Learned in Section D.5.2; forecasting depends on these values.
  • Repressilator drift and volatility params = beta, n, k, gamma, sigma or MLP weights
    Learned in Sections D.6.2 and D.7.2.
  • Gulf of Mexico vortex params = gamma, rv, x0, y0, d, rd, x0d, y0d, sigma
    Learned in Section D.9.2.
  • PBMC MLP weights = weights of MLP with hidden [128,128,128]
    Learned in Section D.10.3.
  • RBF kernel length scale = median pairwise distance in data (median heuristic)
    Kernel choice affects MMD landscape; authors note it can prefer diffuse clouds (GoM).
  • Optimization hyperparameters = learning rate 0.05, epochs 300-500, early stopping on R2
    Chosen by hand; affects convergence, not central claim.
assumptions (5)
  • domain assumption Assumption 2.1: b0 and g0 are Lipschitz and have at most linear growth
    Needed for strong solutions of the SDE (Eq. 1); standard but restricts admissible dynamics.
  • domain assumption Assumption 2.2: bounded second moments of particles and initial distribution
    Needed for finite MMD expectations and unbiased U-statistic estimators.
  • domain assumption Observation times are random draws from h(t) and cells are independent
    Justifies treating state-time pairs as i.i.d. samples, the basis of the joint-distribution objective (Eq. 2, Section 2).
  • domain assumption The chosen parametric family contains (or approximates) the true SDE
    Appendix E shows non-identifiability of drift and volatility from marginals without such a restriction; the method's inference relies on it.
  • standard math Kernel factors as Ky(y,y') delta(t-t') with Ky positive definite
    Underlies Proposition 3.1; many common kernels satisfy this.

how reviews work

0 comments
Cite this review

Pith. "Pith review of Oh SnapMMD! Forecasting Stochastic Dynamics Beyond the Schr\"odinger Bridge's End." pith.science (2026). https://pith.science/paper/ZM4EHQLI

@misc{pith2026250516082,
  author       = {Pith},
  title        = {Pith review of: Oh SnapMMD! Forecasting Stochastic Dynamics Beyond the Schr\"odinger Bridge's End},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/ZM4EHQLI}},
  note         = {Machine review of arXiv:2505.16082}
}
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\"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 $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.

Figures

Figures reproduced from arXiv: 2505.16082 by the authors.

Figure 1
Figure 1. Lotka-Volterra results (Section 4.1). We show 200 samples at each of 10 training times and 1 forecast time (red). Forecast points overlap with the training points at time 0 (blue). showed that the barycenter is the weighted mixture of the empirical distributions: fbary(y) := arg min f X I i=1 wi MMD2 Ky (f(y | ti), ˆf(y;ti)) = 1 PI i=1 wiNi X I i=1 X Ni n=1 wiδY n ti (y), (5) with weights wi as in Eq. (4). Our RKHS-… view at source ↗
Figure 2
Figure 2. Repressilator results: mRNA-only (upper, Section [PITH_FULL_IMAGE:figures/full_fig_p007_2.png] view at source ↗
Figure 3
Figure 3. Gulf of Mexico results (Section 4.4). We show 200 samples at each of 10 training times and 1 forecast time (red). the same parametric model as the data-generating process; see Section D.6 for full results. (2) We use a semiparametric model with a multilayer perceptron; see Section D.7.2 for details. Results. In [PITH_FULL_IMAGE:figures/full_fig_p008_3.png] view at source ↗
Figures from the paper (1 more)
Figure 4
Figure 4. Figure 4: PBMC results (Section 4.5). The axes in every plot are the same three principal com￾ponents, computed over the full data: i.e., 41 time steps of the 30-dimensional gene programs. Leftmost four panels: evolution of the training data at time steps 1, 7, 14, and 20. "Trut…

Discussion (0). Continue with ORCID to comment.

Reference graph

Works this paper leans on

12 extracted references · 9 canonical work pages

  1. [1]

    J., Bengio, Y ., Tong, A., and Neklyu- dov, K

    Atanackovic, L., Zhang, X., Amos, B., Blanchette, M., Lee, L. J., Bengio, Y ., Tong, A., and Neklyu- dov, K. (2024). Meta flow matching: Integrating vector fields on the Wasserstein manifold.arXiv preprint arXiv:2408.14608. Bunne, C., Papaxanthos, L., Krause, A., and Cuturi, M. (2022). Proximal optimal transport modeling of population dynamics. InInternat...

    arXiv 2024

  2. [12]

    A complete characterization of identifiability in more general settings is beyond the scope of this work and constitutes an important direction for future research

    For instance, assuming thatb 0 is a gradient field (i.e.,b 0 =∇Φfor some potentialΦand thatg 0 is constant is known to yield identifiability under suitable conditions (Lavenant et al., 2024; Guan et al., 2024). A complete characterization of identifiability in more general settings is beyond the scope of this work and constitutes an important direction fo...

    work page 2024

  3. [14]

    Since the two baselines cannot make vector fields in correct dimension we did not compare vector field reconstruction

    Our method outperform baseline by a large margin, mostly because the correct account of the missing protein observation. Since the two baselines cannot make vector fields in correct dimension we did not compare vector field reconstruction. Table 14: Evaluation metric for Repressilator forecasting with missing protein observations. Repressilator (with miss...

    work page

  4. [31]

    Chen, T., Liu, G.-H., Tao, M., and Theodorou, E. (2024). Deep momentum multi-marginal Schrödinger bridge.Advances in Neural Information Processing Systems,

    work page 2024

  5. [32]

    P., Kumar, A., Ermon, S., and Poole, B

    Song, Y ., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. (2021). Score- based generative modeling through stochastic differential equations.11th International Confer- ence on Learning Representations. Tong, A., Fatras, K., Malkin, N., Huguet, G., Zhang, Y ., Rector-Brooks, J., Wolf, G., and Bengio, Y . (2024a). Improving and gener...

    arXiv 2021

  6. [36]

    (2025).Statistical Optimal Transport, volume 2364 of Lecture Notes in Mathematics

    Chewi, S., Niles-Weed, J., and Rigollet, P. (2025).Statistical Optimal Transport, volume 2364 of Lecture Notes in Mathematics. Springer Cham. École d’Été de Probabilités de Saint-Flour XLIX –

    work page 2025

  7. [430]

    Pavliotis, G. A. (2016).Stochastic Processes and Applications. Springer. Pavon, M., Trigila, G., and Tabak, E. G. (2021). The data-driven Schrödinger bridge.Communica- tions on Pure and Applied Mathematics, 74(7):1545–1573. Saffman, P. G. (1995).Vortex dynamics. Cambridge university press. Schiebinger, G., Shu, J., Tabaka, M., Cleary, B., Subramanian, V ....

    work page 2016

  8. [1995]

    −γ (Y−y 0)rv (Y−y 0)2 + (Y−y 0)2 1−exp p (Y−y 0)2 + (Y−y 0)2 rv !! +d X−x 0,d rd # dt+σdW x dY=

    — with a constant divergence field. The Lamb-Oseen component captures the swirling, rotational dynam- ics typical of a vortex, while the divergence field is added to account for vertical motion or non- conservative forces that may cause a net expansion or contraction of the flow. In other words, this combined model enables us to represent both the core vo...

    work page 2024

Show all 12 references
  1. [2001]

    We then generate particles following Shen et al. (2025). That is, we took the velocity field in a region where a vortex is observed in June 1st 2024 at 5pm. We then select an initial location near the vortex and uniformly sample 4,400 initial positions within a small radius (0...

    2025

  2. [2019]

    Cohen, S., Arbel, M., and Deisenroth, M. P. (2020). Estimating barycenters of measures in high dimensions.arXiv preprint arXiv:2007.07105. De Bortoli, V ., Thornton, J., Heng, J., and Doucet, A. (2021). Diffusion Schrödinger bridge with applications to score-based generative m...

    2020 arXiv

  3. [2020]

    δ(t−t ′)E y∼f(y|t) y′∼f(y|t ′) [Ky(y,y ′)] # = X t,t′∈T

    combines dynamic optimal transport with continuous normalizing flows (CNFs) to generate continuous-time, nonlinear trajectories from snapshot data. These flows are governed by ODEs rather than SDEs, and incorporate regularization that encourages short, energy-efficient paths. ...

    2024

  4. [2024]

    (2025) proposed generaliza- tions using higher-order derivatives, although these approaches are computationally feasible only in low-dimensional settings

    incorporates momentum into the particles to exploit local information and model multi-marginal dynamics, while Hong et al. (2025) proposed generaliza- tions using higher-order derivatives, although these approaches are computationally feasible only in low-dimensional settings....

    2025

Pith tools

Reviewed August 7, 2026 · model on record in the stance chip above.