Multi-Marginal Stochastic Flow Matching for High-Dimensional Snapshot Data at Irregular Time Points
Authors: Justin Lee
Behnaz Moradijamei
Heman Shakeri
Venue: ICML 2025 poster (2025)
Keywords: Stochastic Flow Matching
Multi-Marginal Optimal Transport
Irregular Time Points
Snapshot Data
Simulation-Free Methods
Score Matching
Generative Models
Single-Cell Data
Primary area: applications->health_medicine
TL;DR: Modeling the evolution of high-dimensional systems from irregular timed sampled snapshots using overlapping windows of marginal distributions and measure-valued splines.

Abstract: Modeling the evolution of high-dimensional systems from limited snapshot observations at irregular time points poses a significant challenge in quantitative biology and related fields. Traditional approaches often rely on dimensionality reduction techniques, which can oversimplify the dynamics and fail to capture critical transient behaviors in non-equilibrium systems. We present Multi-Marginal Stochastic Flow Matching (MMSFM), a novel extension of simulation-free score and flow matching methods to the multi-marginal setting, enabling the alignment of high-dimensional data measured at non-equidistant time points without reducing dimensionality. The use of measure-valued splines enhances robustness to irregular snapshot timing, and score matching prevents overfitting in high-dimensional spaces. 
We validate our framework on several synthetic and benchmark datasets, including gene expression data collected at uneven time points and an image progression task, demonstrating the method's versatility.

