Meta Flow Matching: Integrating Vector Fields on the Wasserstein Manifold

## Abstract

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 depends on the microenvironment of cells specific to each patient.
We propose Meta Flow Matching (MFM), a practical approach to integrating 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 prediction of individual treatment responses on a large scale multi-patient single-cell drug screen dataset.

## 1 Introduction

Understanding the dynamics of many-body problems is a focal challenge across the natural sciences. In the field of cell biology, a 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 trajectories (Armingol et al.,, ; Goodenough and Paul,, ). 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.,, ; Rizvi et al.,, ), in diseased states (Molè et al.,, ; Binnewies et al.,, ; Zeng and Dai,, ; Chung et al.,, ), and in response to perturbations (Ji et al.,, ; Peidli et al.,, ).

Figure: Figure 1: Comparison between *Flow Matching* (FM, Eq. 6), *Conditional Generative Flow Matching* (CGFM, Eq. 9), and *Meta Flow Matching* (MFM, Eq. 17). All methods learn a vector field $v_{t}(\cdot;\omega)$ (or *flow* model) parameterized via a neural network to predict treatment response for unseen populations. Models trained with FM only predict an aggregate response over populations. CGFM incorporates seen population IDs $i$, while MFM jointly learns $v_{t}(\cdot;\omega)$ and the populations embedding model $\varphi(\{x_{0}^{j}\}_{j=1}^{N^{\prime}};\theta)$, allowing generalization to new populations with unseen IDs.
Refer to caption: x1.png

The super-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.,, ). While single-cell sequencing has been used to 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 the inherit 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.,, ; Weinreb et al.,, ; Schiebinger et al.,, ; Tong et al.,, ; Neklyudov et al.,, ; ).
These approaches involve the formalisms of optimal transport (Villani,, ; Peyré and Cuturi,, ), diffusion (De Bortoli et al.,, ), or normalizing flows (Lipman et al.,, ), which allow for learning a 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 a number of different dynamics that are available in the training data.

We propose Meta Flow Matching (MFM) — the amortization of the Flow Matching generative modeling framework (Lipman et al.,, ) 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 a 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.,, ; Peidli et al.,, ).
We evaluate MFM on predicting the response of patient-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.,, ).(^1^11This dataset includes more than $25$ million cells collected for ten patients over $2500$ different experimental/treatment conditions.)
This is a challenging task due to the inherent variation that exist across patients, the variability induced by different treatments, and variation due to 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-specific response to the treatment.

## 2 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.,, ), Flow Matching (Lipman et al.,, ), Stochastic Interpolants (Albergo and Vanden-Eijnden,, ).
A flow is a continuous interpolation between densities
$p_{0}(x_{0})$ and $p_{1}(x_{1})$ in the sample space.
That is, the sample from the intermediate density $p_{t}(x)$ is produced as follows

$$ $\displaystyle x_{t}=~{}$ $\displaystyle f_{t}(x_{0},x_{1}),\;\;(x_{0},x_{1})\sim\pi(x_{0},x_{1})\,,$ (1) $\displaystyle\text{ where }\int dx_{1}\;\pi(x_{0},x_{1})=p_{0}(x_{0})\,,\;\; \int dx_{0}\;\pi(x_{0},x_{1})=p_{1}(x_{1})\,,$ (2) $$

where $f_{t}$ is the time-continuous interpolating function such that $f_{t=0}(x_{0},x_{1})=x_{0}$ and $f_{t=1}(x_{0},x_{1})=x_{1}$ (e.g. linearly between $x_{0}$ and $x_{1}$ with $f_{t}(x_{0},x_{1})=(1-t)\cdot x_{0}+t\cdot x_{1}$); $\pi(x_{0},x_{1})$ is the density of the joint distribution, which is usually taken as a distribution of independent random variables $\pi(x_{0},x_{1})=p_{0}(x_{0})p_{1}(x_{1})$, but can also be generalized to formulate the optimal transport problems (Pooladian et al.,, ; ).
The corresponding density can be defined then as the following expectation over the interpolating samples

$$ $\displaystyle p_{t}(x)=\int dx_{0}dx_{1}\;\pi(x_{0},x_{1})\delta(x-f_{t}(x_{0} ,x_{1}))\,.$ (3) $$

The essential part of 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 $v^{*}_{t}(x)$ as a solution

$$ $\displaystyle\frac{\partial p_{t}(x)}{\partial t}=-\mathopen{}\mathclose{{} \left\langle\nabla_{x},p_{t}(x)v^{*}_{t}(x)}\right\rangle\,,\;\;v^{*}_{t}(\xi) =\frac{1}{p_{t}(\xi)}\mathbb{{E}}_{\pi(x_{0},x_{1})}\mathopen{}\mathclose{{} \left[\delta(f_{t}(x_{0},x_{1})-\xi)\frac{\partial f_{t}(x_{0},x_{1})}{ \partial t}}\right]\,.$ (4) $$

The expectation for $v^{*}_{t}$ is intractable, but we can learn a model $v_{t}(x;\omega)$ to approximate $v^{*}_{t}$ so that we can efficiently generate new samples.
Relying on Eq. 4, one can derive the tractable objective for learning $v^{*}_{t}(x)$, i.e.

$$ $\displaystyle\mathcal{L}_{\rm FM}(\omega)=~{}$ $\displaystyle\int_{0}^{1}dt\;\mathbb{E}_{p_{t}(x)}\mathopen{}\mathclose{{} \left\lVert v^{*}_{t}(x)-v_{t}(x;\omega)}\right\rVert^{2}$ (5) $\displaystyle=~{}$ $\displaystyle\mathbb{E}_{\pi(x_{0},x_{1})}\int_{0}^{1}dt\;\mathopen{} \mathclose{{}\left\lVert\frac{\partial}{\partial t}f_{t}(x_{0},x_{1})-v_{t}(f_ {t}(x_{0},x_{1});\omega)}\right\rVert^{2}+\text{constant}\,.$ (6) $$

