Modeling Complex System Dynamics with Flow Matching Across Time and Conditions
Authors: Martin Rohbeck
Edward De Brouwer
Charlotte Bunne
Jan-Christian Huetter
Anne Biton
Kelvin Y. Chen
Aviv Regev
Romain Lopez
Venue: ICLR 2025 Spotlight (2025)
Keywords: Flow Matching
dynamical systems
Primary area: learning on time series and dynamical systems

Abstract: Modeling the dynamics of complex real-world systems from temporal snapshot data is crucial for understanding phenomena such as gene regulation, climate change, and financial market fluctuations. Researchers have recently proposed a few methods based either on the Schroedinger Bridge or Flow Matching to tackle this problem, but these approaches remain limited in their ability to effectively combine data from multiple time points and different experimental settings. This integration is essential in real-world scenarios where observations from certain combinations of time points and experimental conditions are missing, either because of experimental costs or sensory failure. To address this challenge, we propose a novel method named Multi-Marginal Flow Matching (MMFM). MMFM first constructs a flow using smooth spline-based interpolation across time points and conditions and regresses it with a neural network using the classifier-free guided Flow Matching framework. This framework allows for the sharing of contextual information about the dynamics across multiple trajectories. We demonstrate the effectiveness of our method on both synthetic and real-world datasets, including a recent single-cell genomics data set with around a hundred chemical perturbations across time points. Our results show that MMFM significantly outperforms existing methods at imputing data at missing time points.

Published as a conference paper at ICLR 2025
MODELING COMPLEX SYSTEM DYNAMICS WITH FLOW MATCHING ACROSS TIME AND CONDITIONS
Martin Rohbeck12..41 Edward De Brouwer', Charlotte Bunnel.,, Jan-Christian Huetter', Anne Biton', Kelvin Y. Chen', Aviy Regev!.i, Romain Lopez!
GGenettehh, USA
HHeidelberg University, Germany
German Cancer Research Center, Germany
4 "Eropeaan Molecular Biology Lab, Germany
Stanford University, USA Osaka University, Japan
ABSTRACT
Modeling the dynamics of complex real-world systems from temporal snapshot data is crucial for understanding phenomena such as gene regulation, climate change, and financial market fluctuations. Researchers have recently proposed a few methods based either on the Schrödinger Bridge or Flow Matching to tackle this problem, but these approaches remain limited in their ability to effectively combine data from multiple time points and different experimental settings. This integration is essential in real-world scenarios where observations from certain combinations of time points and experimental conditions are missing, either be- cause experimental costs or sensory failure, and ultimately to uncover general principles shared across different temporal processes. To address this challenge, we propose a novel method named Multi-Marginal Flow Matching (MMFM), MMFM first constructs flow using smooth spline-based interpolation across time points and conditions and regresses it with a neural network using the classifier- free guided Flow Matching framework. This framework allows for the sharing contextual information about the dynamics across multiple trajectories. We demonstrate the effectiveness of our method on both synthetic and real-world datasets, including a recent single-cell genomics data set with around hundred chemical perturbations across time points. Our results show that MMFM signifi- cantly outperforms existing methods at imputing data at missing time points.
INTRODUCTION
Understanding the dynamics governing complex systems is essential for modeling their underly- ing mechanisms as well as for predicting future behavior, such as stock price fluctuations (Gontis et al., 2010), climate change (Franzke et al.. 2013), and gene regulation leading to cell differen- tiation or response to drugs (Yeo et al., 2021 Schiebinger et al. 2019). In many real-world ap- plications, measurements are collected at various time points and under diverse conditions, such as different environmental factors or experimental settings. Additionally, measurements are often unpaired across time and conditions. This may be attributed to the difficulty of matching patients across cohorts (Manton et al.) 2008) in a clinical context or to the destructive nature of single-cell genomics (Ding al. 2022) in modern experimental assays. This results in an extensive yet frag- mented view of the system, necessitating models that can effectively integrate these data.
Such data fragmentation is prevalent in the field of drug development. For example, researchers aim to document the response of cells to various genetic and chemical perturbations in order to under- stand mechanisms of action and predict responses novel conditions (Rood al. 2024). Histor- ically, technological and budgetary constraints limited researchers' ability to generate datasets that
This work conducted during internship Genentech. Correspondence regeva @ eeeevaa@ennecomm gome nomain.lopee@@yeele

---