Multi-Marginal Stochastic Flow Matching for High-Dimensional Snapshot Data at Irregular Time Points
Justin Lee'
Behnaz Moradijamei Heman Shakeri
Abstract
both deterministic and random factors. Developing gener- ative models that accurately represent these dynamics is crucial for simulating cellular behavior and predicting re- sponses in unseen cells under non-steady state dynamics imposed by perturbation agents such as drug stimuli.
Modeling the evolution of high-dimensional sys- tems from limited snapshot observations at irreg- ular time points poses significant challenge in quantitative biology and related fields. Traditional approaches often rely on dimensionality reduction techniques, which can oversimplify the dynam- ies and fail to capture critical transient behaviors in non-equilibrium systems. We present Multi- Marginal Stochastic Flow Matching (MMSFM), a novel extension of simulation-free score and flow matching methods to the multi-marginal setting, enabling the alignment of high-dimensional data measured non-equidistant time points without reducing dimensionality. The use of measure- valued splines enhances robustness to irregular snapshot timing, and score matching prevents overfitting in high-dimensional spaces. We val- idate our framework on several synthetic and benchmark datasets, including gene expression data collected uneven time points and an im- age progression task, demonstrating the method's versatility.'
A common approach to modeling such systems through stochastic differential equations (SDEs), particularly the Langevin equation an Itô SDE (Gardiner, 1985; Risken & Risken, 1996); the evolution of the cellular state X(t) E R can be described by
dx uz dt + g(t) dt)
(I)
where (x) is the drift term representing deterministic dy- namics, gt) is the diffusion coefficient capturing stochas- tic fluctuations, and W(t) is a Wiener process modeling random noise. the population level, the corresponding probability density function p(a,, t) evolves according the Lokker-Plance equation (Risken & Risken, 1996):
@p((z) -V (p:((( u.())) g2) aL
Ap: (2)
where pr shorthand for p(z, t), V. denotes the diver- gence operator, and A is the Laplacian operator.
1. Introduction
In practice we only observe the system through snapshot measurements at discrete, possibly irregular time points to < t < ... < tM providing samples from the marginal distributions = Pt, (z) (Schofield et al., 2023). Therefore, we lack trajectory data that would reveal how individual states evolve between these snapshots due to the destructive nature of single-cell measurements. This raises fundamen- tal question: among the infinitely many stochastic processes that could connect these observed marginals (Weinreb et al., 2018), which one is the most likely?
Understanding cellular responses to perturbations is fun- damental challenge in quantitative biology, with significant implications for fields such as developmental biology, can- research, and drug discovery (Altschuler & Wu, 2010; Saeys et al., 2016). Modeling these responses requires cap- turing complex stochastic dynamics in nigh-dimensional cellular states that evolve over time under the influence of
School of Data Science, University of Virginia, Charlottesville VA. USA. "James Madison University, Harrisonburg VA, USA. Correspondence to: Justin Lee <jgh2xh@ virginia.edu>, Heman Shakeri < virginia.edu...
1.1. Least Action Principle
To address this problem, we turn to the theory of Optimal Transport (OT) (Villani, 2009) which seeks the most ef- ficient way to transform one probability distribution into another. In the simplest case of two marginals Pu and Pr. OT aims to find a transport map T that minimizes the cost
Proceedings of the 4rnd International Conference on Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 hy the author(s).
'Code available Shakeri
httpp://giitee com/

---

Multi-Marginal Stochastic Flow Matching for High-Dimennnn Snapshot Data at Irregular Time Points
functional:
leads us to formulate the problem as finding the drift u. (1) that minimizes the cumulative transport cost and provides the smoothest and most efficient flow connecting the ob- served distributions over time, while ensuring robustness against overfitting. In summary, we require:
min / ((x)]
subject to To = pl. (3)
where Po denotes the pushforward of po under T. Kan- torovich's 's generalized formulation of (3) is a linear program-
ming problem over the set of joint probability distributions, leading to the definition of the Wasserstein-2 distance:
Scalability in High Dimensions: While directly solving high-dimensional transport problems is computation- ally challenging (Benamou Brenier, 2000; Peyré Cuturi, 2019), our approach efficiently approximates the solution. By leveraging advances in simulation-fri score and flow matching methods, we model the high- dimensional stochastic process directly in the ambient space, avoiding dimensionality reduction strategies that could obscure important dynamical features.
min / [|x - vl12 dr dr(x, v), (4)
(200 pI) = =EII(P..P))
where [p P1) is the set of joint distributions with marginals Po and f1- While OT provides a deterministic model based on the principle of least action-finding the shortest path or geodesic in the space of probability distribu- tions-it does not account for the inherent stochasticity of biological systems (Horowitz & Gingrich, 2020). Cells are subject to both extrinsic noise, such as variations in initial conditions and environmental inputs (Hilfinger& & Paulsson, 2011), and intrinsic noise arising from the thermodynamic uncertainty in biochemical reactions (Mitchell & Hoffmann, 2018).
Robustness Against Overfitting: By minimizing the total transport cost across all time intervals, we intro- duce only essential adjustments to match the observed marginals, preventing the model from overfitting to limited observations and ensuring that the inferred dy- namies generalize well beyond the training data.
Insensitivity the Timing of Snapshots: The formu- lation inherently accommodates arbitrary and irregu- lar time points 4, making it robust to the choice of measurement times. By focusing on the minimal ac- tion path that passes through the observed marginals, we capture the system's 's evolution without being con- strained by the timing of data collection.
To incorporate stochasticity and identify the most likely stochastic process connecting the observed marginals, we consider the entropic-reegularized optimal transport prob- lem, particular case of the Schrödinger Bridge Prob- lem (SBP) (Schrödinger, 1931; Léonard, 2014). The SBP seeks the stochastic process that minimally deviates from a matching the
prior-typically a Brownian motion-wiie observed marginals. can be considered a general statistical inference and model improvement methodology in which one updates the probability of hypothesis based the most recent observations while making the fewest possible assumptions beyond the available information (Pavon et al., 2021). This approach aligns with Occam's 's razor principle and aims to find the simplest stochastic process that explains the data with minimal adjustment to our prior belief.
1.2. Literature Review
Direct learning of the sigh-dimensional partial differential equation (2) is computationally prohibitive due the com- plexity of integration and divergence computations in high- dimensional spaces (Benamou & Brenier, 2000; Peyré & Cuturi, 2019). Hence, current approaches typically consider reduced-dimmensionll data representations with gradient- based drifts originating from developmental biology (Wein- reb et al., 2018; Schiebinger et al., 2019) where the focus is primarily on slow time scales and the assumption of low- dimensional manifold dynamics is often useful. In this con- text, dimensionality reduction tools such as t-SNE (Van der Maaten & Hinton, 2008), UMAP (McInnes et al., 2018), and PHATE (Moon et al., 2019) are extensively used to sim- plify the modeling. However, these techniques can obscure critical faster-scale dynamical information, introduce arti- facts (Kiselev et al., 2019), and result in the loss of important biological information in the reduced, folded space.
Extension to Multiple Marginals: Multi- MartinMarginll Opti- mal Transport (MMOT) is natural extension of OT aiming [[(o, Pl: PM.) where IT
to find joint distribution is the set of all joint distributions with marginals {p. Ar minimizing W Pj) for all i # j (Pass, 2015). Of these t, we are interested in the special case where there exists a total order on the time labels t; associated with each Pi- We henceforth use MMOT to refer to this ordered MMOT case. Instances referring to the fully joint MMOT will be explicitly referred to as such.
Extending our least action principle to the MMOT case with arbitrary time points to, fI...., t.. we pose the same question: among all possible stochastic processes that could connect the observed marginal distributions MoM which one is the most probable given our prior knowledge? This
Neural Ordinary Differential Equations (Neural ODEs) have emerged as a powerful tool for modeling continuous- time dynamics and connecting probability measures over time (Chen et al., 2018a). This approach offers an alternative method by parameterizing the time derivative of the hidden
2

---

Multi-Marginal Stochastic Flow Matching for High- Nligh-Dimensional Snapshot Data at Irregular Time Points
state with a neural network, which is trained to approximate the drift term in the Fokker-Planck equation (2). While this method has been successfully applied (Tong et al., 2020; Huguet et al., 2022) for modeling cellular dynamics and trajectory inference, it still operates primarily in reduced- dimensional spaces.
Section 1.1 to ensure robust learning and prevent overfitting.
Our key innovation lies in learning continuous spline measures through overlapping windows of consecutive marginals during training. Specifically, we process over- lapping triplets (Pi. Ai+1, Pi+2) in a rolling fashion, where we demonstrate in Section 2.3 that a window size of two strikes an optimal balance between enforcing smoothness constraints and computational efficiency. This approach enables to capture local dynamics across uneven time in- tervals, maintain consistency between overlapping windows, and generate intermediate states between observed snap- shots, effectively creating a "motion picture" of the system's 's evolution. The overlapping nature of these learned flows ensures robustness against the specific choice of measure- ment times while preserving the high-dimensional structure of the data.
Recent multi-marginn approaches have attempted han- dle multiple time points simultaneously. Chen et al. (2024) developed deep multi-margiinn momentum Schrödinger bridge approach that, while capable of working in high di- mensions, reguires expensive flow integration and memory- intensive caching of trajectories during training. Similarly, Albergo et al. (2024) proposed stochastic interpolants for multi-marginal modeling but still relies on ODE/SDE in- tegration and marginal distributions as supervision signals, which becomes computationally challenging in high dimen- sions. These approaches share common limitations: they either require dimension reduction to handle computational complexity, or they depend on expensive numerical integra- tion and trajectory generation during training.
Although we were unaware during the development of our method, an important concurrent work by Rohbeck et al. (2025) also considers a very similar approach to ours. We include a more in-depth discussion regarding similarities and differences between our approaches in Section 2.3.3.
Alternative approaches using generative models attempt to transform simple distribution an arbitrary target distri- bution. Variational Autoencoders (VAEs) (Kingma et al., 2013) learn an encoder-decoder pair, q(= x) and p(x: 2). such that the decoder can generate x ~ P1 given samples ~ Pa- Generative Adversarial Networks (GANs) (Good- fellow et al., 2014) employ a generator-discriminaat frame- work, where the generator (zz) produces x G(z) with ~ P1 for z ~ Pa. While successful, these methods are limited by the simplicity of the source distribution Po. often chosen to be uniform or normal for analytical convenience. Moreover, these models represent static transformations with no notion of time and cannot generate intermediate states at arbitrary time points, making them unsuitable for modeling dynamic processes where temporal evolution is crucial.
2. Problem Formulation and Methodology
Formally, let 0 = to < t1 < ... < fM = -1 denote a sequence of normalized time points in [0, 1]. and let p be the continuous probability distribution of the system state at time I; in Euclidean space Rdd Our data consists of snapshot samples X,, = {xx(1 ~ p at these time points. The goal is to learn continuous probability path pa (xr) for t E [0, 1] satisfying Pt, p; for all i, describing the evolution of the system over time.
N;
2.1. Dynamic formulation of the Wasserstein distance and Wasserstein Splines
Benamou & Brenier (2000) introduced dynamic formula- tion of the Wasserstein distance, connecting OT with fluid dynamics:
Although diffusion models (Ho et al., 2020; Song & Er- mon, 2019) incorporate a time component by learning denoising Markovian reverse process, their notion of "time" corresponds to a noise schedule rather than physical time. This limitation prevents them from capturing actual tempo- ral dynamies or generating data at arbitrary time points not specified during training.
S [|u(())] pe(x))pp())))))))))< do dt inf Pr.
W:(,vv)-
DE (peua) Po pp, P1
While the fully joint MMOT extends this framework to mul- tiple distributions, computing the MMOT plans becomes computationally challenging in high dimensions. Prior work (Chen et al., 2018b; Benamou et al., 2019) examined the formulation
Our Approach We introduce Multi- Multi-Marginal Stochas- tic Flow Matching (MMSFM) to address these limitations by adapting recent developments in simulation- ffre ap- proaches (Lipman et al., 2023; Tong et al., 2024b) to our setting. These methods learn directly in the ambient space without dimensionality reduction or explicit simula- tion. However, their direct application to our multi- marginal setting requires careful adaptation to principles described in
[
inf N
X
dt,
(5)
termed P-splines by Chewi et al. (2021). Unfortunately, this does not fit our needs because X, here is considered to be
3

---

Multi-Marginal Stochastic Flow Matching for High- Highh-Dimensional Snapshot Data at Irregular Time Points
stochastic process whereas we need a deterministic flow. Moreover, these formulations are still quite computationally expensive given that we need to solve this problem within the training loop. Instead, Chewi et al. (2021) proposed transport splines as a method to efficiently obtain deter- ministic maps that smoothly interpolate between multiple distributions. The key idea to sample points from the dis- tributions P. and apply a Euclidean interpolation algorithm between these points. The specific spline algorithm is left as design choice for the user. Options include the natural cubic spline interpolation which minimizes the integral of the squared acceleration
where (x( is the deterministic component, and V log (x) is the score function of the density (x()). This observation allows us to decouple the learning of the de- terministic drift uf (I) and the score function V log pt (z). Therefore, specifying uf (1) and V log pe pa(x) is sufficient to define the SDE drift (z).. Tong et al. (2024b) proposed the unconditional score and flow matching objective:
((0) -E.-11(011) [[lu((t: 0) (x)]] +
(10)
X(1) 11) V Vlogpa
where (a;; 0) and s,(2; 0) are neural networks approximat- ing the drift and score functions, respectively, and X(t) is a weighting function. However, pc( is unknown and thus di- rectly computing ui (zz) and V log Pe(x) is challenging. To overcome this, Tong et al. (2024b) proposed a conditional formulation of the loss function:
[' [1Aa1l]] dr. Th
inf
(6)
where Yt denotes a curve in space, and the cubic Hermite spline (Hermite & Borchardt, 1878) which represents each interval (z); 21+1) as the third-degree polynomial
(1) =22)² - 312 + 1)x; + (13 - 21 + t)x's
(0) (o.00000)..000)).ppp [1ue(2;))) [
+ (-133 + 36 )o.+1 + -1--11 - P)l'+1
((s(2)) 0) - [logp((<12)11
where To, Iy are the boundary constraints, and z'g- 2' are the derivatives w.r.t. time at those points.
(11)
where z represents conditioning variables, and pe (x|2). In this conditional framework, uff (ar|z) and V log (x)) can be computed analytically estimated empirically based on the conditional distribution pa (2)22) We can reconstruct the learned SDE drift using:
In practice we adopt transport splines using compositions of probabilistic OT plans in place of deterministic OT maps. Let the joint distribution ı be the MMOT plan over (2to Ft,, ...., Ft ). Although the true MMOT plan involves all pairs of marginals, we are interested in the case where there is temporal ordering. This structure allows us to take advantage of a first-order Markov approximation and decompose π into conditional plans
ue (e - =x;;;;)) + g (x;
(12)
and integrate it with given initial conditions zo infer the trajectories that develop from those initial conditions.
M 7 ( se πt (3to Te, II T(E,, ITn.));
(8)
2.3. Learning Overlapping Mini-Flows for Multi-Marginal Data
where π(It, | St. ,) specifies how to transport a point It.-1 ~ Pi-1 to the distribution Pt. We obtain the con- ditional plan as T( v) (xx,))/((x) where p(y) = (2. d) is the marginal distribution of y. By apply- ing the transport spline procedure to batches of vectors (M the conditional plans act alignment operators, allowing us to construct Euclidean splines through optimally coupled points (X;)) M
We aim to train an ODE drift network (z O) and a score network sz( f) to learn an overall flow based on the mini- flows on overlapping (k + 1)-tuples (p:. Pe+1...., Ps+k) for
= 0, 1, M 1 in a rolling window fashion. Because transport splines are ultimately just approximations for the true MMOT, the rolling windows provide a variation of perturbations in the approximated error from any single geodesic spline segment estimate. See Figure for visual representation of the variation of paths in an interval.
2.2. Simulation- Free Score and Flow Matching
We aim to model the stochastic process bridging the multi- ple distributions Pt by learning the underlying dynamics of the system in Equation (1) and the associated Fokker-Plane equation (2). Tong et al. (2024b) introduced a reparameteri-
Theorem 2.1. The gradient of the loss for single interval (ti, f++1) with k overlapping mini-flows is given by
zation of the drift ux(x) as
Volusti =V
ve(=::)
Σ ue( 24
ux(z) ui (z) + (t) V log pr (z).
(9)
+ (1: - (2;
2
4

---

Multi-Marginal Stochastic Flow Matching for High- Biggh-Dimnnional Snapshot Data at Irregular Time Points
where the expectation is laken over q( (erite))) ((1)) and p(æ (tseattten)) 1 E Pc
We refer the reader to Appendix C.I for implementation details.
34 34 1i5--1sss s1x).
In the original source-target distribution pair setting, we sam- ple t ~ (4(0,1) 1). To accommodate our mini- floi-flow method, we could sample t ~ [14,,, 1,+*) for the ith mini-flow. How- ever, this approach is ineffective for training uneven time in- tervals for example f,+1 -t, < f,+2 1142--440 -t:1--1+11100-- -1+1-leading to in- sufficient sampling from the smaller interval. To handle this, we adopt stratified sampling strategy, sampling an equal number of time points from ((42 (i+1). (11++1, l1+2), and so on to ensure balanced training across intervals. Specifi- cally, for a total batch size of b, we sample b/ time points from each interval.
Geometrically, the neural network learns the direction of the overall movement given by the sum of the de vectors. The magnitude of this movement is scaled down by the competing objective (k - ( 0)1?
Corollary 2.2. If all the mini-flow regression signals ur(21 | s s1.) are equal, then the single interval loss recovers the CFM loss (Theorem 2.1 of Lipman et al. (2023)) up to a scalar factor.
We provide proofs in Appendix A.
2.3.2. MINI-FLOW ODE AND SCORE REGRESSION TARGETS
Our method handles overlapping trajectories where (1) * u., (x) for fixed x and & * 1j, accommodat- ing the possibility that trajectories may cross over a point at different times in marginla settings. In practice, we train using mini-batcher of size b.
Theorem 3 of Lipman et al. (2023) and Theorem 2.1 of Tong et al. (2024a) derive the ODE flow regression target for conditional Gaussian probability path Pe 2) = N(x | Pt, af) as
By incorporating stochasticity through score matching, we improve robustness and avoid overfitting in high- dimensional spaces. The score V, log lolloop(()) (x]()) allows the model capture the inherent uncertainty and variabil- ity in the data. Using the identity V. log (x)l)) (apt(zl<)/p.((xz)) we see that the score nudges predic- tions towards more likely regions, thereby implicitly ex- ploring the region around the local per-sample flow. This efficiency allows us to remain in the ambient dimension d and sidestep dimensionality reduction strategies which of- ten introduce information loss and additional complexities into the flow dynamics. Moreover, re-projecting the trajec- tories back into the ambient space introduces undesirable reconstruction artifacts.
mc) =)= at -x- 14
ug (x
(13)
where μ4 and or are respectively the time-varying mean and standard deviation of the flow conditioned 2. The prime notation (') denotes differentiation w.r.t. time t. We set ur Pai+& (t) for a transport spline (+s+*(t) [t,, te+al + Rd constructed through the points in 2 via Euclidean spline interpolation.
For or we consider the case of Brownian bridges with con- stant diffusion g(t) = a and set o, o //t(11 t) along the global time t [0, 1]. Alternatively, we can set oe based on the Brownian bridge of the mini-flow from a = L to b = titk. reparameterizing as a; = a r(f)), where r(t) = t-a. In this case, the derivative of must take into account the reparameterization, yielding oj d0f dr da, 1. Because μt, ot. M4: of can be expressed an- dt de b-o alytically, we can directly compute these quantities and effi- ciently compute the regression target ui using Equation (13). Given our Gaussian probability path, we can easily derive the score regression target as V log (x 2) alternatively -t. for e ~ I0,
1.3.1. TRANSPORT SPLINES SAMPLING OF Z AND STRATIFIED SAMPLING OF L
We sample z from a MMOT plan π using transport splines by first drawing samples Xt,, Xta..... X ~ Pis Pi+1:-.., Pi+k. where each Xx, is batch of i.i.d. sam- ples from P:- Then, compute the MMOT plan given by the first-order Markov approximation (8):
or
i+t-k II jmi+2
T T. (r Rs ( (t.n))
( Bey .).
We summarize our method in Algorithm 1. Once we have the trained networks ve (x 0), (x 0), we can construct the SDE drift ue (2; () using (12), and generate trajectories from given initial conditions zo by an SDE integration with drift uu (z; 0) and diffusion Ø.
The initial plan (Ht,, ( is a standard OT plan w.n.t. the squared Euclidean distance || Zt.eil 1|2 as the cost function. Next, we compute the conditional map ( 1 using
2.3.3. WINDOW SIZE k AND SPLINE ALGORITHM
T(T,,| 1 = 7 a(5e, 11 1531e)) = T(e 11 Iei))
We choose our window size k = 2 based on the properties of our chosen Fuclidean spline algorithm and considerations
x (
J T ( 1 Iee)) daze,
5

---

Multi-Marginal Stochastic Flow Matching for Hligh-iimnnsionn Snapshot Data at Irreguiar Time Points
t3
Linear Piecewise
-2
(2.
Itz: ta. tx)
t1
(5)
It5. It)
-t
Figure 1. Comparison of Euclidean splines on overlapping windows of size k 2, demonstrating the potential for overlapping windows capture variations of paths though the same intervals. From left right: The points interpolate with time lahels fo,..., .,, fg- Natural cubic splines on equidistant time intervals. 3) Monotonic cubic Hermite splines on equidistant time intervals. 4) Natural cubic splines on arbitrary time intervals T (0, 0.05, 0.2, 0.27, 0.86, 0.95, 1). Note the overshooting required to satisfy the continuity of S at ta and 44. 5) Monotonic cubic Hermite splines on arbitrary time intervals.
Algorithm 1 MMSFM Training
more direct path such as from a linear piecewise interpola- tion. Third, while using a larger window can potentially fit spline closer to the linear piecewise interpolation, allowing for smaller windows can better capture a wider variation of paths, increasing the robustness of the learned flow. More- over, 3 control points are sufficient to learn a curvature at the interior control point. Fourth, the specific coefficients describing each piecewise cubic polynomial are efficient to compute, scaling linearly in O(k) with the k + points to interpolate. For window size k and M time points, the overlapping window routine computes (M k)}: splines resulting in a total complexity of O(M - -))a). This is "maximized" when and "minimized" when k plexity of O(M). As added bonus, monotonic cubic Hermite splines are highly insensitive to control points that are not immediate neighbors. Thus, choosing larger win- dow size k: >> 2 does not meaningfully increase the amount of information captured by the spline.
Input: Training data,
Initialize networks U (1;; 0), (2;
Set g(t) 4
Set 0r < 0 / - 1) V
20
Set A(t) g()}2
while training do
for i = 0 to M - k do
Sample mini-batches:
; Xtuss ~ Pi. Pitk
Compute OT plans:
T(x,, 11,4),, 7 (X1142
X ( )....
M/2 for complexity of (M2),
Generate (Xi (Xi......
1 using
k M 1 for com-
Compute Hr (t)+
Set (1 2) <= Me, of)
Sample times é from [tis ti+k]
Sample x ~ (x
of (x u4) + 14
Compute ( 2)
Rohbeck et (2025) considers a similar approach arriving at the same ordered MMOT plan approximation, detailed in their Appendix B. Moreover, Rohbeck al. also splines as the interpolation method allowing for irregular time points, but differ in their selection of natural cubic splines over all the time points (k: = M 1). This contrasts with our choice to use monotonic cubic Hermite splines on rolling windows over = 2. Although natural cubic splines may be more analytically tractable, we reiterate that the monotonic cubic Hermite spline nonetheless provides monotonicity guarantee within each dimension and thus avoids overshoot. This property is especially relevant when exploring the sensitivity of our method to highly irregular time points. We attribute the overshooting behavior in nat- ural cubic splines to the presence neighboring intervals with vastly different lengths. For example, consider the time point sequence 0.2 0.27 0.86 in Figure 1. The two corresponding intervals have length 0.07 and 0.59, nearly 9-fold difference. For the natural cubic spline, any change in
Compute Voog Pe (a 2)
M4 o²
Compute C()) from (11)
Update 9 using (())
end for
end while
Output: Trained networks vz(x; sx
to the running time. We opt to use monotonic cubic Hermite splines instead of natural cubic splines for four main reasons. First, the guaranteed monotonicity of each piecewise cubic polynomial ensures no overshoot within each dimension, thus removing overshooting when computing the ODE flow regression target (13). Second, monotonic cubic Hermite splines by construction do not necessarily have contin- uous second derivative. While this is a desired property for smoother curves (and in fact enforced for natural cubic splines), this condition can restrict the curve from taking
6

---

Multi-Marginal Stochastic Flow Matching for Hligh-Dimensional Snapshot Data at Irregular Time Points
velocity and acceleration along the short interval can happen relatively quickly, but the corresponding change for the long interval must be drawn out. In particular, the continuity of the acceleration constraint does not help in this regard as it prevents the spline from instantaneously re-adjusting veloci- ties as necessary. We include a more in-depth discussion of splines in Appendix B.I.
Table I. Comparison of the inferred distributions generated by MIOFlow and our method using Pairwise and Triplet flows at the held-out time point. For equidistant time points Ti. we hold out is 0.83 and 44 0.67 respectively for the S-shaped and a-shaped data. We do the same for the arbitrary time points T2. holding out ts 0.85 and t4 0.54. We also examine distance metrics averaged across all time points for Ta. From DynGen, we hold out 4 0.25.
3. Results
S-SHAPE (HULD (s.
O-SHAPE (HOLD PATR
MIOFLOW
PAIR
TRIP
MTDFLOW
TRIP
We briefly describe our data and setup below, and also in- clude a more detailed experimental setup description in
W W
2:36
21.54 14.54
3.78
4.54
66.91
5.87
3.86
464.36
14.56
21,06
MMDİG
7.26
2.29
1.47
7.65
3.96
4,26
Appendix D. We summarize our results in Tables and 2.
MMD(M)
66.19
5.24
3.11
463 66
13.92
20.01
W
9.42
2.12
1.62
5.04
8.08
3.79
T2 W
76.82
14.73
3.1. Experimental Setup
MMDIM)
MMDAGI
4.12
2:36
6.45
64.81
14.07
S-SHAPE (ALL)
CHSHAPP (ALL)
We applied our rolling window framework to three synthetic datasets, two RNA gene expression datasets, and an im- age classification dataset. From our framework we use the k 1 (Pairwise, equivalent to SF2M (Tong et al., 2024b)) and k 2 (Triplet) mini-flow settings. We approximate the MMOT plan with transport splines computed on mini- batch OT given the smaller computation cost and asymp- totic convergence properties (Fatras et al., 2020; 2021). We additionally use MIOFlow (Huguet et al., 2022) on the syn- thetic datasets to examine the difference in performance for ambient-space and latent-space models. The initial condi- tions for the generated trajectories are from a held-out set of samples from the source distribution Po. Evaluations all synthetic and RNA gene expression datasets are com- puted by leaving out time point marginal during training and calculating the Wasserstein metrics Wi and W3 (us- ing Euclidean distance as the cost function), the maximum mean discrepancy with mixture kernel (MMD(M)), and the maximum mean discrepancy using a Gaussian kernel (MMD(G) at the left-out time point. For the image dataset, we train using all time point marginals and examine the training stability and loss over epochs.
W W
12,06
3.71
33.95
186.68
257.86
35.01
2400.15
2,5125
MMDÍGI
4.12
2.12
4.62
1.35
MMD(M)
241.01
2168.01
7.76u5
DYNGEN (1OLLD 11)
Wi
0.85
0.74
W
0:98
0..63
0.82
MMIRGI
0:53
0.38
0.22
MMD(M)
0.51
0.19
0.10
(0, 0.27, 0.3, 0.88, 0.98, 1) with neighboring small and large intervals. We introduce To as highly irregular time point set with very small interval of 0.03 followed by a disproportionately large interval of 0.58 in 0.27 + 0.3 + 0.88. Also included is a second miniscule gap of 0.02 in the last interval 0.98 - 1. The DynGen dataset has marginal distributions on equidistant time points T (0, 0.25, 0.75, 1).
We consider gene expression data from the Multiome and CITEseq datasets published as part of a NeurIPS competition (Burkhardt et al., 2022). Measure- ments are taken at T (2,3,4, 7) days, which again nor- malize to T = (0, 0.2, 1). We follow the procedure in (Tong et al., 2024b) and preprocess the data into the first 50 and 100 principal components, along with the top 1000 highly variable genes (Satija et al., 2015; Stuart et al., 2019; Zheng et al., 2017).
Real Data:
Synthetic Data: Our three synthetic datasets are the S- shaped Gaussians, the a-shaped Gaussians, and syn- thetic scRNA dataset generated by the DynGen simula- tor (Cannoodt et al., 2021) which we repurpose from Huguet et al. (2022). The S and a-shaped Gaussians both con- sist of 7 marginal distributions in ?R2 We select these three datasets because S-shaped Gaussians involve learn- ing a flow with changing curvature, the a-shaped Gaus- sians have a cross-over point for some x where the flow un (2) # ue, (x) and i # j. and the DynGen dataset intro- duces a bifurcation. We evaluate both Gaussian datasets on three different time point labels: equidistant time points Ti = (0, 0.17, 0.33, 0.5, 0.67, 0.83, 1), a first set of arbi- trary time points Ta = (0, 00,0.08.00 0.38, 0.42, 0.54, 0.85, 1). and a second set of arbitrary time points Ts =
Finally, we consider a generative perspective and look at the performance of our model on image progression through var- ious classes from the Imagenette dataset (Jeremy Howard), subset of 10 easily classified classes from the ImageNet dataset (Deng et al., 2009). Specifically, we look at the progression: gas pump to golf ball to parachute. For the Pairwise model we set T = (0, 0.25, 1), and for the Triplet model we set T (0. 0.5. 1).
7

---

Multi-Marginal Stochastic Flow Matching for Highh Biggh-Dimensionll Snapshot Data at Irregular Time Points
Figure 2.
Example trajectories for 32 32 pixel image progression through the Imagenette classes (gas pump + golf ball + parachute). Results are generated using our Triplet model with an equidistant time scheme.
between separate branches of the bifurcated flow. We did not explore methods to mitigate this problem and believe this to be an avenue for future work.
Table Comparison of the Pairwise and Triplet methods on the CITEseq and Multiome gene expression datasets. We hold out 0.4 for both datasets.
Remaining in the high-dimensional ambient space without aggressive dimensionality reduction preserves important bi- ological information, leading to more biologically plausible
PCA 50
PCA
100
Hi-VAR
1O00
FAIR
TRIP
FAIR
TRIP
PAIR
TRIP
W
54.18
$3.98
62.85
62,08
50.64
50.71
trajectories. The ability to generate samples arbitrary time points allows us to explore the system's 's behavior beyond
CITESCO
W
3027,28
4036.41
3942,08
2579.84
MMDiG)
MMD(M)
339.20
344.89
345,09
331.72
48.53
49.83
the observed data, potentially identifying critical time win-
W;
61.79
60.92
70.72
70.39
36.15
56.10
MULTI
W
3918.50
3806.89
5037 07
5029.56
3166.01
3160.84
dows where intervention might be most effective. This has
MMD(M)
MMDRG
793.34
0.30
705.21
11,27
656.86
0:25
621.32
0.23
40.71
0.04
40.29
0.04
implications for understanding drug resistance mechanisms
and designing more effective therapeutic strategies.
In all instances, MIOFlow generated idiosyncratic trajec- tories which matched the marginals at the specified time points but performed poorly between those time points. We believe this to be the case because MIOFlow operates in the embedding space generated by a GAB. This structure works very well for trajectories in the embedding space but poses problem when reconstructing the trajectories in the ambi- ent space. The GAE is only trained on the data marginals (p:) 2if at times (4, = which means that data points not specified in the data are effectively the GAE. These duitrf-distribution points arise naturally from generating trajectories which spend time traveling be- tween the data distributions. In addition, we notice that all the reconstructed points exhibit high bias and low variance, tending be bunched very close to each other. This quality perhaps captures the first moment well but not any higher moments.
3.2. Discussion
Learned flows are visualized in Appendix F. Our method consistently outperformed MIOFlow on the interpolation at the held-out time point for the synthetic data. Interestingly, the Pairwise model slightly outperformed the Triplet model for the a-shaped Gaussians on Tı. We believe that in this specific instance, the masked time point corresponded to an interval where the momentum from the prior interval was enough for the Pairwise model to infer the held-out marginal. In contrast, the a-shaped Gaussians on T2 show the Triplet model outperformed the Pairwise model by a significant margin; even MIOFlow generally outperformed the Pairwise model in this instance. This suggests that the Triplet method is more effective for non-equidistant time snapshots especially when capturing complex temporal dy- namics because the variation of flows provided by splines in overlapping windows helps learn the held-out marginal. The success of our methods on To demonstrates the robustness and stability of our approach even when handling arbitrary time points. Looking at the trajectory plots, we can also con- firm that our method able to handle datasets with varying flow curvatures and flow cross-overs.
w.r.t.
We validate our Triplet model's ability to learn flows on high dimensional, noisy data, taken at irregular time points by comparing results from the Pairwise and Triplet models on the CITEseq and Multiome datasets. These two datasets contain high dimensional samples with noise inherent to bio- logical measurements and measurements from non-uniform time intervals. We see the Triplet model successfully out- perform the Pairwise model on inferring the distribution at the held out time point even in these conditions.
The bifurcating flow of DynGen posed challenge for our models: while they outperformed MIOFlow on the metrics, but struggled to handle the bifurcating trajectories. We suspect this behavior to stem from bani- OT because it does not enforce a consistency constraint on the sampling process, resulting in cases where particles are able to jump
On Imagenette, we examine our Triplet model's 's perfor- mance against that of the Pairwise model both by the train-
8

---

Multi-Marginal Stochastic Flow Matching for Hligh-Dimensional Snapshot Data at Irregular Time Points
mensional spaces using the CITEseq and Multiome datasets. Our application to the Imagenette dataset further demon- strates the method's generalizability to varied datasets as well as generative capabilities for sampling from a time point marginal.
Pairwise Triplet
0.5
0.4
0.3
0.2
2
The incorporation of stochasticity through score matching improves robustness and avoids overfitting, enabling the model to generalize new conditions. This work opens new avenues for both generative algorithms and modeling cellular responses to arbitrary, user-defined perturbations, providing a computationally efficient and biologically ac- curate framework capable of handling the complexities of high-dimensional, stochastic biological systems.
50 100
150
50
100
150
Epochs
Figure Loss plots for the Triplet and Pairwise models on the Imagenette dataset. The x-axis represents the epoch (1000 gradient steps). and the y-axis represents the mean loss value for that epoch. Left) Flow loss. Right) Score loss.
Acknowledgment
The authors Justin Lee and Heman Shakeri gratefully ac- knowledge the support provided by the University of Vir- ginia Cancer Center. Heman Shakeri additionally thanks Capital One for their generous support of this research.
ing loss curves and visually. In both cases, we used the model state after training for 160k gradient steps. The Triplet model performed better than the Pairwise model with respect to the flow loss and fairly even on the score matching loss. Both models were able to generate image trajectories which pass through the golf ball distribution, but the Pairwise model exhibited more difficulty than the Triplet model going to the parachute distribution. This suggests the Triplet model exhibited more stable training, further validating our approach.
Impact Statement
We have developed a method to robustly model the evolu- tion over time of a system of variables given snapshot data containing samples of the system state at various arbitrary time points, scalable in both the number of variables and number of snapshots. We apply a multi-marginal optimal transport plan order to couple samples from neighboring snapshots and use the couplings as control points to gener- ate splines overlapping time points which interpolate the system state. The combination of the optimal transport plan and the overlapping splines provide a more robust model less prone to overfitting whilst also minimizing the total work done to transform one snapshot to the next. One par- ticularly important application to help oncologists model cellular behavior and potentially accelerate the development of successful drug treatments.
Finally, we examine performance of the Pairwise and Triplet models on the S and a-shaped datasets using the highly unbalanced time points Ta. Here, we find that the Triplet model greatly outperforms the Pairwise model on the overall learned flow for the S-shaped dataset, but seemingly vastly underperforms for the a-shaped dataset. By taking a closer look at the visualizations of the learned flow, that in both cases, the short-long-short interval pattern poses a significant difficulty, suggesting that arbitrary time points do indeed increase the difficulty of the learning task. In the S-shaped case, the Pairwise model completely fails to learn this trajectory, whereas the Triplet model does better, learning to speed up and then slow down between t2 0.27 to fs 0.3, and ia 0.3 t4 0.88. Unfortunately, neither model was able to converge in the a-shaped case.
see
References
Albergo, M. S., Lindsey, M., Boffi, N. M., and Vanden- Eijnden, E. Multimarginal generative modeling with stochastic interpolants. In 12th International Conference on Learning Representations, ICLR 2024, 2024.
4. Conclusion
We present a novel framework for modeling the dynamics of high-dimensional systems from snapshot data in a multi- marginal setting with non-equidistant time points, while remaining in the high-dimensional space and avoiding the pitfalls of dimensionality reduction. By expanding the liter- ature of Conditional Flow Matching, we have developed a method that learns flows for overlapping triplets, enhancing robustness and stability in multi marrinal settings. We vali- date our method's scalability and ability to learn high di-
Altschuler, S. J. and Wu, L. F. Cellular heterogeneity: do differences make difference? Cell, 141(4):559-563, 2010.
Benamou, J.-D. and Brenier, Y. A computational fluid me- chanics solution the Longe-Kantorovicr mass transfer problem. Numerische Mathematik, 84(3):375-393, 2000.
Benamou, J.-D., Gallouët, T. O., and Vialard, F.-X, Second-

---

Multi-Marginal Stochastic Flow Matching for Hligh-Dimensionn Snapshot Data at Irregular Time Points
order models for optimal transport and cubic splines the Wasserstein space. Foundations of Computational Mathematics, 19:1113-1143, 2019.
Seguy, V., Sutherland, D. Tavenard, R., Tong, A., and Vayer, T. POT: Python Optimal Transport. Journal of Ma- chine Learning Research, 22(78):1-8, 2021. URL http:
/ //jmrr. wwwpg//appr...ee s v22/2222222222 201/200-455 20t1//2mm.............
Burkhardt, D. B., San Juan, B. P., Lock, J. G., Krish- naswamy, S., and Chaffer, C. L. Mapping phenotypic plasticity upon the cancer cell state landscape using mani- fold learning. Cancer Discovery, 12(8):1847-1859, 2022.
Fritsch, F. N. and Carlson, E. Monotone piecewise cubic interpolation. SIAM Journal on Numerical Analysis, 17 (2):238-246, 1980.
Cannoodt, R., Saelens, W., Deconinck, I.., and Saeys, Y. Spearheading future omies analyses using dyngen, a multi-modal simulator of single cells. Nature Communi- cations, 12(1):3942, 2021.
Gardiner, C. W. Handbook of stochastic methods, volume 3.
Springer Berlin, 1985.
Goodfellow, I., Pouget-Abadic J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. Generative adversarial nets. Advances in Nearal Informa- tion Processing Systems, 27, 2014.
Chen, R. T., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. Neural ordinary differential equations. Advances in Neural Information Processing Systems, 31, 2018a.
Harris, C. R., Millman, K. J., van der Walt, S. J., Gommers, R., Virtanen, P., Cournapeau, D., Wieser, E., Taylor, J., Berg. S., Smith, N. J., Kern, R., Picus, M., Hoyer, S., van Kerkwijk, M. H., Brett, M., Haldane, A., del Río, F., Wiebe, M., Peterson, P., Gérard- P., Sheppard. K., Reddy, T., Weckesser, W., Abbasi, H., Gohlke, C., and Oliphant, T. E. Array programming with NumPy. Nature, 585(7825)3357-362 September 2020. doi: 10. 1038//441586-020-26444-99 URL https: /tto::/dooi.oooo 11.1038/s4155 020---020--- 000-00200--2222 2649-22.
Chen, T,, Liu, G.-H., Tao, M., and Theodorou, E. Deep momentum multi marlinla Schrödinger bridge. Advances in Neural Information Processing Systems, 36, 2024.
Chen, Y., Conforti, G., and Georgiou, T. T. Measure-valued spline curves: An optimal transport viewpoint. SIAM Journal on Mathematical Analysis, 5046))5947-5568, 2018b.
Chewi, S., Clancy, J., Le Gouic, T., Rigollet, P., Stepamiants, G., and Stromme, A. Fast and smooth interpolation on Wasserstein space. In International Conference on Artifi- cial Intelligence and Statistics, pp. 3061-3069. PMLR, 2021.
Hermite, M. C. and Borchardt, M. d' d'interpolation de Lagrange. Journal für die reine und angewandte Mathematik (Crelles Journal), 1878(84y:70- 79, 1878.
Sur la formule
Deng, J., Dong, W. Socher, R., L.J, K., and Fei- Fei, I.. ImageNet: A large-scale hierarchical image database. 2009 IEEE Conference on Computer Vi- sion and Pattern Recognition, pp. 248 255, 2009. doi: 10.1 (09/CVV..2009952066
Hilfinger, A. and Paulsson, J. Separating intrinsic from extrinsic fluctuations in dynamic biological systems. Pro- ceedings of the National Academy of Sciences, 108(29): 12167-12172, 2011.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion proba- bilistic models. Advances in Neural Information Process- ing Systems, 33:6840-655, 2020.
Dhariwal, P. and Nichol, A. Diffusion models beat GANs on image synthesis. Advances in neural information processing systems, 34:8780-877 8794, 2021.
Horowitz, J. M. and Gingrich, T. R. Thermodynamic un- certainty relations constrain non-equilibrim fluctuations. Nature Physics, 16(1):15-20, 2020.
Fatras, K., Zine, Y., Flamary, R., Gribonval, R., and Courty, N. Learning with minibatch Wasserstein: asymptotic and gradient properties. In AISTATS 2020-23nd Interna- tional Conference on Artificial Intelligence and Statistics, volume 108, pp. 1 1-20, 2020.
Huguet, G., Magruder, S., Tong, A., Fasina, O., Kuchroo, M., Wolf, G., and Krishnaswamy, S. Manifold inter- polating optimal-transport flows for trajectory inference. Advances in Neural Information Processing Systems, 35: 19705-29777 18, 2022.
Fatras, K., Zine, Y., Majewski, S., Flamary, R., Gribon- val, R., and Courty, N. Minibatch optimal transport distances; analysis and applications. arXiv preprint arXiv:2101.0199,, 2021.
Hunter, J. D. Matplotlib: 2D graphics environment. Com- puting in Science & Engineering, 9(3):90-95, 2007. doi: 10.1109/MCSE.2007.55
Flamary, R., Courty, N., Gramfort, A., Alaya, M. Z., Bois-
bunon, A., Chambon, S., Chapel, L., Corenflos, A., Fatras, K., Fournier, N., Gautheron, L., Gayraud, N. T., Janati, H., Rakotomamonjy, A., Redko, I., Rolet, A., Schutz, A.,
Jeremy Howard. Imagenette. URL htttt:///ithhitl com/ Soan/fastaal/imggeetttee
10

---

Multi-Marginal Stochastic Flow Matching for
Snapshot Data at Irregular Time Points
Kidger, P., Foster, J., Li, X., Oberhauser, H., and Lyons,
Risken, H. and Risken, H. Fokker-Planck equation. Springer, 1996.
T. Neural SDEs infiite-Diimeniionll GANs. Interna- tional Conference on Machine Learning, 2021.
Rohbeck, M., De Brouwer, E., Bunne, C., Huetter, J.-C., Biton, A., Chen, Y., Regev, A., and Lopez, R. Model- ing complex system dynamics with flow matching across time and conditions. In The Thirleenth International Conference on Learning Representations, 2025.
Kingma, D. P., Welling, M., et al. Auto-encoding variational Bayes, 2013.
Kiselev, V. Y., Andrews, T. S., and Hemberg, M. Challenges in unsupervised clustering of single-cell RNA-seq data. Nature Reviews Genetics, 20(5):273 282, 2019.
Saeys, Y., Van Gassen, S., and Lambrecht, B. N. Compu- tational flow cytometry: helping make sense of high- dimensional immunology data. Nature Reviews Immanol- ogy, 16(7):449-44) 462, 2016.
Léonard, C. A survey of the Schrödinger problem and some of connections with optimal transport. Discrete Contin. Dyn. Syst. A, 34(4):1533-1574 2014.
Satija, R., Farrell, J. A., Gennert, D., Schier, F., and Regev, Spatial reconstruction of single-cell gene ex- pression data. Nature Biotechnology, 33(5):995-502 2015.
Li, X., Wong, T.K.. L., Chen, R. T. Q., and Duvenaud, D. Scalable gradients for stochastic differential equations. International Conference on Artificial Intelligence and Statistics, 2020.
Schiebinger, G., Shu, J., Tabaka, M., Cleary, B., Subra- manian, V., Solomon, A., Gould, J., Liu, S., Lin, S., Berube, P., Lee, L.,, Chen, J., Brumbaugh, J., Rigollet, P., Hochedlinger, K., Jaenisch, R., Regev, A., and S. Lan- der, E. Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in repro- gramming. Cell, 176(4):928-943. 2019.
Lipman, Y., Chen, R. T., Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. In 11th International Conference on Learning Representations, ICLR 2023, 2023.
McInnes, L., Healy, J., Saul, N., and Großberger, L. UMAP: Uniform manifold approximation and projection. Journal of Open Source Software, 3(29):861, 2018.
Schofield, M. R., Barker, J., Link, W. A., and Pavanato, H. Estimating population size: The importance of model and estimator choice. Biometrics, 79(4):3803-3811 2023.
Mitchell, S. and Hoffmann, Identifying noise sources governing cell-to-cell variability. Current Opinion in Systems Biology, 8:39 4:39-45,222 2018.
Schrödinger, E. Über die Umkehrung der Naturgesetze. Sitzungsherichte Preuss. Akod. Wiss. Berlin. Phys. Math., 144:144-155 153, 1931.
Moon, K. R., van Dijk, D., Wang, 7., Gigante, S., Burkhardt, D. B., Chen, W. S., Yim, K., Elzen, A. v. d., Hirn, M. J., Coifman, R. R., Ivanova, N. B., Wolf, G., and Krishnaswamy, S. Visualizing structure and tran- sitions in high-dimensional biological data. Nature Biotechnology, (((12):442-1440 2019, doi: 10.1038/ $41587-019-0336-33 URL https: /
Song, Y. and Ermon, S. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems, 32, 2019.
Stuart, T., Butler, A., Hoffman, P., Hafemeister, C., Papalex, E., Mauck, W. M., Hao, Y., Stocckius, M., Smibert, P., and Satija, R. Comprehensive integration of single-cell data. Cell, 177(7):1888-11 1902, 2019.
1038/ / 0//////155800199 0336-3.
Pass, B. Multi-marginal optimal transport: theory and appli- cations. ESAIM: Mathematical Modelling and Numerical Analysis, (9(())177111799 2015.
Tong, A., Huang, J., Wolf, G., Van Dijk, D., and Krish- naswamy, S. TrajectoryNet: A dynamic optimal trans- port network for modeling cellular dynamics. Interna- tional Conference on Machine Learning, pp. 9526 9526-9536. PMLR, 2020.
Paszke, A., Gross, S., Chintala, S., Chanan, G., Yang, E., DeVito, Z., Lin, Z., Desmaison, A., Antiga, L. and Lerer, A. Automatic differentiation in pytorch. 2017.
Tong, A., Fatras, K., Malkin, N., Huguet, G., Zhang, Y., Rector-Brooks, J., Wolf, G., and Bengio, Improv- ing and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research, pp. 1-34, 2024a.
Pavon, M., Trigila, G., and Tabak, E. G. The data-driven Schrödinger bridge. Communications on Pure and Ap- plied Mathematics, 74(7):1555-1573 2021.
Peyré, G. and Cuturi, M. Computational optimal trans- port: With applications to data science. Foundations and Trends@ Machine Learning, 1 115-6))335-6677 2019.
Tong, A. Y., Malkin, N., Fatras, K., Atanackovic, L., Zhang, Y., Huguet, G., Wolf, G., and Bengio, Y. Simulation-fres

---

Multi-Marginal Stochastic Flow Matching for Fligh-Dimensional Snapshot Data at Irregular Time Points
Schrödinger bridges via score and flow matching. In International Conference Artificial Intelligence and Statistics, pp. 1279-1287. PMLR, 2024b.
Van der Maaten, L., and Hinton, G. Visualizing data using t-SNE. Journal of Machine Learning Research, 9(11), 2008.
Villani, C. Optimal transport: old and new, volume 338.
Springer, 2009.
Virtanen, P., Gommers, R., Oliphant, T. E.,, Haberland, M.,
Reddy, T., Cournapeau, D., Burovski, E., Peterson, P., Weckesser, W., Bright, J., van der Walt, S. J., Brett, M., Wilson, J., Millman, K. J., Mayorov, N., Nelson, R. J., Jones, E., Kern, R., Larson, E., Carey, C. J., Polat, I. Feng, Y., Moore, E. W., VanderPlas, J., Laxalde, D., Perktold, J., Cimrman, R., Henriksen, I., Quintero, E. A., Harris, C. R., Archibald, A. M., Ribeiro, A. H., Pedregosa, F. van Mulbregt, P., and SciPy 1.0 Contributors. SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods, 17:261-272, 2020. doi: 10.1038/s41592-01 15922-99-99----622
Weinreb, C., Wolock, S.. Tusi, B. K., Socolovsky, M., and Klein, A. M. Fundamental limits on dynamic inference from single-cell snapshots. Proceedings of the National Academy of Sciences, 115(10)FF2477 2018
Zheng, G. X., Terry, J. M., Belgrader. P., Ryvkin, P., Bent, Z. W., Wilson, R., Ziraldo, S. B., Wheeler, T. D., Mc- Dermott, G. P., Zhu, J., et al. Massively parallel digital transcriptional profiling of single cells. Nature Communi- cations, 8(1):14049, 2017.
12

---

Multi-Marginal Stochastic Flow Matching for Iligh-Dimennionall Snapshot Data at Irregular Time Points
A. Proofs
Proof of Theorem 2.1. The gradient of the loss for a single interval (Li, (441) with overlapping mini-flows is given by
Σ ue(x 24, pla || (I: (1))))le(=::0111
VoLuday VolE (x(2;))
where the expectation taken over q(2, 111 uts f++
and p(x 2t, a a. = Σÿ=d Pe 2n-jat- sie).
We begin by considering the individual losses with respect to the mini-flows
E [1:) ue(x Rtg- [11]
(14)
1=0
where the expectation is taken over (( (ti, 1.+1), and pa(x | zt sit-) We combine this into a single expectation by expanding q to cover 2t, a ix) noting that the expanded window does not affect each individual expectation. Next, we combine all the mini-flow probability paths into p(x Re.k-ptt4 & Ei Pe(x "t because we are functionally only sampling from a single "active" probability path time. The resulting combined expectation is
Liete = E
(): - ue( Re I|
(15)
taken over (iitt (to, (441), and p(x as follows:
-t We derive the rest by considering the inside of the expectation
k-1 ||= 2(vr(=; velx ((-nitt-s+te) + 11 - )11
[[<(x;; 0) ux(x 20c-gite-442) )1(
1=0
[1; - vo(x;0), Σ uo(x (ls-sde-stt))
k-1
+
Σ
|u (2 1l1 + (k - ();
h-1 = =1; 0)11 - ve((;;)) 0).
Σ j=0
ue(x ((c-gite-st))
k-1
+
Σ
u
2t, Stt-)ate-
+ -k= (1) -U
1=0
2
pe((r; 0)
Σ uexx 22, ((1-5
+ (k 1)); U
where U is some constant term of inner products and sums of ux(x) (t- such that E ue(x 2t-ng En= [[ee žt; ytt, 112 U. We utilize the well-known fact that ||a: + y/1² = | + 2x, y) + || for first, third, and fourth equalities. Crucially, in the third equality we introduce U in order "complete the square" of the sum of (a(x =(c-ssh(ste)- We use the bilinearity of the inner product in the second equality. This results in the interval loss of
Cipte =E
(1;
Σ uel Rds =
+ (k: ())la,(;; -U
(16)
13

---

Multi-Marginal Stochastic Flow Matching for Fligh-Dimensional Snapshot Data at Irregular Time Points
Finally, by taking the gradient w.r.t. 0, arrive at Theorem 2.1, noting that U does not depend 0 and therefore VoU 0.
Proof of Corollary 2.2. If all the mini-flow regression signals ue(a (e:-5it(-===) ) are equal, then the single inter- val loss recovers the CFM loss (Theorem 2.1 of Lipman et al. (2023)) up to a scalar factor. Consider the inside the expectation in (16) where all the u. ( (es-site--x)) are equal (denoted simply as us (2 =)):
R (e 2)
& U -11ba<;;001 - (x; 0), E
vr(x; 0) - Σ 1=1
+ (k -
(x z)
5=1
2
+
Σ
U (x
[2)
1)]l0 11)-1)]loxxx 0)|1 U
(k:
j=1
=1/1v (2;; 0)11
(e(2;; 0),
U (
j=1
2
k
+
U (x |z)
-U
1=1
k
= (x; -
ve(a;; 0),
us
(x 2)
/=1
ue =)]1
j=1
(x; 0)11 -2k(ve(=; 0))) k||ua( 1))
<l),,(<;; 0) ue (ox =)|
We "complete the square" for the first and last equality, group like terms the second and fourth equality. and use the definition of in the third equality. Once we reintroduce the expectation, we see that this is proportional to Theorem 2.1 of Lipman et al. (2023):
CLSn = E[blll,((;;)) (x =)] = (x; 0) un (x =)]
(17)
B. Cubic Splines
Cubic splines are a class of piecewise functions interpolating between control points (to: vo), ... (tn, In), taking the form
Sp(t) to <
5(1) =
Sn-1(!) w- S 5 bn
where S; is the cubic polynomial S,(() at(t - 4 + by(t - 42 + ct(t - 4.) + d, for i in O,..., There are 4
coefficients to solve for per equation which results in n equations and 4n unknowns.
B.1. Natural Cubic Splines
Natural cubic splines solve for the above coefficients a,, by. ct, de by applying four conditions. The first requires the spline to interpolate the data points (t,, æ;) such that S(t.) = I; resulting in n + 1 constraints. The second requires S to be continuous at the interior points such that St(t,) = S (t:), resulting in n - 1 constraints. The third and fourth conditions respectively require S' and S" to be continuous for total of 2n 2 constraints. Finally two boundary conditions are added such that (*(oo) (t.) 0. In total, we have constructed system of equations with 4n unknowns and 4n
14

---

Multi-Marginal Stochastic Flow Matching for Hligh-Dimensionn Snapshot Data at Irregular Time Points
Linear Piecewise
Rt (2.
Itz: fax ta)
to
t3
tg
Its. It)
-4 Es
Figure 4. Euclidean splines on overlapping windows size h = 2. using the means of each Gaussian in the S-shaped dataset as the control points. From left right: 1) The points interpolate with time labels to.-...ff: Natural cubic splines on equidistant time intervals. 3) Monotonic cubic Hermite splines on equidistant time intervals. 4) Natural cubic splines on arbitrary time intervals Ta (0, 0.08, 0.38, 0.42, 0.54, 0.85, 1). 5) Monotonic cubic Hermite splines on T2.
Linear Piecewise
its
Its. 1s)
Ita. to)
Ms. 16.
Figure Euclidean splines on overlapping windows size 2, using the means each Gaussian in the a-shaped dataset as the control points. From left to right: 1) The points to interpolate with time labels fo 2) Natural cubic splines equidistant time intervals. 3) Monotonic cubic Hermite splines on equidistant time intervals. 4) Natural cubic splines on arbitrary time intervals T2 (0, 0.08, 0.38, 0.42, 0.54, 0.85, 1). 5) Monotonic cubic Hermite splines T2.
fe.
constraints. Ultimately, this setup constructs a tridiagonal system of equations which is efficiently solvable in (nn) time using single forward and backward pass. Perhaps reasonably, natural cubic splines quite local the influence of neighboring intervals greatly decreases the further away the neighbor is.
B.2. Monotonic Cubic Hermite Splines
Cubic Hermite splines approach the problem differently. Consider single time interval [0, 1] and corresponding points Tp. F1. Let the position of time be given by the following cubic polynomial:
Ft at
bt + ct + d.
Likewise, let me be the velocity of xe at time t, given by
m. = 3at + 2bt +
At f 0 and = 1, we can solve for æo 21, mo, my in terms of b, c, d to get the following system of equations:
Zo d
= a + b + c
mg = c
mi = 3a 26 c.
Solving this system of equations, we get
s(1) (2)3 - 312 + 1)29 + (13 212 + f)mg + (-2(3 + +32)22 + (13 - (2)m
as the polynomial interpolating (0, mg) to (1, 11. All that remains is to specify values for ma and mi In other words, cubic Hermite spline algorithms are defined by how the velocities m, are selected.
15

---

Multi-Marginal Stochastic Flow Matching for High-Dimensioone Snapshot Data at Irregular Time Points
Monotonic cubic Hermite splines set my using the following strategy. Define hy lx+1 lx and dg 2442m#k signs of dg and d do not match either O, then set my = 0. Otherwise, me is given by
If the Ax
W1
+ 02
W1
W22
m
dg-1
where W1 = 2h; + hx-1 and u/2 - h4 + 2h We direct the reader to (Fritsch & Carlson, 1980) for an exact derivation and proof of monotonicity. This formula is also solvable in (()) time, but differs from the natural cubic spline in that it is very local. In fact, only the immediate neighboring data points (4;-1, and T1+1) influence the curve.
C. Implementation Details
Our implementation uses NumPy (Harris al., 2020), SciPy (Virtanen et al., 2020), and PyTorch (Paszke et al., 2017) for the mathematical operations. Plots were generated using Matplotlib (Hunter, 2007). We extended the base CFM framework provided by Tong et al. (2024a;b)
C.1. Optimal Transport Plans
We use the Python Optimal Transport (POT) library (Flamary et al., 2021) to compute all OT plans with ot emd () and the squared Euclidean distance. This method operates on two batches of samples and returns a matrix where the i, j-ht entry (f)
( The superseript (i) denotes the i-th index into batch of samples. We implement the conditional OT plan ( (x2)) by first constructing the marginalization column vector q representing f ( Ftzz )d ((x22), where the i-th entry d; == Σj mij. Finally, we construct the matrix î representing the conditional plan π (Ztzy1 xxz) by clement-wise dividing each column of т by q. Thus, the i-th row of A
Tay indicates the probability of sampling (x 3 -
contains the probabilities of selecting 2tz+ given zt Because the procedure operates by coupling the indexes of samples from the data, we can consider this as an alignment operation on the initial mini-batch Our full sampling procedure for single mini-flow is follows. We set the initial time index to be 0 without loss of generality with respect to any given mini-flow. All samplers sample with replacement.
Algorithm 2 MMOT Sampling
Set batch size 5
Sample batches Xto..... Xr. ~ g(xto).... .... Sample (Xi,; ~ x (Xtos Xea)
(6) ((Ito)...,g(.,,,,, (x,,,) where X., = (0 li=1
for É E 2..k do
Compute the conditional OT plan a *(x,, (it.)) X1. Sample X AXx,, x
end for
Output: Aligned batches Xi,.....
C.2. Neural ODE and SDE Solvers
We use the same setup as Tong et al. (2024b) and use the SDE solver from torchsde (Li et al., 2020; Kidger et al., 2021). We did not adjust any of the default hyperparameters for the solver. For the SDE model, we set the drift to fe(x) t vz(x; 0) 200 (x; 0)) from (12) where v, and s; are respectively the flow and score networks. The diffusion set to constant gr (x) = 0.15.
D. Experimental Setup
D.I. Training Setup
For all experiments except for image generation, we used a MLP with an input layer, two hidden layers of width 64, an output layer, along with SELU activation functions. We optimize these networks using AdamW, We set a = 0.15 for our method and likewise as the noise scale in MIOFlow.
16

---

Multi-Marginal Stochastic Flow Matching for Hligh-Dimensional Snapshot Data at Irregular Time Points
For the S-shaped, a-shaped, and DynGen datasets, we trained for 2500 gradient steps and learning rate of le-4. For the CITEseq and Multiome datasets, trained for 1000 gradient steps and learning rate of 1e-5.
MIOFlow is method to infer "opttmal" trajectories on manifolds which correspond to geodesics. As we do not have access to the underlying manifold itself, the authors propose learning it from data using a GAE such that the encoder á is a mapping from the ambient space to the manifold. More specifically,, the encoder learns an embedding such that the Euclidean distance of two embedded points [Ø(x) - d(y)] matches some geodesic distance G(r, y) based on a diffusion affinity matrix.
Additionally, MIOFlow requires training GAE to embed high-dimensional data into a liver-dimensionll space and to then reconstruct trajectories learned in the embedded space. We define the encoder MLP with three hidden layers of sizes 128, 64, and 32. This encoder outputs an embedding into IR. The decoder has the same architecture but in reverse. We use ReLU as the activation function. The GAE is trained for 1000 gradient steps using the AdamW optimizer.
For Imagenette, we used the same setup as in Lipman et al. (2023) and use the U-Net architecture from Dhariwal & Nichol (2021) as well as the Adam optimizer. We resize all images to 32 x 32 and normalize across all RGB channels with μ = 0.5 and Ø- = 0.5. No other preprocessing done. We match hyperparameters where we can but opt for a smaller experiment in terms of the number of GPUs and the effective batch size per window. Note that for the Triplet model there is only one window but for the Pairwise model there two. We reproduce the hyperparameters below.
Imagenette-32
Channels
256
Depth
3
Channels Multiple
2, 2,
Head
4
Heads Channels
64
Attention Resolution
16, 8
Dropout
0.0
Effective Batch Size per Window
192
GPUs
1
Epochs
250
Iterations
250k
Learning Rate
le-4
Learning Rate Scheduler
Polynomial Decay
Warmup Steps
20k
Table Imagenette experiment hyperparameters. We use the same model architecture and learning rate hyperparameters as Lipman et al. (2023). However, use 250 epochs instead of 200 as arbitrarily set epoch to 1000 gradient steps. This matches the 250k iterations. Moreover, we have a smaller effective batch size per window of 192 and only use one GPU.
D.1.1. SCORE MATCHING IMPLEMENTATION
As noted in Section 2.3.2, we have V Pt 2) - for 1), However, this direct formulation does not protect against numerical instability when og is small. We follow the approach used by Tong et al. (2024b) and take advantage
of the user-defined weighting schedule X(t) to cancel out the division and learn the scaled target g(t)² v log pe(x) 2) based
2
on the Fokker-Planck Equation (2). By rewriting the inside of the expectation of the scaled score loss as
g(t) 2
- ||000 A(1) gt) +00))
x(1)2 [x( 0)
V log
20,
we can see that when setting X(t)
20, 2 the score loss becomes
(x;t)& 0) + ell
c~ c~~1(0,)),
This approach allows us to reconstruct the mini-flow SDE drift as the sum of the mini-flow ODE drift and the scaled score network output:
U (x;; 0) = v. (x; 0) + (2;
(18)
17

---

Multi-Marginal Stochastic Flow Matching for Hligh-Dimensional Snapshot Data at Irregular Time Points
D.2. DynGen
We repurpose the DynGen data used in MIOFlow (Huguet et al., 2022) for our experiments. Notably, the data itself not the raw simulated reads; is preprocessed into 5 dimensions using PHATE (Moon et al., 2019).
PHATE operates as dimensionality reduction scheme aiming to preserve both local and global dependency structures. Local structure is learned first by imposing Pairwise affinities under a Gaussian kernel. Global structure is inferred by propagating the local affinities via diffusion, effectively learning a statistical manifold based on the information geometry. Finally, metric MDS is used as the dimensionality reduction strategy.
We believe that the GAE used in MIOFlow learns the data manifold for the -taAAT--raanffrmme) DynGen dataset especially well given that, by construction, the DynGen dataset does indeed reside manifold equipped with diffusion- baadee metric. This matches the prior belief in MIOFlow that diffusion- based affinities can accurately capture the data manifold.
D.3. CITEseq and Multiome
These datasets were published part of a NeurIPS competition for multimodal single-cell integration (Burkhardt et al., 2022). We present brief overview, and refer the reader to the competition itself for more in-depth descriptions ?. The data is collected from peripheral CD34+ hematopoietic stem and progenitor cells from healthy human donors. The CTTEseq data is measured using 10x Genomics Single Cell Gene Expression with Feature Barcoding technology. The Multiome data is measured using 10x Chromium Single Cell Multiome ATAC + Gene Expression technology.
Technically, both the CIT'seq and Multiome datasets are labeled, with the former about predicting protein levels given gene expressions, and the latter about predicting gene expressions given ATAC-seq peak counts. We are only interested in the gene expression data, so we only use the CTTEseq input data and the Multiome target data. Following Tong et al. (Tong et al., 2024b), we only select cells from the respective datasets from a single donor id 131 76.76.7767 The gene expression data is already library-size normalized and log Ip transformed, so we compute the PCA and top highly variable genes without any further preprocessing step.
D.4. Imagenette
This dataset subset of the well-known ImageNet dataset (Deng et al., 2009) containing a curated collection of 10 easily classifiable classes. The classes present are: tench, English springer, cassette player, chain saw, church, French horn, garbage truck, gas pump, golf ball, and parachute. More details and variants of the dataset can be found at the GitHub page
E. Ablation Studies
We report in Table 4 ablation experiments held-out time points for the S-shaped and a-shaped Gaussians the Pairwise (k- = 1). Triplet (k 2). and "All" (k M 1) models. We test both Ti = 0.10.17, 0.33, 0.5, 0.67, 0.83, 1) and To = (0, 0.08, 0.38, 0.42, 0.54, 0.85, 1). We evaluate the W1 metric on the held-out marginal for these experiments. In general, the Pairwise model performed worst, whereas the Triplet and All models were relatively even, providing experimental validation for minimal performance boosts when 2.
F. Flow Visualizations
We visualize our experiments here. Viewing in color is recommended.
wttele..coo//oomm/eeeeo http:////oo///tttttooooooeeagaoa semmmmmttt mmodal overvvew
18

---

Multi-Marginal Stochastic Flow Matching for Hligh-Dimensionall Snapshot Data at Irregular Time Points
S-shaped Triplet
a-shaped Triplet
Held-out index
Pairwise
All
Pairwise
All
2.431
2.13
2.08
3.38
4.95
5,13
2
2.59t
1.96
2.14
4.38
4.37
4.84
Ti
3
1.44
1.28
1.13
2.72
2.94
2.971
4
2.121
1.95
1.74
3.78
4.54*1
4.53
5
2.36*t
1.83*
2.04
4.17
5.041
4.71
2.35t
1.71
1.48
2.78
3.94
3.94
2
2.65
3.27
3.64
5.741
4.94
4.68
T2
3
2.65t
1.90
1.76
3.37
3.24
3.91
4
0.86
1.09
2.11
8.08*1
3.79*
2.42
5
2.12*1
1.62*
1.59
5.06
5.12
4.54
Table W1 metrics on ablation experiments varying the held-out time point. Entries with indicate values reported Table Entries with indicate the performance out of the Pairwise, Triplet, and All models.
to 0.00
f = 0.08
82 = 0.38
t3 = 0.42
t4 0.54
& 0.85
to = 1.00
Trajs
Figure 6. Trajectories for S-shaped Gaussians using arbitrary time points (0, 0.08, 0.38, 0.42, 0.54, 0.85, 1) and holding time point ts 0.85. Trajectories start from the leftmost point and follow the curve to reach the rightmost point. From left to right: MIOFlow. Pairwise, Triplet.
to = 0.00
t = 0.08
t2 0.38
t3 0.42
t 0.54
t = 0.85
t = 1.00
Trajs
(0, 0.08, 0.38, 0.42, 0.54, 0.85, 1) and holding out time point 0.54. Trajectories start from the upper right and loop around the bottom right. From left right: MIOFlow, Pairwise, Triplet.
Figure Trajectories for a-shaped Gaussians using arbitrary time points T- =
19

---

Multi-Marginal Stochastic Flow Matching for Hlghh Biggh-Dimessionll Snapshot Data at Irregular Time Points
to = 0.00
t1 0.25
t2 0.50
t 0.75
t = 1.00 Trajs
Figure 8. DynGen simulated trajectories. Trajectories start from the leftmost point and quickly bifurcate into the upper and lower right. trajectories are 1R but only the first and second dimensions are shown here. hold t 0.25. From left to right: MIOFlow. Pairwise, Triplet.
to 0.00
t 0.20
t
t 0.30
4 = 0.88
ts = 0.98
1.0000
Trajs
Figure Trajectories for and a-shaped Gaussians predicted by the Pairwise and Triplet models using all time points from (0, 0.2, 0.27, 0.3, 0.88, 0.98, 1). From left to right: 1) Pairwise on S-shaped. 2) Triplet on S-shaped. 3) Pairwise on a-shaped. 4) Triplet a-shaped.
0.00
0.20
0.27
0:00
0:17
0:33
0.50
0.67
0.83
1:00
5 0.30
0.89
0.98
1.00
0:00
0.83
Figure 10. Trajectories for S-shaped Gaussians predicted by the Pairwise and Triplet models using all time points from Ts (0, 0.200,200 0.27, 0.3, 0.88, 0.98, 1). Top row) Pairwise. Bottom row) Triplet.
20

---

Multi-Marginal Stochastic Flow Matching for High- DinhnsiDnnssin Snapshot Data at Irregular Time Points
0.20
022
0.00
0.50
0:57
0.63
1:00
0.30 0.30
0.98
1.00
Figure Trajectories for a-shaped Gaussians predicted by the Pairwise and Triplct models using all time points from To (0, 0.2,0.200 0.27, 0.3, 0.88, 0.98, 1). Top row) Pairwise. Bottom row) Triplet.
21

---

Multi-Marginal Stochastic Flow Matching for Hligh-Dimensione Snapshot Data at Irregular Time Points
Figure 12. Pairwise model on Imagenette-33 32 for the class progression: gas pump golf ball parachute. Time points are from
T == 0.25. 1). Images are generated from a checkpointed model at 160k gradient steps.
22

---

Multi-Marginal Stochastic Flow Matching for Bligh-Dimensiona Snapshot Data at Irregular Time Points
Figure 13. Triplet model on Imagenette-33 32 for the class progression: gas pump golf ball parachute. Time points from
T == 0.5, 1). Images are generated from checkpointed model at 160k gradient steps.
23


Reviews:
[{"review_id": "tNoxdnNy4b", "reviewer": "Reviewer_rnGo", "summary": "This work proposes Multi-Marginal Stochastic Flow Matching (MMSFM) with the goal of training a translation model across multiple snapshots taken at non-equidistant time points. MMSFM builds upon the Flow Matching framework and extends it through measure-valued splines. (Stochastic) Flow Matching can be applied to train a pair-wise flow between all adjacent time point snapshots. However, this approach can struggle to capture the global dynamics across multiple time points and does not generalize well to arbitrary held-out time points. MMSFM enables capturing local dynamics across uneven time intervals and maintaining consistency between overlapping windows.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful feedback.\n\n**Q1)**\n\nWe use the Python Optimal Transport (POT) package for computing the OT plans $\\pi(x_0, x_1)$ where we use $x_0, x_1$ as notational shorthand for $x_{t_\\ell}, x_{t_{\\ell+1}}$. The conditional plans are generated by using probability rules where $\\pi(x_{\\ell+1} | x_\\ell) = \\pi(x_\\ell, x_{\\ell+1}) / q(x_\\ell)$. Specifically, we use the Earth Movers Distance function (`POT.emd()`) with exact matching.\n\nWe manipulate the probability matrices $\\pi$ returned by `POT.emd()` representing $\\pi(x_i, x_{i+1})$ by using the column vector $\\hat{q}$ representing $\\int \\pi(x_i, x_{i+1})dx_{i+1}$. We construct $\\hat{q}$ by summing over the columns of $\\pi$. We can then obtain the conditional plan $\\hat{\\pi} \\gets \\pi(x_{i+1} | x_i) = \\pi / \\hat{q}$.\n\nSampling itself is at simple as using $\\pi$ as a 2D probability table or $\\hat{\\pi}$ as a list of 1D probability vectors where the row index corresponds to the conditioning variable. We also note that the whole sampling procedure technically operates on the indexes of a mini-batch and so we can also consider this as an alignment operation on the initial mini-batch. Further, this procedure adopts Algorithm 1 from [2] using OT plans in place of OT maps.\n\n**Q2)**\n\nAs you sharply note, the \"true\" MMOT plan would correspond to the joint distribution $\\pi$ which considers all pairs of distributions. However, this is not easy to compute and moreover we are specifically interested in processes which do have a temporal ordering. With this in mind, we turn to the framework of joint distributions and apply the chain rule of probability to obtain that $\\pi(x_0, \\dots, x_M) = \\pi(x_0, x_1) \\prod_{i=2}^M \\pi(x_i | x_{< i})$. Finally, by considering a sequential sampling process where $x_i$ must be sampled prior to $x_j$ for $i < j$, we can apply a first-order Markov approximation where $x_i$ only depends on $x_{i-1}$. Thus, the conditional probabilities in the product all reduce to the form $\\pi(x_i | x_{i-1})$ and we recover equation (8). It is in this sense that we call (8) the first-order Markov approximation to the true MMOT plan.\n\n**Q3)**\n\nWe use the same setup as [3] and use the SDE solver `torchsde` from `github.com/google-research/torchsde`. The drift function $f_t(x)$ is set to the SDE drift $u_t(x ; \\theta) = v_t(x ; \\theta) + \\frac{g^2(t)}{2} s_t(x ; \\theta)$ from equation (12) for some learned deterministic flow $v_t$ and score $s_t$. We set the diffusion schedule $g(t) = \\sigma$ for a constant $\\sigma = 0.15$.\n\nWe have not explored any ablation to the SDE solver hyperparameters nor the sensitivity to the number of function evaluations, however we believe these to be avenues for interesting future studies. If referring to differences between trajectories inferred by an SDE and an ODE, neither meaningfully differed from one another because the the same learned flow model $v_t$ is used in both cases.\n\n**Q4)**\n\nYes, MMSFM can work without mini-batch OT. Notice how in our Theorem 2.1 we define the regression signal to be certain $u_t(x | z)$ objects with an expectation over some $q(z)$. OT comes into play when constructing $q(z) \\gets \\pi(z)$ for sampling $z = (x_0, \\dots, x_M)$ and evaluating the expectation. However, the only strict requirement is that $\\pi(z)$ be a coupling distribution such that $q(x_i) = \\int \\pi(x_0, \\dots, x_M)dx_{-i}$. We can just as easily construct $\\pi(z) = \\prod_i^M q(x_i)$ and satisfy the requirement, noting that this $\\pi$ is not the OT plan. Based on prior work such as [1], we focused on the OT coupling and did not evaluate on the independent coupling.\n\n**Comments**\n\nThank you for suggesting [4]. Although we have not applied the methodologies and metrics referenced, we find them very relevant and valuable for future work.\n\n**References**\n\n[1] Tong, A., Malkin, N., Huguet, G., Zhang, Y., Rector-Brooks, J., Fatras, K., Wolf, G., and Bengio, Y. Improving and generalizing flow-based generative models with mini-batch optimal transport. arXiv preprint arXiv:2302.00482, 2023b.\n\n[2] Chewi, S., Clancy, J., Le Gouic, T., Rigollet, P., Stepaniants, G., and Stromme, A. Fast and smooth interpolation on wasserstein space. In International Conference on Artificial Intelligence and Statistics, pp. 3061–3069. PMLR, 2021.\n\n[3] Tong, A., Malkin, N., FATRAS, K., Atanackovic, L., Zhang, Y., Huguet, G., Wolf, G., and Bengio, Y. Simulation-free schrodinger bridges via score and flow matching. In ICML Workshop on New Frontiers in Learning, Control, and Dynamical Systems, 2023a.\n\n[4] Luca Eyring, Dominik Klein, Théo Uscidda, Giovanni Palla, Niki Kilbertus, Zeynep Akata, Fabian Theis. \"Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation\". ICLR 2024.", "claims_and_evidence": "The authors claim that \"the use of measure-valued splines enhances robustness to irregular snapshot timing, and score matching prevents overfitting in high-dimensional spaces.\" While this is somewhat supported by empirical evidence, this overstates the contribution of the paper as the use of the score-matching loss is taken from prior work.\n\nAlso, the authors propose using mini-batch OT, following previous work. However, no ablation of this design choice is done, and it remains unclear how reliant MMSFM is on it.\n\nEmpirically, adding splines with triplets (k=2) compared to pairwise (k=1) shows overall improvements, although not consistent across different experimental settings.", "methods_and_evaluation_criteria": "The evaluation approach is appropriate. The W1/W2 metrics provide standard measures of distribution similarity, and MMD metrics offer complementary insights. The synthetic datasets test specific capabilities (bifurcations, topology changes), and the biological applications demonstrate practical relevance. Biological metrics could further strengthen this experimental section.\n\nEmpirically, the authors compare MMSFM (k=2) to SF2M and MIOFlow. It could be beneficial to include more competing works here, as well as include results for k > 2 to shed further insight into the parameter k. While Appendix D gives one ablation on this, it is not at all discussed in the main text. Additionally, it remains unclear how e.g. k=3 would perform and what tradeoff choosing larger k would give w.r.t performance and computational cost.", "theoretical_claims": "The theoretical foundation draws from established work on optimal transport and measure-valued splines. The paper states that their approach minimizes total action across time points but doesn't prove that the overlapping window strategy actually achieves this global minimization.\n\n Additionally, I am unsure about the author's use of the term MMOT (Multi-marginal Optimal Transport). This term, see e.g. [1], is usually used in a different context with the goal of obtaining a **joint** coupling across N different marginals. This can cause confusion, as this is also the notion that is thoroughly discussed in [2]. As also mentioned by the authors, [2] tackles the multi-marginal OT problem, i.e., learning a mapping between **all pairs** of distributions. On the contrary, MMSFM learns a sequential mapping across time-adjacent distributions, which is not MMOT. Adding a formal definition of MMOT following e.g. [1], and building upon that would help clarify this confusion and position MMSFM better within related work.\n\n[1] Brendan Pass. \"Multi-marginal optimal transport: theory and applications\". 2014.\n\n[2]Michael S. Albergo, Nicholas M. Boffi, Michael Lindsey, Eric Vanden-Eijnden. \"Multimarginal generative modeling with stochastic interpolants\". 2023.", "experimental_designs_or_analyses": "Yes, they do make sense. The synthetic datasets make sense and the biological applications demonstrate practical relevance. However, the work is submitted under \"**Primary Area:** Applications->Health / Medicine\" while having no modality specific metric. Measuring W1/W2 and MMD is fine but I think these experiments would be strengthened through biologically meaningful metrics as e.g. leveraged in [3].\n\n[3] Luca Eyring, Dominik Klein, Théo Uscidda, Giovanni Palla, Niki Kilbertus, Zeynep Akata, Fabian Theis. \"Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation\". ICLR 2024.", "relation_to_broader_scientific_literature": "This paper proposes to improve existing Flow/Bridge Matching literature for mapping across multiple non-equidistant time points. Prior work has shown that these methods are competitive in this problem. Through the addition of measure-valued splines, MMSFM achieves improvements in this specific problem and takes a step toward solving this problem.", "essential_references_not_discussed": "[3] is very closely related to MMSFM and related work as it also tackles unpaired single-cell translation. It improves upon mini-batch OT sampling through unbalanced Optimal Transport. These ideas could in future also be applied to MMSFM and would additionally be a relevant competing method.\n\n[3] Luca Eyring, Dominik Klein, Théo Uscidda, Giovanni Palla, Niki Kilbertus, Zeynep Akata, Fabian Theis. \"Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation\". ICLR 2024.", "other_strengths_and_weaknesses": "Clarity could be improved in Section 2. Specifically, regarding the term MMOT.", "other_comments_or_suggestions": "Formally introduce multi-marginal optimal transport. A toy-example visualization could strengthen the clarity of the work. In general, a concrete description of all used hyperparameters seems to be missing.", "questions_for_authors": "- What solver and hyperparameters do you use to compute the OT plans in Equation 8?\n- What is here referred to as MMOT (Multi-marginal OT) is implemented by computing pair-wise OT plans across adjacent time points. Could the authors extend the sentence \"Then, we compute the MMOT plan given by the first-order Markov approximation\" and explain how this is an MMOT plan?\n- Missing description and ablation of hyperparameters for SDE solver, which solver is used, and how many NFE are used? How sensitive is the trained model to the NFEs?\n- Does MMSFM also work without min-batch OT? How much worse does it perform?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0SglJmPYrm", "reviewer": "Reviewer_cNq2", "summary": "This paper proposes an extension of flow-matching for multi-marginals - i.e. when multiple snapshots are observed, typically over time. The method sample conditioning points from all snapshots using an approximation of the multi-marginal optimal transport map and then fits a spline to these points that is used as the conditional drift. The authors then evaluated their method on synthetic and single-cell data.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful feedback.\n\n**Q1)**\n\nThank you for the clarifying question. By perturbations we mean a system which is not currently at a steady state. For example, a cell system can be perturbed by some drug stimuli. Further, we focus on cases where the perturbation is fixed. Learning a model to handle user-defined, arbitrary perturbation cases (e.g. arbitrary drug types and dosages) is an important avenue for future work.\n\n**Q2)**\n\nThe original flow matching method assumes a linear interpolation function to fill in the gaps for particles at times not given in the training data. This is based on the fact that we effectively do not have access to the true temporal dynamics---if we did, there would be no reason to use a flow matching framework because we could then directly simulate the system for a desired initial condition. As such, we opt for an Occam's razor approach about using simpler dynamics. In our case, this amounts to cubic splines which offer tractable solutions and avoid Runge's phenomenon. As you correctly point out, our method assumes the temporal dynamics to be approximable or model-able using spline interpolations in the data space. These data-space interpolations are then used to learn a vector field generating the probability path $p_t$. This $p_t$ is also a spline, but one over measure space, interpolating between the marginal distributions $\\rho_i$ given in the data.\n\n**Q3)**\n\nOur sampling algorithm is the Transport Spline Interpolation in Algorithm 1 of [2]. However, those authors look at a single interpolating spline whereas we generate multiple overlapping splines in our own Algorithm 1 via a rolling window over the timepoint marginals. Moreover, we require a generative model for evaluations on unseen initial conditions which [2] does not provide because it looks only at interpolations between existing points.\n\n**Comments**\n\nThank you for bringing [1] to our attention as we had not seen it at the time of developing our model. Certainly, there is a good amount of similarity and relevancy:\n\na) We both use splines as interpolants.\n\nb) We both arrive at the same MMOT plan approximation in our work's equation (8). The authors of [1] discuss this in their Appendix B.\n\nc) We both allow for irregular snapshot timings.\n\nHowever, we differ in that:\n\na) We use splines on overlapping triplets whereas [1] looks at splines over the whole sequence.\n\nb) We explore the spline algorithm and opt for monotonic Hermite cubic splines whereas [1] opts for natural cubic splines. The latter is perhaps nicer analytically, but we nonetheless use the former due to practical concerns involving overshoot. For example, natural cubic splines enforce $C^2$ continuity which produce \"smoother\" splines, but this constraint can introduce severe overshooting which we illustrate in Figures (2, 3, 4) found in our Appendix B. Monotonic Hermite cubic splines only enforce $C^1$ continuity which is sufficient in terms of the general Flow Matching framework. This weaker smoothness constraint, along with the monotonicity between control points, avoids the overshooting problem whilst also still allowing for continuous derivatives at each interior point unlike a simple linear interpolation.\n\nc) Related to (b), we explore the sensitivity of our method to highly irregular timepoints. This is in part because we suspect the overshooting behavior in the natural cubic splines to be caused by neighboring short and long time intervals. For example, consider our timepoint set $T_3 = 0, 0.2, 0.27, 0.3, 0.88, 0.98, 1)$ for synthetic data. Notice that the sequence $0.27 \\to 0.3 \\to 0.88$ involves a short interval of 0.03 followed by a long interval of 0.58. This means that any change in velocity and acceleration along the short interval can happen relatively quickly, but the corresponding change for the long interval must be drawn out. The continuity of the acceleration does not help in this regard as it prevents the spline from instantaneously re-adjust velocities.\n\n**References**\n\n[1] Rohbeck, Martin, et al. \"Modeling Complex System Dynamics with Flow Matching Across Time and Conditions.\" The Thirteenth International Conference on Learning Representations, 2025.\n\n[2] Sinho Chewi, Julien Clancy, Thibaut Le Gouic, Philippe Rigollet, George Stepaniants, and Austin Stromme. Fast and smooth interpolation on Wasserstein space. In International Conference on Artificial Intelligence and Statistics, pp. 3061–3069. PMLR, 2021.", "claims_and_evidence": "The authors show in their experiments (both synthetic and single-cell) that the multi-marginal extension is more effective than pairwise flow matching. This demonstrates the added value of their approach.", "methods_and_evaluation_criteria": "The evaluation criteria are clear and make sense (reconstruction of held-out snapshots).", "theoretical_claims": "I did not check all theoretical claims but I did not spot any major issue.", "experimental_designs_or_analyses": "The experimental designs are sound, the authors learnt their method on given snapshots and evaluated on a left out time point, which is typically done in the literature.", "relation_to_broader_scientific_literature": "The paper extends the flow matching framework, and embeds itself in that literature.", "essential_references_not_discussed": "The authors missed an important, albeit very recent, reference that is seemingly very related to their approach [1]. Given the papers are concomitant, I don't expect the authors to compare against it but I encourage them to position their paper with respect to it.\n\nThe paper is also related to [2]. The authors could compare their method against it.\n\n\n[1] Rohbeck, Martin, et al. \"Modeling Complex System Dynamics with Flow Matching Across Time and Conditions.\" The Thirteenth International Conference on Learning Representations, 2025.\n[2] Sinho Chewi, Julien Clancy, Thibaut Le Gouic, Philippe Rigollet, George Stepaniants, and Austin Stromme. Fast and smooth interpolation on Wasserstein space. In International Conference on Artificial Intelligence and Statistics, pp. 3061–3069. PMLR, 2021.", "other_strengths_and_weaknesses": "This paper addresses an important use case of flow matching, where multiple snapshots are observed over time. The authors proposed a sound and effective way to leverage temporal dependencies and showed favorable experimental results.\n\nWeaknesses:\n\n- The authors claim to propose a promising method for single-cell perturbation but I did not see any mechanism to incorporate perturbations nor any experimental result involving perturbations.\n\n- The method heavily relies on the assumption that the temporal dynamics follow a spline. It's not clear whether this assumption is justified in practice.", "other_comments_or_suggestions": "Cfr above.", "questions_for_authors": "1. The authors claim to propose a promising method for single-cell perturbation but I did not see any mechanism to incorporate perturbations nor any experimental result involving perturbations. Could you please clarify the claim ? \n\n2. The method heavily relies on the assumption that the temporal dynamics follow a spline. It's not clear whether this assumption is justified in practice. Can the authors comment on that ? \n\n3. Could the authors include [2]  as a baseline ? \n\n\n\n[2] Sinho Chewi, Julien Clancy, Thibaut Le Gouic, Philippe Rigollet, George Stepaniants, and Austin Stromme. Fast and smooth interpolation on Wasserstein space. In International Conference on Artificial Intelligence and Statistics, pp. 3061–3069. PMLR, 2021.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uKJxKho2m3", "reviewer": "Reviewer_4cti", "summary": "The paper presents Multi Marginal Stochastic Flow Matching Model(abbreviated as MMSFM), which is an extension of simulation-free score and flow matching method. The method enables the alignment of high-dimensional snapshots obtained from non-equidistant time points without reducing dimensionality. MMSFM uses third-degree polynomial to represent the non-equidistant time interval which connects optimally coupled points. Next, MMSFM introduce score matching and flow matching to reconstruct the Fokker-Planck equation while modeling stochastic process which bridges the distributions. As the non-equidistant snapshots are multi-marginal data which consists of overlapping mini-flows, MMSFM proposes a rolling window frame algorithm to deal with overlapping triplets, enhancing robustness and stability of the model. The model has been applied on synthetic datasets , single-cell dataset and other high- dimension real-world datasets, such as COLO58(melanoma single-cell) and CITEseq(Cellular Indexing of Transcriptomes and Epitopes by Sequencing)", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback.\n\n**Q1)**\n\nThis is a very insightful question, as you are correct to hint at the sensitivity of splines to the time-labels in the data.\n\nFor the S-shaped and $\\alpha$-shaped synthetic datasets, we evaluate on 3 different timepoint distributions to see the effect of varying the gaps.\n* Uniform gaps $T_1 = (0, 0.17, 0.33, 0.5, 0.67, 0.83, 1)$, or alternatively $(i/6)\\_{i=0}^6$.\n* One set of arbitrary gaps $T_2 = (0, 0.08, 0.38, 0.42, 0.54, 0.85, 1)$. The smallest gap here is 0.04 from the interval $(0.38 \\to 0.42)$. The largest gap is approximately 0.3 for which there are two intervals: $(0.08 \\to 0.38)$ and $(0.54 \\to 0.85)$.\n* A second set of arbitrary gaps $T_3 = (0, 0.2, 0.27, 0.3, 0.88, 0.98, 1)$. This is an extremely uneven and difficult timepoint set as it contains a disproportionately large interval of 0.58 in $(0.3 \\to 0.88)$. Also included are two very small intervals of 0.02-3 in $(0.27 \\to 0.3)$ and $(0.98 \\to 1)$.\n\nWe can see in Table 1 that for $T_1$ and $T_2$, the Triplet model tends to outperform the Pairwise model. Interestingly, the Pairwise model did outperform the Triplet model on the $\\alpha$-shaped $T_1$ case, however we believe that given the Pairwise model's weakness on $\\alpha$-shaped $T_2$ compared to the Triplet model, the experiments validate the ability for splines to be flexible enough for irregular timepoints.\n\nStill, it is probably a good idea to be cautious. We can see snapshots of the probability path given in the $T_3$ case by plotting how the point cloud moves in time. The relevant figures are in Appendix E, Figures (9, 10). In the S-shaped case (Figure 9) we see that the Triplet model learns the flow but includes shearing effects. Moreover, in the $\\alpha$-shaped case (Figure 10) neither model was able to successfully learn the flow. We believe exploring regulation effects to prevent this behavior can lead to fruitful future work.\n\nThe COLO858 dataset contains 8 timepoints of $T = (0, 0.5, 2, 6, 15, 24, 72, 120)$. We normalize this to $T = (0, 0.004, 0.017, 0.05, 0.125, 0.2, 0.6, 1)$. We can see that the size of the gaps is monotonically increasing. Table 1 again shows the Triplet model to outperform the Pairwise model in this setting with irregular timepoints.\n\nThe Multiome and CITEseq datasets both contain 4 timepoints of $T = (2, 3, 4, 7)$ which we normalize to $T = (0, 0.2, 0.4, 1)$. Likewise, we see that the gap size is monotonically increasing. Table 2 shows that Triplet model marginally outperforms the Pairwise model.\n\nWe see that the time intervals are monotonically increasing for COLO858, Multiome, and CITEseq because typically cell dynamics are most active early into the perturbation and gradually slow down over time. The higher time-fidelity in the early time points reflects this.\n\nOur experimental setup descriptions can be found in section 3.1.\n\n**Q2**\n\nFlow matching methods including ours are ultimately attempts at learning the dynamics driving a system subject to constraints that the marginals at given times must match certain distributions. Moreover, these underlying dynamics are not necessarily unique, implying that there are infinitely many dynamics which solve the system subject to the marginal constraints. It is this unidentifiability which leads prior work to assume linear interpolations and our work to assume spline interpolations.\n\nWe are also beholden to problems relating to signal sampling which are beyond the scope of this work. Consider, for example, that we have data generated from a true underlying process of $f(t) = \\sin(t)$. If we happen to take measurements at $t = 0, \\pi, 2\\pi, \\dots, n\\pi$ then our signal is a constant 0. Based on the data $\\{ (k\\pi, 0) \\}_{k=0}^n$, the simplest reconstruction without any additional assumptions on the data or underlying process is $\\tilde{f}(t) = 0$ but in fact $\\tilde{f} \\neq f$!\n\nWith this in mind, decreasing the number of snapshots is equivalent to decreasing the time-fidelity of our data. As such, this would increase the \"simplicity\" of the dynamics learned by our model. On the other hand, the total length of the process is not as important given that we normalize time to run for $t \\in [0, 1]$.\n\n**Q3**\n\nYes, the Dyngen dataset introduces a bifurcation as can be seen in Appendix E, Figure 7. We repurpose the synthetic dataset from [1], which the authors generated using the Dyngen simulator [2]. We have used the same data files as [1] which defines a realistic dynamic cellular process.\n\n**References**\n\n[1] Huguet, G., Magruder, D. S., Tong, A., Fasina, O., Kuchroo, M., Wolf, G., and Krishnaswamy, S. Manifold interpolating optimal-transport flows for trajectory inference. Advances in neural information processing systems, 35: 29705–29718, 2022.\n\n[2] Cannoodt, R., Saelens, W., Deconinck, L., and Saeys, Y. Spearheading future omics analyses using dyngen, a multi-modal simulator of single cells. Nature Communications, 12(1):3942, 2021.", "claims_and_evidence": "Yes. The paper gives out several theorems to prove the feasibility of score matching and overlapping mini-flow matching. The monotonic cubic Hermite splines introduced to couple different distributions promise the monotonicity of each piece of polynomial and allow a smaller and more flexible window while guaranteeing the computational efficiency and robustness. The experiments on both synthetic datasets and real-world datasets provides a solid foundation for the method proposed in the paper.", "methods_and_evaluation_criteria": "Yes. The algorithm MMSFM combines rolling window, optimal transportation and score matching to cope with high-dimension non-equidistant time interval snapshots. The author choose COLO58 and CITESeq as  experiment datasets, which is sampled as non-equidistant time interval distribution. The model outperforms traditional algorithms such as MIOFlow.", "theoretical_claims": "Yes. The paper gives one theorem and one corollary. The theorem calculates the gradient of the loss for a single interval with overlapping mini-flows and the corollary claims the condition in theorem is a special case of conditional flow matching. The proof of the theorem and corollary is correct and is well shown in the supplementary material.", "experimental_designs_or_analyses": "Yes. The paper conducts the algorithm proposed on three synthetic datasets, one single-cell dataset from COLO858 melanoma cells, and two RNA gene expression datasets. Data is collected at several non-equidistant timepoints and the results are shown in tables and illustrations in the section 3. They both proves the soundness and validity of the experimental designs.", "relation_to_broader_scientific_literature": "The paper combines score and flow matching network with optimal transport and Brownian bridge to solve the multi-marginal problem which only non-equidistant time interval snapshots are provided. The newly-proposed method is one of the contribution of this work . The introduction of rolling window algorithm further improves the performance of the model and the results of experiments on both synthetic datasets and real-world datasets outperforms the traditional ones, which is another contribution of the paper.", "essential_references_not_discussed": "No as far as I know. The paper mentions most of the related work which helps the reader understand the model.", "other_strengths_and_weaknesses": "Strength:\n1. The paper innovatively provides a novel model to solve the non-equidistant time interval snapshots inferring problem. The model based on score and flow matching outperforms traditional methods and increases the interpretability of the model.\n2. The model combines rolling windows with spline measurements to address the issue of non-uniform time points which is inspired by optimal transport and Brownian bridge and shows great potential in dealing with other similar problems.\n3.The experiments cover both synthetic and real-world data, verifying the capability of high-dimensional modeling, which is better than MIOFlow. Also, the visualization of the results is illustrated clearly in both section3 and supplemental materials.\nWeakness:\n1.The paper lacks in analysis of the bifurcation stream problem in Dyngen datasets, it should be added or be discussed in the future works.\n2.The paper forgets to discuss the computational complexity of the rolling window size k in a broader view, which is an inevitable part of the theoretical analysis. The paper shows a M-dependent analysis of the computational complexity, but more analysis can be conducted in the study.", "other_comments_or_suggestions": "Some figure (such as figure 2) should be added into main body of the paper instead of showing in the supplementary material. In section 2.3, the author mentions ‘See Figure 2 for a visual representation of the variation of paths in an interval.’ While the figure 2 is only shown in the supplemental materials, which makes it difficult to read the paper.", "questions_for_authors": "1.The distant between different time intervals in the experiments shows a monotonically increasing property, what will happen if the gap between time intervals is more irregular?\n2.What will happen if less snapshots is given? Is the model robust enough to solve a longer time experiment?\n3. Can you explain the bifurcation of flow in Dyngen datasets and gives a analytical explanation?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "EMdNIkvWev", "reviewer": "Reviewer_7a1W", "summary": "The authors proposed a multimarginal extension of flow matching that are simulation free and can work with high dimension data.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback.\n\n**Q1)**\n\nWe assume that the data space of the marginals is some continuous metric space. In particular, we assume that the data lies in $R^d$ space with the standard Euclidean distance (or distance squared) $d(x, y) = \\| x - y \\|^2$ for all our experiments. Moreover, we also assume that all the true marginals $\\rho_i$ are continuous and that our data is just an empirical realization of the $\\rho_i$ marginals.\n\n**Q2)**\n\nThe numerical computation is done on samples (e.g. computing OT plans with respect to samples from the marginal distributions) because we do not assume access to the true closed form functionals defining the various marginal distributions. The underpinning theory, however, is based on learning a parameterized model which can generate a time-varying probability path $p_t$ that is time-differentiable and is constrained to match the true marginals (up to the fidelity given by the data) at the respective times.\n\nIn terms of trajectory inference, the SDE solvers operate on individual particles. Applying this to a point cloud outputs the trajectory of the empirical distribution, for example seen in Figure 1.\n\n[1] Tong, A., Malkin, N., FATRAS, K., Atanackovic, L., Zhang, Y., Huguet, G., Wolf, G., and Bengio, Y. Simulationfree schrodinger bridges via score and flow matching. In ¨ICML Workshop on New Frontiers in Learning, Control, and Dynamical Systems, 2023a.", "claims_and_evidence": "The theoretical claims are supported by theorem and proofs. I am a bit skeptical on the empirical performance without uncertainty.", "methods_and_evaluation_criteria": "They make more sense if with some uncertainty.", "theoretical_claims": "I did not check in depth, but did not find obvious mistakes.", "experimental_designs_or_analyses": "I checked both simulated and real data. I think the experimental design make sense.", "relation_to_broader_scientific_literature": "Being able to perform scalable trajectory inference is useful in biology, ecology and environmental science. And the paper is targeting at an important problem.", "essential_references_not_discussed": "I do not have one particular -- trajectory inference literature is too large.", "other_strengths_and_weaknesses": "- The triplet idea is very interesting\n- Scales well in dimensions.", "other_comments_or_suggestions": "None", "questions_for_authors": "1) I want to know more is the assumptions on the marginals -- are they discrete/empirical or continuous/population? \n2) related to 1), Tong et al. 2023a used this mixture of Brownian bridge results of SB problem that would produce trajectories pass through one particles at each time point, this seems not inherited in the proposed method. Is this the case and does it imply that the proposed method operates on the continuous marginals not empirical version of it?\n\nTong, A., Malkin, N., FATRAS, K., Atanackovic, L., Zhang, Y., Huguet, G., Wolf, G., and Bengio, Y. Simulationfree schrodinger bridges via score and flow matching. In ¨ICML Workshop on New Frontiers in Learning, Control, and Dynamical Systems, 2023a.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper considers the problem of  modeling high-dimensional system dynamics from irregularly sampled snapshot data. The paper introduces multi-marginal stochastic flow matching, which extends simulation-free score and flow matching methods to the multi-marginal setting. The paper demonstrates the approach on synthetic benchmarks and single-cell gene expression data.

Overall, the reviews for this paper are positive, and the majority of reviewers were in favor of accepting the paper. The paper is well-written and has a strong application-focused motivation, while proposing novel methodology that uses simulation-free approaches to address to multi-marginal problem. 

In the next revision of the paper, we encourage the authors to address Reviewer rnGo's comment about clarifying the term MMOT and other contexts of the term in the literature. Additionally, please incorporate the references discussed with reviewers during the rebuttal and other changes suggested by the reviewers.
Decision: Accept (poster)