Finally, the vector field $v_{t}(\xi,\omega)\approx v_{t}^{*}(\xi)$ defines the push-forward density that approximately matches $p_{t=1}$, i.e. $T_{\#}p_{0}\approx p_{t=1}$, where $T$ is the flow corresponding to vector field $v_{t}(\cdot,\omega)$ with parameters $\omega$.

### 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 (; Rombach et al.,, ) as well as conditioning on other images ().
To learn the conditional generative process with diffusion models, one merely has to pass the conditional variable (sampled jointly with the data point) as 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 is independently introduced in several works
(Zheng et al.,, ; Dao et al.,, ; Isobe et al.,, ) and it operates as follows.
Consider a family of time-continuous densities $p_{t}(x_{t}\,|\,c)$, which corresponds to the distribution of the following random variable

$$ $\displaystyle x_{t}=f_{t}(x_{0},x_{1}),\;\;(x_{0},x_{1})\sim\pi(x_{0},x_{1}\,| \,c)\,.$ (7) $$

For every $c$, the density $p_{t}(x_{t}\,|\,c)$ follows the continuity equation with the following vector field

$$ $\displaystyle v^{*}_{t}(\xi\,|\,c)=\frac{1}{p_{t}(\xi\,|\,c)}\mathbb{{E}}_{\pi (x_{0},x_{1})}\delta(f_{t}(x_{0},x_{1})-\xi)\frac{\partial f_{t}(x_{0},x_{1})} {\partial t}\,,$ (8) $$

which depends on $c$.
Thus, the training objective of the conditional model becomes

$$ $\displaystyle\mathcal{L}_{CGFM}(\omega)=~{}$ $\displaystyle\mathbb{E}_{p(c)}\mathbb{E}_{\pi(x_{0},x_{1}\,|\,c)}\int_{0}^{1} dt\;\mathopen{}\mathclose{{}\left\lVert\frac{\partial}{\partial t}f_{t}(x_{0}, x_{1})-v_{t}(f_{t}(x_{0},x_{1})\,|\,c;\omega)}\right\rVert^{2}\,,$ (9) $$

where, compared to the original Flow Matching formulation, we first have to sample $c$, then produce the samples from $p_{t}(x_{t}\,|\,c)$ and pass $c$ as input to the parametric model of the vector field.

### 2.3 Modeling Process in Natural Sciences as Vector Fields on the Wasserstein Manifold

Figure: Figure 2: Illustration of flow matching methods on the 2-Wasserstein manifold, $\mathcal{P}_{2}(\mathcal{X})$, depicted as a two-dimensional sphere. *Flow Matching* learns the tangent vectors to a single curve on the manifold. *Conditional* generation corresponds to learning a finite set of curves on the manifold, e.g. classes $c_{1}$ and $c_{2}$ on the plot. *Meta Flow Matching* learns to integrate a vector field on $\mathcal{P}_{2}(\mathcal{X})$, i.e. for every starting density $p_{0}$ Meta Flow Matching defines a push-forward measure that integrates along the underlying vector field.
Refer to caption: x2.png

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 a Graph Neural Network (GNN).

To provide the reader with the necessary intuition, we are going to use the geometric formalism developed by Otto, ().
That is, time-dependent densities $p_{t}(x_{t})$ define absolutely-continuous curves on the 2-Wasserstein space of distributions $\mathcal{P}_{2}(\mathcal{X})$ (Ambrosio et al.,, ).
The tangent space of this manifold is defined by the gradient flows $\mathcal{S}_{t}=\mathopen{}\mathclose{{}\left\{\nabla s_{t}\,\middle|\,s_{t}:
\mathcal{X}\to\mathbb{R}}\right\}$ on the state space $\mathcal{X}$.
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 as a gradient flow (Neklyudov et al.,, ).

Under the geometric formalism of the 2-Wasserstein manifold, Flow Matching can be considered as learning the tangent vectors $v_{t}(\cdot)$ along the density curve $p_{t}(x_{t})$ defined by the sampling process in Eq. 2 (see the left panel in Fig. 2).
Furthermore, the conditional generation processes $p_{t}(x_{t}\,|\,c)$ would be represented as a finite set of curves if $c$ is discrete (e.g. class-conditional generation of images) or as a family of curves if $c$ is continuous (see the middle panel in Fig. 2).

Finally, one can define a vector field on the 2-Wasserstein manifold via the continuity equation with the vector field $v_{t}(x,p_{t}(x))$ on the state space $\mathcal{X}$ that depends on the current density $p_{t}(x)$ or its derivatives.
Below we give two examples of processes defined as vector fields on the 2-Wasserstein manifold.

###### Example 1 (Mean-field limit of interacting particles) .

Consider a system of interacting particles, where the velocity of the particle at point $x$ interacting with the particle at point $y$ is defined as $k(x,y):\mathbb{R}^{d}\times\mathbb{R}^{d}\to\mathbb{R}^{d}$.
In the limit of the infinite number of particles one can describe their state using the density function $p_{t}(x)$.
Then the change of the density is described by the following continuity equation

$$ $\displaystyle\frac{dx}{dt}=\mathbb{E}_{p_{t}(y)}k(x,y),\;\;\frac{\partial p_{t }(x)}{\partial t}=-\mathopen{}\mathclose{{}\left\langle\nabla_{x},p_{t}(x) \mathbb{E}_{p_{t}(y)}k(x,y)}\right\rangle\,,$ (10) $$

which is the first-order analog of the Vlasov equation (Jabin and Wang,, ).

###### Example 2 (Diffusion) .

Even when the physical particles evolve independently in nature, the deterministic vector field model might be dependent on 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 in the density-dependent vector field when written as a continuity equation, i.e.

$$ $\displaystyle\frac{\partial p_{t}(x)}{\partial t}=\frac{1}{2}\Delta_{x}p_{t}(x )=-\mathopen{}\mathclose{{}\left\langle\nabla_{x},p_{t}(x)\mathopen{} \mathclose{{}\left(-\frac{1}{2}\nabla_{x}\log p_{t}(x)}\right)}\right\rangle \implies\frac{dx}{dt}=-\frac{1}{2}\nabla_{x}\log p_{t}(x)\,.$ (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

$$ $\displaystyle\frac{\partial p_{t}(x)}{\partial t}=-\mathopen{}\mathclose{{} \left\langle\nabla_{x},p_{t}(x)v_{t}\mathopen{}\mathclose{{}\left(x,p_{t}} \right)}\right\rangle\,,\;\;p_{t=0}(x)=p_{0}(x)\,,$ (12) $$

where the state-space vector field $v_{t}\mathopen{}\mathclose{{}\left(x,p_{t}}\right)$ depends on the density $p_{t}$.

The dependency might vary across models, e.g. in Example 1 the vector field can be modeled as an application of a kernel to the density function, while in Example 2 the vector field depends only on the local value of the density and its derivative.

###### Example 1 (Mean-field limit of interacting particles) .

###### Example 2 (Diffusion) .

## 3 Meta Flow Matching

In this paper, we propose the amortization of the Flow Matching framework over the marginal distributions.
Our model is based on the outstanding ability of the Flow Matching framework to learn the push-forward map for any joint distribution $\pi(x_{0},x_{1})$ given empirically.
For the given joint $\pi(x_{0},x_{1})$, we denote the solution of the Flow Matching optimization problem as follows

$$ $\displaystyle v_{t}^{*}(\cdot,\pi)=\operatorname*{argmin}_{v_{t}}\mathcal{L}_{ GFM}(v_{t}(\cdot),\pi(x_{0},x_{1}))\,.$ (13) $$

Analogous to amortized optimization (Chen et al.,, ; Amos et al.,, ), we aim to learn the model that outputs the solution of Eq. 13 based on the input data sampled from $\pi$, i.e.

$$ $\displaystyle v_{t}(\cdot,\varphi(\pi))=v_{t}^{*}(\cdot,\pi)\,,$ (14) $$

where $\varphi(\pi)$ is the embedding model of $\pi$ and the joint density $\pi(\cdot\,|\,c)$ is generated using some unknown measure of the conditional variables $c\sim p(c)$.

### 3.1 Integrating Vector Fields on the Wasserstein Manifold via Meta Flow Matching

Consider the dataset of joint populations $\mathcal{D}=\{(\pi(x_{0},x_{1}\,|\,i))\}_{i}$, where, to simplify the notation, we associate every $i$-th population with its density $\pi(\cdot\,|\,i)$ 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 $p_{0}(x_{0}\,|\,i)=\int dx_{1}\;\pi(x_{0},x_{1}\,|\,i)$ are sampled from some unknown distribution that can be parameterized with a large enough number of parameters (ii) the endpoint marginals $p_{1}(x_{1}\,|\,i)=\int dx_{0}\;\pi(x_{0},x_{1}\,|\,i)$ 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 $\mathcal{D}$ using the conditional version of the Flow Matching algorithm (see Section 2.2) where the population index $i$ plays the role of the conditional variable.
However, obviously, such a model will not generalize beyond the considered data $\mathcal{D}$ and unseen indices $i$.
We illustrate this empirically in Section 5.

To be able to generalize to previously unseen populations, we propose learning the density-dependent vector field motivated by Eq. 12.
That is, we propose to use an embedding function $\varphi:\mathcal{P}_{2}(\mathcal{X})\to\mathbb{R}^{m}$ to embed the starting marginal density $p_{0}$, which we then input into the vector field model and minimize the following objective over $\omega$

$$ $\displaystyle\mathcal{L}_{\text{MFM}}(\omega;\varphi)=~{}$ $\displaystyle\mathbb{E}_{i\sim\mathcal{D}}\mathbb{E}_{\pi(x_{0},x_{1}\,|\,i)} \int_{0}^{1}dt\;\mathopen{}\mathclose{{}\left\lVert\frac{\partial}{\partial t} f_{t}(x_{0},x_{1})-v_{t}(f_{t}(x_{0},x_{1})\,|\,\varphi(p_{0});\omega)}\right \rVert^{2}\,.$ (15) $$

Note that the initial density $p_{0}$ is enough to predict the push-forward density $p_{1}$ since the Cauchy problem for Eq. 12 has a unique solution.
The embedding function $\varphi(p_{0})$ can take different forms, e.g. it can be the density value $\varphi(p_{0})=p_{0}(\cdot)$, which is then used inside the vector field model to evaluate at the current point (analogous to Example 2); a kernel density estimator (analogous to Example 1); or a parametric model taking the samples from this density as an input.

###### Proposition 1 .

Meta Flow Matching recovers the Conditional Generation via Flow Matching when the conditional dependence of the marginals $p_{0}(x_{0}\,|\,c)=\int dx_{1}\pi(x_{0},x_{1}\,|\,c)$ and $p_{1}(x_{1}\,|\,c)=\int dx_{0}\pi(x_{0},x_{1}\,|\,c)$ and the distribution $p(c)$ are known, i.e.
there exist $\varphi:\mathcal{P}_{2}(\mathcal{X})\to\mathbb{R}^{m}$ such that $\mathcal{L}_{MFM}(\omega)=\mathcal{L}_{CGFM}(\omega)$.

###### Proof.

Indeed, sampling from the dataset $i\sim\mathcal{D}$ becomes sampling of the conditional variable $c\sim p(c)$ and the embedding function becomes $\varphi(p_{0}(\cdot\,|\,c))=c$.
∎

Furthermore, for the parametric family of the embedding models $\varphi(p_{t},\theta)$, we show that the parameters $\theta$ can be estimated by minimizing the objective in Eq. 15 in the joint optimization with the vector field parameters $\omega$.
We formalize this statement in the following theorem.

###### Theorem 1 .

Consider a dataset of populations $\mathcal{D}=\{(\pi(x_{0},x_{1}\,|\,i))\}_{i}$ generated from some unknown conditional model $\pi(x_{0},x_{1}\,|\,c)p(c)$.
Then the following objective

$$ $\displaystyle\mathcal{L}(\omega,\theta)=~{}$ $\displaystyle\mathbb{E}_{p(c)}\int_{0}^{1}dt\;\mathbb{E}_{p_{t}(x_{t}\,|\,c)} \mathopen{}\mathclose{{}\left\lVert v^{*}_{t}(x_{t}\,|\,c)-v_{t}(x_{t}\,|\, \varphi(p_{0},\theta),\omega)}\right\rVert^{2}$ (16) $$

is equivalent to the Meta Flow Matching objective

$$ $\displaystyle\mathcal{L}_{\text{MFM}}(\omega,\theta)=~{}$ $\displaystyle\mathbb{E}_{i\sim\mathcal{D}}\mathbb{E}_{\pi(x_{0},x_{1}\,|\,i)} \int_{0}^{1}dt\;\mathopen{}\mathclose{{}\left\lVert\frac{\partial}{\partial t} f_{t}(x_{0},x_{1})-v_{t}(f_{t}(x_{0},x_{1})\,|\,\varphi(p_{0},\theta);\omega)} \right\rVert^{2}$ (17) $$

up to an additive constant.

###### Proof.

We postpone the proof to Appendix A.
∎

###### Proposition 1 .

###### Proof.

###### Theorem 1 .

###### Proof.

### 3.2 Learning Population Embeddings via Graph Neural Networks (GNNs)

Figure: Algorithm 1 Meta Flow Matching (training)

Figure: Algorithm 2 Meta Flow Matching (sampling)

In many applications, the populations $\mathcal{D}=\{(\pi(x_{0},x_{1}\,|\,i))\}_{i=1}^{N}$ are given as empirical distributions, i.e. they are represented as samples from some unknown density $\pi$

$$ $\displaystyle\{(x_{0}^{j},x_{1}^{j})\}_{j=1}^{N_{i}}\,,\;\;(x_{0}^{j},x_{1}^{j })\sim\pi(x_{0},x_{1}\,|\,i)\,,$ (18) $$

where $N_{i}$ is the size of the $i$-th population.
For instance, for the diffusion process considered in Example 2, the samples from $\pi(x_{0},x_{1}\,|\,i)$ can be generated by generating some marginal $p_{1}(x_{1}\,|\,i)$ and then adding the Gaussian random variable to the samples $x_{1}^{j}$.
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 a parametric model $\varphi(p_{0},\theta)$, i.e.

$$ $\displaystyle\varphi(p_{0},\theta)=\varphi\mathopen{}\mathclose{{}\left(\{x_{0 }^{j}\}_{j=1}^{N_{i}},\theta}\right)\,,\;\;(x_{0}^{j},x_{1}^{j})\sim\pi(x_{0}, x_{1}\,|\,i)\,.$ (19) $$

For this purpose, we employ GNNs, which recently have been successfully applied for simulation of complicated many-body problems in physics (Sanchez-Gonzalez et al.,, ).
To embed a population $\{x_{0}^{j}\}_{j=1}^{N_{i}}$, we create a k-nearest neighbour graph $G_{i}$ based on the metric in the state-space $\mathcal{X}$, input it into a GNN, which consists of several message-passing iterations (Gilmer et al.,, ) and the final average-pooling across nodes to produce the embedding vector.
Finally, we update the parameters of 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 1 and sampling in Algorithm 2.

## 4 Related Work

#### Meta – amortizing learning over distributions

The meta-learning of probability measures was previously studied by Amos et al., () where they demonstrate 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 push-forward map without embedding the second marginal and without restricting ourselves to Input-Convex Neural Network (ICNN) by Amos et al., () and optimal transport maps. MFM is *meta* in the same sense of amortizing optimal transport (or in our case flow) problems over multiple input distributions. Hence, we follow the naming convention of *Meta Optimal Transport* (Amos et al.,, ). This is different but related to how *meta* is used in the meta learning setting, which amortizes over learning problems (Hospedales et al.,, ; Achille et al.,, ).

#### Generative modeling for single cells

Single cell data has expanded to encompass multiple modalities of data profiling cell state and activities (Frangieh et al.,, ; ). 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.,, ; Lopez et al.,, ; Hetzel et al.,, ). Orthogonal non-negative matrix factorization (oNMF) has also been used for dimensionality reduction combined with mixture models for cell state prediction (Chen et al.,, ). Other approaches have tried to use Flow Matching (FM)  (; ; Neklyudov et al.,, ) or similar approaches such as the Monge gap (Uscidda and Cuturi,, ) to predict cell trajectories. Currently, the state of the art method uses the principle of Optimal Transport (OT) to predict cell trajectories with ICNN (Makkuva et al.,, ; ). What determines the significance of the method is its capability in generalizing out of distribution to a new population of cells, which may be from different culture or individuals. As of this time, our method is the only method that takes inter-cellular interactions into account.

#### Generative modeling for physical processes

The closest approach to ours is the prediction of the many-body interactions in physics (Sanchez-Gonzalez et al.,, ) 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.
Neklyudov et al., () 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 to unseen data.
Campbell et al., () 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.,, ; Verma et al.,, ) 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 task population prediction, we consider two experimental settings. (i) A synthetic experiment with well defined coupled populations, and (ii) experiments on a publicly available single-cell dataset consisting of populations from patient dependent treatment response trials.
To quantify model performance, we consider three distributional distances metrics: the 1-Wasserstein distance (${\mathcal{W}}_{1}$), 2-Wasserstein (${\mathcal{W}}_{2}$) distance, and the radial basis kernel maximum-mean-discrepancy (MMD) distance (Gretton et al.,, ).
We parameterize all vector field models $v_{t}(\cdot\,|\,\varphi(p_{0});\omega)$ using a Multi-Layer Perceptron (MLP).
For MFM, we additionally parameterize $\varphi(p_{t};\theta,k)$ using a Graph Convolutional Network (GCN) with a $k$-nearest neighbour graph edge pooling layer.
We include details regarding model hyperparameters, training/optimization, and implementation in Appendix B and Section B.2.
The results for all the models are averaged over three random seeds.

