Meta Flow Matching: Integrating Vector Fields on the Wasserstein Manifold
Authors: Lazar Atanackovic
Xi Zhang
Brandon Amos
Mathieu Blanchette
Leo J Lee
Yoshua Bengio
Alexander Tong
Kirill Neklyudov
Venue: ICLR 2025 Poster (2025)
Keywords: Flow matching
Dynamics
Cell dynamics
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We propose a new framework, Meta Flow Matching, for integrating vector fields on the Wasserstein manifold of probability densities and show the use case to improve prediction of patient specific treatment responses.

Abstract: Numerous biological and physical processes can be modeled as systems of interacting entities evolving continuously over time, e.g. the dynamics of communicating cells or physical particles. Learning the dynamics of such systems is essential for predicting the temporal evolution of populations across novel samples and unseen environments. Flow-based models allow for learning these dynamics at the population level - they model the evolution of the entire distribution of samples. However, current flow-based models are limited to a single initial population and a set of predefined conditions which describe different dynamics. We argue that multiple processes in natural sciences have to be represented as vector fields on the Wasserstein manifold of probability densities. That is, the change of the population at any moment in time depends on the population itself due to the interactions between samples. In particular, this is crucial for personalized medicine where the development of diseases and their respective treatment response depend on the microenvironment of cells specific to each patient. We propose *Meta Flow Matching* (MFM), a practical approach to integrate along these vector fields on the Wasserstein manifold by amortizing the flow model over the initial populations. Namely, we embed the population of samples using a Graph Neural Network (GNN) and use these embeddings to train a Flow Matching model. This gives MFM the ability to generalize over the initial distributions, unlike previously proposed methods. We demonstrate the ability of MFM to improve the prediction of individual treatment responses on a large-scale multi-patient single-cell drug screen dataset.

Published as a conference paper at ICLR 2025
META FLOW MATCHING: INTEGRATING VECTOR FIELDS ON THE WASSERSTEIN MANIFOLD
Lazar Atanackovic!.2+
Xi Zhang 3.4
Brandon Amos
Mathieu Blanchette
Leo J. Lee!.2
Yoshua Bengio
Alexander Tong
Kirill Neklyudov
University of Toronto
Vector Institute
Milaa Quebec Al Institute
MMGGlll University "Meta
Université Montréal
"CIAAR Fellow
ABSTRACT
Numerous biological and physical processes can be modeled as systems of interact- ing entities evolving continuously over time, e-g. the dynamics of communicating cells or physical particles. Learning the dynamics of such systems is essential for predicting the temporal evolution of populations across novel samples and unseen environments. Flow-based models allow for learning these dynamics at the population level they model the evolution of the entire distribution of samples. However, current flow-based models are limited to a single initial population and a set of predefined conditions which describe different dynamics. We argue that multiple processes in natural sciences have to be represented vector fields the Wasserstein manifold of probability densities. That is, the change of the population at any moment in time depends on the population itself due to the interactions between samples. In particular, this is crucial for personalized medicine where the development of diseases and their respective treatment response depend on the mi- croenviromment of cells specific to each patient. We propose Meta Flow Matching (MFM), a practical approach integrate along these vector fields on the Wasser- stein manifold by amortizing the flow model over the initial populations. Namely, we embed the population of samples using Graph Neural Network (GNN) and use these embeddings to train a Flow Matching model. This gives MFM the ability to generalize over the initial distributions, unlike previously proposed methods. We demonstrate the ability of MFM to improve the prediction of individual treatment responses on a large-scale multi calli--patent single-cell drug screen dataset.
INTRODUCTION
Understanding the dynamics of many-body problems is a focal challenge across the natural sciences. In the field of cell biology, central focus is the understanding of the dynamic processes that cells undergo in response to their environment, and in particular their response and interaction with other cells. Cells communicate with one other in close proximity using cell signaling. exerting influence over each other's 's trajectories (Armingol et al., 2020; Goodenough and Paul, 2009). This signaling presents an obstacle for modeling due to the complex nature of intercellular regulation, but is essential for understanding and eventually controlling cell dynamics during development (Gulati et al., 2020; Rizvi et al., 2017), ), in diseased states (Molè et al., 2021: Binnewies et al., 2018; Zeng and Dai, 2019; Chung et al., 2017), and in response to perturbations (Ji et al., 2021; Peidli et al., 2024).
The exponential decrease of sequencing costs and advances in microfluidics has enabled the rapid advancement of single-cell sequencing and related technologies over the past decade (Svensson et al., 2018). While single-cell sequencing has been used great effect to understand the heterogeneity in cell systems, it is also destructive, making longitudinal measurements extremely difficult. Hence, learning dynamical system models of cells while also capturing their inherent heterogeneity and stochasticity of cellular systems remains a central challenge in biology.
Instead, most existing approaches model cell dynamics at the population level (Hashimoto et al., 2016; Weinreb et al., 2018; Schiebinger al., 2019; Tong et al., 2020; Neklyudov et al., 2022; Bunne et al., 2023). These approaches involve the formalisms of optimal transport (Villani, 2009; Peyré and Cuturi, 2019), diffusion (De Bortoli et al., 2021), normalizing flows (Lipman et al., 2022), which
"Jiint first authorship. Correspondence to: 1. aaaacckoviii iewaaiiiiieiiiiiimmmm atoronto.caa Our code is available at: https: / / thg.. com aazaratan
/etaattnnmmmaaaaaaaaaeaeeaaaaa /eeta- fattaaaffowwwmmhhh---

---

Published as a conference paper at ICLR 2025
(a)
(b)
Flow
FM
MFM
Tree
Figure 1: Illustration Meta Flow Matching (MFM, Eq. (17)). (a) Comparison between Flow Matching (FM, Eq. (6)) and MFM. (b) Depiction of differences between MFM and FM generated predictions. Given a point 2t. a vector field (flow) model trained with MFM can generate different points è1 for different initial distributions Po (represented by red, green, and purple). FM trained models can only predict aggregate response over populations (shown in gray). FM best can incorporate known (seen) conditional information available in the training data, denoted as Conditional Generative Flow Matching (CGFM, Eq. (9). In contrast, MFM jointly learns population embedding model (p(o) and a vector field vs allowing generalization to unseen populations.
learn map between empirical measures. While these methods are able to model the dynamics of the population, they are fundamentally limited in that they model the evolution of cells as independent particles evolving according to the learned model of a global vector field. Furthermore, these models can be trained to match any given set of measures, but they are restricted to modeling of a single population and can at best condition on number of different dynamics that are in the training data.
We propose Meta Flow Matching (MFM) the amortization of the Flow Matching generative modeling framework (Lipman et al., 2022) over the input measures, which (i) takes into account the interaction of particles (instead of modelling them independently) and (ii) allows for generalization of the learned model to previously unseen input populations. In practice, our method can be used to predict the time-evolution of distributions from a given dataset of the time-evolved examples. Namely, we assume that the collected data undergoes universal developmental process, which depends only on the population itself as in the setting of the interacting particles or communicating cells. Under this assumption, we learn the vector field model that takes samples from the initial distribution as input and defines the push-forward map on the sample-space that maps the initial distribution to the final distribution (see Fig. 1).
We showcase the utility of our approach on two applications. To illustrate the intuition of the proposed method, we first test MFM on a synthetic task of "letter denoising". We show that MFM is able to generalize the denoising process to unseen letter silhouettes, whereas the standard FM approach cannot. Next, we explore how MFM can be applied to model single-cell perturbation data (Ji et al., 2021; Peidli et al., 2024). We evaluate MFM predicting the response of patient-derive derived cells to chemotherapy treatments in a recently published large scale single-cell drug screening dataset where there are known to be patient-specific responses (Ramos Zapatero et al., 2023).' This is challenging task due to the inherent variation that exist across patients, the variability induced by different treatments, and variation due local cell compositions. Addressing this problem can lead to better prediction of tumor growth and therapeutic response to cancer treatment. We demonstrate that MFM can successfully predict the development of cell populations on replicated experiments, and most importantly, that it generalizes to previously unseen patients, thus, capturing the patient-speciff response to the treatment.
BACKGROUND
2.1 GENERATIVE MODELING VIA FLOW MATCHING
Flow Matching is an approach to generative modeling recently proposed independently in different works: Rectified Flows (Liu et al., 2022b). Flow Matching (Lipman et al., 2022), Stochastic Inter- polants (Albergo and Vanden-Eijndee 2022). flow is continuous interpolation between densities Pa (r(o)) and pı (z,) in the sample space. That is, the sample from the intermediate density Pr (xr) is produced as follows
fe 21): (zo, 21) ~ T T2o, 21);
(1)
where [ dæi m( = pu po()) / dao 21) = (11)),
(2)
TThis dataset includes more than 25 million cells collected for ten patients over 2500 different experimen- tal/treatment conditions.
2

---