Published as a conference paper at ICLR 2025
captured cell profiles either across time for single condition across multiple conditions but a single time point. Existing methods reflect these limitations, either modeling conditions without
temporal components Yu & Welch 2022] Lotfollahi et al. [023] Piran al. 2024 Bunne et al.
2022 2023) or modeling dynamics without interventions (Schulz et al. 2012 Tong et al. 2020,
Wang al. 2023). Recent advances in screening technologies enable cost-effective, large-scale measurement of cell populations over time under multiple conditions. This advancement motivates the development of new modeling strategies that can effectively leverage this combined complex- ity. Moreover, even with such advances, the sheer scale of biology and chemical space makes it prohibitive experimentally characterize the entire temporal and perturbation space. When human patient data are concerned, these abilities are further diminished. Thus, better models are essential to generalize to unseen lab experiments and into human biology.
To leverage the rich and complex data now available and deliver better models, we propose framing the problem of learning population dynamics as modeling the transport of the probability distribu- tion of cellular states across time and conditions. This approach allows us to capture changes over time and across conditions, while accommodating unpaired samples. Recent advances in genera- tive modeling, such as diffusion models (Ho et al. 2020) and Flow Matching (FM) (Lipman et al., 2023). constitute effective methods for learning mappings between arbitrary distributions. This makes them particularly suitable for our task. However, despite the popularity of these approaches, they have been hitherto limited to transporting between two marginal distributions. Although multi- ple time points could be naively modeled by decomposing the problem into series of two-marginal problems, such an approach would have significant limitations. It would fail to (1) leverage depen- dencies between the dynamics of multiple (more than two) related conditions (i.e., similar conditions are assumed to follow similar dynamics), and (2) incorporate prior knowledge or constraints on the long-range dynamics of the system.
To address the aforementioned challenges, we introduce Multi-Marginal Flow Matching (MMFM), a general Flow Matching framework that learns system dynamics from populations measured at multiple time points and under various conditions. Notably, our model uniquely leverages dynamic dependencies between conditions and enables both interpolation and extrapolation to unobserved time points and/or conditions. Through extensive experiments, we demonstrate MMFM's superior performance and generality compared to existing methods. We apply our MMFM to a real-world dataset of immune cells subjected to each of hundred kinase inhibitor perturbations measured over four time points and to a dataset of air pollutants concentrations. Empirically, our method successfully predicts cellular trajectories and air pollutants concentrations for unobserved conditions and time points.
2 BACKGROUND
In this section, we introduce the concepts of vector fields and probability density paths and then pro- vide an overview of the Flow Matching (FM) framework. Unless mentioned otherwise, X' denotes an Fuclidean space (i.e., X = R.d).
Vector Fields We model the underlying dynamics of a system using a time-dependen vector field [0, 1] x X + X. This vector field induces time-depente dependent diffcomorphic map X, called flow, 6: 00, 1] x X + XX defined by the Ordinary Differential Equation (ODE)
d
(x) = ue (x))),
(1)
dt
with the initial condition (b((x) = x. Given an initial probability distribution Po : X R+ the flow reshapes the distribution Po into a probability density path (z) defined via the furh-forward equation (xr) = [o)).) The resulting density for each time t is then obtained via the change of variable formula.
Flow Matching The original FM framework for generative modeling (Lipman et al. 2023 con- siders two marginal distributions: the source distribution po (e.g., Gaussian random noise) and the target distribution P1 (e.g., a set of images), both with support contained in X. The goal is to ap- proximate the target vector field u, that generates probability path p. (x) satisfying the boundary
We use the subscript notation for the time parameter, e.g., (z).
2

---

Published as a conference paper at ICLR 2025
conditions Pa and P1 with trainable vector field vo(. :0) parameterized by A. This leads to the Flow Matching objective:
(FM( (x || 8) un (x))2)
with t sampled from the uniform distribution t ~~ 24 1).. Given an approximation of the target vector field uy (z). one can easily sample from the distribution pl given samples from Po (or the other way around), as well as predict dynamics of individual samples over time. However, the objective function in Equation i2is generally intractable, as both p. and ", are unknown. Lipman et al. therefore introduced Conditional Flow Matching (CFM), tractable yet equivalent objective that aims to approximate a conditional vector field us z))
CCFM = E, ae,, [[])(r; us(a =)1(
(3)
The conditioning variable z and conditional probability paths pr (a =) are chosen such that the marginals match the boundary distributions Po and P1. In general, one chooses 2 as pair of samples from the source and target distributions (Zo, 21) according with marginals Po and p1- Remarkably, (0) and (O) are equivalent objectives as they have identical gradients with respect to (Theorem of Lipman et al. (2023).
joint distribution g(2) π (æg. 21)
Conditional Probability Paths
A common choice for the form of conditional probability paths
(r 2) = N( (=)))
(4)
where ; [0, 1 X' and [0, 1]] R+ represent the dime-dependent mean and standard deviation Gaussian distribution. Since the number of potential vector fields inducing a desired path is infinite, a reasonable approach select simple vector field, e.g., one that results flow of the form:
(x- (o-mo(z) (*= (~))(2) Go(2)
(1 12) (=() + o1(2))
(5)
The above functional form leads to a unique inducing vector field (Theorem in Lipman et al. 2023 defined as
(x [e] 2) d((2) or(2) (z, mc(2)) + p1(2).
Given this parameterization, different (z) and or (z) are possible, as long as the marginals coincide with the boundary conditions (i.e., Po(x =))(())d po(z) and f Pı(z =)q(=)dz P1 (x:). The most natural is to consider a small variance function at the boundaries with an interpolation function for (z),, such that po(=) = 3o and μi (z) FT. For instance, Tong al. (2024) consid- Tong ered a linear interpolation for the mean (=((2) = Lxg + (1 ()ær, with a constant (small) variance (z1). for their Independent-CEF. (I-CFM) method.
(2(2) > 0, and q(2) =
Optimal Transport Tong al. 2024 highlighted that sampling points 2g and xy independently (i.e., x1) = Po (o ( [z1) could lead to training instabilities and poor performance. Their improved method, OT-CFM, instead sampled using the optimal transport coupling between Po and Pr, rr" (1), and resulted in stabilized training and better performance.
CONDITION-AWAR MULTI-MARGINAL FLOW MATCHING
We first extend the Flow Matching framework to multiple temporal snapshots, and then present a condition-aware variant that leverages information across various marginals and contexts.
Problem Statement Generalizing the Flow Matching framework that considers two distributions
Po and P1 only, we assume the data sampled from a collection of distributions {pr, (r c)
[C], k [K + 1}} over C different conditions and K + 1 different sampling times. Given these observations, our goal is to learn an approximation of a target vector field u,(x c) that generates a probability path pa (xx |c) that satisfies the boundary conditions l, 4,...,, l for all conditions.
3

---

Conditional Probability Paths A key feature of MMFM is the specification of Gaussian prob- ability density path (r 2) that goes through all the time points (zo...., .K). In contrast to classical FM, the MMFM framework can naturally incorporate prior knowledge about the system's Published as a conference paper at ICLR 2025
No
Multi-Marginal Flow Matching
Figure I: Schematic overview of the data and our modeling approach. Left: Data observed across multiple time points and varying conditions, originating from a common source distribution (white), with "?" marking an unmeasured point of interest. Middle: Learning process. We train neural network to regress the derivative (arrww at points r (red) sampled around the interpolation path (equipped with time-de -eneeednddent variance, dashed gray lines). Right: Estimated trajectories after learning the underlying vector field allowing the generation of paths for novel samples, imputing missing time points, and generalizing unseen domains.
3.1 MULTI-MARGINAL FLOw MATCHING
We consider a collection of K + discrete time stamps (4.....,4k) Without loss of generality, we consider that the time stamps are distinct, ordered, and scaled so that 0 to tt .... < tk 1. We don't t impose any constraints on the spacing between consecutive time points, This flexibility accommodates scenarios where data collection occurs at irregular intervals, which often the case in real-world applications. At each discrete time point i,, we observe samples from a data density Pt, valued in X.
Flow Matching Objective We define = == (zo..... Rk) 2 X and for now, define
as the data distribution with independent samples from each marginal. q admits as density g(z) 11 (z&). Applying Equation 3 with the augmented vector leads to the Multi Multi-Marginal Flow
Matching (MMFM) objective:
[MMFM(() -AAna((0)BBarnng))e)eeep (app
[|o, [1,( - ux(x =>]3
for any probability density path Pr (æ 2) and conditional vector field uz 2). This objective has the same form Equation but the density ((2) and the probability density path are defined differently (over multiple time steps), resulting in a different target conditional vector field u, As demonstrated below, this remains valid surrogate objective function for Flow Matching across multiple time points (the proof appears in Appendix(A) (A}.
Proposition 1. Assuming that pe (x) O for all X and [0, 1], then, up to constant independent of 0, LFM and LMMFM are equal. Hence, for all values of the parameters 0:
[oLru(O) = (oLMMFM(O))
dynamics over multiple time steps. For physical systems, a meaningful prior for he is defined as an interpolating path of minimum energy or minimal curvature:
[x(2) = arg
min TEH*((to.tc))
J.
(1)(())
dt s.t. vg = (tx) for I (O....,K) (9)
where (Lo, (x)) denotes the set of functions whose first derivative is absolutely continuous [to, *x] and admit a weak second derivative. The celebrated Holladay's theorem (Holladay 1957) demonstrates that is the natural cubic spline interpolation between the points { +x) coefficients can be computed efficiently since they only require solving tri-diagonal system of linear equations.
J whose
Unlike the typical constant variance a, employed in CFM, we also use a time-dependent variance function defined piece-wise on each time interval as
4(1e+1 1)(1 44) (x(2)
(10)
(t++1 22
4

---

Published as a conference paper at ICLR 2025
for t E [Bk: fx+1l. Intuitively, this variance function adds noise to the vector field estimation in- between observed time points and is a crucial component for sharing information across conditions. Derivations for the temporal derivatives of (z) and (2 are provided in Appendix D
Optimal Transport Following the enhancements proposed by Tong et al. (2024) in the case with two marginals, we propose to sample data points across marginals using a joint distribu- tion, computed the solution of multi-marginal optimal transport (MMOT) problem (Pass) 2015): ((2) = = (Tg...., Ik). Notably, when the cost structure is pairwise additive for each pair of marginals, the MMOT problem reduces to set of K independent OT problems, such that
I (1x,0_))))) where ("E, Tk+1) is the solution of the OT problem be-
T T"......., Tx)
Ih== (xy]
tween Pt, and Pta+: (the proof is provided in Appendix B). Throughout this work, assume such a pairwise additive cost structure. We pre-compute the solutions of the OT problems on the training data using the squared Euclidean distance When the dataset is too large, OT can be approximated using mini-batches, TTog et al. 2024), although this was not necessary our experiments.
cost function and use the multi- multi-marggial optimal
transport coupling for q in Equation]] 7
Relationship with CFM MMFM is generalization of CFM in the following sense. When the conditional vector field is defined with a piece-wise linear interpolation function for the mean and constant variance, and in addition, the neural network U, is defined piece- piece-wise on [to, fx), the MMFM problem is equivalent to solving K distinct CFM problems between consecutive time points.
Let us assume that for all 2 = (Io...., E 2, (c(2) is defined as the piecewise linear function going through all the points of 2, and 0x(2) = a is constant for all 2. Additionally, lel us assume that the vector field is learned separately on each time interval: ve(2; 0) = EE= (2(2;; 0x) (tr (tn,))))( Then, the MMFM problem is equivalent to solving K sepa- rate CFM problems between each pair of consecutive marginals.
Proposition 2.
The proof appears in Appendix A This result highlights key advantage of MMFM over multiple pairwise CFMs when the system's s behavior exhibits similarities across time periods (or conditions, below). In such scenarios, MMFM can utilize a single set of parameters model the dynamics across all time points (or, different conditions. below). This approach allows MMFM to leverage common patterns in the data that would be treated independently in separate pairwise CFM prob- lems, leading to more efficient and robust modeling, and deeper understanding of the underlying phenomena (c.g., biological commonalities).
3.2 CONDITIONAL MMOT-BASED MALLTIMARGINAL FLOW MATCHING
We now extend our framework to consider cases where each observation z from time point ty is associated with a condition {1...., C}. We assume the condition information is provided for each observation. We use the notation of ~ (2g c) to indicate the data density for condition at time point t4. Additionally, we define 2 (=) ......... +x)). Importantly,
model can handle
incomplete data scenarios, where not all conditions are observed at all time points. This extension enables us capture and analyze condition-specific dynamies, such the response of cells to different drugs in a chemical screen.
We leverage the conditional information in our architecture by (1) modifying the coupling ((2) to be condition specific, and (2) extending to architecture of the vector field neural network i/ to use the condition as input.
For the coupling, we use condition-wise multi- marginal optimal transport coupling:
I T (of xh+1)
= :.
(11)
l Pta (1f. c)
where " is the optimal transport coupling between pt, (of c) and Pta+a (+ c). and the second equality follows from the assumption of pairwise additive structure of the cost function (Section 3.1
and proof in Appendix B)
For the vector field neural network, we integrate the condition an additional input to the neural model, modifying it to v (z, c; 8). This neural network takes as input a condition index c that is used to internally retrieve learnable condition embedding (details about the architecture are provided in
5

---

Published as a conference paper at ICLR 2025
Appendix D). Importantly, we learn set of shared weights 8 across all conditions, which enables the model to generalize across different conditions and potentially improve performance on conditions with limited data.
To train the conditional model, we employ classifier-free guidance as proposed by Zheng et al. (023). The objective for the condition-informed version of MMFM, termed C-MMFM, is defindd as:
CCMMMMM Be..0--Br(rs,.coCCCC((1... poops ve (1 b)e begz; u 25)
(12)
where Pu represents the probability of switching to the unconditional model, and eæ represents the null conditioning. The null conditioning Ca is an additional condition index that retrieves a learnable embedding corresponding to condition-inded indition-indeeeeneent model. When the Bernoulli variable b is 1, the model uses a conditiionll-indeeeneen model. Otherwise, the conditional model is used. This approach allows information sharing across conditions via joint training of conditional 2n (2, c; 0) and unconditional ey (2, Ca: () models.
At inference time, we use a weighted combination of the two vector fields (details about the sampling appear in Appendix C
(1, a w. 0) (1 w) (ar, cw; 0) + wee o;
(13)
and therefore can model the trajectory of sample zo in condition e es:
Rr = xn + / (r G w. O)dl.
(14)
The parameters Pu and are selected based on the value of the objective function on held-out validation data.
4 RELATED WORK
Flow Matching Flow-based generative models have gained significant attention in recent years due to their ability to efficiently model generative processes for complex data distributions. Flow Matching, introduced by Lipman al. 2023), proposed a simulation-free approach to train neural ODEs for generative modeling computationally attractive alternative to maximum likelihood training of continuous normalizing flows (Chen et al. 2018)). Building upon this, Tong et al. (2024) extended FM by (1) allowing flows between arbitrary distributions (I-CFM), (2) proposing OT- CFM coupling samples using optimal transport theory (linking FM to dynamic optimal transport), and (3) introducing Schrödinger Bridge Conditional FM (SB-CFM), linking FM to Schrödinger Bridges. Kapusniak et al. (2024) improved FM by estimating the data manifold to adjust the linear interpolation path and ensuring it remains within high-density regions of the data. However, none of these methods takes into account multiple (more than two) time points.
Label-Guided Flow Matching Zheng et al. (023): ; Dao et al. (2023); Isobe et al. (2024) first integrated conditional information into FM. Guided Flows (Zheng et al. 2023) extended the FM framework to incorporate conditional information, allowing for more precise control over the gen- eration process. At the same time, [ao et al. (2023) modeled the flows space, offering improved computational efficiency. Meta Flow Matching (Atanackovic et al.) 2024) extends traditional FM by modeling the flow over populations, egg., a set of cells from a patient, by embedding it using Graph Neural Networks. As with other Flow Matching studies, none of these methods consider multiple time points.
jointly learned latent
Fast and Smooth Interpolation The approach in Chewi et al. 2021 shares similarities with MMFM in its use of OT principles to guide the interpolation process but differs in its specific im- plementation by directly using splines rather than neural network- based Flow Matching and cannot handle multiple conditions. Therefore, it also has the disadvantage of not being able to make infer- ences test time using learned dynamics from other conditions.

---

Published as a conference paper at ICLR 2025
Multi Marginal Schrödinger Bridges In contrast to Flow Matching, Schrödinger Bridge (SB) models, which are based on Stochastic Differential Equations (SDEs) and therefore predominantly estimated using Diffusion Models, have already been extended to make use of multiple marginals. The Deep Multi-Marginal Momentum Schrödinger Bridge (DMSB), introduced by Chen et al. (2023), extends the classical SB problem to scenarios with multiple marginal distributions by ad- justing the Bregman iteration approach.
Modeling of biological single-cell time series data Understanding cellular mechanisms from biological time series data or pseudo time series data is a popular research area, as it helps understand fundamental dynamic processes in cells, from differentiation during development
pathogenesis
(Schiebinger et al. 2019 Weinreb et al. 2020, Setty ct al. 2019; Yeo et al. 2021 Qiu et al. [022]
Farrell et al., 2023] Bergen et al., 2020 Tong al.] 2020 Zhang al. 2024). Advancements in
single-cell profilin have enabled researchers to capture temporal snapshots of individual cell states across cell populations, from which computational methods can infer temporal trajectories and the underlying gene regulatory networks that control them. Despite the fact that temporal processes both span across multiple time points and are often shared across conditions, models to date have focused on either a single condition monitored along multiple time points (through modeling consecutive pairs) or on a series of conditions measured in a single time point.
EXPERIMENTS
We assessed the performance of MMFM using synthetic data well single-cell RNA-seq dataset where each of multiple conditions (perturbations) was measured along multiple time points. We compared the performance of MMFM to that of several other methods, first considering methods explicitly designed to incorporate multiple time points. We considered FSI (Chewi et al. (021). a method interpolating data densities using OT and natural cubic splines, and as an ablation study, we assessed the performance of a variant of MMFM with a piece-wise linear interpolation function (L-MMFM). We also sought report the performance of DMSB (Chen et al. 2023) but did not see meaningful outputs after a long training time and, therefore. were not able to draw conclusions. We therefore considered more direct applications of the FM framework. applying OT-CFM (Tong et al.] 2024) across consecutive pairs of marginals, and applied separately per condition (PCFM). Fi- nally, we also compared to an OT-CFM model that would ignore all intermediate time points (CFM). We used three metries for evaluation: Mcan-Soe Seuun-Suared-aaae Error (MSE), Maximum-Mee Discrepancy (MMD), and Wasserstein-2 (W2) distance. The MSE provides a straightforward comparison of clus- ter centers, while MMD and W2 offer more sophisticated assessments of distributional differences by considering higher-order moments and geometric properties. To ensure a fair comparison, we provided the condition-specti specific OT couplings input to all models. Additional details about these experiments, including hyperparameter search grids and hold-out-based validation strategies, are provided in Appendix D Considerations about computational complexity are given in Appendix]] !]]
5.1 SYNTHETIC EXPERIMENTS
To evaluate MMFM's 's ability to estimate condition-specific vector fields from time course data, we created two-dimensional synthetic datasets with a known ground truth vector field. Data were gen- erated for 12 conditions, defined as Cm m for m E (1....,1 12} using the vector field
Uy (2 cm) =
( + 1)7/4 cos/44t))(5
(15)
0, Cm 2
xo(m)= =
4
We sampled from the time points and conditions described in Table 1 An overview of the data and the actual underlying trajectories appear in Figure AA.. We generated 50 data points per condition and time point according to the initial value problem above and included an additive Gaussian noise with variance o? This scenario is similar to that in single-cell genomics, where the experimental assay is destructive, so each cell (sample) can be measured at most once. We held-out time point t = 0.15 during training for hyperparameter tuning of all FM-based models (but left it as input for the FSI model).
7

---

Published as a conference paper at ICLR 2025
Table I: Sampling times and conditions for the synthetic experiments.
0
0.1
0.15
0.3
0.5
0.7
0.9
1.0
V
v
C5 Rest
V
Note: v indicates that samples from that combination available in the training set.
A.
Vector iPhose Diagram) with Training Data
B.
Predicted
C.
Predicted Trajectory MMFM
3
6.5
25
33
213
225
Figure 2: Assessment on synthetic data. (A) Ground truth vector field, which incorporates training samples for ten distinct conditions, each represented by different color. (B,C) Interpolated paths generated by FSI (B) and MMFM (C) illustrated for conditions E [3, 5, 12).
For evaluation, we compared deviations from the ground truth the mean level, estimated with the Mean-Square Error (MSE), and at the distribution level, estimated by the Maximum Mean Dis- crepancy (MMD) between predicted samples and samples from the ground-truth vector field at sgudddstanttspaced time points (from t = 0 to = 1, by increments of At = 0.05).
The experimental results (Table ( demonstrate the superior performance of MMFM on this time course data, as it most effectively predicts trajectories for conditions and cs. where measurements are sparse and irregularly spaced. While FSI struggles to capture the right dynamics for the inferred trajectories (Figure 2) and systematically underestimated the number of inflection points, MMFM accurately adjusted the trajectories and aligned them with the curvature profile of other conditions (visualization of interpolation quality for other conditions and methods appear in Appendix E
FSI and PCFM both cannot leverage information from other conditions, leading to incorrect trajec- tories and much higher errors. By contrast, CFM and MMFM have access multiple conditions. However, CFM only considers the first and last time points, resulting higher loss than MMFM. Finally, L-MMFM performs substantially worse than MMFM, suggesting that the cubic splines effective models for the dynamics of this dataset. Further results are provided in Appendix F where we also discuss the extrapolation properties of MMFM.
Together, these results validate that MMFM is beneficial in scenarios where one has access to mea- surements taken different time points and across different conditions.
Table 2: Results on synthetic data. We report the mean square error (MSE) and maximum mean diserepancy (MMD), with lower values indicating better performance. Mcans and standard deviations are computed over 21 time points. Best-performing models are highlighted in bold.
Cm 3
MSE J
MMD J
MSE |
MMD 1
ESI
0.72 + 0.60
0.42 0.37
0.72 + 0.93
0.33 # 0.45
CFM
0.46
0.25 + 0.15
0.93 + 0.52
0.60 + 0.33
PCFM
0
0.45 0.35
0.78
0.43 # 0.44
L-MMFM (ours)
0.39 # 0.25
0.20
0.36
0.17 0.19
MMFM (ours)
0.13 + 0.08
0.04 0.02
0.22 0.15
0.12 0.08
8

---

Published as a conference paper at ICLR 2025
points
Drug Labels
Drug Category
JANSTAT
More
No
OTher
4000 Mad
Weic
KAAIUMAPI
MOVI UIMAPY
NOVI UHARP
Figure 3: UMAP embeddings of single-cell RNA sequencing data visualizing drug responses over time. (A)
Cells colored by time point (Oh, 24h, 48h, 72h). highlighting the temporal progression of cellular states. (B)
The same embedding colored by individual drugs (93 distinct labels), showcasing response heterogeneity. (C)
Cells colored by drug categories, highlighting shared response patterns within drug classes. UMAP coordinates
were derived using scVI (single-cell Variational Inference) (Lopez et al. 2018
5.2 APPLICATION TO SINGLE-CELL PERTURBATION SCREENING DATA
We assessed the performance of MMFM on a world single-cell perturbation screen dataset. The dataset records single-cell gene expression (18, 250 genes) profiles of T cells treated with kinase in- hibitors and undergoing activation, measured at four time points (Oh, 24h, 48h, and 72h). The dataset includes 93 distinct inhibitors, each used at varying concentrations (100 nM, 1 μM, and 10 μM), as well as negative control conditions (vehicle and non-activation). In each condition (combination of inhibitor, concentration and time point), hundreds of cells have been profiled. We treated the non- activation data as the & = Oh time point for all conditions. To focus on the most significant effects, we filtered the data and retained M 123 distinct treatments, each representing a unique combina- tion of compound and concentration (Figure 3p. We then randomly selected 60 of those treatments for analysis. For evaluation purposes, we withheld ten non-overlapping random treatments for each of the three time points. Standard single-cell data processing pipelines were used to normalize and scale the features, with the first 25 principal components used as input features (details regarding data preparation appear in Appendix HH..
Because the ground-truth vector field is unknown in this setting, we evaluated the capacity of the methods to generalize to held-out time points. We thus assessed model performance at the mean level, estimated with the Mean-Square Error (MSE), and at the distribution level, estimated by the Wasserstein distance W2 between predicted and hold-out samples. The training was conducted with five random seeds, using one held-out treatment to select the best model for testing. For this benchmark, the PCFM baseline was applied once using data from all conditions, as it otherwise required training 150 individual models, which was prohibitively conecconu.ing.
Table 3: Results for the drug response imputation task. We report the mcan-square error (MSE) between the predicted and actual distributions' means and Wasserstein distance (W2), where in both cases lower is better. Means and standard deviations are computed over five folds. Best best-performing models are highlighted in bold.
24h
48h
72h
MSE
W2
MSE 4
W2
MSE
W₂
FSI
07.77 + 01.49
56.13 = 13.44
04.52 01.14
38.93 # 08.56
06.62 # 01.65
78.46 = 23.48
CFM
09.10 01.59
60.86 14.67
04.92
33.28 + 13.53
PCFM
08.88 01.58
56.68 = 12.12
01.21
34.85 + 08.01
06.07 = 01.32
61.83 16.74
L-MMFM (ours)
08.05 +: 01.29
50.82 # 10.79
03.56 01.50
26.07 +: 08.62
06.30 + 01.92
47.97 19.69
MMFM (ours)
07.54 00.99
51.86 + 07.84
03.38 + 01.34
26.01 07.72
04.92 01.61
37.71 12.24
Our experimental results (Table 3b demonstrate the superior performance of MMFM at imputing missing data across multiple time points. Interestingly, MMFM and L-MMFM have similar per- formance for the f = 24h and t = 48h, but MMFM significantly outperforms L-MMFM at the t = 72h time point. This suggests that the cubic spline is a more desirable prior for this modeling problem. In contrast, the reliance of PCFM on linear paths may introduce a bias towards straight- line trajectories, limiting its flexibility in capturing complex dynamics. Overall, MMFM demonstrates robust performance in handling complex datasets with missing temporal information, emphasizing its potential for diverse applications in generative modeling.

---

Published as a conference paper at ICLR 2025
5.3 APPLICATION TO BEIJING AIR QUALITY DATA
To further study how well MMFM generalizes, especially under irregular sampling over time, applied to the Beijing multi-site air quality data set (Chen 2017). This dataset comprises hourly air pollutant data from 12 air-quality monitoring sites across Beijing. We focused on PM2.5 data, which measures the density of particulate matter smaller than 2.5 micrometers, an important air pollution indicator, between January 2015 and February 2017. The 12 monitoring sites were modeled as the condition (c E {1..... 1. ). For each site, we grouped together the measurements collected over the same month, resulting in 26 temporal snapshots (t E {1...., 263). We selected different months as training data to simulate irregular sampling, with the interval between successive snapshots varying between 1 and 7 months. For 9 out of 12 stations, we selected 50% of the measurements, i.e. 13 months. For the other three stations, we selected only 7, 6 and 7 months as training data to simulate missing sensor data. These three stations are represented by the conditions = 4, c = 7 and 10. We evaluated our method on all months that were not part of the training data set. This experiment represents a common challenge in sensor networks, where some measurements are lost or corrupted and must be imputed using data from other sensors. The results (Table 4 indicate that both MMFM versions (cubic spline and linear paths) show strong performance compared to FSI, CFM, and PCFM when it comes to interpolating the measurements. L-MMFM models show the best performance across all evaluated methods, as measured by MSE, MMFM is in second place in three out of four scenarios. Further details and evaluation metries can be found in Appendix G
Table 4: Results on the Beijing data. We report the maximum mean discrepancy (MMD), with lower values indicating better performance. Best-performing models are highlighted bold.
7
c 10
Rest
FSI
1.70 + 0.40
1.89 # 0.16
1.92
0 0.10
1.47 + 0.54
CFM
1.16
1.78 0.35
1.85 0.21
1.83 + 0.30
PCFM
1.61
0.36
0.15
1.50 3:0+000 0.50
L-MMFM (ours)
1.33
1.34 x
1.24
1.33 i 0.59
MMFM (ours)
1.52
0.49
1.32 = 0.50
1.52 + 0.51
6 DISCUSSION
We introduced Multi-Marginal Flow Matching (MMFM), a novel method for modeling complex system dynamics from temporal snapshot data across multiple conditions. We demonstrated that MMFM effectively combines data from various time points and experimental settings, outperform- ing existing methods in imputing missing time points. Our approach showed particular strength in scenarios with sparse or irregularly spaced measurements, leveraging information across conditions.
An important parameter in the design of MMFM is the time-varying variance function o, as well as the sharing of parameters across conditions. An interesting direction for future work would be to add set of parameters to the probability density path used in Flow Matching and provide a principled way to learn those parameters while simultaneously solving the FM problem. For example, one could use latent attentive neural processes (Kim et al. 2019) to construct meta-mode models across conditions. This approach could alleviate the need to find the right set of hyperparameters, ensuring that the conditioning network has sufficient, but not excessive, influence parameters. Another promising avenue for constructing flows on real-world data is to explore more complex cost functions for multi-marginal optimal transport, (e.g., the energy landscape-based loss functions for biological applications, Appendix D-
the model
MMFM opens up new modeling opportunities for single-cell RNA-seq datasets. For example, so-called cellular velocities can be estimated from data either using nascent (unspliced) RNA, or metabolic labeling (La Manno et al.) 2018: Bergen et al., 2020). Despite being inherently noisy, this additional information could be valuable addition to enhance and steer the synthetic gradient flows based on interpolation paths, as previously demonstrated by Tong et al.) (2020). Further- more, constraining the model to satisfy specific stochastic differential equations, such as Ornstein- Uhlenbeck processes [Wang et al., 2022), could help bring interpretability, as well as mechanistic insights. Indeed, such models have analytical steady-state solutions and also accommodate for in- terventions (Rohbeck et al. 2024) (e.g., modeling gene knock-outs). Lastly, biomedical data often includes various modalities providing multiple measurements expression and chromatin accessibility). Access to these multi view datasets can allow for modeling coupled trajectories across different spaces (Somnath et al. 2023).
the same target (e.g.. a cell's gene
10

---

Representation Learning and Published as a conference paper at ICLR 2025
REFERENCES
Lazar Atanackovic, Xi Zhang, Brandon Amos, Mathieu Blanchette, Leo Lee, Yoshua Bengio, Alexander Tong, and Kirill Neklyudov. Meta flow matching: Integrating vector fields on the Wasserstein manifold. In ICML Workshop on Generative Modeling, 2024.
Volker Bergen, Marius Lange, Stefan Peidli, F Alexander Wolf, and Fabian Theis. Generalizing RNA velocity to transient cell states through dynamical modeling. Nature Biotechnology, pp- 1408 1408-1414, 2020.
Charlotte Bunne, Andreas Krause, and Mareo Cuturi. Supervised training of conditional Monge maps. In Advances in Neural Information Processing Systems, volume 35, 2022.
Charlotte Bunne, Stefan G Stark, Gabriele Gut, Jacobo Sarabia Del Castillo, Mitch Levesque, Kjong-Vate Van Lehmann, Lucas Pelkmans, Andreas Krause, and Gunnar Rätsch. Learning single- cell perturbation responses using neural optimal transport. Nature Methods, 20(11)1179-11788, 2023.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. Advances in Neural Information Processing Systems, 31, 2018.
Song Chen. Beijing Multi-Site Air Quality. UCI Machine Learning Repository, 2017. DOI: http:://doo.og//022444//CCC
Tianrong Chen, Guan-horng Liu, Molei Tao, and Evangelos A Theodorou. Deep multi-marginal momentum Schrödinger bridge. In Advances in Neural Information Processing Systems, pp- 77058-57086, 2023.
Sinho Chewi, Julien Clancy, Thibaut Le Gouic, Philippe Rigollet, George Stepaniants, and Austin Stromme. Fast and smooth interpolation Wasserstein space. In International Conference Artificial Intelligence and Statistics, pp. 3061-3069. PMLR, 2021.
Quan Dao, Hao Phung, Binh Nguyen, and Anh Tran. Flow matching in latent space. arXiv, 2023,
Jun Ding, Nadav Sharon, and Ziv Bar-Joseph. Temporal modelling using single-cell transcriptomics. Nature Reviews Genetics, 23(6):355-368 2022.
Pavel Dvurechensky, Alexander Gasnikov, and Alexey Kroshnin. Computational optimal transport: Complexity by accelerated gradient descent is better than by sinkhorn's 's algorithm. Interna-
tional conference machine learning, pp. 1367-1376. PMLR, 2018.
Spencer Farrell, Madhav Mani, and Sidhartha Goyal. Inferring single-cell transcriptomic dynamics with structured latent gene expression dynamics. Cell Reports Methods, 3(9), 2023.
Christian Franzke, Terence J Kane, Judith Berner, Paul D Williams, and Valerio Lucarini. Stochastic climate theory and modeling. Wiley Interdisciplinary Reviews: Climate Change, (1):63-78, 2015.
Gontis, Ruseckas, and A Kononovičius. A long-range memory stochastic model of the return in financial markets. Physica A: Staristical Mechanics and its Applications, 189(1):100-10))) 2010.
Isabel Haasler, Axel Ringh, Yongxin Chen, and Johan Karlsson. Multimarginal optimal transport with a tree-structured cost and the Schrödinger bridge problem. SIAM Journal on Control and Optimization, (9(4))2428-244 2021.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6853 6851, 2020.
C. Holladay. A smoothest curve approximation. Mathematical Tables and Other Aids to Compu- tation, 11 (60):2333-23 1957.
Noboru Isobe, Masanori Koyama, Jinzhe Zhang, Kohei Hayashi, and Kenji Fukumizu. Extended flow matching: a method conditional generation with generalized continuity equation. arXiv, 2024.
11

---

Published as a conference paper at ICLR 2025
Kacper Kapusniak, Peter Potaptchik, Teodora Reu, Leo Zhang, Alexander Tong, Michael Bronstein, Avishek Joey Bose, and Francesco Di Giovanni. Metrie flow matching for smooth interpolations on the data manifold. In Advances in Neural Information Processing Systems, 2024.
Hyunjik Kim, Andriy Mnih, Jonathan Schwarz, Marta Gamelo,, Ali Eslami, Dan Rosenbaum, Oriol Vinyals, and Yee Whye Teh. Attentive neural processes. In International Conference on Learning Representations, 2019.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. In International Conference on Learning Representations, 2015.
Gioele La Manno, Ruslan Soldatov, Amit Zeisel, Emelie Braun, Hannah Hochgerner, Viktor Petukhov, Katja Lidschreiber, Maria Kastriti, Peter Lönnerberg, Alessandro Furlan, al. Rna velocity of single cells. Nature, 560(7719):494 2018.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. In International Conference on Learning Representations, 2023.
Romain Lopez, Jeffrey Regier, Michael Cole, Michael Jordan, and Nir Yosef. Deep generative modeling for single-cell transcriptomics. Nature Methods, 15(12): 2018.
Mohammad Lotfollahi, Anna Klimovskaia Susmelj, Carlo De Donno, Leon Hetzel, Yuge Ji, Igna- cio L Ibarra, Sanjay R Srivatsan, Mohsen Naghipourfar, Riza M Daza, Beth Martin, et al. Predict- ing cellular responses to complex perturbations in throuhhroughput screens. Molecular systems biology, 19(6)e1 1517, 2023.
Kenneth G Manton, XiLiaag Gu, and Gene R Lowrimore. Cohort changes in active life expectancy in the us elderly population: Experience from the 1982-2004 national long-term care survey, The Journals of Gerontology Series B: Psychological Sciences and Social Sciences, 63(5):526--S281, 2008
Brendan Pass. Multi-marginal optimal transport: theory and applications. ESAIM: Mathematical Modelling and Numerical Analysis, 49(6):1711-1799) 2015.
7oe Piran, Niv Cohen, Yedid Hoshen, and Mor Nitzan. Disentanglement of single-cell data with Biolord. Nature Biotechnology, pp. 1-6, 2024.
Xiaojie Qiu, Yan Zhang, Jorge D Martin-Rufine Chen Weng. Shayan Hosseinzadeh, Dian Yang. Angela N Pogson, Marco Y Hein, Kyung Hoi Joseph Min, Li Wang, al. Mapping transcriptomic vector fields of single cells. Cell, 185(4)6600-71,, 2022
Martin Rohbeck, Brian Clarke, Katharina Mikulik, Alexandra Pettet, Oliver Stegle, and Kai Ueltzhöffer. Bicycle: Intervention- based causal discovery with cycles. In Causal Learning and Reasoning, pp. 209-242. PMLR, 2024.
Jennifer E. Rood, Anna Hupalowska, and Aviv Regev. Toward a foundation model of causal cell and tissue biology with a perturbation cell and tissue atlas. Cell, 187(17)445204455, 2024.
Geoffrey Schicbinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Arych Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. Optimal- analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell, 176(4):928-943, 2019.
Marcel H Schulz, William E Devanny, Anthony Gitter, Shan Zhong, Jason Ernst, and Ziv Bar- Joseph. Drem 2.0: Improved reconstruction of dynamic regulatory networks from time-series expression data. BMC systems biology, 6:1-9, 2012.
Manu Setty, Vaidotas Kiseliovas, Jacob Levine, Adam Gayoso, Linas Mazutis, and Dana Er. Characterization of cell fate probabilities in single-cell data with Palantir. Nature Biotechnology, 37(4):451-460, 2019.
Vignesh Ram Somnath, Matteo Pariset, Ya-Ping Hsieh, Maria Rodriguez Martinez, Andreas Krause, and Charlotte Bunne. Aligned diffusion Schrödinger bridges. In Uncertainty in Artificial Intelli- gence, pp. 1985-1995. PMLR, 2023.
12

---

Published as a conference paper at ICLR 2025
Alexander Tong, Jessie Huang, Guy Wolf, David van Dijk, and Smita Krishnaswamy. Trajecto- rynet: A dynamic optimal transport network for modeling cellular dynamics. In International Conference Machine Learning, 2020.
Alexander Tong, Kilian Fatras, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research, 2024.
Konstantin Tsiolkovsky. Reactive flying machines. Izdatel'stve Akademii Nauk SSSR: Moscow,
Russia, 1954.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances Neural Infor- mation Processing Systems, pp. 5998-6008 2017.
Lingfei Wang. Nikolaos Trasanidis, Ting Wu, Guanlan Dong, Michael Hu, Daniel E Bauer, and Luca Pinello. Dictys: dynamic gene regulatory network dissects developmental continuum with single-cell multiomics. Nature Methods, 20(9):1368-133 1378, 2023.
Caleb Weinreh, Alejo Rodriguez-Frattielli, Fernando D Camargo, and Allon M Klein. Lineage tracing on transcriptional landscapes links state to fate during diffcrentiation. Science, 367(6479): caaw3381, 2020.
Grace Hui Ting Yeo, Sachit Saksena, and David K Gifford. Generative modeling of single- cell time series with prescient enables prediction of cell trajectories with interventions. Nature communications, 12(1):3222, 2021.
Hengshi Yu and Joshua D Welch. Perturbnet predicts single-cell responses to unseen chemical and genetic perturbations. BioRxiv, pp. 2022-07, 2022.
Jiaqi Zhang, Erica Larschan, Jeremy Bigness, and Ritambhara Singh. scnode: generative model for temporal single cell transcriptomic data prediction. Bioinformatics, 40: 10:146-154, 2024.
Qinqing Zheng, Matt Le, Neta Shaul, Yaron Lipman, Aditya Grover, and Ricky T. Q. Chen. Guided flows for generative modeling and decision making. arXiv, 2023.
13

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS AND DISCLOSURES OF FUNDING
The authors would like to thank Maciek Wiatrak for fruitful discussions.
Disclosures: This work done while Martin Rohbeck Lopez, Edward De Brouwer, Charlotte Bunne, Anne Biton, Jan-Christian Huetter and Aviv Regev employees Genentech and/or have equity in Roche. Aviv Regev is co-founder and equity holder of Celsius Therapeutics and an equity holder in Immunitas. She was an SAB member of ThermoFisher Scientific, Syros Pharmaceuticals, Neogene Therapeutics, and Asimov until July 31, 2020, and has been an employee of Genentech since August 1, 2020.
intern Genentech. Romain
CODE AVAILABILITY STATEMENT
The code to reproduce the figures and tables, as well as to run the model and generate the simulated data, be found at ggthuu.ccom/GGeeeetch/MMMMMM
DATA AVAILABILITY STATEMENT
Input data used for the experiments in this manuscript are currently undergoing publication as separate manuscript. They will become publicly available upon acceptance.
APPENDICES
The following list gives an overview of the structure of the Appendix:
In Appendix A present the proofs for the theoretical results behind our Multi-Marginal
Flow Matching (MMFM) framework.
we go into the details of Multi-Marginal Optimal Transport (MMOT), explaining how it reduces to sequence of pairwise Optimal Transport problems in our specific case.
In Appendix B
In Appendix
C present the algorithm for sampling from the COT-MMFM model.
In Appendix
outline model training and evaluation approach, including hyper-
parameter tuning strategies.
In Appendix we provide additional experimental details and visualizations for our syn- thetic data experiments.
In Appendix we showcase an additional synthetic experiment focused on the extrapola- tion capabilities of our MMFM model.
In Appendix
f provide further details preprocessing and training-validationn validation-test splits on the Beijing weather data set.
Appendix H describes the preprocessing steps and treatment selection criteria for the per- turbation screening data used in our real-world experiments.
In Appendix ! we offer deeper motivation for using the minimum curvature prior proxy for energy minimization along a trajectory.
Finally, in Appendix ! including the optimal transport step.
provide detailed analysis the complexity of the method
A PROOF OF PROPOSITIONS
Proposition 1. Assuming that pr()) > O for all x E X and E [0, 1}. then, up to constant independent of 0, Lew and LUMPEM are equal. Hence, for all values of the paramelers 0:
voLru(0) = V LoM
(8)
Proof. Following Lipman et al. 2023 we begin by constructing the target probability path (x) from the conditional probability path p. (x: z). Given particular draw from the sample distribution
14

---

Published as a conference paper at ICLR 2025
(Zu..... ................................ IK), design the conditional probability path pz 2) such that it satisfies the following conditions:
Wk E {O..... ,,K,,,,,, K}, Pre (x 2) = Normal (ax,, of1))
(16)
where o > O is small so that Pt, (x: =) is well concentrated around Ty- Marginalizing against the distribution g(z), we obtain the marginal probability path:
p((x) Pc(21 2)g(2)dz.
(17)
For each time point tRs the marginal probability Pta closely approximates the k-th marginal:
/
26 gt (2k))da's -- Qt (Ix),
Ptz (x)
(18)
where f is the density of the d-dimensional isotropic Gaussian distribution.
Similarly to the concept of marginal probability paths, we can also introduce the marginal vector field (assuming (r()) 0 for all æ E X):
Po(x 2)
u (x(x) ua(x 2) pe(x)
(())dz,
(19)
where (1 =) is the conditional vector field that generates pr (3: 2). This is guaranteed by applying the proof of Theorem from Lipman et al. (003)..
With this proper setup, the remainder of the proof follows closely the one from Theorem in Lipman et al. 2023). We reformulate it below for completeness.
We assume that q, Pe z) decreasing to sufficient speed ||ar|| and that Ut, Vav; are bounded. We start to derive from the left-hand side:
(eL = Volpu(a)) 11 (2; 0) ur(x)112
0)11 ((m 0), (); + [[xx(())1)
Vo [|ua(1; 0]12 (xx (), (), u,(x))
and rewrite the right-hand side:
VoLMMFN(()) (eEuttt)mmt ux, 2)]
t VoBa(t))
[<; 0) (1v 0), ux 2)) + |/ue( =))p²)
- ((()) Vo [1: 0) (xv 0). u. 2)}
where we drop all terms independent of 0 (*) and use the bilinearity of the 2-norm (t). First, we show that the first terms of both equations are equal.
/ [1v. 0)
(a)) lo (z; 0)]
/ [|. 0)| (f pe( ==)(())dz)
/ [1: p))
(x 2)4 22d22
E (e (1]
where we use Fubini's theorem to change the order of the integrals (*), and the law of total expecta- tion (t).
15

---

Published as a conference paper at ICLR 2025
Second, show the equality of the latter parts. To get started, we plug the definition from Equa- tion 19 into the term of interest:
K { /ul(s
z)pr(a =)q(=)dz
(x) (1 f), u, (x())
De (o((o))))
Pe(x))dz.
(20)
(e(x)
Thanks to the linearity of the integral, we may take the inner integral outside of the inner product:
(a)(s) (cz 0), un (x)) = / (2; ue(x z)pr(2 x))(=)) pilz))daz..
(21)
(((x)
Employing the linearity of the inner product, we may rearrange the terms as:
(x) (ve(2; ux = f (va( 0), u4(x) =)) (<=)pe(x =)dedx.
(22)
Finally, using Fubini's theorem to change the order the integrals and the law of total expectation, recognize the following term:
Ep.(x) (re( 0), ue (x)) (111)))) ue(x =));
(23)
which concludes our proof.
Let us assume that for all = (Ko.....,KK) E Z, pr(z) is defined as the piecewise linear function going through all the points of 2, and oz(2) = a is conslant for all 2. Additionally, let us assume that the vector field is learned separately on each time interval: ve(z; 0) = >h=n (e 0g)) ( (te)e Then, the MMFM problem is equivalent to solving K sepa- rate CFM problems between each pair of consecutive marginals.
Proposition 2.
Proof. We remind the reader of the expression of the MMFM optimization problem:
MMMFM (0) = F,, a (xxp) [1u 0) u 2)13
(24)
We first notice that this time integral from 0 to may be broken up into each the segments in between observed marginal distributions:
K
(MMFM(()
=
Σ IN.
(25)
k=
where we define the integral Ig as;
16+1
[1o( 0) ue(x 1=) dt.
(26)
The main goal of this proof is to see whether the integrand each 1ş depends only on information from the marginals It, and Itot and can be related to Flow Matching problem.
Because (z(2) is the piecewise interpolation between all the points 2 = (Zo...., σK) at time .o,....,t).. it has the closed-form expression:
K
1 1- Lu (II+1 - 22)) - ((eett)))
(x(2)=
2k
(27)
t++1 tt
k=
where (t)
(t) is the indicator function, which equals when E [tx, tx-1) and 0 otherwise.
Using these, we may plug the expressions for (z(2)) and (2)2) into the formula for the conditional vector field:
Jb+1 2k
(x 2)
1 ((antte)))
(28)
t tx
Finally, plugging in the piece-wise formula for vr, this is enough to identify that:
12+1
1g =
Bywatee 2.2 (wl 2144) x1)))
(e((2;0) (x)
Fk-1 - 22
dt.
(29)
t-- 1x
16

---

Published as a conference paper at ICLR 2025
where is the result of marginalizing with respect to all variables, except 24 and 2k++, and Pe(x myy Øx+1) a probability path linearly interpolating the means between 24 and T and with constant variance. Therefore, it is CFM problem between two data distributions, as in Tong et al. (2024).
Because each Ig does not depend on other marginals than time k and k+ 1, and the neural network 2r (2; 0x) is only learned on the interval [1x. 1.. then optimizing the sum each I becomes equivalent optimizing each I independently. Therefore, the MMFM problem is equivalent K separate CFM problems.
B MULTI-MARGINAL OPTIMAL TRANSPORT
Multi-marginn optimal transport (MMOT) Passs 2015) extends the classical Optimal Transport (OT) problem to multiple probability measures, aiming to find coupling that minimizes the total cost associated with transporting mass between multiple marginals.
The optimization problem for marginals can be expressed considering potential couplings between measures pi..... M4 supported on the same metric space (X, 122). For the cost function e: X + R. the MMOT objective is:
inf =EI(p.....kk)) / (I1,,..., TK) (T ..., RK),
(30)
where IT vK) is the set of probability measures over X with marginals equal to MKK. For K 2, this reduces to the classical two-marginal optimal transport problem.
Potentially, this could be a great opportunity to constrain the matching across multiple time points be physically realistic, for example, by minimizing the curvature of the transport across time. In the case of cellular biology, one might expect that cells strive to maintain homeostasis and minimize changes across time.
For the sake of simplicity, we focus on pairings between consecutive marginals (with respect to time), and define cost function that linear across pairs of consecutive spaces. We demonstrate that this approach reduces the MMOT problem to sequence of pairwise OT problems. We refer the reader (Haasler al. 2021) for the general case.
In our case, the MMOT problem becomes:
K-11 M (IR, m841)
inf PETT
Σ
dr IIK.
(31)
XK
which, after permutation of the integral and the discrete sum, yields:
Σ EJ
K-1
inf =E Mg)
_Jxxx
xx+1)
(32)
Let us now denote by (Ik, Tx41) the solution of the OT problem with cost ch for each k e
(1,...,-- --1.. We consider the composite density:
II T (+1,
n (x1
FK
(33)
II
2 (x(2x)
Because मे. is an optimal transport plan for each of the & elements in the sum above, it must be that for all æ E II(p.......
K-1 Σ E
K-1
Σ JJ xxx k=
I4+1) &
(
(x41) Fx+1)-
(34)
we can show that (1) the marginals * correspond to the intended marginals, that is n' e II -- -.p.....--... (PK) and (2) the pairwise marginals of =" correspond to the solution of the pairwise
17

---

Published as a conference paper at ICLR 2025
OT, then the lower bound in Equation 34 is tight. Consequently, T" solves the MMOT problem; therefore, the solution of MMOT can be achieved by solving a sequence of OT problems. These steps of the proof rely on the specific assumption about the form of the marginals.
We present the example for K 3 without loss of generality (the main proof can be written by induction on K). We have that:
"1 T2 (22,
(01, 22,
(35)
Pa (2(22))
In that case, the pairwise marginal of ' (x"2 Hz) is evaluated as:
(I22 v3) (22, (2(
// Ti (x), (21, w2)de
ग (r2)
(36)
where the second equality holds because nं is optimal transport map, have that it satisfies the marginal condition. Together, this verifies the equality between the pairwise marginals of n² and the solutions of pairwise OT by the symmetry of the roles played by 21, 22, and æz. Because the pairwise marginals are the OT transport plan, we must have that π II(,....... MK), which completes the proof.
21,
C METHOD DETAILS
We present the detail for the sampling of trajectories from the MMFM model in Algorithm | To generate test trajectories from the vector fields, we employed the fourth order Runge-Kutte method (rk4). Note that the guidance strength w can be greater than this formulation.
Algorithm Pseudocode: Sampling from COT-MMFM
Require: Trained Uy (z, c; 0), guidance parameter w, condition e, no. of steps T, sample Tz=0
1: Set: t= 1/T
2: Define: 0, (zx c; 0) = (1 w) v,(x, 0; + w (z(2, c)
3: for k to T do
4: Øk+1 = ODDEIneerrator((2,, 0,
5: end for
6: return 2T
D MODEL TRAINING AND EVALUATION
To ensure fair comparison across all models, constructed the data sets for each model based on the need for validation data and conducted each experiment at least three times with varying seeds. For models necessitating hyperparameter tuning, employed hold-out train-van wrain-validation split of the samples, selecting the model that achieved the best performance on the validation set. In contrast, for FSI, where hyperparameter tuning was not required, we used the entire data set for training. including the hold-out part. Tab. provides an overview of the hyperparameters tuned for each MMFM model.
Architecture of uz(2, c; 0). Given observation æ, condition and time f, first compute a temporal embedding for using a sinusoidal timestep embeddings (Vaswani et al.] 2017). We then retrieve condition embedding from library of learnable condition embeddings, using the index c. The input x is transformed into a representation via a three-layer feedforward network using LeakyReLUU activations. Subsequently, the embeddings of xr and are concatenated and processed through an additional lhree-layer neural network. The model is optimized using the Adam optimizer (Kingma & Ba 2015).
Mean and Variance Function The definition of (2( can be found in Equation and repre- sents the natural cubic spline-basee mean-function of the flow. Hence, the mean function (z(2) is piecewise third-order polynomial function with C2 continuity, meaning its values and first and
18

---

Published as a conference paper at ICLR 2025
second derivatives continuous at interval boundaries. For each interval [Ex: tx+1l. the spline can be expressed as:
p(l,z) = oa (L- (2- 1x))² + by 4x)² + cs(1 - (x) + da.
where dks du, CR. and da are coefficients specific to that interval. Therefore, the first derivative is:
(1, 2) 30q (1 1x,) 2b 1g) + cj.
The spline and its derivative can be efficiently computed by solving a tridiagonal system of linear equations. For the variance function (z),, we considered two options: (1) Constant values and therefore ((()) = 0 or (2) if (=)) is as defined in Equation 10 then we have (= = 4. 11 Both of these formulas are also efficient to compute.
(fx=1-6,)
Data Preprocessing We precomputed optimal couplings between samples of consecutive time steps in all synthetic experiments. This approach eliminated the need for mini mani-batching in each iteration, thereby enhancing computational efficiency.
Table 5: Hyperparameters for model training.
(*) Applicable all variations of Flow Matching models dis-
cussed in the paper.
Model
Hyperparameters
Values/Range
FSI MMFM
learning rate
[le-2, le-3, lc-4]
Pu latent dimensions (x,t,c)
[0.0, 0.1, 0.2, 0.3]
[16, 32, 64, 128, 256]
flow variance
[0.01, 0.1, 1, adaptive]
guidance w
[ for (l...., 10) {20, 30)
We provide distinct approaches for integrating conditional information into the neural network model. The first approach was used throughout this study. This second approach is particularly ad- vantageous when dealing with ordinal condition types, such as drug dosage, allowing incorporate relationships between embeddings into the model.
Learnable Embeddings: This approach learns separate embeddings for each condition and the null condition Ø. While yielding superior results compared to the second approach (see below), does not allow for interr/extrapolation to new conditions. However, when data for a previously unseen condition becomes available, the model can be fine-tuned by fixing the weights of uy (ar,,, e; 0) and re-training only the new embedding. This approach generally necessitates separate hold-out distribution to select the most promising model for test scenarios.
Fixed Embeddings: This approach either (i) concatenates the value of the condition directly to the input vectors z and t or (ii) before concatenating uses another neural network to compute em- beddings for each condition, which are then concatenated to the input, e.g., e(c) I - R". The encoder weights then learned during training, allowing the model generalize to unseen conditions without model retraining.
Optimal Transport Implementation As shown in Appendix B the multi-marginal optimal trans- port map decomposes as a product on two-marginal transport maps T (T&, T++1). Practically, we solve the above problem for each (t4, (+11) and using the ot.mdd2 function provided in the Python package POT.
E EXPERIMENTAL DETAILS ON SYNTHETIC DATA
The following figures illustrate the condition -paiiif trajectory prediction of MMFM, L-MMFM, and the FSI model. Fig. 5 indicates that linear interpolation generally leads to straighter paths. However, the model can correct the paths for C3 and c5. Fig. shows that FSI able to match all marginals perfectly, which leads to vastly incorrect trajectories for 0g and c5-
19

---

Published as a conference paper at ICLR 2025
Conditon
Cunditon
Figure 4: Visualization of MMFM interpolation results across all experimental conditions. The ground truth vector field is shown in blue, with predicted test sample trajectories black. Training samples are displayed the background using transparent, colored markers.
Condition
Conditton
Conditton
Conditton
Condition
Condition
Cenittion.e7
Condition see
Contition eee
Conditen asto
Confition eit
Condition o112
Figure 5: Visualization of L-MMPM interpolation results across all experimental conditions. The ground truth vector field is shown in blue, with predicted sample trajectories in black. Training samples displayed in the background using transparent, colored markers.
Condition pet
Conditton o*2
Condition o+3
Condition oet
Condition o=t
Condition end
Candiion.cd ea7
Condition get
Canditin coe
Condito 2310
Comfition estt
Conditoo es12
Figure 6: Visualization of FSI interpolation results across all experimental conditions. The ground truth vector field is shown in blue. with predicted test sample trajectories in black. Training samples are displayed in the background using transparent, colored markers.
Condition
Conditio
Condition
Conditio
Condition
Figure 7: Visualization of PCEM interpolation results across all experimental conditions, The ground truth vector field is shown in blue, with predicted test sample trajectorics in black. Training samples are displayed in the background using transparent, colored markers.
20

---

Published as a conference paper at ICLR 2025
E. EXTRAPOLATION To UNSEEN STARTING POINTS
We evaluate the extrapolation capabilities of our model by extending the estimated vector fields unseen regions for fixed conditions, illustrated in Fig. 8 The results demonstrate that the underlying vector fields exhibit smooth behavior in the neighborhoods surrounding the interpolated paths (highlighted in blue). However, when initializing trajectories from unexplored regions, i.c., regions without training data being provided to the model beforehand. specifically where < O, we observe a tendency for the trajectories overshoot beyond we note that the trajectories are influenced by the probability path derived from the spline, which was constrained to only five time points and followed a linear trajectory, and therefore divides the vector field into two regions (above and below the linear, horizontal line). This highlights the impact of the interpolation method on the model's extrapolation performance.
Furthermore, for condition ca+
Prodided Tejectieis ln Coudition
Prolicted Pegesbeics for Caniitin
Prolicied Tegecbucio for Canillin emi55
-5
Figure 8: Visualization of MMFM inter-/extrapolation using unseen starting points across three conditions, The ground truth vector field shown in blue, overlaid with the predicted test sample trajectories.
F ADDITIONAL SYNTHETIC EXPERIMENT: EXTRAPOLATION TO UNSEEN TIME POINTS
Next, we evaluated the extrapolation capabilities of our MMFM model to unseen time points for some conditions. Therefore, generated data samples from sinusoidal curve at time points t = 0, 0.25,, 0.5, 0.75, 1 across conditions с = 1,..., 10. For condition 5, the model was provided with measurements only samples. The findings demonstrate that MMFM effectively extrapolates the time course data for Cs by leveraging a smooth vector field constructed based on the other conditions. In contrast, as expected, FSI fails to predict future behavior accurately despite having access to an additional data point at 0.55 during training.
0, 0.25, 0.5. Figure 9 shows the interpolated trajectories for test
A. True Vector Fieed (Phase Diagram) with Training Data
B.
Predicted Trajectory FSI
C.
Predicted Trajectory MMFM
2.0
7.0
9.0
220
10.0
Figure 9: Overview of extrapolation time points. (A) Visualization of training samples, (B) predicted trajec- tories FSI, and (C) predicted trajectories MMFM. Trajectories are shown for cı, Co Cio. Here ea only measured until 0.5.
G DETAILS ABOUT THE BEIJIN WEATHER DATA
Data Preprocessing The data was downloaded from archive.ics.uci.edu We selected the PM2.5 concentration the dependent variable for analysis. Next, we grouped the data by month and computed the mean PM2.5 values. For training, used the subset of months shown in Table and Table 7 (months enumerated
21

---

Published as a conference paper at ICLR 2025
from 0 (January 2015) to 25 (February 2017)). For testing, used all other time points that were not part of the training. For the MMFM methods, we used month 11 for condition e₁ for validation and provided this data as additional training data to FSI.
Table 6: Sampling times (month 0- 12) and conditions for the Beijing weather experiment.
0
2
3
4
5
7
8 9 10
11 12
C4
V
V
47
V
V
Rest
Note: / indicates that samples from that combination are available the training set.
Table 7: Sampling times (month 13-25) and conditions for the Beijing weather experiment.
13
14
15
16
17
18
19
20
21
22
23
24
25
C4
V
V
C7
V
Rest
Note: v indicates that samples from that combination are available in the training set.
Table 8: Results the Beijing data. We report the Mean-Squared Error (MSE), with lower values indicating better performance. Best-performing models are highlighted in bold.
c=
10
Rest
FSI
29.54 30.99
45.31
34.05
$3.42 + 35.53
17.20 # 20.11
CFM
28.18 + 18.60
20.22 12.43
28.79 + 18.46
26,86 : 16.86
PCFM
18.47 # 16.42
6,87 5.37
35.18 + 23.27
12.77 # 11.10
L.MMFM (ours)
9.333
6.33 #7.68
9.03 9.00
9.21 8.90
MMFM (ours)
11.43 +
13.74 14.44
17.07
17.02 =18.72
PM2.5 Seasonal Patterns by Station
PM125 (pg.mm.)
150
100
2015-01
2016-02
2016-07
2015-10
2016-01
2016-01
2016-07
2016-10
2017-01
(Year-Meuih))
Station
Aotishongzin
Dingling
Guanyuau
Husizou
Shmayi
Wanliu Wansheuxigong
Changping
Dongsi
Gucheng
Nongzhanguan
Tantan
Figure 10: Overview of the preprocessed Beijing data set.
H DETAILS ABOUT THE PERTURBATION SCREENING DATA
Treatment selection We implemented a two-stage approach to eliminate treatments with minimal effects. First, we trained scVI model (Lopez et al., 2018) on the full count dataset using 50 latent features and negative binomial likelihood. We then calculated the energy distance between the control group and each treatment the latent space, keeping treatments with an energy distance above certain threshold, here 1. As described in the main text, randomly selected subset of treatments and sampled the hold-out data for our experiments also at random. Non-stimulated control and vehicle data were also preserved for the experiments.
22

---

Published as a conference paper at ICLR 2025
Visualization of MMFM interpolation results across all stations. The vertical bars indicate the training time points.
Figure 11:
Figure 12: Visualization L-MMFM interpolation results across stations. The vertical bars indicate the training time points.
Figure 13: Visualization of FSI interpolation results across stations. The vertical bars indicate the training time points.
Data Preprocessing To train the models, the mRNA count data was normalized for library size to remove technical noise and log(x + -)-traanformedd stabilize the variance and reduce the impact of extreme values. Following normalization and transformation, we applied Principal Component Analysis (PCA) to reduce the dimensionality the 18,50 08.500-dimensional gene feature space. We retained the top 25 principal components, capturing most of the data's 's variance.
I PRIORS FOR TRAJECTORIES DYNAMICS
In this section, further motivate the minimum curvature prior proxy for energy minimization along a trajectory. We consider the trajectories of a particle of initial mass m in an inertial frame, following path 7 : [o, 1] + X. We assume constant velocity over the trajectory, i.e., (t)ll)) uo- From the law of conservation of momentum (no extemal forces), we have that the particle should shed mass according to the Tsiolkovsky rocket equation Tsiolkovsky 1954p)
d7((1) = (v,(t) (t)) dm(()
-m
(37)
dt
dt
ve(() dm(t) dl.
(38)
23

---

Published as a conference paper at ICLR 2025
where (t) stands for the output velocity of the expelled mass and ve (t) is the output velocity in the particle reference frame. Rearranging the above equation we obtain:
d dllo(m(t)) -)")))
(39)
This equation shows that ve (1) and (!) are parallel vectors with different magnitudes. The above system of equations is thus equivalent to the following:
dllog(m(t)) dt
1/0e (t)ll2 --
(40)
Assuming a constant relative for the output velocity lIve (4)ll2 ve, and assuming potential function in the system, we finally get the difference in log- engrryyrgy of the particle between time t = 0 and 1 as
m(1) m(0)
---- -
(og(Ei) log( Eo) = log
(r)
(41)
The logarithm being concave function, we obtain that minimizing the loss of energy along a path Y amounts to minimizing the cumulative curvature along the path.
COMPUTATIONAL COMPLEXITY
Compared to CFM, the complexity of our method differs by two parameters: the number of snap- shots K + and the number of conditions C. Indeed, while CFM only considers the matching between two distributions, single OT solution needs be computed. In the MMFM case, C K OT problems need to be solved. The complexity of our method is thus linear in the number of time points and in the number of conditions.
Optimal transport solvers typically show a quadratic dependence on the number of data points, al- though recent methods can marginally improve upon this bound Dvurechensky et al. (2018). When feasible, pre-computed the OT couplings before training. However, when the datasets grow larger, we rely on mini-batchin hani-batching and solve OT couplings over mini mani-batches during training, follow- ing the approach of Tong et al. (024)). We note that experiments all fall in the first scenario and OT-couplings could be computed before training. Given the successful applications using mini- batch OT approximations in Tong al. (2024), believe that MMFM with minibatch OT will demonstrate strong empirical performance as well.
We assessed the computational complexity of MMFM with respect to the number of time points, conditions, and total number of samples in a new benchmark experiment. We precomputed the solution to the OT problem. Our baseline scenario uses K + time points, = 12 conditions and N = 50 samples following the setup of our synthetic experiment in Section 5.1, i.e. we generate N samples for each condition and each time point, resulting in N + 1) C' samples. We then varied these parameters independently to study their impact on running time. We kept all other hyperparameters and training details identical the original experiment, executed on GPU and results were averaged across three seeds. The computational time required for solving the Optimal Transport problem is indicated in parentheses beneath the run time, Results are shown Table 9
24

---

Published as a conference paper at ICLR 2025
No. Samples (N)
Time [s]
No. Cond. (C)
Time [s]
No. Timep. (K + 1)
Time [s]
50
27.45
(0.08)
12
26.98
7
26.12
100
48.99
(0.09)
(0.09)
(0.10)
24
49.62
14
26.20
500
232,86
(0.12)
(0.19)
(2.11)
100
195.07
28
27.73
1000
461.31
(0.72)
(0.67)
(9.58)
Table 9: Comparison of run time (and pre-computing the Optimal Transport solution, given in parentheses) with respect different sample sizes, number of conditions, and time points. Left: We set C 12, K + vary the sample size N. Middle: We set = 50,
and vary the number conditions C. Right: set 50, C 12 and vary the number time points 1.
25


Reviews:
[{"review_id": "QfgDXBlWih", "reviewer": "Reviewer_19bk", "summary": "This paper aims to model the dynamics of complex real-world systems across time and conditions using flow matching. Current methods, such as Flow Matching, remain limited by the inability to effectively capture information across multiple time points and experimental conditions. This paper aims to alleviate these limitations with a novel method they refer to as  Multi-Marginal Flow Matching, which leverages the Flow Matching framework and optimal transport theory to train a neural network to regress onto the flow. The authors evaluate the method’s effectiveness on synthetic and real-world data, demonstrating superior performance to existing baselines, measured by qualitative and quantitative metrics.", "questions": "In the proof of proposition 1, the final (*), I am confused how Fubini’s theorem applies here to switch the order of the inner product and integral. \nI may have a misunderstanding: is it necessary for conditions {c} to be given by the data. If they are missing, how may one proceed? In many real word systems, for example, conditions or categories may be latent and unobserved.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "This paper introduces a novel method which expands on previous methodology in a parsimonious manner.\nThe work is relevant to the research community, addressing the problem of modeling complex system dynamics. \nThe approach is simple and intuitive, with sound theoretical backing. The empirical evidence they provide is strong but limited. \nThe experiments are well-designed and include sufficient baselines. The authors appear to know the literature very well and do a good job providing citations throughout. \nThe paper is well-written and organized and easy to follow. I particularly appreciate how the authors clearly provide the minimally sufficient information in the main body of the paper and detail the appendix to complement.", "weaknesses": "The authors rigorously compare their method to multiple existing methods. However, they only do so on one real-world rna-seq dataset and one synthetic dataset. In the introduction of the paper, the authors give a multitude of real-world settings where their method may be considered useful. While the authors clearly motivate their method, I am concerned by the rigor with which they evaluate their method empirically. Are there additional datasets that could further demonstrate the effectiveness of the method, such as some of the genomics datasets used in Tong et al. (2024)? If so, inclusion would greatly increase this paper's contribution. The authors address this weakness in their rebuttal, providing additional experiments supporting their method's usefulness in practice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "j0SW6X96WQ", "reviewer": "Reviewer_FqP7", "summary": "This paper introduces Multi-Marginal Flow Matching (MMFM), a method for modeling the dynamics of biological single-cell time series data. The proposed approach extends the Flow Matching (FM) framework by incorporating multiple time points and conditions, thereby learning the transport of probability distributions for cellular states across these different scenarios. MMFM leverages techniques like smooth spline-based interpolation and neural networks with classifier-free guidance to model changes over time and conditions. Through experiments on synthetic and single-cell RNA-seq datasets, the authors show MMFM's superior performance in imputing missing time points compared to other baseline methods.", "questions": "- Some generative models designed for imputation task can be also adapted for extrapolation/forecasting, such as CSDI[1]. Can the MMFM or CSD method be extended to handle **extrapolation** for sequences beyond the given time points?\n\n- The current model appears to rely on temporal snapshots at specific intervals. Can **MMFM** handle **non-uniformly sampled** time points? If so, how is this addressed within the interpolation or flow matching framework?\n\n[1]: Tashiro, Yusuke, et al. \"Csdi: Conditional score-based diffusion models for probabilistic time series imputation.\" Advances in Neural Information Processing Systems 34 (2021): 24804-24816.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The motivation for using MMFM to model cell population dynamics across multiple steps and conditions is well-founded, as it directly addresses challenges inherent in biological experiments, such as fragmented and incomplete data. The multi-step transport approach is impressive in efficiently capturing the system's temporal and conditional complexities.\n\n- The method extends (Conditional) Flow Matching (FM) to the *Multi-Marginal** setting, which is a natural and well-justified progression for addressing biological systems that require modeling across multiple time points. This approach enables the framework to leverage inter-dependencies across multiple conditions, significantly enhancing generalizability.\n\n - The use of **optimal transport** and **Schrödinger Bridge** for improving the integration of multiple marginals is insightful, resulting in improved model robustness. Compared to **piece-wise linear interpolation**, the introduction of cubic splines for smoother interpolation significantly boosts the performance and efficiency of the model, which is well demonstrated in their experiments.", "weaknesses": "- The title, **\"Modeling Complex System Dynamics,\"** might be an over-claim relative to the paper's focus, which is primarily on **biological single-cell dynamics**. Typically, \"Complex System Dynamics\" also refers to broader scenarios, such as physical phenomena in fluid dynamics, climate systems, or even economic time series. Since the proposed method appears promising and generalizable, demonstrating its applicability to tasks in broader domains would further strengthen this claim.\n\n- In the Multi-Marginal Conditional Flow Matching (MMCF), the authors pre-compute the solution for the multi-marginal optimal transport, which can be time-consuming. While they mention that **mini-batches** can be used to handle large datasets, further discussion on the **scalability** of the approach would be valuable. In particular, it would help if the authors could expand on the tradeoff between the **performance** and **efficiency** when comparing **L-MMFM** and **MMCF**, as OT solvers tend to become computationally expensive with increased dimensionality.\n\n- Although MMFM is demonstrated effectively on biological single-cell data, there is no extensive discussion about its generalization capabilities in handling diverse temporal datasets beyond biology. Considering the potential to model other types of time series (e.g., physical or environmental data), the paper would benefit from a discussion about such possible applications and any inherent limitations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "fzCy5TzCrm", "reviewer": "Reviewer_Xe9p", "summary": "The paper introduces a Multi-Marginal Flow Matching (MMFM) method to model the dynamics of complex systems from fragmented temporal snapshot data. MMFM constructs a flow using smooth spline-based interpolation across time points and experimental conditions, and regresses this flow with a neural network using a classifier-free guided flow matching framework. The paper demonstrates the effectiveness of MMFM on both synthetic and real-world datasets, including a single-cell genomics dataset with around a hundred chemical perturbations, showing the MMFM’s power in imputing data at missing time points.", "questions": "Why is the condition c multiplied by a scalar in Eq.(11) possible, as c is only the identifier of the condition?\nWhat does the y-axis represent in Figure 2?\nHow is optimal transport specifically implemented in MMFM?\nThe MMFM loss function (7) and the CFM loss function (3) seem to be the same, what is the purpose of introducing the MMFM loss function?\nFrom which literature do the formulas used in the synthetic dataset originate, What is the purpose of studying this synthetic dataset?\nWhat is the exact definition of c_{\\phi}?\nWhat is the exact definition of \\mu_{t}^{‘}(x) and \\sigma_{t}^{‘}(x)?\nWhat is the exact network architecture of \\bar{v}_t(.)?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "-The use of smooth spline-based interpolation is a plausible approach for constructing the flow across different time points.\n-The authors evaluate MMFM on a real-world dataset involving cellular responses to drugs, demonstrating the method's applicability to practical problems in biology and medicine.", "weaknesses": "Lack of Clear Problem Definition: The paper lacks a clear and concise problem definition. There is no centralized notation or formal problem statement, making it difficult to follow the methodology and understand the specific goals of the research.\nIncomplete Methodological Details: The method section is incomplete and lacks crucial details. For example, the loss function used for training is not fully explained, and key components such as mean and variance formulas are missing. Additionally, the specific implementation of optimal transport is not described, leaving readers uncertain about how the method works in practice.\nUnclear Experimental Setup: The experimental setup is poorly documented. The origin of the formulas used in the synthetic dataset is not explained, and the y-axis labels in Figure 2 are missing, making it difficult to interpret the results. For the real-world dataset on cellular responses to drugs, the criteria for evaluating performance are not clearly defined, making it hard to assess the effectiveness of MMFM.\nLimited Baseline Comparisons: The performance comparison with baseline methods is inadequate. The paper mentions that MMFM outperforms existing methods but does not provide detailed comparisons or statistical significance tests, which are necessary to validate the claims.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (Spotlight)