### 5.1 Synthetic Experiment

Figure: Figure 3: Synthetic letters experiment visualizations. Examples of model-generated samples for synthetic letters from the source distribution ($t=0$) to predicted target distribution ($t=1$). See Fig. 5 in Appendix D for a larger set of examples.
Refer to caption: x5.png

**Table 1: Results of the synthetic letters experiment for population prediction on seen train populations and unseen test populations. We use ${}^{\text{w/}}\mathcal{N}$ to denote models trained with source distributions sampled via $x_{0}\sim\mathcal{N}(\mathbf{0},\mathbf{1})$ (All MFM models use noisy letters source distributions denoted by $p_{0}$ as input to the population embedding model $\varphi_{\theta}(p_{0})$). We report the 1-Wasserstein (${\mathcal{W}}_{1}$), 2-Wasserstein (${\mathcal{W}}_{2}$), and the maximum-mean-discrepancy (MMD) distributional distances. We consider 4 settings for MFM with varying $k$. Letter populations that comprise test set (X and Y) are entirely unseen during training.**
|  | Train | Test (X’s) | Test (Y’s) |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | ${\mathcal{W}}_{1}$ | ${\mathcal{W}}_{2}$ | MMD ($\times 10^{-3}$) | ${\mathcal{W}}_{1}$ | ${\mathcal{W}}_{2}$ | MMD ($\times 10^{-3}$) | ${\mathcal{W}}_{1}$ | ${\mathcal{W}}_{2}$ | MMD ($\times 10^{-3}$) |
| FM | $0.209\pm 0.000$ | $0.277\pm 0.000$ | $2.54\pm 0.00$ | $0.234\pm 0.000$ | $0.309\pm 0.000$ | $2.45\pm 0.00$ | $0.238\pm 0.000$ | $0.316\pm 0.000$ | $3.32\pm 0.01$ |
| FM${}^{\text{w/}}\mathcal{N}$ | $0.806\pm 0.000$ | $0.960\pm 0.000$ | $31.68\pm 0.00$ | $0.764\pm 0.000$ | $0.931\pm 0.000$ | $25.04\pm 0.00$ | $1.030\pm 0.000$ | $1.228\pm 0.000$ | $45.36\pm 0.00$ |
| CGFM | $0.090\pm 0.000$ | $0.113\pm 0.000$ | $0.25\pm 0.00$ | $0.334\pm 0.000$ | $0.407\pm 0.000$ | $5.55\pm 0.00$ | $0.327\pm 0.000$ | $0.405\pm 0.000$ | $6.85\pm 0.00$ |
| CGFM${}^{\text{w/}}\mathcal{N}$ | $0.156\pm 0.025$ | $0.201\pm 0.027$ | $1.02\pm 0.39$ | $0.849\pm 0.004$ | $0.993\pm 0.003$ | $35.08\pm 0.75$ | $1.062\pm 0.011$ | $1.229\pm 0.010$ | $55.66\pm 0.76$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=0$) | $0.148\pm 0.003$ | $0.195\pm 0.010$ | $0.94\pm 0.11$ | $0.347\pm 0.011$ | $0.431\pm 0.012$ | $6.47\pm 0.44$ | $0.402\pm 0.011$ | $0.485\pm 0.010$ | $10.92\pm 0.18$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=1$) | $0.154\pm 0.004$ | $0.208\pm 0.010$ | $0.91\pm 0.01$ | $0.349\pm 0.023$ | $0.433\pm 0.023$ | $6.53\pm 0.52$ | $0.391\pm 0.035$ | $0.477\pm 0.041$ | $10.71\pm 1.86$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=10$) | $0.151\pm 0.013$ | $0.197\pm 0.015$ | $0.94\pm 0.15$ | $0.343\pm 0.020$ | $0.427\pm 0.019$ | $6.38\pm 0.67$ | $0.413\pm 0.018$ | $0.502\pm 0.024$ | $11.93\pm 1.14$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=50$) | $0.174\pm 0.005$ | $0.232\pm 0.006$ | $1.40\pm 0.13$ | $0.363\pm 0.010$ | $0.449\pm 0.013$ | $7.46\pm 0.44$ | $0.446\pm 0.021$ | $0.536\pm 0.028$ | $13.40\pm 0.23$ |
| MFM ($k=0$) | $\textbf{0.081}\pm\textbf{0.003}$ | $\textbf{0.100}\pm\textbf{0.004}$ | $\textbf{0.16}\pm\textbf{0.06}$ | $0.202\pm 0.002$ | $0.249\pm 0.003$ | $2.29\pm 0.05$ | $0.218\pm 0.001$ | $0.262\pm 0.002$ | $3.79\pm 0.11$ |
| MFM ($k=1$) | $0.082\pm 0.001$ | $0.101\pm 0.002$ | $\textbf{0.16}\pm\textbf{0.01}$ | $0.205\pm 0.008$ | $0.251\pm 0.008$ | $2.38\pm 0.22$ | $0.215\pm 0.006$ | $0.258\pm 0.007$ | $3.78\pm 0.25$ |
| MFM ($k=10$) | $0.088\pm 0.002$ | $0.109\pm 0.003$ | $0.21\pm 0.01$ | $\textbf{0.201}\pm\textbf{0.006}$ | $\textbf{0.248}\pm\textbf{0.006}$ | $2.20\pm 0.15$ | $0.208\pm 0.003$ | $0.252\pm 0.002$ | $3.55\pm 0.06$ |
| MFM ($k=50$) | $0.092\pm 0.004$ | $0.116\pm 0.004$ | $0.25\pm 0.06$ | $0.206\pm 0.008$ | $0.257\pm 0.008$ | $\textbf{2.18}\pm\textbf{0.25}$ | $\textbf{0.204}\pm\textbf{0.005}$ | $\textbf{0.249}\pm\textbf{0.006}$ | $\textbf{3.14}\pm\textbf{0.18}$ |

We curate a synthetic dataset of the joint distributions $\{(p_{0}(x_{0},\,|\,i),p_{1}(x_{1}\,|\,i))\}_{i=1}^{N}$ by simulating a diffusion process applied to a set of pre-defined target distributions $p_{1}(x_{1}\,|\,i)$ for $i=1,\dots,N$.
To get a paired population $p_{0}(x_{0}\,|\,i)$ we simulate the forward diffusion process without drift $x_{0}\sim\mathcal{N}(x_{1},\sigma)$.
After this setup, for reasonable values of $\sigma$, we assume that one can reverse the diffusion process and learn the push-forward map from $p_{0}(x_{0}\,|\,i)$ to $p_{1}(x_{1}\,|\,i)$ for every index $i$.
For this task, given the $i$-th population index we denote $p_{0}(x_{0}\,|\,i)$ as the *source* population $p_{1}(x_{1}\,|\,i)$ as the $i$-th *target* population.

To construct $p_{1}(x_{1}\,|\,i)$, we discretize samples from a defined silhouette; e.g. an image of a character, where $i$ indexes the respective character.
We use upper case letters as the silhouette and generate the corresponding samples $x_{1}\sim p_{1}(x_{1}\,|\,i)$ from the uniform distribution over the silhouette and run the diffusion process for samples $x_{1}$ to acquire $x_{0}$.
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 $k$ for the GCN population embedding model $\varphi(p_{0};\theta,k)$. When $k=0$, $\varphi(p_{t};\theta,k)$ becomes identical to the DeepSets model (Zaheer et al.,, ).
We compare MFM to FM and CGFM.
We repeat this experiment for models trained with source distributions sampled from a a standard normal $x_{0}\sim\mathcal{N}(\mathbf{0},\mathbf{1})$. We label these models as FM${}^{\text{w/}}\mathcal{N}$, CGFM${}^{\text{w/}}\mathcal{N}$, and MFM${}^{\text{w/}}\mathcal{N}$. Here, MFM${}^{\text{w/}}\mathcal{N}$ still takes the original $p_{0}$ as input to the population embedding model $\varphi(p_{0};\theta,k)$, while $v_{t}(\cdot;\omega)$ uses $x_{0}\sim\mathcal{N}(\mathbf{0},\mathbf{1})$ as the source distribution.

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 generalize to the test conditions.
For the training data, 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 to the unseen condition in the test dataset.
CGFM indicates when the model can fit the training data and demonstrates that the test data is substantially different and cannot be generated by the same model.
For MFM, we expect to 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 fails to adequately learn to sample from $p_{1}(x_{1}\,|\,i)$ in the training set, while CGFM is able to effectively sample from $p_{1}(x_{1}\,|\,i)$ in the training set.
As expected, CGFM fits the training data, but fails 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.
FM fails to fit the train data and fails to generalize under the test conditions.
Interestingly, although MFM performs better for certain values of $k$ versus others, overall performance does not vary significantly for the range considered.

### 5.2 Experiments on Organoid Drug-screen Data

Figure: Figure 4: Organoid drug-screen dataset overview. *Left*: a given replica consists of a control distribution $p_{0}$ and corresponding treatment response distribution $p_{1}$ for treatment condition $c_{i}$. *Right*: train and test data splits for replica (top) and patients (bottom) splits, restively. For each experiment there are 11 treatments, 10 patients and 3 culture conditions.
Refer to caption: x8.png

#### Data

For experiments on biological data, we use the organoid drug-screen dataset from Ramos Zapatero et al., (). 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 experimental condition. Populations from each patient are treated with 11 different drug treatments of varying dose concentrations.(^2^22We consider only the highest dosage and leave exploration of dose-dependent response to future work.)
We use the term *replicate* to define control-treatment population pairs, $p_{0}(x_{0}\,|\,c_{i})$ and $p_{1}(x_{1}\,|\,c_{i})$, respectively (see Fig. 4-left). In each patient, cell population 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 averaged over Fibroblast, PDO, and PDOF cultures and results for the individual cultures (this is reported in Appendix D).