Published as a conference paper at ICLR 2025
where ft is the time- interpolating function such that fi=o 21) To and (2o. 21) = 21 (e.g. linearly between 2g and 21 with fz 21) = (1 - t) Io + t 21); (2o, 21) is the density of the joint distribution, which is usually taken distribution of inde- pendent random variables ([n, x1 ) = po (x1). but can also be generalized to formulate the optimal transport problems (Pooladian al., 2023; Tong et al., 2024a). The corresponding density can be defined then as the following expectation over the interpolating samples
Pc(x)
drgdrı π (zo. z1)8(x fa 21))
The essential part Flow Matching is the continuity equation that describes the change of this density through the vector field on the state space, which admits vector field 2tz (a) solution
(p((x) (z) - (V (1()) (z))
E (x( ((f x1) - 0 (
(E(() = Pr (E)
(4)
aL
01
The expectation for of is intractable, but we can model us (x;ww) to approximate of so that we can efficiently generate new samples. Relying on Eq. (4), one can derive the tractable objective for leamnnng uf (z), i.e.
(FM(W)= = L. di (p)(a))lo((()) - (x)))(())))))) (x;
(5)
-En(so.R1) (
oor Si x1) ve ve x1 x1):)) + constant.
(6)
Finally, the vector field (,(E, Re 1z defines the push- forward density that approximately matches Pt=t, i.e. ToPo 2e Pt=l, where T is the flow corresponding vector field va (-(,,) with parameters
2.2 CONDITIONAL GENERATIVE MODELING VIA FLOW MATCHING
Conditional image generation is one of the most common applications of generative models nowadays; it includes conditioning on the text prompts (Saharia et al., 2022b; Rombach al., 2022) well as conditioning on other images (Saharia et al., 2022a). To learn the conditional generative process with diffusion models, one merely has to pass the conditional variable (sampled jointly with the data point) an additional input to the parametric model of the vector field. The same applies for the Flow Matching framework.
Conditional Generative Modeling via Flow Matching independently introduced several works (Zheng et al.. 2023; Dao et al., 2023; Isobe et al., 2024) and it operates as follows. Consider family of time-continuous densities p. ( c), which corresponds to the distribution of the following random variable
~ (2 x1 (c. x; x1( (1)
(7)
For every the density Pt (x, c) follows the continuity equation with the following vector field
) (ff( 21 -) (f
(c) Pr(Ele) c)
(8)
Øt
which depends on Thus, the training objective of the conditional model becomes
LCGFM M (w) Em(o..r (c) L' dt
fi(eo.
01) us (f(,
0t
where, compared to the original Flow Matching formulation, we first have to sample c, then produce the samples from (1 c) and pass input to the parametric model of the vector field.
2.3 MODELING PROCESS IN NATURAL SCIENCES AS VECTOR FIELDS ON THE WASSERSTEIN MANIFOLD
We argue that numerous biological and physical processes cannot be modeled via the vector field propagating the population samples independently. Thus, we propose to model these processes as families of conditional vector fields where we amortize the conditional variable by embedding the population via Graph Neural Network (GNN).
To provide the reader with the necessary intuition, we are going to use the geometric formalism developed by Otto (2001) (see the complete discussion in Appendix A). That is, time-dependent
3

---

Published as a conference paper at ICLR 2025
Flow Matching
Conditional Flow Matching
Vector Field Pa (2(2)
Ps(A)
Ps(x)
P((
Pto1
Po-1
Pe-s(Hlay)
sfles)
prafftes))
Figure 2: Illustration of flow matching methods the 2-Wasserstein manifold, (2()), depicted two- dimensional sphere. Flow Matching learns the tangent vectors single generation corresponds to learning Flow Matching learns to integrate a vector field on (X),, i.c. for every starting density MFM defines a push-forward measure that integrates along the underlying vector field.
the manifold. Conditional
set curves
manifold, e.g. classes e₁ and e2 on the plot. Meta
curves on the 2- 2-Wasserstein space of distributions (2))) (see (Ambrosio et al., 2008), Chapter 8). The tangent space of this manifold is defined by the gradient flows S, &Vor 8t X - R on the state space X, to be precise, its L²( X)) closure Vuu E (2t) (see (Ambrosio et al., 2008), eq. (8.0.2)). In the Flow Matching context, we are going to refer to the tangent vectors as vector fields since one can always project the vector field onto the tangent space by parameterizing it gradient flow (Neklyudov et al., 2022).
densities pe(x;) define conoluuely-continnoee
Under the geometric formalism of the 2- 2-Wasserstein manifold, Flow Matching can be considered as learning the tangent vectors vy (-) along the density curve pr(x); defined by the sampling process in Eq. (2) (see the left panel in Fig. 2). Furthermore, the conditional generation processes pa (z c) would be represented as a finite set of curves if is discrete (e.g. class-conditional generation of images) a family of curves if is continuous (see the middle panel in Fig. 2).
Finally, one can define a vector field on the 2- Wasserssteinn manifold via the continuity equation with the vector field (I, Pt (x)) on the state space X that depends on the current density Pe(x) or its derivatives. Below we give two examples of processes defined as vector fields on the 2- Wassersteii manifold.
Example aelean-fiele limit of interacting particles). Consider system of interacting par- ticles, where the velocity of the particle at point x interacting with the particle at point is defined k(x,y) y) Id Rd - -Rd In the limit of the infinite number of particles one can describe their state using the density function Pe(x)). Then the change of the density is described by the following continuity equation
dr Epecy v).
Öp((x) = -V,, (c,
=
(10)
dt
Øt
which is the first-order analog of the Vlasov equation (Jabin and Wang, 2016).
Example 2 (Diffusion). Even when the physical particles evolve independently in nature, the deterministic vector field model might be dependent the current density of the population. For instance, for the diffusion process, the change of the density is described by the Fokker- Planck equation, which results the density deepededd vector field when written continuity equation. i.e.
Öp((x) Aepi(a) = Øt 2
(aan))--vvor -VV. - V2
-v (o))))((x). dx -
-1
1
(11)
Motivated by the examples above, we argue that using the information about the current or the initial density is crucial for the modeling of time-evolution of densities in natural processes, to capture this type of dependency one can model the change of the density as the following Cauchy problem
@p)(x) - -Va,, 0t.
pr=o( =po(x).
(12)
4

---

Published as a conference paper at ICLR 2025
where the state-space vector field (z(2, pt) depends on the density Pt-
The dependeney might vary across models, e-g. in Example 1 the vector field can be modeled as an application of a kernel the density function, while in Example the vector field depends only on the local value of the density and its derivative.
3 META FLOW MATCHING
In this paper, we propose the amortization of the Flow Matching framework over the marginal distributions. Our model based on the outstanding ability of the Flow Matching framework to learn the push-forward map for any joint distribution π ( z1) given empirically, For the given π (Io, 21), we denote the solution of the Flow Matching optimization problem as follows
of (+ π) arg min LGFM (v)()), ((xo,
(13)
Analogous to amortized optimization (Chen et al., 2022; Amos et al., 2023), we aim to learn the model that outputs the solution of Eq. (13) based on the input data sampled from i... (14)
(e (T)) = x (:;
where (Tz) is the embedding model of and the joint density x(: (c) is generated using some unknown measure of the conditional variables c ~ p(c).
3.1 INTEGRATING VECTOR FIELDS ON THE WASSERSTEIN MANIFOLD VIA META FLOW MATCHING
Consider the dataset of joint populations we associate every i-th population with its density T[- ()) and the conditioning variable here is the index of this population in the dataset. We make the following assumptions regarding the ground truth sampling process (i) we assume that the starting marginals Po( ()) dx1 π ( () are sampled from some unknown distribution that can be parameterized with a large enough number of
(< 21 [ where, to simplify the notation,
parameters (ii) the endpoint marginals pi (æt ()) f T xı () are obtained as push-forward densities solving the Cauchy problem in Eq. (12), (iii) there exists unique solution to this Cauchy problem.
One can learn a joint model of all the processes from the dataset D using the conditional version of the Flow Matching algorithm (see Section 2.2) where the population index plays the role of the conditional variable. However, obviously, such model will not generalize beyond the considered data D and unseen indices i. We illustrate this empirically in Section 5.
To be able generalize to previously unseen populations, we propose learning the density vector field motivated by Eq. (12). That is, we propose to use an embedding function şo P22(X)) - Tnm to embed the starting marginal density Po. which we then input into the vector field model and minimize the following objective over
(.' du
② 0i [ 21) (e an) ((p))::))
(15) Note that the initial density is enough to predict the push-forward density since the Cauchy problem for Eq. (12) has unique solution. The embedding function ((p)) can take different forms, e.g. it can be the density value p(Pb) -Po(-), which is then used inside the vector field model to evaluate the current point (analogous Example 2); a kernel density estimator (analogous to Example I): or parametric model taking the samples from this density as an input.
LMEM (w:; sp) EL (4)
Proposition 1. Meta Flow Matching recovers the Conditional Generation via Flow Match- ing when the conditional dependence of the marginals Pu(zol c) = J dx (xg. 21 (c) and p1 (2 c) dx I1 c) and the distribution p(c) are known, i.e. there exist (2)) + 12th such that (MFM(W) = LCGFM (w).
Proof. Indeed, sampling from the dataset ~ DD becomes sampling of the conditional variable c ~ p(c) and the embedding function becomes ((po(- (())(-lc))==
Furthermore, for the parametric family of the embedding models (p(,, 0). we show that the parameters i can be estimated by minimizing the objective in Eq. (15) in the joint optimization with the vector field parameters w. We formalize this statement in the following theorem.
5

---

Published as a conference paper at ICLR 2025
Algorithm 1: Meta Flow Matching (training)
Input dataset of populations {(T #1 i), ()) and treatments c¹, and parametric models for the velocity, vz(:;w), and population embedding sp(:;))
for training iterations do i i~ ((,,(1,))1) // sample batch of n populations ids
(1 =( 12) ( ()M))))( // sample N; particles for every population x
fe(zh. 24) + -(1 -') + t/x
10(0) + 40 ((=d) 1/ embed popslation (rg) For CGFM h t i, FM h O.
(MFM 0) - Er th Σÿ
di f( x3) ues x3) (t)))
d
w' + Updato(w, CWCMMM (w, 0)) // evaluate new parametera of the flow model 0' Update(O, Vo VoLMMM(W;0)) 0))
avaluate new parameters of the embedding modal C w', 0 + 0' // update both models
return vx(;w"),
Theorem 1. Consider dataset of populations D = {(T(xo, x1 | generated some unknown conditional model π (Zo. F1 c)p(c). Then the following objective
C(w, 0) = =Fppte [-= [dt Eppce oylligg
ue (are (c) vx (p0o 0), w)ll
(16)
is equivalent to the Meta Flow Matching objective
② LMPMrww, 0) Ei Dex 11) S dt ae (10o (0 ():) up to an additive constant.
(17)
Proof. We postpone the proof to Appendix B.
3.2 LEARNING POPULATION EMBEDDINGS VIA GRAPH NEURAL NETWORKS (GNNS)
In many applications, the populations 'D ((=(10, I1 ()) are given as empirical distributions, i.e. they are represented as samples from some unknown density æ
Algorithm 2: Meta Flow Matching (sampling)
Input:
initial population {= treatment condition e¹. models v=(===)) and p(=;
((*)) - (1) ~ π(zo, x1 [i),
h = p (m))) // urbed the population क So (a( h ch wd)t + z // ODE sclver return predicted population 1N'
(18)
where N; is the size of the i-th population. For instance, for the diffusion process con- sidered in Example 2, the samples from π (No, 21 () can be generated by generating some marginal pi (.n ()) and then adding the Gaussian random variable to the samples z/- We use this model in our synthetic experiments in Section 5.1.
Since the only available information about the populations is samples, we propose learning the embedding of populations via parametric model sppo, 0), i.e.
(po,, 0) = sp (==))))))) (1 a+) n (wo, .1 (t)
(19)
For this purpose, we employ GNNs, which recently have been successfully applied for simulation of complicated many-body problems in physics Sannhez-Gonzalee et al., 2020). To embed population [ N. 1= we create k-nearest neighbour graph G; based the metric in the state-space X. input into a GNN, which consists of several message-passingg iterations (Gilmer et al., 2017) and the final average-pooling across nodes to produce the embedding vector. Finally, update the parameters of

---

Published as a conference paper at ICLR 2025
the GNN jointly with the parameters of the vector field to minimize the loss function in Eq. (17). We show pseudo-code for training in Algorithm and sampling in Algorithm 2.
4 RELATED WORK
Meta - amortizing learning over distributions The meta-learning of probability measures was previously studied by Amos al. (2022) where they demonstrated that the prediction of the optimal transport paths can be efficiently amortized over the input marginal measures. The main difference with our approach is that we are trying to learn the mush-forwart map without embedding the second marginal and without restricting ourselves to Input-Conver Neural Network (ICNN) by Amos et al. (2017) and optimal transport maps. MFM is meta in the same sense of amortizing optimal transport (or in case flow) problems over multiple input distributions. Hence, we follow the naming convention of Meta Optimal Transport (Amos et al., 2022). This is different but related to how meta is used in the meta learning setting, which amortizes over learning problems (Hospedales et al., 2021; Achille et al., 2019).
Generative modeling for single cells Single cell data has expanded to encompass multiple modali- ties of data profiling cell state and activities (Frangich et al., 2021; Bunne et al., 2023). Single-cell data presents multiple challenges in terms of noise, non-time resolved, and high dimension, and generative models have been used to counter those problems. Autoencoder has been used to embed and extrapolate data Out Of Distribution (OOD) with its latent state dimension (Lotfollahi et al., 2019; Lopez et al., 2018; Hetzel et al., 2022). Orthogonal non-negative matrix factorization (oNMF) has also been used for dimensionality reduction combined with mixture models for cell state prediction (Chen et al., 2020). Other approaches have tried to use Flow Matching (FM) (Tong et al., 2024b;a; Neklyudov et al., 2023) or similar approaches such as the Monge gap (Uscidda and Cuturi, 2023) to predict cell trajectories. Currently, the state of the art method uses the principle of Optimal Transport (OT) to predict cell trajectories (Makkuva et al., 2020; Bunne et al., 2023). These methods are based on input convex neural network (ICNN) architectures and can generalize out of distribution to new cells. As of this time, our method is the only method that takes inter-cellular interactions into account and learns embeddings of entire cell populations to generalize across unseen distributions.
Generative modeling for physical processes The closest approach to ours is the prediction of the many-boue body interactions in physics Saanche--Gozzlle et al., 2020) via GNNs. However, the problem there is very different since these models use the information about the individual trajectories of samples, which are not available for the single-cell prediction. Liu et al. (2022a; 2024) consider a generalized form of Schrödinger bridges also with interacting terms, but do not consider the generalization unseen distributions. Neklyudov et al. (2022) consider learning the vector field for any continuous time-evolution of a probability measure, however, their method is restricted to single curves and do not consider generalization unseen data. Campbell al. (2024) use input and space independent conditions for applications in co-protein design, but similarly, their method cannot generalize to novel distributions. Finally, the weather/climate forecast models generating the next state conditioned on the previous one (Price et al., 2023; Verma et al., 2024) are similar approaches to ours but operating on a much finer time resolution.
5 EXPERIMENTS
To show the effectiveness of MFM to generalize under previously unseen populations for the pop- ulation prediction task, we consider two experimental settings. (i) A synthetic experiment with well defined coupled populations, and (ii) experiments on publicly available single-cell dataset consisting of populations from patient dependent treatment response trials. We parameterize all vector field models vy(: (p(p)); w) using Multi-Layer Perceptron (MLP). For MFM, we additionally parameterize p(p;; 0, k) using a Graph Convolutional Network (GCN) with a k--earest neighbor graph edge pooling layer. We include details regarding model hyperparameters, training/opiimization, and implementation in Appendix D and Appendix D.2. We report results over 3 random seeds.
5.1 SYNTHETIC EXPERIMENT
Synthetic data. We curate a synthetic dataset of the joint distributions (Po ()). (r 10)) by simulating diffusion process applied to set of pre-defined target distributions P1 (#1 [i) for i= 1,..., N. To get a paired population Po (xo i) we simulate the forward diffusion process without drift 2o Ma(1, a). After this setup, for reasonable values of a, we assume that can reverse the
7

---

Published as a conference paper at ICLR 2025
source
tw0,50
target
source
tm0.50
t=1.00
target
t=1.00
target
Train
Y
YY
Y
Y
Test
FM
CGFM
MFM
Figure 3: Synthetic letters experiment visualizations. Examples of muoel--generate samples from the source distribution (L 0) to predicted target distribution 1). See Fig. Appendix F for further examples.
Table 1: Results of the synthetic letters experiment for population prediction on seen train populations and unseen test populations. We report the 1-Wasserstein (Wi) 2- 2-Wasscrttein (W22), and the maximum mean- discrepancy (MMD) distributional distances. We consider settings denote models that use Gaussian source distributions sampled via 2g ~ NO0,, 1). but maintains the noisy letters source populations
MFM with varying k. We use " to
as imput to the population embedding model (po: 0, k). Populations that comprise test sets (X and Y) are entirely unseen during training. We include an ablation the effect of changing the numher of training populations (Fig. 5) and results when using OT couplings between samples (Table 4) in Appendix F.
Trols
Teel
Hy(l]
MMD
11111]
Hy(0)
MMD [1
MMD (i)
RM
0.205 0.000
0.277
0.200
251 0.00
01:31 : 0.000
0.300 4 -.100)
2:15 0.03
0.235 - 0.200
0.316 à 0.007
332--0.1 0.01
FM
0,505 8.000
0,050
2.0000
31.05:2 2: 8.00
0.764 2: 0.000
0,381
=: 0,000
25.04 2: (1B)
1,030 = 0.200
1.228 0.002
46.36
-U20
OGPM
0.000
0.113 10:00
0.25 0.00
0.31 0.000
0.00P 0.000
5.55 0.00
0.327
0.000
0.405 0.003
CGPMPAA
0.156 0.025
0.221 0.027
1 0.39
0.88 0.003
0.906
1 0.003
55.08 35.08..0.75
1.002 0.011
1.229 0.000
55.65 0.78
0.019 0,400 0.01.
0:013
MPM
0.174 ::
01232 0.000
0.000
7
14.40=
MFM
(ou5))
0
0,003
0,100 =: 8,004
0.16 0196
0.312 5: 0.82
0.240 :: 0.003
2.29 0.05
1,215 2 0.201
0,202 #: 0.0022
0.79
MEM
(our))
0.002
0.101
2: 0:0000
0.16 000
0.312 2: 0.00
0.251
= 0.005
2:35 1122
[1215
2
0,264
0.007
3.75 -1.25
MFM
0.02
0.150 0:000
0:21 0.0
0,086
6:348 0.006
n.15
0.236
0.00
0.362
0.002
0.08
MPM
(MU5)
0.004
0.136:50:00 0.304
0:55 0.05
0.23% 0.088
0:357
0,004
3 0.25
0.204
9,005
0.349
0.001
3.110.11
diffusion process and learn the push-forward map from Po [) to Pı (x1 (t) for every index i. For
this task, given the i-th population index we denote po i) as the source population pr (Ia (1) as the i-th target population.
To construct p1 (x i), discretize samples from defined silhouette; e.g. image of a character, where i indexes the respective character. We use upper case letters as the silhouette and generate the corresponding samples x1 Pı (x1 i) from the uniform distribution over the silhouette and run the diffusion process for samples 21 to acquire Ty- We construct the training data using 10 random orientations of 24 letters. We construct the test data by using 10 random orientations of "X" and "Y". Test data populations of "X" and "Y" are entirely unseen during training.
We train FM, CGFM and 4 variants of MFM of varying & for the GCN population embedding model (p(o 0, f:). When k 0, (p(,; 0, k) becomes identical the DeepSets model (Zaheer et al., 2017). We compare MFM to FM and CGFM. We repeat this experiment for models trained with source distributions sampled from a standard normal 20 AO(0,) 1). We label these models FM N, CGFM and MFM" M. Here, MFM" ''' still takes the original Po as input to the population embedding model (p(o; 0, ke), while (e(; w)) uses To ~ ~N(00 1) as the source distribution.
MFM generalizes to populations from unseen letter silhouettes. FM does not have access to conditional information; hence will only learn an aggregated lens of the distribution dynamics and will not be able to fit the training data, and consequently won't 't generalize to the test conditions. For the training data, the CGFM vector field model takes in the distribution index i as a one-hot input condition. On the test set, since none of these indices is present, we input the normalized constant vector, which averages the learned embeddings of the indices. Because of this, CGFM will fit the training data, however, will not be able to generalize the unseen condition in the test dataset. CGFM indicates when the model can fit the training data and demonstrates that the test data substantially different and cannot be generated by the same model. For MFM, we expect both fit the training data and generalize to unseen distributional conditions.
We report results for the synthetic experiment in Fig. 3 and Table 1. We observe that indeed FM struggles adequately learn sample from pı (wa i) in the training set, while CGFM is able to
8

---

Published as a conference paper at ICLR 2025
effectively sample from P1 (21 [a) in the training set. As expected, CGFM fits the training data, but struggles to generalize beyond its set of training conditions. In contrast, we see that MFM is able to both fit the training data while also generalizing to the unseen test distributions. Interestingly, although MFM performs better for certain values of & versus others, performance does not vary significantly for the range considered. In Fig. 7 (Appendix F.1) we plot the letter population embeddings.
5.2 EXPERIMENTS ON ORGANOID DRUG-SCREEN DATA
Organoid drug-screen data. For ex- periments on biological data, use the organoid drug-screen dataset from
Treatment Conditioning
Replica split
Ramos Zapatero et al. (2023). This dataset is a single-cell mass-cytometry dataset collected over 10 patients. Somewhat unique to this dataset, unlike many prior perturbation-screen datasets which have a single control population, this dataset has matched controls to each experimen- tal condition. Populations from each pa- tient are treated with 11 different drug treat- ments of varying dose concentrations." 2 We use the term replicate to define control- treatment population pairs. Po (z c4) and
Patlent split
samples (k-))
Figure 4: Organoid drug-screen dataset overview. (Leff) given replica consists of control distribution po and cor- responding treatment response distribution p1 for treatment condition c₁. (Right) train and test data splits for replicates (top) and patients (bottom). Experiments are conducted for 11 treatments, 10 patients, 3 culture conditions, and repeated (replicated) numerous times, resulting in a dataset of many control and treated population pairs.
Pr (I1 G): respectively (see Fig. 4-left). After pre-processing, we acquire total of 927 population pairs. In each patient, cell popula- tion are categorized into 3 cell cultures: (i) cancer associated Fibroblasts, (ii) patient-derived organoid cancer cells (PDO), and (iii) patient-derived organoid cancer cells co-cultured fibroblasts (PDOF). We report results for the individual cultures, i.e. Fibroblast, PDO, and PDOF cultures, in Appendix F.
Pre-processing and data splits. We fil- ter each cell population to contain at least 1000 cells and use 43 bio-markers. We con- struct two data splits for the organoid drug- screen dataset (see Fig. 4-right). Replicate split; here we consider 2 settings, (i) leave- out replicates evenly across all patients for testing (replica-1, depicted in Fig. 4 see Table 5 for results on this split), and (ii) leaving out a batch of replicates from a sin- gle patient while including the remaining replicas from the same patient for training (replica-2, results in Table 2). The second setting is the Patients split (results in Ta- ble 3); here we leave-out replicates fully in one patients- - in this setting, we are testing the ability of the model to generalize popu- lation prediction of treatment response for unseen patients. We do this for 3 different patients and report results across these in- dependent splits. Further details regarding the organoid drug-screen dataset, data pre-processing, and data splits are provided in Appendix D.2.
Table 2: Experimental results on the organoid drug-screen dataset for population prediction of treatment response over replicates. We use "T to denote models that incorporate OT couplings between source and target distributions. Note, the ICNN model operates in the OT regime. We use " to denote models that use Gaussian source distributions sampled via æp ~ (O, 1). but maintain the control populations po as input to the population embedding model pppn; 0, k). We denote the best non-OT method with bold and the best OT- based method with underline.
19:13)
Wall)
MMD cooen (1)
-(1)
4.513
4.055
19.11
0.76
dus pu + Rt.p))
6.222
6.540
74.90
0.876
FM
1.30 0.00%
<.501
0.111
0.67
0.505 0.01
CGFM
1.443
0.015
0.041
1.00
0.066
MFM
(nun)
4.200 0.007
4.380
0.50
8.918
JJR2
MBM
(non)
4.210 :: 0,050
1.300
0.088
11:90
1.30
0
0000
MFMiaz
4.214
0.01T
L.3 0.000
0.75
0,915 0000
MPM
(mun)
4,100
0,093
4,269
0,164
1.58 5,96
0.917 + 0.00
FMP
7.114
7.404 2: 0.056
64.97
8.75
0,613 2: 0.006
CO
7:155
0:915
1:300 01177
70.78
4.67
0.437 + 0.10
MFM, "X (nun)
4.177 0.991
4.355 0.016
10.53
0.50
0.911 + 0.001
MI "W" (oun)
0.055
0.067
1,55
0,912 D.ICI
MFM "W
=158==0.00
1.4G
0.912 0.012
4.000
4.34 0.000
9.32 1.221:0.33
0,915 D.00S
FMPOT
4.000
4.397 0.001
0..76::00722
0.910
D.0N5
CO
KNN
4 0.035
4.865 0.038
0.55
0,584 0.002
Baselines. For the organoid drug- trug-screen experiments, we also consider an ICNN model as a baseline as well as the OT variations of FM and CGFM. The ICNN is equivalent to the architecture used in CellOT (Bunne et al., 2023); a method for learning cell specific response to treatments. The ICNN (and likewise CellOT) counterparts our OMT model in that it does not take the population index i as a condition. CGFM takes in the population identity index i as a one-hot input condition. Unlike CGFM, MFM does not require knowledge of population identities. We additionally include two model-agnostic baselines: (i) the distributional distances between the target and source distributions
We consider only the highest dosage and leave exploration of dose-dependent response to future work.

---

Published as a conference paper at ICLR 2025
d(po. Pı). and (ii) the distributional distances between the target and the source distributions shifted by the average mean of target distributions (A1)) of the training data d(po Ho + P1).
Predicting treatment response across replicates. We show results for generalization across replicates for the replica-2 split in Table 2. In Appendix F, we include an extended evaluation as well as report results for separate cell cultures in Table 8 and Table 9. We observe that MFM outperforms all baselines on left-out replicas. Although this is in general a simpler generalization problem compared to the left-out patient populations setting, this result suggests that there is arguably sufficient biological heterogeneity across replicas for MFM to learn meaningful embeddings.
Predicting treatment response across patients. We show results for general- ization across patients in Table We ob- serve that MFM (w/o OT) outperforms all non-OT baseline methods as well as the ICNN (which learns OT couplings), while yielding competitive performance to the FM*OT and CGFM OT. All
Table 3: Experimental results on the organoid drug-screen dataset for population prediction of treatment response across left-out patient populations. We report the mean and standard deviations across metries computed over 3 different patient splits. We denote the best non-O1 method with bold and the best OT- based method with underline.
Mill)
111
d(po.P))
4.175 0.115
4.303
12:11 10.07.12.07
0.902
d(u - Ps th Pi.
6.158
+ 0.230
6.235
0.229
77.74
10.72
0.3022
0.006
methods yield generally equal degrees of variation across 3 patient splits, with
FM
4.171 0.
4.315
0.142
10.95 1: 1.98
0.807
0. 0.025
CGRM
4,180 4: 0.088
4.321
60.19
11.57 = 0.96
119114
MFM
(ours)
4.1.55 + +00094
4.268
0..128
10.18 + 1.28
0.918
0.007
comparable variation relative to the triv- ial baseline. See Appendix F.1 for an
MFML
(ours)
4.112 + 0.086
4.243
0.122
9:50 + 0.99
0.927
0.008
MFM
(ours)
4.067 1: 0.122
4,218
0,160
9.26 1: 1.56
0.026
0.007
MPM4-100
4.112
4.243
0.931 0,002
analysis of population embeddings.
4,054 0.352
4.189
0.194
9:43 2:40
COEMOOOT
4,087 1. 0.129
4.217
0.165
9.83 2.03
0.924
0.009
Through the biological and synthetic ex- periments, we have shown that MFM is able generalize to unseen distributions. The implication of
IONN
4.157 + 0.368
4.282 = 0.213
11.18
2.01
0.904 0.005
our results suggest that MFM can learn population dynamics in unseen environments. In biological contexts, like the one we have shown in this work, this result indicates that we can improve learning of population dynamics, of treatment response or any arbitrary perturbation, in new/unseen patients. This works towards a model where it is possible to predict and design an individualized treatment regimen for each patient based on their individual characteristics and tumor microenvironment.
6 CONCLUSION
Our paper highlights the significance of modeling dynamics based on the entire distribution. While flow-based models offer promising avenue for learning dynamics at the population level, they were previously restricted to a single initial population and predefined (known) conditions. In this paper, we introduce Meta Flow Matching (MFM) as a practical solution to address these limitations. By integrating along vector fields of the Wasserstein manifold, MFM allows for a more comprehensive model of dynamical systems with interacting particles. Crucially, MFM leverages graph neural networks to embed the initial population, enabling the model to generalize over various initial distributions. MFM opens up new possibilities for understanding complex phenomena that emerge from interacting systems in biological and physical systems. In practice, we demonstrate that MFM learns meaningful embeddings of single-cell populations along with the developmental model of these populations. Moreover, our empirical study demonstrates the possibility of modeling patient-specific response to treatments via the meta-learning. We discuss broader impacts of this work in Appendix C.
Limitations & Future work. In this work, we focused on exploring the amortization of learning over the space of distributions using flow matching. We argue this is a more natural model for many biological systems. However, there are many other aspects of modeling biological systems that we did not consider. In particular we did not consider extensions to the manifold setting (Huguet et al., 2022: 2023), unbalanced optimal transport (Benamou, 2003; Yang and Uhler, 2019; Chizat et al., 2018), aligned (Somnath et al., 2023; Liu et al., 2023), or stochastic settings (Bunne et al., 2023; Koshizuka and Sato. 2022) in this work. We observed that OT improves performance of FM and CGFM in the biological settings which we considered. As such, we view investigating the use of OT in MFM as a natural future direction. Since the focus of this work was to investigate the effect of conditioning flow models on initial distributions, we leave this exploration for future work. We also note that our framework can be extended beyond flow matching. One of the central innovations in this work is conditioning vector fields on representations of entire distributions. Hence, this can be easily extended to other training regimes. Furthermore, we proposed the joint training of parameters and 0 for the vector field vt ()) and population embedding model (p(:; f) using the loss defined in Eq. (17). Although this is a novel approach to v, (-(p)p((: ();:: training, there remains room to explore the pre-training and/or use of different losses for the population embedding model.
10

---

Published as a conference paper at ICLR 2025
REPRODUCIBILITY STATEMENT
To ensure reproducibility of our findings and results we submitted our source code as supplementary materials. In addition, describe the mathematical details of our method throughout Section 3 and provide pseudo-code to reproduce training and inference in Algorithm 1 and Algorithm 2. We provide details regarding parameterization of models, optimization, and implementation in Appendix D.4 and Appendix E. For our datasets, we include details regarding the synthetic letters dataset and how to construct in it Section 5.1 and Appendix D.I, and details for data download and data processing for the Organoid drug-screen data in Section 5.2 and Appendix D.2. Information regarding the datasets that we consider this work can also be found in the source code. For our theoretical contributions, we state our assumptions throughout the text in Section 2 and Section and provide detailed proofs for Proposition (in the main text) and Theorem (in Appendix B).
ACKNOWLEDGMENT
The authors acknowledge funding from UNIQUE, CIFAR, NSERC, Intel, and Samsung. The research was enabled in part by by the Province of Ontario and companies sponsoring the Vector Institute (http:
aapeeaii/ppae aitpartners////eeeeeaeaaaaaaas )) the computational resources provided by the wawaaancccccann.mma ca), Mila (https:/ /mmmmmla la. quebec), and NVIDIA. KN is supported by IVADO.
Digital Research Alliance of Canada
REFERENCES
Achille, A., Lam. M., Tewari, R., Ravichandran, A., Maji, S., Fowlkes, C. C., Soatto, S., and Perona, P. (2019). Task2vec: Task embedding for meta meeaallerrnnn.. In Proceedings of the IEEE/CVF international conference on computer vision, pages 6430-6433
Albergo, M. S. and Vanden-Eijnden, E. (2022). Building normalizing flows with stochastic inter- polants. arviv preprint 15555
Ambrosio, L., Gigli, N., and Savaré, G. (2008). Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media.
Amos, B., Cohen, S., Luise, G., and Redko, I. (2022). Meta optimal transport. arXiv preprint arXiv:2206.05262.
Amos, B. et al. (2023). Tutorial on amortized optimization. Foundations and Trends@ in Machine Learning, 16(5):592-732.
Amos, B., Xu, L., and Kolter, J. Z. (2017). Input convex neural networks. In International conference on machine learning, pages 146-155. PMLR.
Armingol, E., Officer, A., Harismendy, O., and Lewis, N. E. (2020). Deciphering cell cell interactions and communication from gene expression. Nature Reviews Genetics, 22(2):71-88)
Benamou, J.-D. (2003). Numerical resolution of "unbalanced" mass transport problem. ESAIM: Mathematical Modelling and Numerical Analysis, 37(5):851-868.
Binnewies, M., Roberts, E. W., Kersten, K., Chan, V., Fearon, D. F., Merad, M., Coussens, L. M., Gabrilovich, D. I., Ostrand Ostrand-Rosenbergg S., Hedrick, C. C., Vonderheide. R. H.. Pittet, M. J., Jain, R. K., Zou, W., Howcroft, T. K., Woodhouse, E. C., Weinberg, R. A., and Krummel, M. F. (2018). Understanding the tumor immune microenvironment (time) for effective therapy. Nature Medicine, 24(5):541-555
Bunne, C., Stark, S. G., Gut, G., Del Castillo, J. S., Levesque, M., Lehmann, K.-V., Pelkmans, L., Krause, A., and Rätsch, G. (2023). Learning single-cell perturbation responses using neural optimal transport. Nature Methods, 17511)1155-.
Campbell, A., Yim, J., Barzilay, R., Rainforth, T., and Jaakkola, T. (2024). Generative flows discrete state-spaces: Enabling multimodal flows with applications to protein co-design. arXiv preprint xxXiv:242..04997
11

---

Published as a conference paper at ICLR 2025
Chen, S., Rivaud, P., Park, J. H., Tsou, T., Charles, E., Haliburton, J. R.. Pichiorri, F., and Thomson, M. (2020). Dissecting heterogeneous cell populations across drug and disease conditions with popalign. Proceedings of the National Academy of Sciences, 117/46)28784-2879)
Chen, T., Chen, X., Chen, W., Heaton, H.,, Liu, Wang, Z., and Yin, W. (2022). Learning to optimize: A primer and benchmark. Journal of Machine Learning Research, 23(18)):
Chizat, L., Peyré, G., Schmitzer, B.. and Vialard, F.-.. (2018). Unbalanced optimal transport: Dynamic and kantorovich formulations. Journal o Functional Analysis, 274(11):3000-312
Chung, W., Eum, H. H., Lee, H.-O., Lee, K.-M.,, Lee, H.-B., Kim, K.-T., Ryu, H. S., Kim. S., Lee, J. E., Park, Y. H., Kan, Z., Han, W., and Park, W.-Y (2017). Single-cell maasee enables comprehensive tumour and immune cell profiling in primary breast cancer. Nature Communications, 8(1).
Dao, Q., Phung, H., Nguyen, B., and Tran, A. (2023). Flow matching in latent space. arXiv preprint arXiv:2307.08698
De Bortoli, V., Thornoon, J., Heng, J., and Doucet, A. (2021). Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems, 34:17695-17709.
Frangieh, C. J., Melms, J. C., Thakore, P. I., Geiger-Schuller, K. R., Ho, P., Luoma, A. M., Cleary, B., Jerby-Amon, L., Malu, S., Cuoco, M. S., Zhao, M., Ager, C. R., Rogava, M., Hovey, L., Rotem, A., Bernatchez, C., Wucherpfennig, K. W. Johnson, B. E., Rozenblatt- D., Regev, A., and Izar, B. (2021). Multimodal pooled perturb-cite-see seq screens in patient models define mechanisms of cancer immune evasion. Nature Genetics, 53(3):332 341.
O., Schadendorf,
Gilmer, J., Schoenholz, S. S., Riley, F., Vinyals, O., and Dahl, G. (2017). Neural message passing for quantum chemistry. In International conference on machine learning, pages 1263-1272. PMLR.
Goodenough, D. A. and Paul, D. L. (2009). Gap junctions. Cold Spring Harb Perspect Biol, 1(1)a002576.
Gulati, G. S., Sikandar, S. S., Wesche, D. J., Manjunath, A., Bharadwaj. A., Berger, M. J., Ilagan, F., Kuo, A. H., Hsieh, W., Cai, S., Zabala, M., Scheeren, A., Lobo, N. A., Qian, D., Yu, F. B., Dirbas, F. M., Clarke, M. F,, and Newman, A. M. (2020). Single-cell transcriptional diversity is hallmark of developmental potential. Science, 36766476:
Hashimoto, T. B., Gifford, D. K., and Jaakkola, T. S. (2016). Learning population-level diffusions with generative recurrent networks. Proceedings of the 33rd International Conference on Machine Learning, pages 2417-2426.
Hetzel, L., Boehm, S., Kilbertus, N., Günnemann, S., Lotfollahi, M., and Theis, F. (2022). Predicting cellular responses novel drug perturbations at single-cell resolution. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A., editors, Advances in Neural Information Processing Systems, volume 35, pages 2671 1-26222. Curran Associates, Inc.
Hospedales, T., Antoniou, A., Micaelli, P., and Storkey, A. (2021). Meta-learning in neural networks: A survey. IEEE transactions on pattern analysis and machine intelligence, 44(9)511495516)
Huguet, G., Magruder, D. S., Tong, A., Fasina, O., Kuchroo, M., Wolf, G., and Krishnaswamy, (2022). Manifold interpolating optimal-transport flows for trajectory inference.
Huguet, G., Tong, A., Zapatero, M. R., Wolf, G., and Krishnaswamy, S. (2023). Geodesic sinkhorn: Optimal transport for high-dimensional datasets. In IEEE MLSP.
Isobe, N., Koyama, M., Hayashi, K., and Fukumizu. K. (2024). Extended flow matching: a method of conditional generation with generalized continuity equation. arXiv preprint arXiv:2402. 188399
Jabin, P.-E. and Wang, Z. (2016). Mean field limit and propagation of chaos for vlasov systems with bounded forces. Journal Functional Analysis, 27 (7)
12

---

Published as a conference paper at ICLR 2025
Ji, Y., Lotfollahi, M., Wolf, F. A., and Theis, F. J. (2021). Machine learning for perturbational single-cell omics. Cell Systems, 12(6))522-537,
Koshizuka, T. and Sato, I. (2022). Neural lagrangian schrödingerdinger bridge: Diffusion modeling for population dynamics. arXiv preprint crXiv:2204.04853
Lipman, Y., Chen, T., Ben-Hamu, H., Nickel, M., and Le, M. (2022). Flow matching for generative modeling. arXiv preprint arXiv:2210.02747.
Liu, G.-H., Chen, T., So, O., and Theodorou, E. (2022a). Deep generalized schrödinger bridge. Advances in Neural Information Processing Systems, 35:9374-9388
Liu. G.-H., Lipman. Y., Nickel, M.. Karrer, B., Theodorou, E. A., and Chen, Q. (2024). Generalized schrödinger bridge matching. International Conference on Learning Representations.
Liu, G.-H., Vahdat, A., Huang, D.-A., Theodorou, A., Nie, W., and Anandkumar, A. (2023). Psb: Image-to-image schröddiner bridge. In ICMI..
Liu, X., Gong, C., and Liu, Q (2022b). Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003.
Lopez, R., Regier, J., Cole, M. B., Jordan, M. I., and Yosef, N. (2018). Deep generative modeling for single-cell transcriptomies. Nature Methods, 15(12):1053-1058.
Lotfollahi, M., Wolf, F. A., and Theis, J. (2019). segen predicts single-cell perturbation responses. Nature Methods, 16(8):715-72..
Makkuva, A. V., Taghvaei, A., Oh, S., and Lee, J. D. (2020). Optimal transport mapping via input convex neural networks. ICML.
Molè, M. A.. Coorens, T. H. H., Shahbazi, M. N., Weberling, A., Weatherbee, B. A. T., Gantner, C. W., Sancho-Serra, C., Richardson, L., Drinkwater, A., Syed. N., Engley, S.. Snell, P., Christie, L., Elder, K., Campbell, A., Fishel, S., Behjati, S., Vento-Tormo, R., and Zernicka-Goetz, M. (2021). A single cell characterisation of human embryogenesis identifies pluripotency transitions and putative anterior hypoblast centre. Nature Communications, 12(1).
Neklyudov, K., Brekelmans, R., Tong, A., Atanackovic, L., Liu, Q. and Makhzani, A. (2023). com- putational framework for solving wasserstein lagrangian flows. arXiv preprint arXiv:2310.10641
Neklyudov, K., Severo, D., and Makhzani, A. (2022). Action matching: A variational method for learning stochastic dynamics from samples.
Otto, F. (2001). The geometry of dissipative evolution equations: the porous medium equation.
Peidli, S., Green, T. D., Shen, C., Gross, T., Min, J., Garda, S., Yuan, B., Schumacher, L. J., Taylor- King, J. P., Marks, D. S., et al. (2024). scperturb: harmonized single-cell perturbation data. Nature Methods, pages 1-10.
Peyré, G. and Cuturi, M. (2019). Computational Optimal Transport. arXiv:1803.0m
Pooladian, A-A.,, Ben-Hamu, H., Domingo-Enrich, C., Amos, Lipman, Y., and Chen, R. (2023). Multisample flow matching: Straightening flows with minibatch couplings. arXiv preprint 11714:2304.14772.
Price, I., Manchez-Gonzaler A., Alet, F., Ewalds, T., El-Kadi, A., Stott, J., Mohamed, S., Battaglia, P., Lam, R., and Willson, M. (2023). Geneast: Diffusion- ensemble forecasting for medium- range weather. arXiv preprini arXiv:2312.1572
Ramos Zapatero, M., Tong, A., Opzoomer, J. W., O'Sullivan, R., Cardoso Rodriguez, F., Sufi, J., Vlckova, P.. Nattress, C., Qin, X., Claus, J., Hochhauser, D.,, Krishnaswamy, S., and Tape, C. J. (2023). Trellis tree-based analysis reveals stromal regulation of patient-derived organoid drug responses. Cell, 186(25):5606556192)
13

---

Published as a conference paper at ICLR 2025
Rizvi, A. H., Camara, P. G., Kandror, E. K., Roberts, J., Schieren, I., Maniatis, T., and Rabadan, R. (2017). Single-cell topological ma-seq analysis reveals insights into cellular differentiation and development. Nature Biotechnology. 35(6):551-560
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. (2022). High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10684-10695.
Saharia, C., Chan, W., Chang. H., Lee, C.,, Ho, J., Salimans, T., Fleet, D., and Norouzi, M. (2022a). Palette: Image-to-image diffusion models. In ACM SIGGRAPH 2022 conference proceedings, pages 1-10.
Saharia, C., Chan, W., Saxena, S., L., Whang. J., Denton, E. L.,, Ghasemipour, K., Gontijo Lopes, R., Karagol Ayan, B., Salimans, T., et al. (2022b). Photorealistic text-to-image diffusion models with deep language understanding. Advances in neural information processing systems, 35:36479- 36494.
Sanchez- Sanchez-Gonzale Gonzalez, A., Godwin, J., Pfaff, T., Ying, R., Leskovec, J., and Battaglia, P. (2020). Learning to simulate complex physics with graph networks. International conference on machine learning, pages 8459-8468 PMLR.
Schiebinger, G., Shu, J., Tabaka, M., Cleary, B., Subramanian, V., Solomon, A., Gould, J., Liu, S., Lin, S., Berube, P., al. (2019). Optimal-transpon analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell, 176(4):928-9) 943.
Somnath, R., Pariset, M., Hsieh, Y.-P., Martinez, M. R., Krause, A., and Bunne, C. (2023). Aligned diffusion schröodinger bridges. In UAI.
Svensson, V., Vento-Tormo, R., and Teichmann, S. A. (2018). Exponential scaling single-cell rna-seq in the past decade. Nature protocols, (((4):599-044.
Tong, A., Fatras, K., Malkin, N., Huguet, G., Zhang. Y., Rector- Rector-Brooks, J., Wolf, G., and Bengio, Y. (2024a). Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research. Expert Certification.
Tong, A., Huang, J.. Wolf, G., Van Dijk, D., and Krishnaswamy, S. (2020). Trajectorynet: A dynamic optimal transport network for modeling cellular dynamics. In International conference on machine learning, pages 9526-9536. PMLR.
Tong, A. Y., Malkin, N., Fatras, K., Atanackovic, L., Zhang, Y., Huguet, G., Wolf, G., and Bengio, Y. (2024b). Simulation-frr schrödinger bridges via score and flow matching. In International Conference on Artificial Intelligence and Statistics, pages 1279-1287, PMLR.
Uscidda, T. and Cuturi, M. (2023). The monge gap: A regularizer to learn all transport maps. In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J., editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages
PMLR.
Verma, Y., Heinonen, M., and Garg, V. (2024). Climode: Climate and weather forecasting with physics-informed neural odes. arXiv preprint 200.:2404..022..
Villani, C. (2009). Optimal transport: old and new, volume 338. Springer.
Weinreb, C., Wolock, S., Tusi, B. K., Socolovsky, M., and Klein, A. M. (2018). Fundamental limits on dynamic inference from single-cell snapshots. 115(10)E2467-E2277
Yang, K. D. and Uhler, C. (2019). Scalable unbalanced optimal transport using generative adversarial networks. In 7th International Conference on Learning Representations, page 20.
Zaheer, M., Kottur, S., Ravanbakhsh, S., Poczos, B., Salakhutdinov, R., and Smola, A. (2017). Deep sets. Advances in neural information processing systems, 30.
Zeng. T. and Dai, H. (2019). Single-cell ma sequencing- computational analysis to describe disease heterogeneity. Frontiers in Genetics, 10.
Zheng, Q., Le, M., Shaul, N., Lipman, Y., Grover, A., and Chen, R. T. (2023). Guided flows for generative modeling and decision making. arXiv preprint arXiv:2311. 13443.
14

---

Published as a conference paper at ICLR 2025
A DEFINITIONS
All the densities we consider are the densities of the probability measures from the following set
P(X): J du(x)) 5)² +o0 for some I E x
P2(x)
EEX}.
(20)
where P(X) is the family of all Borel probability measures on A'. Thus, one can introduce 2- Wasserstein distance between measures in P2(A')) i.e.
W (po. M1) = min So dt || [|
Öp:
[mo.
+ (V. pret) = =
(21)
at
where used the Sntamou-Breniee formulation of optimal transport (Benamou, 2003). That is, the possible changes of the density are restricted by the continuity equation
@p((x) (V (z))
(22)
where the define the divergence as
d (v,(x)
V v())
(23)
dr,
Note that (-:-,)) define metric on P2(X'), which allows to introduce absolutely-contimmuus curves and the following tangent space
Tan, P2(X) = [vs: & (X)} (1x:2) Vp E P2(2')
(24)
These developments lead to the celebrated Riemannian interpretation of the Wasserrstei distance (-..) .) on P2(X)) proposed by (Otto, 2001) and rigorously studied and presented in (Ambrosio al., 2008).
B PROOF OF THEOREM
Theorem 1. Consider a dataset of populations' D = [[T( xı (i)}} generated from some unknown conditional model (( 21 c)p(c). Then the following objective
L(w, 0) = Expe) {-² di Eppi le)llv (ate (c) - vi (x p(po, w)l
(16)
is equivalent to the Meta Flow Matching objective
LM 0) DE (4) S(' di ② 01 (1( ue(e( - ():0)
(17)
up to an additive constant.
Proof. The loss function
C(w. 0) = Ep(c) /'' dt Em (oll - vx(xe (pp, (]?
(25)
2Ep(o) / dtdr (pr(=lc)u (a c), (x pp(r. 0):
(26)
+ Ep(c) /'' dt Epe(rs (x p(p:, 0), 1
(27)
/'' di Ep,(os le (ax |c)]
+ Ep(c)
(28)
The last term does not depend on 0, the second term we can estimate, for the first term, we use the formula for the (from Eq. (8))
1.(El (E) c) ((( -)) (1 -
(29)
Øt.
15

---

Published as a conference paper at ICLR 2025
Thus, the loss is equivalent (up to constant)
/ dt 21) x1)) p(pe. ():w))
C) = - 28 25)() Eatoomm
(30)
Ōt
+ E (10 (c)
dt (l:( pppp 0) w)1
(31)
(zu, .1)
+ Ex B (x
dt
(32)
Ot
-Eeple) Ex(2n.21
di
S (1
(1)
us (f ((r,, 0);w)
(33)
Note that in the final expression we do not need access to the probabilistic model of p(c) if the joints π (2o, 21 c) are already sampled in the data Thus, we have
L(w, 0) Eo-ple) Bx( B (o) S dt on f1(00, (e(F((2, ( 0):w)
(34)
E EENDBr(RO,TI (0) L dt
0t f( x1) - ve(f((zo, pppe, 0):w)
(35)
LMMM (w, 0).
(36)
0
C BROADER IMPACTS
This paper is primarily theoretical and methodological contribution with little societal impact. MFM can be used to better model dynamical systems of interacting particles and in particular cellular systems, Better modeling of cellular systems can potentially be used for the development of malicious biological agents. However, we do not see this significant risk at this time.
D EXPERIMENTAL DETAILS
D.1 SYNTHETIC LETTERS DATA
The synthetic letters dataset (used for Fig. and Table contains 240 train populations a test populations. Each population contains roughly between 750 and 2700 samples in this dataset.
D.2 ORGANOID DRUG-SCREEN DATA
The organoid drug-screen dataset consists of patient derived organoids (PDOs) from 10 different patients (Ramos Zapatero et al., 2023). For each patient, experiments are replicated and per- formed over varrying conditions. This results in significant amount of coupled populations pairs (po (zolc;), pi (c)]c))), where we use c; to denote the treatment condition corresponding to respective population pair. After pre-processing, we acquire total of 927 replicates (or coupled population pairs). In the replica-1 split, we use 713 populations for training, 111 left-out population for vali- dation, and 103 left-out populations for testing. In the replica-2 22 split, training, 33 left-out populations for validation, and 33 left-out populations for testing. We use the replica splits to set and select reasonable hyperparameters for MFM and the baseline models.
861 populations for
For the patients split, consider 3 different patient splits where independently leave out all populations from either patient 21, patient 27, and patients 75. Organoid PDO-21 was found to present interesting chemoprotective response (High CAF Protection) to cancer treatments, and hence was selected by Ramos Zapatero et al. (2023) for further study and data collection. For the same reason select PDO-21 the unseen test patient for the patients splir experiments. We additionally consider populations from PDO-75 as a third left-out patient split. In the patients splits, we only split data into train and test sets, resulting in 839 population pairs for training and 88 population pairs for test evaluation for PDO-21 and PDO-27 splits. For PDO-75, the train split contain 830 train population pairs and 97 test population pairs.
PPDOs are cultures of cell populations which derived from patient cells.
16

---

Published as a conference paper at ICLR 2025
D.3 METRICS
We use four metrics to evaluate the quality of generated empirical distribution compared to an empirical ground truth distribution p. In this section we detail the calculation and interpretation of each metric. For letters, evaluate over all samples, while for the orgranoid drug drug-screen dataset we evaluate over 5000 cells. If a target population contains less than 5000 cells, we evaluate using sample size corresponding the size of the entire target population.
1- I-Wasserstein (Wi and 2-Wasserstein (W2). We evaluate the Wasserstein distance for 1 and a 2. Specifically for { = p and {2 = p we calculate the a- ---asssrsttinn distance with respect to the Euclidean distance as
1/a
/ | 10-vlie 2
Wa inf (ETI(A)))
(37)
where 11(0,p) denotes the set of all joint probability measures with marginals p. p.
Maximum mean discrepancy (MMD). We evaluate an estimate for MMD between empirical using the radial basis function (RBF) kernel k(, (7) as:
distributions { 1 and {2x 1
1 n
n
MMD(p. p) n
Σ
+ ΣΣ kxI,
Σ n²
k(ūs, (38)
n
We estimate MMD for = 2, 0.1, 0.01, 0.005 and report the average.
Coefficient determination (1²). We adopt the r2 metric, overall average correlation coefficient, from Bunne et al. (2023) used to compare predictions and observations. This metric is computed as follows. First, given two arrays X E Ra and Y E Ruxd, where n is the number of samples and is the number of features, we can compute the correlation
E xxx)((; NY)
corr( (X,))
(39)
VE - -xx, - -xx) Eis My [Y, my))
where yx and Jy are the empirical means of X and Y. We compute correlation across all points for X (predicted samples) and Y (target samples) individually
Σ corr co, X,),
Σg = =orr Yj).
(40)
Then, we consider the two arrays
S Vi j, S = Eイ < j.
(41)
The overall average correlation coefficient is reported 22 corr s
(42)
We compute a 2 for every population in the evaluation sets and report the average y2 over this set. We take this evaluation directly from Bunne et al. (2023), which is typical evaluation of single-cell prediction methods.
D.4 MODEL ARCHITECTURES AND HYPERPARAMETERS
ICNN The ICNN baseline was constructed with two networks ICNN network f(x) and g(z), with non-negative leaky ReLU activation layers. f(x) is used minimize the transport distance and g(x) used to transport from source to target. It has four hidden units with width of 64, and a latent
dimension of 50. Both networks uses Adam optimizer (Ir 10-4, 31 0.5, 92 0.9). g(x) is trained with an inner iteration of 10 for every iteration (xx is trained.
Vector Field Models All vector field models iy are parameterized with linear layers of 512 hidden units and SELU activation functions. For the synthetic experiments, we use hidden layers. For the biological experiments, we use hidden layers and skip connections across every layer. We found that this setup worked well for the biological experiments. The FM vector field model additionally takes conditional input for the one-hot treatment encoding. CGFM takes the conditional input for the one-hot treatment conditions as well as a one-hot encoding for the population index condition i. The MFM vector field model takes population embedding conditions, that is output from the GCN, as input, as well as the treatment one-hot encoding. All vector field models use temporal embeddings for time and positional embeddings for the input samples. We did not sweep the size of this embeddings space and found that a temporal embedding and positional embeddings sizes 128 worked sufficiently well.
17

---

Published as a conference paper at ICLR 2025
MPM
-X-
CGFM
FM
Train
Test (X's)
Test (Y's)
0.3
0.1
2'
2'
2
2
2
2
2
2
2
2
2
2
2
2
populations
train populations
populations
Train
Test (X's)
Test (Y's)
Pereeteera M
0.3
0.1
26
2'
2
2'
2
26
2
2
2
2
2
24
2
2
train populstions
train populations
populations
Train
Test (X's)
Test (Y's)
0.006
X
0.004
0.002
25
2'
2
2
2
2
2.5
2
2
24 2
2
train populations
train populations
train populations
Figure 5: Synthetic letters ablation over number training populations. Here we fix the test (X's and Y's) with 10 random rotations (same across each experiment). We then ablate the number populations used for training FM (red), CGFM (green), and MFM (blue). by changing the number of random pnations/oriemtationn used for each letter silhouette. We observe that for MFM the distributional errors on the test sets consistently decrease
increase the number of training populations. In contrast, since FM and CGFM cannot generalize across novel populations, this increase in training populations does not lead to an overall improved performance on the test populations. For a large number of training populations, CGFM exhibits exhaustive memory requirements since it requires one-hot encodings as input conditions to denote the population index.
Graph Neural Network We considered GCN model that consists of k-nearest neighbour graph edge pooling layer and graph convolution layers with 512 hidden units. For the synthetic experiment we found that 3 GCN layers to work well. while for the biological experiments we found GCN layers to perform well. The final GCN model layer outputs embedding representation the Synthetic experiment, we found that d = 64 performed well, and d = 128 performed well for the biological experiments. We normalize and project embeddings onto hyper-sphere, and find that this normalization helps improve training. Additionally, the GCN takes a one-hot cell-type encoding (encoding for Fibroblast cells PDO cells) for the control populations po- This may be beneficial for PDOF populations where both Fibroblast cells and PDO cells are present. However, it is important to note that labeling which cells are Fibroblasts versus PDOs withing the PDOF cultures is difficult and noisy in itself, hence such a cell-type condition may yield no additive informmtioo/pperfrmmnncc gain.
Rd For
Optimization We use the Adam optimizer with a learning rate of 0.0001 for all Flow-matching models (FM, CGFM, MFM). We also used the Adam optimizer with learning rate of 0.0001 for the GCN model. To train the MFM (FM+GCN) models, we alternate between updating the vector field model parameters w and the GCN model parameters 0. We alternate between updating the respective model parameters every gradient step. For the syntehtic experiment, FM and CGFM model were trained for 12000 epochs, while MFM models were trained for 24000 epochs, with a population batch size of 10 and sample batch size of 700. Due to the altermating optimization, the MFM vector field model receives half as many updates compared to its counterparts (FM and CGFM). Therefore, training for the double the epochs is necessary for fair comparison.
18

---

Published as a conference paper at ICLR 2025
Influence of GCN parameter is on performance. We observe that clear single selection for k yields the best performance across all tasks on the single-cell experiments. For the replicates split, k 0 on average performs better than & >> O, whereas the patients split, the opposite is true, 100 performs best on average. This is possibly since the patients split forms a more difficult generalization problem (more diversity between training populations and test populations), and hence it is more difficult to over-fit during training with higher k.
The hyperparameters stated in this section were selected from brief and small grid search sweeps. We did not conduct any thorough hyperparameter optimization.
E IMPLEMENTATION DETAILS
We implement all our experiments using PyTTrc and Py Torch Geometric. All experiments were conducted HPC cluster primarily NVIDIA Tesla T4 16GB and A40 48GB GPUs. Each individual seed experiment run required only GPU. Depending on the model (FM vs CGFM vs MFM) Each synthetic experiment ran between 6-24 hours and each real-data experiment ran between 12-36 hours. All experiments took approximately 500 GPU hours in aggregate.
19

---

Published as a conference paper at ICLR 2025
F
EXTENDED RESULTS
Soarce
1-050
T-1.00
target
SOUPCE
t-0.50
1w1.00
target
SOUFCE
t-0.50
T=100
tagget
PBO
Train
3
E
TE
EE
L
T
150.50
X
Test (X'S)
X
target
A
Test (Y's)
Y
L
FM
CGFM
MFM
Figure 6: Model-generane generated samples for synthetic letters from the source 0) target (t 1) distributions.
20

---

Published as a conference paper at ICLR 2025
Table 4: Comparison of models trained using independent coupling versus OT couplings on the synthetic letters setting. We report the 1- I-Wasserstein (W1). 2- Wasserstein (W2), and the maximum mean disan--ndsrpancy (MMD) distributional distances. We use OT to denote models that incorporate OT couplings between source and target distributions (Tong al., 2024a; Pooladian al., 2023). We report results for MFM (k = 50). Here we report results for models trained with a population batch size of and sample hatch size of 700. We observe using OT couplings in this setting does not provide a significant change in performance.
Train
Text (Y'o
Wall)
MMD
1. 183
(.(-))
Wall
MMD -.te- (1)
Wi(
Welc]
MMD extonts (U)
EM
02200 a
0.279 1: 0.000
2.548 0.000
0.201
0.000
0.307 =: 0.000
2,422 2 00000
0.236 a 0.000
0.315
0.000
0,305 2: 00000
EM-lOT
0:210
0.000
0.277 2: 0.001
2.555 = 0.003
0.254
5: 0.0001
0.30 0.000
2.458 20003
0.239
0001
0.318
0.001
3,365 2: 0055
OGMM
0,893
0.000
0.318 2 6,000
-122220000 0.000
0-321
2.0000
0.310
0.000
5.125 2:000
0.309
0000
0.388
0.000
6,384 1: (0000)
CGPM-OOT
0,893
0.000
0.110 2: 0.000
0.215 - 0.036
0.314
8: 0.002
0.000
: 0.002
4,871 0.004
0.307 a: 0.001
0.588
0.002
6,021 :: 0.08
MRM
0.100=00.000
0.127
+: 0.002
0.322 0.002
RJ083
0.238
0.006
1.742
0.190
0.082
0.003
1,473 0.120
MiM"or
0.100:::00000
0.1333 a: 0.010
0.400 0:330
0.189
1.006
0,230
0.007
1,706 : 0.139
0.183 c 0.006
0:200 01000
2,522 a 0.220)
Table 5: Experimental results on the organoid drug-sereen dataset for population prediction of treatment response across replica-/ populations. Results shown in this table are for dim = 43. We use "NV to denote models that use Gaussian source distributions sampled via zo ~ . (O, 1), but maintain the control populations po as input to the population embedding model (Po: 0, k). We further report d(pr p4o + (1,P1). ) baseline that uses the means from each true treated populations p1-
Train
Test
Wy(-)
MMD
(1 (L)
(1(1)
W.(4)
Wa([]
MMD 10 31 (-)
((1)
Alp..p.)
4.182
4.254
10.84
0.925
4.100
4,222
10.21
0.931
Alpo Ho + Pc.
6:089
6,074
69.71
0.925
5.986
6.058
68.00
0.931
dip- Pt
3.915
3.952
1.97
0.925
TM
3,925 1. 0.019
4.041
0.023
3.76.1 0.20
0.952 1 0.007
3,961 1. 0.006
4.060 1 0.042
5.90 0.25
0.041
0.010
OGFM
3364 a: 0.064
3,900
(1,000
5.16
0
0.361 =: 0.006
1.NT
=: 0.063
4.211
= 111166
8.84 4: 0.75
0,333
0.0M5
MFM (6 0)
3.874 0.015
3,973
0,020
3:37 +
0.967 + 0.003
3,580 + 0.009
3,990 + 00.011
1.88+-0.16
0.055
0.002
MFM (6-
3.896 + 0.021
4.000
0.021
3.82 0.12
0.961 + 0.001
3.890
+ 0.013
4.012
- 0.011
6.13+00.8
0.053 0.001
MIM 500
3,888 0.008
3:901
0.000
3.50 0.41
0,953 1 0.001
3,900 1 0.038
4.013
0.054
5.06 0.22
0.054 0.005
MPM 2000
4,006 0.010
4.008
0.00G
4.05
0.002
0.008
4,000
0,300
5.19 0..05
0.957 0.000
FMC
6,008 0.007
7.181 0.003
57.70
1 0.005
6.972 0.022
7:244
60.30
0:812
OGFM
1.187 0.008
4.340 0.000
8.60:1:0.3 0.30
0.936 +: 0.002
6.852 +: 0.045
7.114 = 0.044
71.21 5:711
0.895 0.016
WV 0)
3,940 0.022
4.047
0.023
3.91 +-0.18
0.930
+ 0.006
3.896 + 0.026
4/002 - 0.090
4.35 0.18
0.950 0.005
MFM (k- - 10)
3,976 +: 0.044
4.066 =: (1,04)
4.52 0.42
0.961 + 0.002
3.943 + 0.082
4.051 = 0.1154
5,28 + 0.25
0,972 0.001
MPMPW
1.968
4.075
01014
4,36
0.001
3.9144
4.041 11.108
01.35
0.304
0.00K
MFMPW (k- 100)
3.937
3,96
0.001
0.053
0,953
ICNN
4,286 à: 0.018 4.313 2: 0.112
38,60
0.897 =: 0.031
4.191 0.310 4.313 =: 0.112
37,90 2: 2.84
0,897
Table 6: Experimental results on the organoid drug-sereen dataset for population prediction of treatment response across replica-2 populations. Results shown in this table are for dim- 43.
Train
Text
Wy(c)
MMD (x (()
-(t)
(9(())
MMD -. 7 (.)
P(t)
plpu.p))
4,218
8:GG
0,125
4.518
4.1855
19.34
0.876
d/p- p..p
6.002
6.083
63.50
0.929
6.222
6.346
74:00
0.876
4/2- Pt.
3,887
3.962
0.929
FM
1,985
0.054
4.315
0,067
1.64 0.48
0,333
0.014
4,340 :: 0.078
4.564
11.11
10,00 2: 01.77
0.865
0.004
CGEM
3,882 0.019
3:959
0.020
3,16 0.59
0,952
1 0.0014
4.448 1. 0.083
4.621
0.041
17.00 1.00
0,999 0.008
MFM
3,905
0.006
0.25
0,058
0.001
4,200
0.067
4.540
0.002
12.84
0.50
0 0JN2
MPM
3.80
0.005
4.216
0.060
4.596
11.00 2.56
0.017
2.012
4.008
0,0022
4.20
0.955 0.000
4.214 0.017
4,306
12.099 0.75
0.916
MEM 2000
3.884 0.009
3,986
0,044
3.77
0.001
4.100
1 0.693
4.269
0.104
1.58
0.917 0.000
PM
6,892 & 0.027
7.364 2: 0,003
$7.03 4: 1.000
0.633 1: 0.001
7.114 1: 0.100
7.404 =: 11186
64.97 3.79
0.613 à 0.008
CGIMPNN
4.315 1 0.077
4.480
0,881
11.51 1 1:56
0.918 1 0,004
7.135
1. 0.045
1,300
1 0.007
79.78 1. 4.67
0.637 1 0.010
MEMPWN (k- 0)
3,982 d: 0.014
4.086
01015
6.044 0.36
0.951
= 0.002
4.177 =: 0.042
4.550
= 0.114
30.03 2: 0.50
0,911 0.001
WPM
4.0016
0,012
5.13
0.001
4.156 =: 0.065
4.324
1,63
0,912
0.003
MEM 0
3,082
0.018
4.005
0.016
4.74
0.951
0.002
4.155 1. 0.009
4.324
1.45
0.912
0.002
MPM (k- 100)
4.004 0.012
4.110
0.014
5.10.1..44
0.43
0.949
1 0.002
4.166
1. 0.001
4.541 1: 0.003
9.52 11 0.53
0.015 0.005
ICNN
4.308 L 0.034
4.413 1. 0.008
7.07 0.11
0.929 1 0.006
4. 4.488 1. 0.035
4.065 1. 0.038
17.60.00.55
0.884 L 0.002
Table 7: Experimental results on the organoid drug-sereen dataset for population prediction of treatment response across left-out patient populations. Results shown in this table are for dim = 43. We report the mean and standard deviations across metrics computed over 3 patient splits. We denote the best non-OT method with bold and the best OT- based method with underline.
Train
Wb(-))
MMD ouor (4)
(())
W3(4)
MMD ccse (1)
d(pu.P)) PL)
4.128 à 0.009
4.254
4 0.076
9:44=122 1.21
0.907 a 0.007
4.175 à 0.135
4.803 L: 0.174
4:011
0,502 & 0,008
App- p0 + P..
6.007 + 0.027
6.065
0.034
68 06:30-0.00 0.00
0.957
0.007
6.158 + 0.259
6.235 + 0.229
77.74 + 10.72
0,002 + 0.006
4/2- Pt)
3.905 1 0.042
3.987 0.001
2.01-0.22
0.907
1. 0.007
FM
3,950 1. 0.055
4:087
0.056
4.62 1.64
0.996
L. 0.006
4.171 1 0.100
43315
0.144
10.95 1. 1.98
0.897 1. 0.023
CGFM
3.808 : 0.0344
3.954
0.017
3.61= 11.00
0.955
0.000
6.180
+ 0.888
4.321
+ 0.119
11.57
+ 0.56
0.3114 + 0.008
MFM (k - 0f
1,800 2: 0.064
3.573
0.1776
3.50 2.00--.77
0.060
0.005
6.135 2 0.1194
4.268
= 0.328
11.188 4: 1.28
1133. 0.00T
MFM (1 10)
3.833 0.862
3.963
0.067
3.55
0.92
0.568
0.805
6.112 + 0.086
4.245 +-0.1
0.00 + 0.50
0.025 0.008
MPM
3.863
3,560
0.053
3.64 0.64
0.006 0.966
4,067 1. 0.122
4.218
0.160
1.56
0.926 0.007
MPM 100)
0.065
3,550
3.60 0.91
0.966
4.112
4.244
0.386
0.951 0JI02
IM"OT
3,896 0.056
3,961
3.76
0.007
6.064 0.152
4.189 1 0.194
2.49
0,002
OG
3.783 0.040
0.057
2.38 0.50
0.074
0.004
4.087
1 0.129
4.217
4: 0.165
0.55 1. 2.03
0.024 0.009
JCNN
4,384 2: 0.477
4.508 0.518
3.57 7.80
0.914 1: 0.0822
4.157
4 0.168
4.282
1: 012111
11.18 4: 2.51
0.304 à: 0.005
21

---

Published as a conference paper at ICLR 2025
Table 8: Extended experimental results the organoid drug-screen dataset for population prediction of treatment response across replicas- populations. Results shown this table are for dim 43. Here we show results for the individual cell cultures: Fibroblasts, patient derived organoids (PDOs), and patient derived organoids with Fibroblasts (PDOFs), respectively. We consider 4 settings for MFM with varying nearest-neighhours parameter (k= 0, 10, 50, 100).
Fibroblasts
Irain
Test
Wi(
WelL)
MMD
((1)
Wi(1) (1)
Wel(1
MMD
TM
3,768 0.000
5.838
0.047
2.171:0.10
0.957
0.0020
3,750 1: 0.018
5,816
0.055
2.20 0.11
0.048 0.014
FM
6.701
51.55 5:11
0.15
48.95 3.60
0.007
CGEM
3.674
0.86
0.001
3.732
3.797
0.929
OGEM
3.671 00011
3.725
0.011
0.06.1...7
0,980
0.001
5.844 0.192
6.148
0.204
1..00
0.895
0.007
ICNN
4.075 0.000
4.145
0.008
4.47:1.0.0 0.00
0.025
1 0.006
3,784
0.007
5,848 0.007
3.05 1: 0.07
0.052 0.005
MFM W 0K 0)
3.745 1: 0.004
3.808
0.045
1.731:0.27
0.968 1 0.010
3,730 1: 0.006
5.792
0.045
2.02 0:229
0.067
0.011
MEM V
3.718 0.025
3.774
0.023
1.8++0.22
0.976 + 0.002
3.719 + 0.025
9:777
+ 0.023
2.15 0.22
0.973 0.000
MEMPW 1t-
3.716
3.771
1.55
1 0.001
3,717
3.773
0.033
2.13100.44
0.976 0.001
MEMMW 0t. 100)
3,712
1.37
0.002
3.708
0.011
3.764
3
1.43
0,027
3.719
0.0221
3.7844 00J022
1,60
0.001
3,080
3.7AH
0..5
0.975
MFM (E= 50)
3.664 0.042
3. 0.043
1.66-+-0.55
0.971
0.002
3.681
+: 0.009
3.722
0.041
1.84 0.23
0.975 0.001
MPM I8- 1000
3.676 4 0.085
3.732
0.006
1.56 0.27
0.972 1 0.003
3.674 1 0.001
3.731 2 0.002
2.05 0.20
0.976 1 0.001
Train
Mi(()
Wy(l)
((1)
Wi(1)
Wh(t]
MMD
3:907
7,354
7.578
0.100
60.00.1.3.0
0.585 1 0.000
7.461 1: 0.001
7.602
0.094
72:40 1: 5:39
0.592
0.000
OGEM
1.792 0.061
3.806
0,063
1.27 4: 0.39
0.979
1: 0.005
4002
=: 0.086
4.181 = 1.100
7.75 2: 1.45
0,000 0.011
CGPMM
3.46 RU832
3,807
0,037
1.19
0.29
0.986
+ 0.002
7.672
+ 0.140
7:509
- 0.144
94.96
7:53
0.002
0.022
ICNN
4,555 0.008
4.635
0.007
11.55 0.11
0.901
1 0.005
4,152 1 0.014
4.261
0.013
8.66
0.27
0.028 1 0.001
W 0)
3.788 0.047
3.851
0.033
1.50
0.20
0.382 + 0.002
3.741
0.042
3,829
0.047
2,39 0.23
0,978
0.002
WF
1.76=00.55
0.001
3.785
0.40
0.979
MEM (k-
1.794
3,558
0.1134
1.71 0.50
0.985
0.0012
3,775
0.022
3,888
0,980 0.002
MPM (k- 100)
3.783 11855L000 0.018
3.845
0,018
0.19 1.53
0.984 1 0.002
3,749 1. 0.627
3.835
0.028
0.24 2.59
0,981 0.002
MEM (K- 01
3.843 0.056
3.916
0,006
2.65 0.81
0.971 1. 0.002
3.704 1: 01005
3,891 =: 0.0074
0,34 1.00
0,974
0.002
MFM (k- 100
3.852 0.039
3.502
0.046
2.80++0.15
0.972 + 0.005
3.781 + 0.013
9:879
= 0.00.
3.45 + 0.52
0.976
0.002
MFM (K- 500
3.864 0.036
3.924
0.003
2,51 1. 0.01
0.973
1 0.005
3,701
1. 0.029
3.560
0.025
3.35
0.14
0.074 0.006
MFM (E 100)
3.005 0.082
3.968
0.088
3.87 5: 1.81
0.071 3: 0.001
3.785
3: 0.025
3.552 : 0.009
3.35
5: 0.45
0.076
0.005
PIMIES
Train
We(l)
Wy(c)
MMD (())
(11)
W:(1)
315(3
MMD :uw (-)
FM
4.117 0.088
4.287
0.007
5.75 0.63
0.953 a: 0.001
4.328 4: 0.057
4.314=10. 0.064
12.11 0.76
0,911
0.07
6,696 0.005
6.521
0.095
52.12.1..30
0.722 1 0.005
6,906 1. 0.073
7.154
1 0.073
30.75 3.16
0.696 0.009
CGEM
4,180 1 1.130
4.385
0.145
6.35.1.2.55 2.13
0.933 1 0.020
4.467 1. 0.080
4.653
1 0.083
15.51
1.000
0.005 0.001
OG
5.150 0.046
5.499
0.010
20.70 1:21
0.843
:: 0.005
7.050 :: 0.147
7.285
0.1
78.06 6:63
0,700 0.020
ICNN
4.034
4.582
0.006
5.09
0.002
4.552
0.005
4.735
0.006
17.02
0,858 0.005
WFMW
4.287
8-41+-0.00
0.011
4.215 0.016
4.585
0.34
0.905 0.002
MEM (k-
4.372
0.048
9.70
0.923
0.007
4.326
0.062
4.501
0.057
10.86 0.00
0.904
W 000
4.356 0.036
4.566 0,040
9.72 0.088
0.922
= 0.00NG
4.291
+ 0.014
4,460
1.18
30.14 0.00
0.305 0.001
WFM 200)
4.200 0.026
4,005
(1,033
8.55 0.36
01332 à: 0,0017
4.251 :: 0.071
4,423
11,180
9.44 2: 1.51
0,905 0.003
MFM (k 0)
4.111 8J816
4.281
0,023
G.04 0.26
0.954 :: 0,009
4.184 :: 0.623
4,359 :: 0.022
8,86 2: 0.64
0,015
0L003
MFM (L- 10)
4.158 0.028
1.332 0.072
6.97 0.50
0.946 + 0.006
4.255 + 0.019
4.418 - 0 0.022
9.80 + 0.51
0.914 0.002
MFM (L- -500
4.155 0.055
4.330 0.015
6.80 0.98
0.943 + 0.002
4.245 + 0.044
4.429 - 0.055
10.00 0..3
0.912 0.001
MIM 100)
4.157 0.048
4.305 0.068
6.75 0.42
0,945 1 0.002
4.256 L 0.001
4.412 0.004
10.18 0.10
0.918 0.005
22

---

Published as a conference paper at ICLR 2025
Table 9: Extended experimental results the organoid drug-screen dataset for population prediction of treatment response across replicas-2 populations. Results shown this table are for dim 43. Here we show results for the individual cell cultures: Fibroblasts, patient derived organoids (PDOs), and patient derived organoids with Fibroblasts (PDOFs), respectively. We consider 4 settings for MFM with varying nearest-neighhours parameter (k= 0, 10, 50, 100).
Fibroblasts
Irain
Test
Wi(
We(L)
MMD (MMD...c-"111
((1)
Wi(1)
Wer(1
MMD
-11
TM
3.604 0.088
5.775 0.023
1.24::0.:00
0,951
0.000
3,606 2: 06.
5,805
0.225)
3.67 1: 1.46
0.867 0.070
FM
0.120
0.002
64.40
1.51
0.90
0.003
0.011
3.584
0.975
OGEM
3,683 0.111
3.718 0.0011
0.05 1.0.1
0,982 0.001
5.813 1. 0.572
6.136
0.398
48.44 5.188
0.853 0.026
ICNN
4.054 1: 0.03
4.124
2: 0.033
4.43
0.025
1 0.007
3.477
1: 0.001
5.534
0.053
3.67
0.11
0.009 0.004
WFM W 0K 0)
3,671 0.021
3.726
0.022
1.171:0.17
0.079 1: 0.001
3.505 1: 0.000
5.555
0.051
1.88
0.11
0.075
0.001
MEMM W. (k
3.601 0.015
3.748 + 0.014
1.18+-0.14
0.978 + 0.000
3.526 + 0.224
5:562
0.0 0.022
180++00.5
0.972 0.001
MEMPW it-
3,678 0.017
3.733 0.017
1.15 0.12
1 0.001
3,315 0.029
3.567
0.029
0.63
0.917 0.002
MEMM 0t. 100)
3,706
0.000
1,90
0,075
0.013
1.20
0.020
0.081
MPM
4,024 0.025
3.67% 01/026
1.04 0.11
0.000
3.451 + 0.007
3:504
0.50
0J001
MFM (E= 50)
3.639 +: 0.016
3.604
0.006
1.71
0.979
0.001
3.443 0.635
3,497
2.80 1.85
0.981 0.001
MPM I8- 1000
3.654 à: 0.085
3.712
0.005
1.56 0.20
0.978 1 0.000
3,480 1 0.027
3.534 = 0.0029
3,32 11 0.49
0,980 1 0.001
Train
Mi(()
Wy(l)
((1)
Wi(1)
MMD owwo
0.025
4.174
0,131
0.031
7,109 0.028
7:351
0.000
63.42:1
0.625
1 0.008
7.007
1: 0.212
7:900
0.182
77.52
0.519 0.002
OGEM
1.750 0.018
3.815
0.017
1.77 0.71
0.977
:: 0.0004
4.255
4: 0.146
4:50
= 11171
9.50 3.14
0,933 0.018
CGPMM
3.288 0J824
3.787
0.026
1.14 0.16
0.985
0.000
6.355
+ 0.309
5:334
- 02288
114.46
14.17
0.508
0.016
ICNN
4,476 1 0.000
4.58
0.041
11.17
0.05
0,906
1 0.007
1.415
1 0.048
1.300
1 0.051
11.18
1. 0.05
0.896
1 0.001
W 0)
3.775 0.012
3
0.013
1.75
0.06
0.979 + 0.001
3.855
0.075
3:941
0.087
3.73 0.79
0.978
0.003
WF
0.000
0.978
MEM (k-
8,810
3,563
0,122)
1.03
0.980
0.001
3,857 =: 0.059
3,938
0.37
0.981
MPM (k- 100)
3,800 11100-1000 0.008
3.563
0.010
0.12 2,05
0.979 1 0.001
3,882 1. 0.008
3,973
0.009
3,80 1. 0.25
0.979 0.000
MEM (K- 01
3.797 & 0.020
3,864 0,020
2.40 0.32
0.971 1. 0.001
3,894 1: 0.026
3.582 =: 01,037
3,87 0.53
0,976
0.002
MFM (k- 100
3.792 0.041
3,363
0.049
2.20 + 0.47
0.975
+ 0.005
3.955 + 0.085
4.066
= 0.11.
4.80 + 2.00
0.976
0.001
MFM (K- 500
3.796
0.014
5.864
0.018
2.50 0.22
0.972
1 0.001
4.055 1. 0.009
4.153
1 0.116
6.40 2.19
0.071 0.005
MFM (E = 100)
3.708 1 0.007
3.866
0.041
2.54 0.28
0.070 3 0.002
3,660 1: 0.004
3.963 = 0.113
3.40 1: 0.76
0.075 0.001
PIMIES
Train
We(1)
Wy(c)
MMD ovie (()
(11)
W:(()
315(3)
MMD :ww (-)
FM
4.133 0.062
4.521 0.176
8.21
1: 0.55
0.919 a 100014
5,300
4: 0.112
5.712
0.113
30.45
3,30
0.798 a: 0.002
6.778 0.046
7.015
0,042
$2.69
2.63
0.708 1 0.004
6,741
1. 0.177
8:985
1 0.163
$2.98
5.35
0.713 0.014
CGEM
4.81 1 0.065
4.511
0.078
6.1
1.33
0,806
1 0.001
5.674
1. 0.061
5.945
1 0.052
$7.15 0.30
0.790
0.009
OG
5.549 0-206
5:906
0.214
32.44 5.58
0,357
: 0.014
7.258
+ 0.107
7:501
0.212
80.45
5:00
0:713 0.005
ICNN
4,353
4.504
0,007
5.60
0,956
0.003
5.573
0.049
5.943
0.056
37.96
1,58
0.798
0.004
WEM
4.500
4.791
0.043
0.005
5.175
5.567
26.17 2:20
0.779 0.005
MEM (k-
4.028 0.033
4.761
0.045
12.31
0.853 0.003
5.005
23.34
0.786 0.005
WF 000
4.008 0.0266
4.5281
010000
11.15 0.34
0.894 = 0.000
5,088 =: 0.121
5.468 0.135
20.43 8.35
0.776 0.005
WFM 200)
0.042 4.508
4.781
(1,142
1.34 12
0,891 à 0,006
5,082 :: 0.007
5.461 =: 0,0147
20.16
0.50
0,792 0.015
MFM (k 0f
4.203 2: 0.015
4.451 HIJKS
8.34 0.54
0.325 :: 0.002
5.283 d: 0.017
5,558 =: 11,129
30.60 2: 0.77
U.797 0.005
MFM (L- 10)
4.273 0.050
4.474 0.041
8.35 0.83
0.915 + 0.000
5.244 + 168
5.606 + 0. 175
28.71 + 5.10
0.700 0.016
MFM (L- -500
4.271 8J838
4.466 0.010
8.40 0.40
0.921 + 0.001
5.165
+ 152
5.538
- 0. 1444
36.07 2,46
0.796 0.000
MIM 100)
4,200 0.079
4,351 0.095
7,20
0 1 0.004
4,950
0.199
5.321
0.219
5.10 20.06
0.796 0.013
23

---

Published as a conference paper at ICLR 2025
F.1 ANALYSIS OF POPULATION EMBEDDINGS
Here, we analyze the embedding space of the population embedding model (Pp; 0) for the synthetic letters dataset and the organoid drug-screen dataset. To do this, we compute embeddings = (((())))))) )===;))))))))))) 0) for each initial population. For letters we consider 200 random rotations per letter and for the organoid drug-screen dataset we consider all population pairs (Po, P1). We then project the embeddings into 2-dimensional space using uniform manifold approximation and projection (UMAP) to visualize model embeddings. We compute pairwise Euclidean distances of the samples in the projected embedding space and report these distances visually using heat maps. See Fig. 7.
On the letters dataset, the population embedding model learns to embed similar letter close silhouettes together. For example. ''' and 'V', "T" and 'Y' (even though 'Y' silhouettes are never seen during training), and more. Note, that in this dataset, letter silhouette populations po are both corrupted with noise and rotated randomly. We conduct the same analysis for the organoid drug-screen dataset. We observe that some sense ((Po: 8) groups populations from particular patients in generally consistent matter to that found in Ramos Zapatero et al. (2023), Specifically, we observe patient who are chemosensitive (PDO 21, 75, 23, 27) cluster together, away from patients who are chemorefractory (PDO 5, 11, 141, 216). This illustrates the embedding are able to capture patient drug response characteristics from untreated patient sample po-
Leter
Q LL
15
20-311
F40-2
-35
-27
-20
UMAP1
-23
paee 5
pae 102
ppo--11
pco-14t
22
Figure 7: Analysis of population embeddings from (Ppo; 0). We provide UMAP visualization of population embeddings space (top) and pairwise distances of population embeddings computed in 2-dimensional space (bottom). (Lefl) analysis of letter population embeddings plotted for 200 random rotations of each noisy letter silhouette. (Right) analysis of patient population embeddings plotted for all control populations.
F.2 ANALYSIS OF PREDICTION QUALITY ON ORGANOID DRUG-SCREER DATA
In this section, we analyze the predictions produced by vector field models trained with MFM. For this analysis, look model predictions for the three respective left-out test patient splits: PDO-21, PDO-27, PD0-75. For each control population in the patient split (specifically, for the test patient in the given split), predict the respective treatment response. We then subset the data look
24

---

Published as a conference paper at ICLR 2025
populations in the PDOF culture and take the mean of observations in each population (for both predicted and target populations). We project the means the predicted and target populations into 2D space using principal component analysis (PCA) across all. We fit PCA using means from source, target, and predicted samples. Lastly, we plot the target and predicted samples separately. We show the result of this analysis in Fig. 8.
We observe that for all three test patients the general structure is preserved, with the treatment, Oxaliplatin (Green), being the furthest away both for the target and ground truth datasets. This is because (as shown in the original paper by Ramos Zapatero al. (2023)) Oxaliplatin has large effect on these cancer cells for the PDOF subset of PDOs. We see from Fig. 8 that this more pronounced for PDO-21 and PDO-27 than for PDO-75. Overall, this reflects the conclusions drawn from Figure in the original dataset paper Ramos Zapatero et al. (2023). In this way, we are able to draw similar conclusions from the predicted populations as from the data.
pdo-21 Target
pdo-21 Prediction
Treatments
Oxaliplatin
SN:38
LGK-744
PCA2
PCA
S-FU
SN-38 S-FU
Cetux
-1
-1
S-FU Cetax
S-FU
-3
pdo-27 Target
pdo-27 Prediction
Treatments
0.2
Dxaliplatin
1.5
0.0
SN-38 VX-970
1.0
PCA2
-0.2
VX-970
-0.4
Celux
-0.6
S-FU
-0.5
SN-38 + Cetux
-0.8
SN-38 i S-FU Cotux
-1.0
PCA
PCA
pdo-75 Target
pdo-75 Prediction
Treatments
Oxaliplatin
2.5
SN-38
2.0
LGK-974
PCA
1.5
S-FU
0.5
1.0
0.0
0.5
SN-38 + S-FU
SN.38
-0.5
0.0
SIFU Cetux
SN-38 S-FU Cetux
-1.0
-0.5
-2
-1
1
PCA
Figure Analysis of treatment--ppecffic response prediction. We plot population means PCA space for target populations (Left), predicted populations (Middle), and the respective treatment identifier legend (Right).
25


Reviews:
[{"review_id": "cRhPU4lWzR", "reviewer": "Reviewer_jcAc", "summary": "This paper introduces Meta Flow Matching (MFM), a novel approach for modeling the evolution of systems consisting of interacting samples/populations. Unlike previous flow-based models, MFM can generalize to unseen initial populations by using a Graph Neural Network to embed the population and amortizing the flow matching model over these embeddings  The authors demonstrate MFM's effectiveness on an interesting synthetic letter denoising task and a large-scale single-cell organoid screening dataset, showing its ability to predict patient-specific responses to treatments.", "questions": "- In Eq.17, shouldn't the condition `c` be part of `\\phi` or `v_t` as well?\n- I am wondering in what cases I should *not* use MFM to model the change of a distribution using a flow? I can see that in Exp. 5.1 and 5.2, the samples have strong dependencies (e.g., to create the silhouettes), but in case there are no dependencies, would MFM also model this well?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "- Novel approach: The paper proposes a new method for integrating vector fields using the full of probability densities - allowing the modeling of the evolution of entire distributions of samples (which is important for many biological and physical processes). I also believe this amortization effect is especially helpful in small-sample regimes.\n- Theoretical foundation: The paper provides a solid theoretical basis for the proposed method, including the connections to existing approaches like conditional generative flow matching.\n- Readability: The paper is *very* well written and easy to read and understand (although it's covering a complex topic)", "weaknesses": "- Honestly, no weaknesses come to my eye. I've been working on Flows for a quite some time now, and have to say this is a good and well-executed idea.\n-  Figure 1 (c) is not that helpful in my opinion.\n\nTypos:\n- \"a a standard\" (l420)\n- \"of of model\" (l464)", "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": "39OQTPMb3F", "reviewer": "Reviewer_Kj6P", "summary": "This article deals with the problem of learning many-body dynamics via flow matching. The main objective is to model the interaction between components (e.g. particles) rather than modeling the entities separately. The resulting methodology allows to correctly model the evolution of distributions from the time evolution of observations of sample populations.", "questions": "All weaknesses as stated above are my main questions.", "limitations": "", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 1, "contribution": 3, "strengths": "-- The problem is known to have plagued science for a long time. Therefore, new approaches and solutions are important contributions.\n-- The idea of modeling the fields in the Wasserstein setting is intriguing, and intuitively meaningful.  \n-- The authors relate their approach with mean-field limits and diffusion processes. This is certainly something of interest.", "weaknesses": "-- First and foremost, the article is written in a very confusing way. It seems that the narrative is scattered, and this really makes it difficult to follow. \n-- The problem seems to boil down to embedding populations using graph neural networks and then flow matching in Wasserstein manifolds. \n-- It is unclear what a Wasserstein manifold is to start with. It clearly should not be a finite dimensional manifold, as depicted in their Figure 2. This is an infinite dimensional manifold, but the atlases are not clearly specified. For instance, a Banach manifold is a well understood object. In this case, the authors do not bother to introduce what they are talking about. P_2(X) is not generally even a Banach manifold. My understanding is that P_2(X) has a geometric structure only in some relatively special cases (Otto calculus), and usually people talk about metric over P_2(X), and the notion of curvature in Wasserstein spaces is often not the curvature of a manifold in the traditional sense of differential geometry. How is the tangent space defined in this context? In section 2.3 it seems that the author suggest that using the Wasserstein setup is rather important. As a consequence, I think it should not be left to the interpretation of the reader, but it should be clearly explained (at least defined). \n-- In general, there is a terrible lack of definitions. Even the theorem, which should be a formalization of the work, is unclear because the definitions are either absent, or scattered throughout the article. At least an appendix with all definitions should be added. \n-- The results do not seem overall particularly good. The table included in the text seems to show good results, but those in the appendix seem to show that the model is not particularly outperforming. One might not care much about this if the paper were a very well written piece of theoretical work, but as explained above it does not seem to be that case either.", "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": "tv8xz9dI0f", "reviewer": "Reviewer_Y1RT", "summary": "The paper introduces an extension to Flow Matching (FM) called Meta Flow Matching (MFM) that efficiently accommodates generalizing to varying base distributions by featurizing the base distribution with a nearest-neighbor graph neural network (GNN).\n\nMore specifically, FM is a recent popular framework for generative modeling that works by first matching up a target distribution with a base distribution (e.g., a standard normal distribution, coupled to the target distribution via an independent coupling) and then training a vector field to match the resulting particle vector fields via minimizing the Mean Squared Error (MSE). A natural extension of this is Conditional Flow Matching (CFM) where the vector fields can be conditioned on some covariate, such as a generative prompt for image generation. The authors argue that there are compelling reasons to extend the problem setup to include varying base distributions. First, naturally occurring dynamics such as interacting particles or diffusion equations have dynamics that are explicitly distribution dependent and thus their particle versions cannot be modeled by vector fields that only depend on the particle location alone. Second, problems in molecular biology naturally lend themselves to considering different base distributions, such as predicting the drug response for different baseline cell distributions per patient.\n\nTherefore, the authors suggest expanding the FM formulation with a distributional embedding that captures the base distribution and that is used as input to the vector field neural network. To obtain the embedding, they first calculate a k-nearest-neighbor graph on the base distribution and then embed it with a GNN, leading to the proposed MFM formulation. They show that MFM can be trained in the same manner as FM and show that it outperforms baseline architectures in experimental results on synthetic data (denoising rotated letter profiles) and on real data (mass-cytometry cell profiles on patient-derived cultures under drug treatments).", "questions": "I am recapitulating some of the problems outlined in \"Weaknesses\" above as questions to the authors here, as well as some other things that were not clear to me:\n\n* How are uncertainties on your results calculated? If they are not derived from independent splits/seeds, I would suggest you do that.\n* Why did you not benchmark MFM on the [dataset consider in the CellOT manuscript](https://onlinelibrary.wiley.com/doi/epdf/10.1111/exd.12683)?\n* How is CFM set up on the patient data set? I assume you condition on the drug, while ignoring the base distribution/patient identity. Is this correct, or are you also feeding the patient identity as condition? If the latter, I would recommend the former.\n* Can you add the following uninformed baselines to the experiments in Sections 5.1 and 5.2?\n    * The base distribution itself, i.e., distance of the unperturbed patient cells to the perturbed cells without applying any model.\n    * The base distribution shifted by a constant vector that is the mean over all perturbed cells (e.g., those in the train dataset)\n* Are there ways to better understand the overall quality of the predictions? Are there downstream conclusions of the paper that introduced that can or cannot be reproduced by the various predictive algorithms? Could you plot some of the predicted distributions just as they did in Figure 3?\n* I am surprised by the dynamic range of the Wasserstein distances vs MMD. For example, in Table 2, ICNN on the patient holdout achieves MMD of 74.00 and $\\mathcal{W}_2$ of 4.681 vs MFM (k=100) with 8.96 and 4.269, respectively. Is there any intuition about this discrepancy in dynamic range?\n* How is the $r^2$ calculated here? The model does not actually predict an output for a single target cell, but a distribution instead, so I am confused as to how this is done.\n* Would you consider running experiments with Optimal Transport matching instead of independent coupling as considered in [Tong et al (2023)](https://arxiv.org/abs/2302.00482)?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "* The manuscript is overall well-written and organized.\n* Flow matching is a popular generative modeling framework and is of significant interest to the ML community. I like the authors' motivation of the distribution dependence through the continuity equation and interacting particle systems. Their idea to solve this problem is simple and elegant and deserves exploration.", "weaknesses": "* **Empirical results are not fully convincing and need strengthening:** The presentation of the empirical results leave me slightly puzzled, as explained in the following.\n    * First of all, the authors consider [ICNN/CellOT](https://www.nature.com/articles/s41592-023-01969-x) as a baseline that can accommodate varying baseline distributions and conditioning. I think the introduction to ICNN presented in in Section 4 is slightly misleading: in lines 348-350, it is sometimes not clear what \"the method\" refers to. ICNN *can* generalize to new distributions, but the novelty in MFM is to take additional interactions into account.\n    * Since ICNN is conceptually the closest competitor to MFM, why did the authors not benchmark their method on the [dataset consider in that manuscript](https://onlinelibrary.wiley.com/doi/epdf/10.1111/exd.12683)? The chosen dataset in Section 5.2 seems well suited for the task, but it leaves me wondering why they swapped out datasets. Did their method not give them the desired results on previously considered datasets?\n    * The authors often claim that one method, in particular FM, \"fails to fit the train data\" (e.g., lines 472-473). This is a strong statement given that FM often surprisingly scores second-best out of the baseline methods. This makes me wonder about the underlying dataset and the employed metrics. At a minimum, I suggest reporting $\\mathcal{W}_1$, $\\mathcal{W}_2$ and MMD for the following very simple uninformed baselines to account for trivial underlying biological phenomena:\n        * The base distribution itself, i.e., distance of the unperturbed patient cells to the perturbed cells without applying any model.\n        * The base distribution shifted by a constant vector that is the mean over all perturbed cells (e.g., those in the train dataset)\n    * Further, are there ways to better understand the overall quality of the predictions? Are there downstream conclusions of the dataset paper that can or cannot be reproduced by the various predictive algorithms? Can the authors plot some of the predicted distributions just as they did in Figure 3?\n    * The generalizability of the results might be compromised by the authors' picking specific patients for the patient holdout setup (see Section C.2). Can the experiment be performed over several splits?\n    * In the same vein, I do not understand how the authors arrive at the error estimates for their experimental results if only one split is considered. I am afraid these estimates could be vastly underestimating the actual uncertainty. This is important since MFM sometimes only shows a small edge compared to FM.\n* **Independent matching may be supoptimal for all considered methods:** [Tong et al (2023)](https://arxiv.org/abs/2302.00482) propose to train flow matching by coupling base and target distribution via Optimal Transport instead of an independent coupling. This is a simple tweak that could be applied to all considered FM methods, i.e., FM, CFM, and MFM. I am curious if this would improve results across the board or specifically help FM and CFM because those methods currently have no way of accounting for different source distributions. Potentially, this could be a much simpler fix than MFM.\n* **(minor)** The text contains quite a few typos. I suggest the authors do another round of proof-reading for the final version.", "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": "rlYJXDSGwj", "reviewer": "Reviewer_kKjh", "summary": "- The paper proposes a new diagram of conditional flow matching models named meta flow matching.\n- The high-level understanding is: meta flow matching represents the conditional input in a new way -- by embedding the information of the whole population (or saying distributions).\n- In math, this can be formulated as constructing the vector field in the Wasserstein manifold.\n- The authors then provide the technical construction of the model, including objective functions and parametrization.\n- Numerical experiments are conducted on synthetic and organoid drug-screen datasets.\n- I am happy to see this conceptual novelty in the approach.\n- I didn't thoroughly check the technical part in the paper while I didn't see anything odd.\n- I am a little concerned on certain claims and the effectiveness of the method, both conceptually and numerically.", "questions": "See Strengths.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "I am listing both strengths (+), weakness (-) and questions (?) in the following to make the review more logistic.\n\n- (+) The meta flow matching diagram is clearly novel. The specific new part is the population/distribution embedding that the model can distinguish and evolve differently in different population, for the same particle.\n- (?) Question on ODE: Then what is the particle-level ODE corresponding to the PDE (12)? Is it similar to FP-equation that dx/dt = vt(x, pt)? Can the author prove it or give the detailed reference?\n- (?) Can you method be extent to SDE case?\n- (?) Analysis of population embeddings: To learn the population embeddings, I believe you should have enough example of populations otherwise the generalization is less promised. I would like to see a detailed analysis of population embeddings for a sanity check in both experiments.\n- (-) Reduced number of data points: Following the last point, for the organoid drug-screen experiment where only contains 10 patients, does it mean there are only 10 populations (data point for patient population embeddings)? It is out of expectation that model can capture the right population embeddings from 10 data points.\n- (?) Why replicate-split in the organoid drug-screen experiment is meaningful? In my understanding the conclusion drawn from different replicates usually would be consistent (that they has similar patterns).\n- (?) I also do not quite understand why other baselines under-perform in the replicate-split since if I understand correctly, generalization would not be a big challenge here.\n- (?) Can you provide a random baseline for your experiments? Random guess baseline and random initialized model baseline.\n- (-) An extremely important reference of \"Deep Generalized Schrödinger Bridge\" is missed here. Conceptually this paper should be sth ahead of yours and need to be discussed in details. Population embedding is also adopted in the paper in a simpler way.\n- (?) To further illustrate the effectviese, I think comparison with the numbers and experiments in \"Deep Generalized Schrödinger Bridge\" is necessary.\n- (?) Comparison with the number and experiment in \"Learning single-cell perturbation responses using neural optimal transport\" is also necessary in my eye. This is a classical paper and one of your competitor method.\n- (?) Complexity: Can you provide a complexity analysis here? Your method need additional embedding of the whole population and I think the scalability performance needed to be shown somehow (w.r.t. the size of population) which is very useful to the community for method selection.\n- (?) How do you construct the graph in the organoid drug-screen experiment? How different graph construction impact your model results?", "weaknesses": "See Strengths.", "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 (Poster)