#### Pre-processing and data splits

We filter each cell population to contain at least 1000 cells and consider 43 bio-markers.
We consider two data splits for the organoid drug-screen dataset (see Fig. 4-right). (1) *Replicate split*; here we leave-out replicates evenly across all patients for testing. (2) *Patients split*; here we leave-out replicates fully in one patients – in this setting, we are testing the ability of of model to generalize population prediction of treatment response for unseen patients.
Further details regarding data pre-processing and data splits are provided in Section B.2.

For the organoid drug-screen experiments, we consider an ICNN architecture in as an additional baseline alongside FM and CGFM. The ICNN model is based on CellOT (); a method for learning cell specific response to treatments. The ICNN (and likewise CellOT) counterparts our FM model in that it does not take the population index $i$ as a condition. Therefore, it will neither be able to fit the training data, nor generalize.

#### Predicting treatment response across replicates

We show results for generalization across replicates in Table 2. As expected, we observe that CGFM fits the training data, but does not generalize to the test replicates. With this, we can observe that the FM and ICNN models fail to fit the train data, relative to CGFM, and also fail to generalize. In general, MFM ($k=0$) performs best on generalization to unseen replicates. We include results reported for the separate cell cultures in Table 4 in Appendix D.

#### Predicting treatment response across patients

We show results for generalization across patients in Table 3. Similar to the replicates data setting, we observe that CGFM fits the training data, but does not generalize to the test replicates. Likewise, the FM and ICNN models fail to fit the train data, relative to CGFM, and also fail to generalize. In general, MFM ($k=100$) performs best on generalization to unseen patients. We include results reported for the separate cell cultures in Table 5 in Appendix D.

We observe that no clear single selection for $k$ yields the best performance across all tasks on the single-cell experiments. We observe that for the *replicates* split, $k=0$ on average performs better than $k>0$. Whereas on the *patients* split, the opposite is true, $k=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$.

Through the biological and synthetic experiments, we have shown that MFM is able to generalize to unseen distributions/populations. The implication of 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 learn 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.

**Table 2: Experimental results on the organoid drug-screen dataset for population prediction of treatment response across *replicate* populations. Results shown in this table are for $\text{{dim}}=\textbf{43}$. We use ${}^{\text{w/}}\mathcal{N}$ to denote models trained with source distributions sampled via $x_{0}\sim\mathcal{N}(\mathbf{0},\mathbf{1})$. We report results results for the 1-Wasserstein (${\mathcal{W}}_{1}$), 2-Wasserstein (${\mathcal{W}}_{2}$), the maximum-mean-discrepancy (MMD), and mean-feature squared correlation coefficient ($r^{2}$) distributional metrics. We consider $4$ settings for MFM with varying nearest-neighbours parameter ($k=0,10,50,100$). For extended results, see Table 4.**
|  | Train | Test |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | ${\mathcal{W}}_{1}(\downarrow)$ | ${\mathcal{W}}_{2}(\downarrow)$ | MMD ($\times 10^{-3}$) $(\downarrow)$ | $r^{2}(\uparrow)$ | ${\mathcal{W}}_{1}(\downarrow)$ | ${\mathcal{W}}_{2}(\downarrow)$ | MMD ($\times 10^{-3}$) $(\downarrow)$ | $r^{2}(\uparrow)$ |
| FM | $3.925\pm 0.019$ | $4.041\pm 0.023$ | $3.76\pm 0.26$ | $0.952\pm 0.007$ | $3.961\pm 0.036$ | $4.089\pm 0.042$ | $5.90\pm 0.25$ | $0.941\pm 0.010$ |
| FM${}^{\text{w/}}\mathcal{N}$ | $6.908\pm 0.037$ | $7.181\pm 0.033$ | $57.70\pm 0.75$ | $0.639\pm 0.005$ | $6.972\pm 0.022$ | $7.244\pm 0.022$ | $60.39\pm 0.98$ | $0.642\pm 0.007$ |
| CGFM | $\textbf{3.864}\pm\textbf{0.064}$ | $3.975\pm 0.069$ | $\textbf{3.16}\pm\textbf{0.89}$ | $0.964\pm 0.006$ | $4.087\pm 0.063$ | $4.211\pm 0.066$ | $8.84\pm 0.75$ | $0.938\pm 0.006$ |
| CGFM${}^{\text{w/}}\mathcal{N}$ | $4.187\pm 0.008$ | $4.340\pm 0.009$ | $8.69\pm 0.50$ | $0.936\pm 0.002$ | $6.852\pm 0.045$ | $7.114\pm 0.044$ | $71.24\pm 3.71$ | $0.666\pm 0.016$ |
| ICNN | $4.286\pm 0.018$ | $4.313\pm 0.112$ | $38.6\pm 0.212$ | $0.897\pm 0.031$ | $4.194\pm 0.110$ | $4.313\pm 0.112$ | $37.9\pm 2.84$ | $0.897\pm 0.008$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=0$) | $3.940\pm 0.022$ | $4.047\pm 0.023$ | $3.91\pm 0.18$ | $0.959\pm 0.006$ | $3.896\pm 0.026$ | $4.002\pm 0.030$ | $\textbf{4.35}\pm\textbf{0.18}$ | $0.950\pm 0.005$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=10$) | $3.976\pm 0.044$ | $4.086\pm 0.049$ | $4.52\pm 0.42$ | $0.961\pm 0.002$ | $3.943\pm 0.032$ | $4.051\pm 0.034$ | $5.28\pm 0.25$ | $0.952\pm 0.001$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=50$) | $3.968\pm 0.013$ | $4.075\pm 0.014$ | $4.36\pm 0.44$ | $0.961\pm 0.002$ | $3.934\pm 0.007$ | $4.041\pm 0.008$ | $4.99\pm 0.35$ | $0.954\pm 0.000$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=100$) | $3.937\pm 0.014$ | $4.040\pm 0.015$ | $3.94\pm 0.00$ | $0.963\pm 0.001$ | $3.908\pm 0.030$ | $4.011\pm 0.033$ | $4.68\pm 0.52$ | $0.953\pm 0.002$ |
| MFM ($k=0$) | $3.874\pm 0.015$ | $\textbf{3.973}\pm\textbf{0.020}$ | $3.37\pm 0.14$ | $\textbf{0.967}\pm\textbf{0.003}$ | $\textbf{3.880}\pm\textbf{0.009}$ | $\textbf{3.990}\pm\textbf{0.011}$ | $4.68\pm 0.16$ | $0.955\pm 0.002$ |
| MFM ($k=10$) | $3.896\pm 0.021$ | $4.000\pm 0.021$ | $3.82\pm 0.12$ | $0.964\pm 0.001$ | $3.899\pm 0.013$ | $4.012\pm 0.011$ | $5.13\pm 0.48$ | $0.955\pm 0.001$ |
| MFM ($k=50$) | $3.888\pm 0.038$ | $3.991\pm 0.030$ | $3.59\pm 0.41$ | $0.963\pm 0.001$ | $3.900\pm 0.038$ | $4.013\pm 0.034$ | $5.06\pm 0.22$ | $0.954\pm 0.003$ |
| MFM ($k=100$) | $3.906\pm 0.010$ | $4.008\pm 0.005$ | $4.05\pm 0.38$ | $0.964\pm 0.002$ | $3.898\pm 0.008$ | $4.009\pm 0.009$ | $5.19\pm 0.05$ | $\textbf{0.957}\pm\textbf{0.000}$ |

**Table 3: Experimental results on the organoid drug-screen dataset for population prediction of treatment response across *patient* populations. Results shown in this table are for $\text{{dim}}=\textbf{43}$. For extended results, see Table 5.**
|  | Train | Test |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | ${\mathcal{W}}_{1}(\downarrow)$ | ${\mathcal{W}}_{2}(\downarrow)$ | MMD ($\times 10^{-3}$) $(\downarrow)$ | $r^{2}(\uparrow)$ | ${\mathcal{W}}_{1}(\downarrow)$ | ${\mathcal{W}}_{2}(\downarrow)$ | MMD ($\times 10^{-3}$) $(\downarrow)$ | $r^{2}(\uparrow)$ |
| FM | $3.985\pm 0.054$ | $4.115\pm 0.067$ | $4.64\pm 0.43$ | $0.938\pm 0.014$ | $4.340\pm 0.078$ | $4.564\pm 0.111$ | $13.00\pm 0.67$ | $0.865\pm 0.034$ |
| FM${}^{\text{w/}}\mathcal{N}$ | $6.892\pm 0.027$ | $7.164\pm 0.033$ | $57.03\pm 1.00$ | $0.655\pm 0.003$ | $7.114\pm 0.100$ | $7.404\pm 0.086$ | $64.97\pm 3.79$ | $0.613\pm 0.008$ |
| CGFM | $\textbf{3.882}\pm\textbf{0.019}$ | $3.999\pm 0.020$ | $\textbf{3.16}\pm\textbf{0.59}$ | $0.952\pm 0.004$ | $4.443\pm 0.033$ | $4.621\pm 0.041$ | $17.00\pm 1.03$ | $0.899\pm 0.008$ |
| CGFM${}^{\text{w/}}\mathcal{N}$ | $4.313\pm 0.077$ | $4.480\pm 0.081$ | $11.51\pm 1.96$ | $0.918\pm 0.004$ | $7.135\pm 0.045$ | $7.390\pm 0.037$ | $79.78\pm 4.67$ | $0.637\pm 0.010$ |
| ICNN | $4.289\pm 0.020$ | $4.382\pm 0.021$ | $37.0\pm 2.84$ | $0.913\pm 0.003$ | $4.525\pm 0.051$ | $4.681\pm 0.054$ | $74.00\pm 0.57$ | $0.862\pm 0.127$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=0$) | $3.982\pm 0.014$ | $4.095\pm 0.015$ | $5.04\pm 0.36$ | $0.951\pm 0.002$ | $4.177\pm 0.042$ | $4.355\pm 0.048$ | $10.53\pm 0.59$ | $0.911\pm 0.001$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=10$) | $4.006\pm 0.008$ | $4.119\pm 0.012$ | $5.13\pm 0.30$ | $0.948\pm 0.001$ | $4.156\pm 0.065$ | $4.324\pm 0.067$ | $9.58\pm 1.63$ | $0.912\pm 0.003$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=50$) | $3.982\pm 0.018$ | $4.095\pm 0.016$ | $4.74\pm 0.21$ | $0.951\pm 0.002$ | $4.153\pm 0.069$ | $4.324\pm 0.070$ | $9.63\pm 1.45$ | $0.912\pm 0.002$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=100$) | $4.004\pm 0.012$ | $4.119\pm 0.014$ | $5.19\pm 0.43$ | $0.949\pm 0.002$ | $4.166\pm 0.001$ | $4.341\pm 0.003$ | $9.52\pm 0.33$ | $0.915\pm 0.005$ |
| MFM ($k=0$) | $3.905\pm 0.005$ | $4.012\pm 0.006$ | $4.18\pm 0.25$ | $\textbf{0.958}\pm\textbf{0.001}$ | $4.209\pm 0.007$ | $4.380\pm 0.012$ | $12.34\pm 0.50$ | $\textbf{0.918}\pm\textbf{0.002}$ |
| MFM ($k=10$) | $3.896\pm 0.033$ | $4.005\pm 0.036$ | $3.89\pm 0.44$ | $0.957\pm 0.005$ | $4.216\pm 0.090$ | $4.395\pm 0.098$ | $11.99\pm 2.36$ | $0.917\pm 0.005$ |
| MFM ($k=50$) | $3.902\pm 0.018$ | $4.008\pm 0.022$ | $4.20\pm 0.17$ | $\textbf{0.958}\pm\textbf{0.000}$ | $4.214\pm 0.017$ | $4.396\pm 0.020$ | $12.09\pm 0.75$ | $0.916\pm 0.002$ |
| MFM ($k=100$) | $3.884\pm 0.039$ | $\textbf{3.986}\pm\textbf{0.044}$ | $3.77\pm 0.49$ | $0.955\pm 0.001$ | $\textbf{4.100}\pm\textbf{0.093}$ | $\textbf{4.269}\pm\textbf{0.104}$ | $\textbf{8.96}\pm\textbf{1.88}$ | $0.917\pm 0.004$ |

## 6 Conclusion

Our paper highlights the significance of modeling dynamics based on the entire distribution.
While flow-based models offer a 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.

#### Limitations

In this work we explored empirically the effect of conditioning the learned flow on the initial distribution. 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.,, , ), unbalanced optimal transport (Benamou,, ; Yang and Uhler,, ; Chizat et al.,, ), aligned (Somnath et al.,, ; Liu et al.,, ), or stochastic settings (; Koshizuka and Sato,, ) in this work.

#### Future Work

In this work, we focused on exploring the amortization of learning over the space of distributions using flow matching. However, 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 $\omega$ and $\theta$ for the vector field $v_{t}(\cdot;\omega)$ and population embedding model $\varphi(\cdot;\theta)$ using the loss defined in Eq. 17. Although this is a novel approach to $v_{t}(\cdot|\varphi(\cdot;\theta);\omega)$ training, there remains room to explore the pre-training and/or use of different losses for the population embedding model.

## 7 Broader Impacts

This paper is primarily a 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 as a significant risk at this time.

## Acknowledgments

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://vectorinstitute.ai/partners/](http://vectorinstitute.ai/partners/)), the computational resources provided by the Digital Research Alliance of Canada ([https://alliancecan.ca](https://alliancecan.ca)), Mila ([https://mila.quebec](https://mila.quebec)), and NVIDIA.
KN is supported by IVADO.

## Appendix A Proof of Theorem 1

See 1

###### Proof.

The loss function

$$ $\displaystyle\mathcal{L}(\omega,\theta)=~{}$ $\displaystyle\mathbb{E}_{p(c)}\int_{0}^{1}dt\;\mathbb{E}_{p_{t}(x_{t}\,|\,c)} \mathopen{}\mathclose{{}\left\lVert v^{*}_{t}(x_{t}\,|\,c)-v_{t}(x_{t}\,|\, \varphi(p_{t},\theta);\omega)}\right\rVert^{2}$ (20) $\displaystyle=~{}$ $\displaystyle-2\mathbb{E}_{p(c)}\int dtdx\;\mathopen{}\mathclose{{}\left \langle p_{t}(x\,|\,c)v^{*}_{t}(x\,|\,c),v_{t}(x\,|\,\varphi(p_{t},\theta); \omega)}\right\rangle+$ (21) $\displaystyle+\mathbb{E}_{p(c)}\int_{0}^{1}dt\;\mathbb{E}_{p_{t}(x_{t}\,|\,c)} \mathopen{}\mathclose{{}\left\lVert v_{t}(x_{t}\,|\,\varphi(p_{t},\theta), \omega)}\right\rVert^{2}+$ (22) $\displaystyle+\mathbb{E}_{p(c)}\int_{0}^{1}dt\;\mathbb{E}_{p_{t}(x_{t}\,|\,c)} \mathopen{}\mathclose{{}\left\lVert v_{t}^{*}(x_{t}\,|\,c)}\right\rVert^{2}\,.$ (23) $$

The last term does not depend on $\theta$, the second term we can estimate, for the first term, we use the formula for the (from Eq. 8)

$$ $\displaystyle p_{t}(\xi\,|\,c)v^{*}_{t}(\xi\,|\,c)=\mathbb{E}_{\pi(x_{0},x_{1} )}\delta(f_{t}(x_{0},x_{1})-\xi)\frac{\partial f_{t}(x_{0},x_{1})}{\partial t}\,.$ (24) $$

Thus, the loss is equivalent (up to a constant) to

$$ $\displaystyle\mathcal{L}(\omega,\theta)=~{}$ $\displaystyle-2\mathbb{E}_{p(c)}\mathbb{E}_{\pi(x_{0},x_{1}\,|\,c)}\int dt\; \mathopen{}\mathclose{{}\left\langle\frac{\partial f_{t}(x_{0},x_{1})}{ \partial t},v_{t}(f_{t}(x_{0},x_{1})\,|\,\varphi(p_{t},\theta);\omega)}\right\rangle+$ (25) $\displaystyle+\mathbb{E}_{p(c)}\mathbb{E}_{\pi(x_{0},x_{1}\,|\,c)}\int_{0}^{1} dt\;\mathopen{}\mathclose{{}\left\lVert v_{t}(f_{t}(x_{0},x_{1})\,|\,\varphi(p _{t},\theta),\omega)}\right\rVert^{2}\pm$ (26) $\displaystyle\pm\mathbb{E}_{p(c)}\mathbb{E}_{\pi(x_{0},x_{1}\,|\,c)}\int_{0}^{ 1}dt\;\mathopen{}\mathclose{{}\left\lVert\frac{\partial f_{t}(x_{0},x_{1})}{ \partial t}}\right\rVert^{2}$ (27) $\displaystyle=~{}$ $\displaystyle\mathbb{E}_{c\sim p(c)}\mathbb{E}_{\pi(x_{0},x_{1}\,|\,c)}\int_{0 }^{1}dt\;\mathopen{}\mathclose{{}\left\lVert\frac{\partial}{\partial t}f_{t}(x _{0},x_{1})-v_{t}(f_{t}(x_{0},x_{1})\,|\,\varphi(p_{t},\theta);\omega)}\right \rVert^{2}\,.$ (28) $$

Note that in the final expression we do not need access to the probabilistic model of $p(c)$ if the joints $\pi(x_{0},x_{1}\,|\,c)$ are already sampled in the data $\mathcal{D}$.
Thus, we have

$$ $\displaystyle\mathcal{L}(\omega,\theta)=~{}$ $\displaystyle\mathbb{E}_{c\sim p(c)}\mathbb{E}_{\pi(x_{0},x_{1}\,|\,c)}\int_{0 }^{1}dt\;\mathopen{}\mathclose{{}\left\lVert\frac{\partial}{\partial t}f_{t}(x _{0},x_{1})-v_{t}(f_{t}(x_{0},x_{1})\,|\,\varphi(p_{t},\theta);\omega)}\right \rVert^{2}$ (29) $\displaystyle=~{}$ $\displaystyle\mathbb{E}_{i\sim\mathcal{D}}\mathbb{E}_{\pi(x_{0},x_{1}\,|\,i)} \int_{0}^{1}dt\;\mathopen{}\mathclose{{}\left\lVert\frac{\partial}{\partial t} f_{t}(x_{0},x_{1})-v_{t}(f_{t}(x_{0},x_{1})\,|\,\varphi(p_{t},\theta);\omega)} \right\rVert^{2}$ (30) $\displaystyle=~{}$ $\displaystyle\mathcal{L}_{\text{MFM}}(\omega,\theta)\,.$ (31) $$

∎

###### Proof.

## Appendix B Experimental Details

### B.1 Synthetic letters data

The synthetic letters dataset contains 242 train populations a 10 test populations. Each population contains roughly between 750 and 2700 samples in this dataset.

### B.2 Organoid drug-screen data

The *organoid drug-screen* dataset consists of patient derived organoids (PDOs) from 10 different patients (Ramos Zapatero et al.,, ).(^3^33PDOs are cultures of cell populations which are derived from patient cells.)
After pre-processing, we use a total of 927 replicates (or coupled populations). In the replicates split, we use 713 populations for training and 103 left-out populations for testing. In the patients split, we use 861 populations for training and 33 left-out populations for testing. Organoid PDO-21 was found to present an interesting chemoprotective response (High CAF Protection) to cancer treatments, and hence was selected by Ramos Zapatero et al., () for further study and data collection. For the same reason we select PDO-21 as the unseen test patient for the patients split experiments. We also use patient 27 (PDO-27) as a validation set to select reasonable hyperparameters for MFM and the baseline models as it is representative of a different but similarly sized patient group (Low CAF Protection).

### B.3 Model architectures and hyperparameters

#### ICNN

The ICNN baseline was constructed with two networks ICNN network $f(x)$ and $g(x)$, with non-negative leaky ReLU activation layers. $f(x)$ is used to minimize the transport distance and $g(x)$ is 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 ($lr=10^{-4}$, $\beta_{1}=0.5$, $\beta_{2}=0.9$). $g(x)$ is trained with an inner iteration of $10$ for every iteration $f(x)$ is trained.

#### Vector Field Models

All vector field models $v_{t}$ are parameterized with linear layers of 512 hidden units and SELU activation functions.
For the synthetic experiments, we use 4 hidden layers. For the biological experiments, we use 7 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 a 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 of $128$ worked sufficiently well.

#### Graph Neural Network

We considered a GCN model that consists of a $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 2 GCN layers to perform well. The final GCN model layer outputs an embedding representation $e\in\mathbb{R}^{d}$. For 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 a hyper-sphere, and find that this normalization helps improve training. Additionally, the GCN takes a one-hot cell-type encoding (encoding for Fibroblast cells or PDO cells) for the control populations $p_{0}$. 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 information/performance gain.

#### 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 a 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 $\omega$ and the GCN model parameters $\theta$. We alternate between updating the respective model parameters every gradient step. For the syntehtic experiment, FM and CGFM model were trained for 2000 epochs, while MFM models were trained for 4000 epochs. Due to the alternating 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.

The hyperparameters stated in this section were selected from brief and small grid search sweeps. We did not conduct any thorough hyperparameter optimization.

## Appendix C Implementation Details

We implement all our experiments using PyTorch and PyTorch Geometric. All experiments were conducted on a HPC cluster primarily on NVIDIA Tesla T4 16GB and A40 48GB GPUs. Each individual seed experiment run required only 1 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.

## Appendix D Extended Results

Figure: Figure 5: Model-generated samples for synthetic letters from the source ($t=0$) to target ($t=1$) distributions.
Refer to caption: x9.png

**Table 4: Extended experimental results on the organoid drug-screen dataset for population prediction of treatment response across *replicate* populations. Results shown in this table are for $\text{{dim}}=\textbf{43}$. Here we show results for the individual cell cultures: Fibroblasts, patient derived organoids (PDOs), and patient derived organoids with Fibroblasts (PDOFs), respectively. We use ${}^{\text{w/}}\mathcal{N}$ to denote models trained with source distributions sampled via $x_{0}\sim\mathcal{N}(\mathbf{0},\mathbf{1})$. We report results results for the 1-Wasserstein (${\mathcal{W}}_{1}$), 2-Wasserstein (${\mathcal{W}}_{2}$), the maximum-mean-discrepancy (MMD), and mean-feature squared correlation coefficient ($r^{2}$) distributional metrics. We consider $4$ settings for MFM with varying nearest-neighbours parameter ($k=0,10,50,100$).**
|  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Fibroblasts |  |  |  |  |  |  |  |  |
|  | Train | Test |  |  |  |  |  |  |
|  | ${\mathcal{W}}_{1}(\downarrow)$ | ${\mathcal{W}}_{2}(\downarrow)$ | MMD ($\times 10^{-3}$) $(\downarrow)$ | $r^{2}(\uparrow)$ | ${\mathcal{W}}_{1}(\downarrow)$ | ${\mathcal{W}}_{2}(\downarrow)$ | MMD ($\times 10^{-3}$) $(\downarrow)$ | $r^{2}(\uparrow)$ |
| FM | $3.748\pm 0.030$ | $3.838\pm 0.047$ | $2.17\pm 0.19$ | $0.937\pm 0.020$ | $3.730\pm 0.018$ | $3.816\pm 0.031$ | $2.29\pm 0.11$ | $0.948\pm 0.014$ |
| FM${}^{\text{w/}}\mathcal{N}$ | $6.701\pm 0.093$ | $7.045\pm 0.094$ | $51.55\pm 3.16$ | $0.621\pm 0.007$ | $6.548\pm 0.115$ | $6.885\pm 0.113$ | $48.95\pm 3.69$ | $0.647\pm 0.007$ |
| CGFM | $3.620\pm 0.004$ | $3.674\pm 0.004$ | $0.86\pm 0.07$ | $0.979\pm 0.001$ | $3.732\pm 0.014$ | $3.797\pm 0.015$ | $3.26\pm 0.07$ | $0.959\pm 0.004$ |
| CGFM${}^{\text{w/}}\mathcal{N}$ | $3.671\pm 0.011$ | $3.725\pm 0.011$ | $0.96\pm 0.07$ | $0.980\pm 0.001$ | $5.844\pm 0.192$ | $6.148\pm 0.204$ | $39.80\pm 3.07$ | $0.695\pm 0.007$ |
| ICNN | $4.057\pm 0.067$ | $4.127\pm 0.070$ | $10.84\pm 6.70$ | $0.889\pm 0.008$ | $3.828\pm 0.094$ | $3.893\pm 0.097$ | $8.06\pm 6.38$ | $0.917\pm 0.009$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=0$) | $3.745\pm 0.034$ | $3.808\pm 0.045$ | $1.73\pm 0.27$ | $0.968\pm 0.010$ | $3.730\pm 0.036$ | $3.792\pm 0.045$ | $2.02\pm 0.29$ | $0.967\pm 0.011$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=10$) | $3.718\pm 0.023$ | $3.774\pm 0.023$ | $1.48\pm 0.22$ | $0.976\pm 0.002$ | $3.719\pm 0.023$ | $3.777\pm 0.023$ | $2.15\pm 0.25$ | $0.973\pm 0.000$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=50$) | $3.716\pm 0.027$ | $3.771\pm 0.027$ | $1.55\pm 0.32$ | $0.976\pm 0.001$ | $3.717\pm 0.032$ | $3.773\pm 0.031$ | $2.13\pm 0.45$ | $0.976\pm 0.001$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=100$) | $3.712\pm 0.013$ | $3.767\pm 0.012$ | $1.37\pm 0.17$ | $0.977\pm 0.002$ | $3.708\pm 0.011$ | $3.764\pm 0.011$ | $1.81\pm 0.29$ | $0.975\pm 0.002$ |
| MFM ($k=0$) | $3.667\pm 0.027$ | $3.722\pm 0.028$ | $1.43\pm 0.18$ | $0.975\pm 0.002$ | $3.661\pm 0.027$ | $3.719\pm 0.027$ | $1.74\pm 0.08$ | $0.976\pm 0.001$ |
| MFM ($k=10$) | $3.679\pm 0.032$ | $3.734\pm 0.032$ | $1.69\pm 0.34$ | $0.974\pm 0.001$ | $3.680\pm 0.030$ | $3.738\pm 0.030$ | $2.04\pm 0.28$ | $0.975\pm 0.001$ |
| MFM ($k=50$) | $3.664\pm 0.042$ | $3.720\pm 0.043$ | $1.46\pm 0.25$ | $0.974\pm 0.002$ | $3.664\pm 0.039$ | $3.722\pm 0.041$ | $1.84\pm 0.23$ | $0.975\pm 0.001$ |
| MFM ($k=100$) | $3.676\pm 0.005$ | $3.732\pm 0.006$ | $1.56\pm 0.27$ | $0.972\pm 0.003$ | $3.674\pm 0.001$ | $3.731\pm 0.002$ | $2.05\pm 0.20$ | $0.976\pm 0.001$ |
|  |  |  |  |  |  |  |  |  |
| PDOs |  |  |  |  |  |  |  |  |
|  | Train | Test |  |  |  |  |  |  |
|  | ${\mathcal{W}}_{1}(\downarrow)$ | ${\mathcal{W}}_{2}(\downarrow)$ | MMD ($\times 10^{-3}$) $(\downarrow)$ | $r^{2}(\uparrow)$ | ${\mathcal{W}}_{1}(\downarrow)$ | ${\mathcal{W}}_{2}(\downarrow)$ | MMD ($\times 10^{-3}$) $(\downarrow)$ | $r^{2}(\uparrow)$ |
| FM | $3.910\pm 0.037$ | $3.998\pm 0.042$ | $3.37\pm 0.25$ | $0.964\pm 0.009$ | $3.825\pm 0.062$ | $3.937\pm 0.072$ | $3.31\pm 0.34$ | $0.965\pm 0.012$ |
| FM${}^{\text{w/}}\mathcal{N}$ | $7.334\pm 0.101$ | $7.578\pm 0.103$ | $69.40\pm 3.18$ | $0.585\pm 0.009$ | $7.461\pm 0.091$ | $7.692\pm 0.094$ | $72.49\pm 3.39$ | $0.592\pm 0.009$ |
| CGFM | $3.792\pm 0.061$ | $3.866\pm 0.063$ | $2.27\pm 0.59$ | $0.979\pm 0.003$ | $4.062\pm 0.096$ | $4.181\pm 0.103$ | $7.75\pm 1.45$ | $0.950\pm 0.011$ |
| CGFM${}^{\text{w/}}\mathcal{N}$ | $3.746\pm 0.032$ | $3.807\pm 0.037$ | $1.19\pm 0.29$ | $0.986\pm 0.002$ | $7.672\pm 0.149$ | $7.909\pm 0.144$ | $94.96\pm 7.33$ | $0.602\pm 0.022$ |
| ICNN | $4.532\pm 0.057$ | $4.634\pm 0.059$ | $41.67\pm 32.45$ | $0.874\pm 0.009$ | $4.225\pm 0.110$ | $4.336\pm 0.114$ | $18.09\pm 15.51$ | $0.896\pm 0.010$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=0$) | $3.788\pm 0.047$ | $3.851\pm 0.053$ | $1.59\pm 0.20$ | $0.982\pm 0.002$ | $3.741\pm 0.042$ | $3.829\pm 0.047$ | $2.39\pm 0.23$ | $0.978\pm 0.002$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=10$) | $3.822\pm 0.063$ | $3.891\pm 0.071$ | $1.76\pm 0.54$ | $0.983\pm 0.001$ | $3.785\pm 0.028$ | $3.875\pm 0.029$ | $2.82\pm 0.49$ | $0.979\pm 0.002$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=50$) | $3.794\pm 0.033$ | $3.858\pm 0.034$ | $1.71\pm 0.50$ | $0.985\pm 0.001$ | $3.775\pm 0.022$ | $3.868\pm 0.020$ | $2.92\pm 0.41$ | $0.980\pm 0.002$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=100$) | $3.783\pm 0.018$ | $3.845\pm 0.018$ | $1.53\pm 0.19$ | $0.984\pm 0.002$ | $3.749\pm 0.027$ | $3.835\pm 0.028$ | $2.59\pm 0.24$ | $0.981\pm 0.002$ |
| MFM ($k=0$) | $3.843\pm 0.056$ | $3.916\pm 0.066$ | $2.65\pm 0.81$ | $0.971\pm 0.002$ | $3.794\pm 0.065$ | $3.891\pm 0.074$ | $3.34\pm 1.03$ | $0.974\pm 0.002$ |
| MFM ($k=10$) | $3.852\pm 0.039$ | $3.932\pm 0.045$ | $2.80\pm 0.19$ | $0.972\pm 0.003$ | $3.781\pm 0.013$ | $3.879\pm 0.015$ | $3.45\pm 0.52$ | $0.978\pm 0.002$ |
| MFM ($k=50$) | $3.844\pm 0.036$ | $3.924\pm 0.033$ | $2.51\pm 0.01$ | $0.973\pm 0.006$ | $3.791\pm 0.029$ | $3.889\pm 0.026$ | $3.33\pm 0.14$ | $0.974\pm 0.006$ |
| MFM ($k=100$) | $3.905\pm 0.082$ | $3.988\pm 0.088$ | $3.87\pm 1.81$ | $0.974\pm 0.001$ | $3.783\pm 0.025$ | $3.882\pm 0.029$ | $3.33\pm 0.45$ | $0.976\pm 0.003$ |
|  |  |  |  |  |  |  |  |  |
| PDOFs |  |  |  |  |  |  |  |  |
|  | Train | Test |  |  |  |  |  |  |
|  | ${\mathcal{W}}_{1}(\downarrow)$ | ${\mathcal{W}}_{2}(\downarrow)$ | MMD ($\times 10^{-3}$) $(\downarrow)$ | $r^{2}(\uparrow)$ | ${\mathcal{W}}_{1}(\downarrow)$ | ${\mathcal{W}}_{2}(\downarrow)$ | MMD ($\times 10^{-3}$) $(\downarrow)$ | $r^{2}(\uparrow)$ |
| FM | $4.117\pm 0.006$ | $4.287\pm 0.007$ | $5.75\pm 0.63$ | $0.953\pm 0.007$ | $4.328\pm 0.057$ | $4.514\pm 0.064$ | $12.11\pm 0.76$ | $0.911\pm 0.017$ |
| FM${}^{\text{w/}}\mathcal{N}$ | $6.686\pm 0.093$ | $6.921\pm 0.095$ | $52.12\pm 2.39$ | $0.722\pm 0.006$ | $6.906\pm 0.073$ | $7.154\pm 0.073$ | $59.73\pm 3.16$ | $0.686\pm 0.009$ |
| CGFM | $4.180\pm 0.130$ | $4.385\pm 0.145$ | $6.35\pm 2.13$ | $0.933\pm 0.020$ | $4.467\pm 0.080$ | $4.653\pm 0.083$ | $15.51\pm 0.93$ | $0.906\pm 0.004$ |
| CGFM${}^{\text{w/}}\mathcal{N}$ | $5.150\pm 0.046$ | $5.499\pm 0.049$ | $23.79\pm 1.21$ | $0.843\pm 0.005$ | $7.039\pm 0.147$ | $7.285\pm 0.153$ | $78.96\pm 6.83$ | $0.700\pm 0.020$ |
| ICNN | $4.419\pm 0.075$ | $4.564\pm 0.078$ | $19.04\pm 14.16$ | $0.936\pm 0.003$ | $4.607\pm 0.093$ | $4.788\pm 0.094$ | $46.62\pm 51.98$ | $0.881\pm 0.004$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=0$) | $4.287\pm 0.019$ | $4.482\pm 0.015$ | $8.41\pm 0.93$ | $0.928\pm 0.011$ | $4.218\pm 0.016$ | $4.385\pm 0.016$ | $8.62\pm 0.24$ | $0.905\pm 0.002$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=10$) | $4.372\pm 0.045$ | $4.582\pm 0.048$ | $9.70\pm 1.33$ | $0.923\pm 0.007$ | $4.326\pm 0.052$ | $4.501\pm 0.057$ | $10.86\pm 0.60$ | $0.904\pm 0.004$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=50$) | $4.356\pm 0.036$ | $4.566\pm 0.040$ | $9.72\pm 0.68$ | $0.922\pm 0.006$ | $4.291\pm 0.014$ | $4.469\pm 0.018$ | $10.14\pm 0.60$ | $0.905\pm 0.001$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=100$) | $4.295\pm 0.026$ | $4.495\pm 0.033$ | $8.55\pm 0.56$ | $0.932\pm 0.007$ | $4.251\pm 0.071$ | $4.423\pm 0.080$ | $9.44\pm 1.51$ | $0.905\pm 0.003$ |
| MFM ($k=0$) | $4.111\pm 0.016$ | $4.281\pm 0.023$ | $6.04\pm 0.26$ | $0.954\pm 0.009$ | $4.184\pm 0.022$ | $4.359\pm 0.022$ | $8.96\pm 0.64$ | $0.915\pm 0.003$ |
| MFM ($k=10$) | $4.158\pm 0.028$ | $4.332\pm 0.029$ | $6.97\pm 0.80$ | $0.946\pm 0.006$ | $4.235\pm 0.019$ | $4.418\pm 0.022$ | $9.89\pm 0.81$ | $0.914\pm 0.002$ |
| MFM ($k=50$) | $4.155\pm 0.035$ | $4.330\pm 0.015$ | $6.80\pm 0.98$ | $0.943\pm 0.002$ | $4.245\pm 0.044$ | $4.429\pm 0.035$ | $10.00\pm 0.28$ | $0.912\pm 0.001$ |
| MFM ($k=100$) | $4.137\pm 0.048$ | $4.305\pm 0.068$ | $6.73\pm 0.42$ | $0.945\pm 0.002$ | $4.236\pm 0.001$ | $4.412\pm 0.004$ | $10.18\pm 0.10$ | $0.918\pm 0.005$ |

**Table 5: Extended experimental results on the organoid drug-screen dataset for population prediction of treatment response across *patients* populations. Results shown in this table are for $\text{{dim}}=\textbf{43}$. Here we show results for the individual cell cultures: Fibroblasts, patient derived organoids (PDOs), and patient derived organoids with Fibroblasts (PDOFs), respectively. We use ${}^{\text{w/}}\mathcal{N}$ to denote models trained with source distributions sampled via $x_{0}\sim\mathcal{N}(\mathbf{0},\mathbf{1})$. We report results results for the 1-Wasserstein (${\mathcal{W}}_{1}$), 2-Wasserstein (${\mathcal{W}}_{2}$), the maximum-mean-discrepancy (MMD), and mean-feature squared correlation coefficient ($r^{2}$) distributional metrics. We consider $4$ settings for MFM with varying nearest-neighbours parameter ($k=0,10,50,100$).**
|  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Fibroblasts |  |  |  |  |  |  |  |  |
|  | Train | Test |  |  |  |  |  |  |
|  | ${\mathcal{W}}_{1}(\downarrow)$ | ${\mathcal{W}}_{2}(\downarrow)$ | MMD ($\times 10^{-3}$) $(\downarrow)$ | $r^{2}(\uparrow)$ | ${\mathcal{W}}_{1}(\downarrow)$ | ${\mathcal{W}}_{2}(\downarrow)$ | MMD ($\times 10^{-3}$) $(\downarrow)$ | $r^{2}(\uparrow)$ |
| FM | $3.694\pm 0.018$ | $3.775\pm 0.023$ | $1.74\pm 0.20$ | $0.951\pm 0.007$ | $3.646\pm 0.163$ | $3.805\pm 0.231$ | $3.67\pm 1.46$ | $0.867\pm 0.070$ |
| FM${}^{\text{w/}}\mathcal{N}$ | $6.790\pm 0.109$ | $7.124\pm 0.120$ | $54.97\pm 1.68$ | $0.661\pm 0.005$ | $6.905\pm 0.092$ | $7.252\pm 0.100$ | $64.40\pm 1.31$ | $0.607\pm 0.012$ |
| CGFM | $3.616\pm 0.031$ | $3.671\pm 0.032$ | $0.90\pm 0.13$ | $0.980\pm 0.003$ | $3.530\pm 0.011$ | $3.584\pm 0.010$ | $4.37\pm 0.55$ | $0.975\pm 0.003$ |
| CGFM${}^{\text{w/}}\mathcal{N}$ | $3.663\pm 0.011$ | $3.718\pm 0.011$ | $0.93\pm 0.13$ | $0.982\pm 0.001$ | $5.813\pm 0.372$ | $6.136\pm 0.398$ | $44.44\pm 8.18$ | $0.688\pm 0.026$ |
| ICNN | $4.166\pm 0.171$ | $4.238\pm 0.174$ | $10.73\pm 6.87$ | $0.859\pm 0.055$ | $3.525\pm 0.048$ | $3.583\pm 0.050$ | $7.00\pm 4.74$ | $0.930\pm 0.023$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=0$) | $3.671\pm 0.021$ | $3.726\pm 0.022$ | $1.17\pm 0.17$ | $0.979\pm 0.001$ | $3.503\pm 0.030$ | $3.555\pm 0.031$ | $1.68\pm 0.11$ | $0.975\pm 0.001$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=10$) | $3.691\pm 0.015$ | $3.748\pm 0.014$ | $1.18\pm 0.14$ | $0.978\pm 0.000$ | $3.526\pm 0.024$ | $3.582\pm 0.022$ | $1.86\pm 0.36$ | $0.972\pm 0.001$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=50$) | $3.678\pm 0.017$ | $3.733\pm 0.017$ | $1.15\pm 0.12$ | $0.979\pm 0.001$ | $3.515\pm 0.029$ | $3.567\pm 0.029$ | $2.07\pm 0.63$ | $0.977\pm 0.002$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=100$) | $3.706\pm 0.020$ | $3.764\pm 0.022$ | $1.44\pm 0.08$ | $0.977\pm 0.001$ | $3.536\pm 0.036$ | $3.590\pm 0.037$ | $1.90\pm 0.23$ | $0.975\pm 0.003$ |
| MFM ($k=0$) | $3.626\pm 0.012$ | $3.682\pm 0.013$ | $1.20\pm 0.04$ | $0.979\pm 0.000$ | $3.451\pm 0.020$ | $3.505\pm 0.021$ | $2.55\pm 0.42$ | $0.981\pm 0.002$ |
| MFM ($k=10$) | $3.624\pm 0.025$ | $3.678\pm 0.026$ | $1.04\pm 0.17$ | $0.981\pm 0.000$ | $3.451\pm 0.037$ | $3.504\pm 0.038$ | $2.47\pm 0.56$ | $0.982\pm 0.001$ |
| MFM ($k=50$) | $3.639\pm 0.016$ | $3.694\pm 0.016$ | $1.71\pm 0.20$ | $0.979\pm 0.001$ | $3.443\pm 0.035$ | $3.497\pm 0.037$ | $2.89\pm 1.85$ | $0.981\pm 0.001$ |
| MFM ($k=100$) | $3.654\pm 0.005$ | $3.712\pm 0.005$ | $1.56\pm 0.20$ | $0.978\pm 0.000$ | $3.480\pm 0.027$ | $3.534\pm 0.029$ | $3.32\pm 0.49$ | $0.980\pm 0.001$ |
|  |  |  |  |  |  |  |  |  |
| PDOs |  |  |  |  |  |  |  |  |
|  | Train | Test |  |  |  |  |  |  |
|  | ${\mathcal{W}}_{1}(\downarrow)$ | ${\mathcal{W}}_{2}(\downarrow)$ | MMD ($\times 10^{-3}$) $(\downarrow)$ | $r^{2}(\uparrow)$ | ${\mathcal{W}}_{1}(\downarrow)$ | ${\mathcal{W}}_{2}(\downarrow)$ | MMD ($\times 10^{-3}$) $(\downarrow)$ | $r^{2}(\uparrow)$ |
| FM | $3.948\pm 0.092$ | $4.050\pm 0.116$ | $3.96\pm 0.68$ | $0.942\pm 0.023$ | $4.043\pm 0.100$ | $4.174\pm 0.131$ | $4.90\pm 0.50$ | $0.931\pm 0.033$ |
| FM${}^{\text{w/}}\mathcal{N}$ | $7.109\pm 0.028$ | $7.351\pm 0.030$ | $63.42\pm 1.97$ | $0.626\pm 0.008$ | $7.697\pm 0.212$ | $7.976\pm 0.182$ | $77.52\pm 7.33$ | $0.519\pm 0.002$ |
| CGFM | $3.750\pm 0.013$ | $3.815\pm 0.017$ | $1.77\pm 0.71$ | $0.977\pm 0.004$ | $4.225\pm 0.146$ | $4.333\pm 0.171$ | $9.50\pm 3.14$ | $0.933\pm 0.016$ |
| CGFM${}^{\text{w/}}\mathcal{N}$ | $3.728\pm 0.024$ | $3.787\pm 0.026$ | $1.14\pm 0.16$ | $0.985\pm 0.000$ | $8.333\pm 0.309$ | $8.534\pm 0.298$ | $114.45\pm 14.17$ | $0.508\pm 0.016$ |
| ICNN | $4.561\pm 0.138$ | $4.666\pm 0.139$ | $34.71\pm 28.68$ | $0.876\pm 0.027$ | $4.456\pm 0.058$ | $4.564\pm 0.062$ | $24.51\pm 22.00$ | $0.869\pm 0.008$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=0$) | $3.775\pm 0.012$ | $3.837\pm 0.013$ | $1.75\pm 0.05$ | $0.979\pm 0.001$ | $3.855\pm 0.075$ | $3.944\pm 0.087$ | $3.73\pm 0.79$ | $0.978\pm 0.003$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=10$) | $3.799\pm 0.009$ | $3.860\pm 0.009$ | $1.89\pm 0.26$ | $0.978\pm 0.001$ | $3.877\pm 0.042$ | $3.961\pm 0.047$ | $3.53\pm 0.30$ | $0.978\pm 0.003$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=50$) | $3.800\pm 0.023$ | $3.863\pm 0.023$ | $1.93\pm 0.27$ | $0.980\pm 0.001$ | $3.857\pm 0.059$ | $3.938\pm 0.066$ | $3.40\pm 0.77$ | $0.981\pm 0.001$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=100$) | $3.800\pm 0.008$ | $3.863\pm 0.010$ | $2.05\pm 0.12$ | $0.979\pm 0.001$ | $3.882\pm 0.008$ | $3.973\pm 0.009$ | $3.80\pm 0.26$ | $0.979\pm 0.000$ |
| MFM ($k=0$) | $3.797\pm 0.020$ | $3.864\pm 0.020$ | $2.40\pm 0.32$ | $0.971\pm 0.001$ | $3.894\pm 0.026$ | $3.982\pm 0.037$ | $3.87\pm 0.83$ | $0.976\pm 0.002$ |
| MFM ($k=10$) | $3.792\pm 0.041$ | $3.863\pm 0.049$ | $2.29\pm 0.47$ | $0.975\pm 0.005$ | $3.953\pm 0.095$ | $4.055\pm 0.115$ | $4.80\pm 2.00$ | $0.978\pm 0.001$ |
| MFM ($k=50$) | $3.796\pm 0.014$ | $3.864\pm 0.016$ | $2.50\pm 0.22$ | $0.972\pm 0.001$ | $4.035\pm 0.099$ | $4.153\pm 0.116$ | $6.40\pm 2.19$ | $0.971\pm 0.004$ |
| MFM ($k=100$) | $3.798\pm 0.037$ | $3.866\pm 0.041$ | $2.54\pm 0.28$ | $0.970\pm 0.002$ | $3.869\pm 0.094$ | $3.953\pm 0.113$ | $3.49\pm 0.76$ | $0.975\pm 0.001$ |
|  |  |  |  |  |  |  |  |  |
| PDOFs |  |  |  |  |  |  |  |  |
|  | Train | Test |  |  |  |  |  |  |
|  | ${\mathcal{W}}_{1}(\downarrow)$ | ${\mathcal{W}}_{2}(\downarrow)$ | MMD ($\times 10^{-3}$) $(\downarrow)$ | $r^{2}(\uparrow)$ | ${\mathcal{W}}_{1}(\downarrow)$ | ${\mathcal{W}}_{2}(\downarrow)$ | MMD ($\times 10^{-3}$) $(\downarrow)$ | $r^{2}(\uparrow)$ |
| FM | $4.313\pm 0.062$ | $4.521\pm 0.076$ | $8.21\pm 0.55$ | $0.919\pm 0.014$ | $5.330\pm 0.112$ | $5.712\pm 0.113$ | $30.45\pm 3.30$ | $0.798\pm 0.002$ |
| FM${}^{\text{w/}}\mathcal{N}$ | $6.778\pm 0.046$ | $7.015\pm 0.042$ | $52.69\pm 2.63$ | $0.708\pm 0.004$ | $6.741\pm 0.177$ | $6.985\pm 0.163$ | $52.98\pm 5.35$ | $0.713\pm 0.014$ |
| CGFM | $4.281\pm 0.065$ | $4.511\pm 0.078$ | $6.81\pm 1.33$ | $0.898\pm 0.011$ | $5.574\pm 0.051$ | $5.945\pm 0.052$ | $37.13\pm 0.40$ | $0.790\pm 0.009$ |
| CGFM${}^{\text{w/}}\mathcal{N}$ | $5.549\pm 0.206$ | $5.936\pm 0.214$ | $32.44\pm 5.58$ | $0.787\pm 0.014$ | $7.258\pm 0.197$ | $7.501\pm 0.212$ | $80.45\pm 8.02$ | $0.713\pm 0.005$ |
| ICNN | $4.483\pm 0.148$ | $4.624\pm 0.146$ | $17.08\pm 13.75$ | $0.932\pm 0.020$ | $5.585\pm 0.065$ | $5.950\pm 0.068$ | $107.00\pm 124.58$ | $0.791\pm 0.006$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=0$) | $4.500\pm 0.039$ | $4.721\pm 0.043$ | $12.21\pm 1.07$ | $0.896\pm 0.005$ | $5.173\pm 0.078$ | $5.567\pm 0.077$ | $26.17\pm 2.29$ | $0.779\pm 0.003$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=10$) | $4.528\pm 0.033$ | $4.751\pm 0.045$ | $12.31\pm 0.73$ | $0.888\pm 0.003$ | $5.065\pm 0.169$ | $5.431\pm 0.172$ | $23.34\pm 4.74$ | $0.786\pm 0.005$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=50$) | $4.468\pm 0.026$ | $4.690\pm 0.030$ | $11.15\pm 0.34$ | $0.894\pm 0.006$ | $5.088\pm 0.121$ | $5.468\pm 0.115$ | $23.43\pm 3.15$ | $0.778\pm 0.005$ |
| MFM${}^{\text{w/}}\mathcal{N}$ ($k=100$) | $4.508\pm 0.042$ | $4.731\pm 0.042$ | $12.09\pm 1.34$ | $0.891\pm 0.006$ | $5.082\pm 0.037$ | $5.461\pm 0.047$ | $23.16\pm 0.90$ | $0.792\pm 0.015$ |
| MFM ($k=0$) | $4.293\pm 0.005$ | $4.491\pm 0.005$ | $8.94\pm 0.54$ | $0.925\pm 0.002$ | $5.283\pm 0.017$ | $5.653\pm 0.029$ | $30.60\pm 0.77$ | $0.797\pm 0.005$ |
| MFM ($k=10$) | $4.273\pm 0.039$ | $4.474\pm 0.041$ | $8.35\pm 0.83$ | $0.915\pm 0.009$ | $5.244\pm 0.168$ | $5.626\pm 0.175$ | $28.71\pm 5.10$ | $0.790\pm 0.016$ |
| MFM ($k=50$) | $4.271\pm 0.038$ | $4.466\pm 0.049$ | $8.40\pm 0.40$ | $0.924\pm 0.001$ | $5.165\pm 0.132$ | $5.538\pm 0.140$ | $26.97\pm 2.46$ | $0.796\pm 0.009$ |
| MFM ($k=100$) | $4.200\pm 0.079$ | $4.381\pm 0.095$ | $7.20\pm 1.13$ | $0.918\pm 0.004$ | $4.950\pm 0.199$ | $5.321\pm 0.219$ | $20.08\pm 5.10$ | $0.798\pm 0.013$ |

## Peer Reviews

confidence: 4
contribution: 4
initial_rating: 8
presentation: 4
questions: - In Eq.17, shouldn't the condition `c` be part of `\phi` or `v_t` as well?
- 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?
soundness: 3
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.
- Theoretical foundation: The paper provides a solid theoretical basis for the proposed method, including the connections to existing approaches like conditional generative flow matching.
- Readability: The paper is *very* well written and easy to read and understand (although it's covering a complex topic)
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.
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.
-  Figure 1 (c) is not that helpful in my opinion.

Typos:
- "a a standard" (l420)
- "of of model" (l464)

---

confidence: 2
contribution: 3
initial_rating: 5
presentation: 1
questions: All weaknesses as stated above are my main questions.
soundness: 2
strengths: -- The problem is known to have plagued science for a long time. Therefore, new approaches and solutions are important contributions.
-- The idea of modeling the fields in the Wasserstein setting is intriguing, and intuitively meaningful.  
-- The authors relate their approach with mean-field limits and diffusion processes. This is certainly something of interest.
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.
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. 
-- The problem seems to boil down to embedding populations using graph neural networks and then flow matching in Wasserstein manifolds. 
-- 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). 
-- 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. 
-- 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.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
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:

* How are uncertainties on your results calculated? If they are not derived from independent splits/seeds, I would suggest you do that.
* Why did you not benchmark MFM on the [dataset consider in the CellOT manuscript](https://onlinelibrary.wiley.com/doi/epdf/10.1111/exd.12683)?
* 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.
* Can you add the following uninformed baselines to the experiments in Sections 5.1 and 5.2?
    * The base distribution itself, i.e., distance of the unperturbed patient cells to the perturbed cells without applying any model.
    * The base distribution shifted by a constant vector that is the mean over all perturbed cells (e.g., those in the train dataset)
* 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?
* 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?
* 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.
* 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)?
soundness: 2
strengths: * The manuscript is overall well-written and organized.
* 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.
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).

More 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.

Therefore, 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).
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.
    * 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.
    * 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?
    * 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:
        * The base distribution itself, i.e., distance of the unperturbed patient cells to the perturbed cells without applying any model.
        * The base distribution shifted by a constant vector that is the mean over all perturbed cells (e.g., those in the train dataset)
    * 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?
    * 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?
    * 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.
* **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.
* **(minor)** The text contains quite a few typos. I suggest the authors do another round of proof-reading for the final version.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: See Strengths.
soundness: 3
strengths: I am listing both strengths (+), weakness (-) and questions (?) in the following to make the review more logistic.

- (+) 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.
- (?) 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?
- (?) Can you method be extent to SDE case?
- (?) 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.
- (-) 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.
- (?) 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).
- (?) 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.
- (?) Can you provide a random baseline for your experiments? Random guess baseline and random initialized model baseline.
- (-) 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.
- (?) To further illustrate the effectviese, I think comparison with the numbers and experiments in "Deep Generalized Schrödinger Bridge" is necessary.
- (?) 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.
- (?) 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.
- (?) How do you construct the graph in the organoid drug-screen experiment? How different graph construction impact your model results?
summary: - The paper proposes a new diagram of conditional flow matching models named meta flow matching.
- 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).
- In math, this can be formulated as constructing the vector field in the Wasserstein manifold.
- The authors then provide the technical construction of the model, including objective functions and parametrization.
- Numerical experiments are conducted on synthetic and organoid drug-screen datasets.
- I am happy to see this conceptual novelty in the approach.
- I didn't thoroughly check the technical part in the paper while I didn't see anything odd.
- I am a little concerned on certain claims and the effectiveness of the method, both conceptually and numerically.
weaknesses: See Strengths.

Meta-review: The paper proposes a 'meta flow matching approach' for modelling a family of distributions, conditioned on the population index i; 
instead of previous approaches (conditional generative flow matching) instead of embedding a sequence i the meta-flow matching approach embeds to population (i.e. the set of points) into the conditioning vector for a vector field. This is very neat and mimics, for example, mean field approaches. The paper is well-written. The applicability is shown for small-dimensional distributions, but overall, this is a good work. Some of the concerns are the actual graph neural network design for particular applications and correct comparison to the baselines, but it also has been A.

Decision: Accept (Poster)