Meta Flow Matching: Integrating Vector Fields on the Wasserstein Manifold

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 18:30:10

Abstract: Numerous biological and physical processes can be modeled as systems of interacting samples evolving continuously over time, e.g. the dynamics of communicating cells or physical particles. 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 treatments depend 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 Meta Flow Matching the ability to generalize over the initial distributions unlike previously proposed methods. Finally, we demonstrate the ability of MFM to improve prediction of individual treatment responses on a large scale multi-patient single-cell drug screen dataset.

Main paper: 
 introduction: 
 Understanding the dynamics of many-body problems is a central 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., 2020; Goodenough and Paul, 2009) . This signaling presents an obstacle for modeling, but is essential for understanding and eventually controlling cell dynamics during development (Gulati et al., 2020; Rizvi et al., 2017) , in diseased states (Molè et al., 2021; Binnewies et al., 2018; Zeng and Dai, 2019; Chung et al., 2017) , and in response to perturbations (Ji et al., 2021; Peidli et al., 2024) .The 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. While single-cell sequencing has been used to great effect to understand the heterogeneity in cell systems, they are also destructive, making longitudinal measurements extremely difficult. Instead, most approaches model cell dynamics at the population level (Hashimoto et al., 2016; Weinreb et al., 2018; Schiebinger et al., 2019; Tong et al., 2020; Neklyudov et al., 2022; Bunne et al., 2023a) . These approaches involve the formalisms of optimal transport (Villani, 2009; Peyré and Cuturi, 2019) and Submitted to 38th Conference on Neural Information Processing Systems (NeurIPS 2024). Do not distribute.generative modeling (De Bortoli et al., 2021; Lipman et al., 2023) methods, 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 a shared dynamical system. 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 is available in the training data.To address this we propose Meta Flow Matching (MFM) -the amortization of the Flow Matching generative modeling framework (Lipman et al., 2023) over the input measures. In practice, our method can be used to predict the time-evolution of distributions from a given dataset of the timeevolved 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.We showcase the utility of our approach on two applications. We first explore Meta Flow Matching on a synthetic task of denoising letters. We show that MFM is able to generalize the denoising process to letters in unseen orientations where a standard flow matching approach cannot. Next, we explore how MFM can be applied to model single-cell perturbation data (Ji et al., 2021; Peidli et al., 2024) .We evaluate MFM 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., 2023) . This dataset includes more than 25M cells collected over ten patients under 2500 conditions. This is a challenging task due to the variance over multiple patients, treatments applied and the local cell compositions, but it can be used to study the tumor micro-environment (TME), thought to be essential in circumventing chemoresistance. We demonstrate that Meta Flow Matching can successfully predict the development of cell populations on replicated experiments, and, most importantly, it generalizes to previously unseen patients, thus, capturing the patient-specific response to the treatment. 
generative modeling via flow matching: 
 Flow Matching is an approach to generative modeling recently proposed independently in different works: Rectified Flows (Liu et al., 2022) , Flow Matching (Lipman et al., 2023) , Stochastic Interpolants (Albergo and Vanden-Eijnden, 2022) . It assumes 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 t ) is produced as follows x t = f t (x 0 , x 1 ), (x 0 , x 1 ) ∼ π(x 0 , x 1 ) ,where dx 1 π(x 0 , x 1 ) = p 0 (x 0 ) , dx 0 π(x 0 , x 1 ) = p 1 (x 1 ) ,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 withf t (x 0 , x 1 ) = (1 -t) • x 0 + t • x 1 ); π(x 0 , x 1 )is the density of the joint distribution, which is usually taken as a distribution of independent random variables π(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., 2023; Tong et al., 2024) . The corresponding density can be defined then as the following expectationEQUATIONThe 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 solutionEQUATIONConditional Flow Matching Vector Field on P2(X )Figure 1 : Illustration of flow matching methods on the 2-Wasserstein manifold, P2(X ), depicted as a twodimensional 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 c1 and c2 on the plot. Meta Flow Matching learns to integrate a vector field on P2(X ), i.e. for every starting density p0 Meta Flow Matching defines a push-forward measure that integrates along the underlying vector field.Relying on this formula, one can derive the tractable objective for learning v * t (x), i.e.EQUATIONFinally, the vector field v t (ξ, ω) ≈ v * t (ξ) defines the push-forward density that approximately matches p t=1 , i.e. T # p 0 ≈ p t=1 , where T is the flow corresponding to vector field v t (•, ω) with parameters ω.Conditional image generation is one of the most common applications of generative models nowadays;it includes conditioning on the text prompts (Saharia et al., 2022b; Rombach et al., 2022) as well as conditioning on other images (Saharia et al., 2022a) . To learn the conditional generative process with diffusion models, one merely has to pass the conditional variable (sampled jointly with the data point) 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., 2023; Dao et al., 2023; Isobe et al., 2024) 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 variableEQUATIONFor every c, the density p t (x t | c) follows the continuity equation with the following vector fieldEQUATIONwhich depends on c. Thus, the training objective of the conditional model becomesEQUATIONwhere, 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.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: We discuss limitations in Appendix D.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.• The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.• The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.• The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.• The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.• If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.• While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
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 π(x 0 , x 1 ) given empirically. For the given joint π(x 0 , x 1 ), we denote the solution of the Flow Matching optimization problem as followsEQUATIONAnalogously to the amortized optimization (Chen et al., 2022; Amos et al., 2023) , we aim to learn the model that outputs the solution of Eq. ( 10) based on the input data sampled from π, i.e.EQUATIONwhere φ(π) is the embedding model of π and the joint density π(• | c) is generated using some unknown measure of the conditional variables c ∼ p(c).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 (2001) . That is, time-dependent densities p t (x t ) define absolutely-continuous curves on the 2-Wasserstein space of distributions P 2 (X ) (Ambrosio et al., 2008) . The tangent space of this manifold is defined by the gradient flows S t = {∇s t | s t : X → R} on the state space 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., 2022) .Under the geometric formalism of the 2-Wasserstein manifold, Flow Matching can be considered as learning the tangent vectors v t (•) along the density curve p t (x t ) defined by the sampling process in Eq. (2) (see the left panel in Fig. 1 ). 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. 1 ).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 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). In the limit of the infinite number of interacting particles one can describe their state with the density function p t (x). Consider the interaction according to the first order dynamics with the velocity k(x, y) : R d × R d → R d of the particles at point x that interact with the particles at point y. Then the change of the density is described by the following continuity equationEQUATION)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.EQUATION)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 problemEQUATIONwhere the state-space vector field v t (x, p t ) 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.Consider the dataset of joint populations D = {(π(x 0 , x 1 | i))} i , where, to simplify the notation, we associate every i-th population with its density π(• | 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 (x0 | i) = dx 1 π(x 0 , x 1 | i) aresampled from some unknown distribution that can be parameterized with a large enough number of parameters (ii) the endpoint marginals p 1 (x 1 | i) = dx 0 π(x 0 , x 1 | i) are obtained as push-forward densities solving the Cauchy problem in Eq. ( 14), (iii) there exists unique solution to this Cauchy problem.One can learn a joint model of all the processes from the dataset D using the conditional version of the Flow Matching algorithm (see Section 2.2) where the population index i plays the role of the conditional variable. However, obviously, such a model will not generalize beyond the considered data 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. ( 14). That is, we propose to use an embedding function φ : P 2 (X ) → 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 ωEQUATION)Note that the initial density p 0 is enough to predict the push-forward density p 1 since the Cauchy problem for Eq. ( 14) has a unique solution. The embedding function φ(p 0 ) can take different forms, e.g. it can be the density value φ(p 0 ) = p 0 (•), 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 marginalsp 0 (x 0 | c) = dx 1 π(x 0 , x 1 | c) and p 1 (x 1 | c) = dx 0 π(x 0 , x 1 | c) and the distribution p(c) are known, i.e. there exist φ :P 2 (X ) → R m such that L M F M (ω) = L CGF M (ω).Proof. Indeed, sampling from the dataset i ∼ D becomes sampling of the conditional variable c ∼ p(c) and the embedding function becomesφ(p 0 (• | c)) = c.Furthermore, for the parametric family of the embedding models φ(p t , θ), we show that the parameters θ can be estimated by minimizing the objective in Eq. ( 15) in the joint optimization with the vector field parameters ω. We formalize this statement in the following theorem.Theorem 1. Consider a dataset of populationsD = {(π(x 0 , x 1 | i))} i generated from some unknown conditional model π(x 0 , x 1 | c)p(c).Then the following objectiveL(ω, θ) = E p(c) 1 0 dt E pt(xt | c) ∥v * t (x t | c) -v t (x t | φ(p 0 , θ), ω)∥ 2 (16)is equivalent to the Meta Flow Matching objectiveL MFM (ω, θ) = E i∼D E π(x0,x1 | i) 1 0 dt ∂ ∂t f t (x 0 , x 1 ) -v t (f t (x 0 , x 1 ) | φ(p 0 , θ); ω) 2 (17)up to an additive constant.Proof. We postpone the proof to Appendix A.In many applications, the populations D = {(π(x 0 , x 1 | i))} N i=1 are given as empirical distributions, i.e. they are represented as samples from some unknown density πEQUATIONwhere N i is the size of the i-th population. For instance, for the diffusion process considered in Example 2, the samples from π(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 j 1 . 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 φ(p 0 , θ), i.e.EQUATIONFor this purpose, we employ GNNs, which recently have been successfully applied for simulation of complicated many-body problems in physics (Sanchez-Gonzalez et al., 2020) . To embed a population{x j 0 } Ni j=1, we create a k-nearest neighbour graph G i based on the metric in the state-space X , input it into a GNN, which consists of several message-passing iterations (Gilmer et al., 2017) 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).Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)• The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean.• It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.• For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).• If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
related work: 
 The meta-learning of probability measures was previously studied by Amos et al. (2022) 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.Generative modeling for single cells. Single cell data has expanded to encompass multiple modalities of data profiling cell state and activities (Frangieh et al., 2021; Bunne et al., 2023b) . Single-cell data presents multiple challenges in terms of noise, non-time resolved, and high dimension, and generative models have been used to counter those problems. Autoencoder has been used to embed and extrapolate data Out Of Distribution (OOD) with its latent state dimension (Lotfollahi et al., 2019; Lopez et al., 2018; Hetzel et al., 2022) . Orthogonal non-negative matrix factorization (oNMF) has also been used for dimensionality reduction combined with mixture models for cell state prediction (Chen et al., 2020) . Other approaches have tried to use Flow Matching (FM) (Tong et al., 2023 (Tong et al., , 2024;; Neklyudov et al., 2023) or similar approaches such as the Monge gap (Uscidda and Cuturi, 2023) to predict cell trajectories. Currently, the state of the art method uses the principle of Optimal Transport (OT) to predict cell trajectories with Input Convex Neural Network (ICNN) (Makkuva et al., 2020; Bunne et al., 2023b) . 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., 2020) via GNNs. However, the problem there is very different since these models use the information about the individual trajectories of samples, which are not available for the single-cell prediction. Neklyudov et al. (2022) consider learning the vector field for any continuous time-evolution of a probability measure, however, their method is restricted to single curves and do not consider generalization to unseen data. Finally, the weather/climate forecast models generating the next state conditioned on the previous one (Price et al., 2023; Verma et al., 2024) are similar approaches to ours but operating on a much finer time resolution. 
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 (W 1 ), 2-Wasserstein (W 2 ) distance, and the radial basis kernel maximum-mean-discrepancy (MMD) distance (Gretton et al., 2012) . We parameterize all vector field models v t (• | φ(p 0 ); ω) using a Multi-Layer Perceptron (MLP). For MFM, we additionally parameterize φ(p t ; θ, k) using a Graph Convolutional Network Table 1 : Results of the synthetic letters experiment for population prediction on seen train populations and unseen test populations. We report the the 1-Wasserstein (W1), 2-Wasserstein (W2), and the maximum-meandiscrepancy (MMD) distributional distances. We consider 4 settings for MFM with varying k.Train Test The results for all the models are averaged over three random seeds.W 1 W 2 MMD (×10 -3 ) W 1 W 2 MMD (×10 -3 ) FM 0.We curate a synthetic dataset of the joint distributions {(p 0 (x 0 , | i), p 1 (x 1 | i))} N i=1 by simulating a diffusion process applied to a set of pre-defined target distributions p 1 (x 1 | i) for i = 1, . . . , N . To get a paired population p 0 (x 0 | i) we simulate the forward diffusion process without drift x 0 ∼ N (x 1 , σ).After this setup, for reasonable values of σ, 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 ∼ 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, while only using the upright orientation for the remaining letters "X" and "Y". We construct the test data by using 10 random orientations of "X" and "Y" (validation and test, respectively) that differ from the upright orientations of the same letters in the training data. We do this to simplify the generalization task -the model will see the shapes of "X" and "Y" during training, but the same letters under different orientations remain unseen.We train FM, CGFM and 4 variants of MFM of varying k for the GCN population embedding model φ(p t ; θ, k). When k = 0, φ(p t ; θ, k) becomes identical to the DeepSets model (Zaheer et al., 2017) .We compare MFM to Flow-Matching (FM) and Conditional Generation via Flow-Matching (CGFM).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. Note that the CGFM can be viewed as an idealized model for the train data since it gets perfect information regarding the population conditions. We use CGFM to assess if other models are fitting the data. For MFM, we expect to both fit the training data and generalize to unseen distributional conditions.In Fig. 2 , we observe that indeed FM fails to adequately learn to sample from p 1 (x 1 | i) in the training set, and likewise fails to generalize, while CGFM is able to effectively sample from p 1 (x 1 | i) in the training set, but fails to generalize. We report results for the synthetic experiment in Table 1 .As expected, CGFM fits the training data, however, fails to generalize beyond its set of training conditions. In contrast, we see that MFM is able to both fit the training data (approaching the performance of CGFM) 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.Data. For experiments on biological data, we use the organoid drug-screen dataset from Ramos Zapatero et al. (2023) . This dataset is a single-cell mass-cytometry dataset collected over 10 patients.Somewhat unique to this dataset, unlike many prior perturbation-screen datasets which have a single control population, this dataset has matched controls to each experimental condition. Populations from each patient are treated with 11 different drug treatments of varying dose concentrations. 1 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. 3-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/PDOF cultures and results for the individual cultures (this is reported in Appendix F).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. 3-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. In both settings, we normalize the data and embed it into a lower dimensional principle components (PC) representation. We do this to reduce the dimensionality of the data and to extract the relevant information from the 43 bio-markers (features) of the ambient space. We train and evaluate all models in the PC space. For all organoid drug-screen dataset experiments we use PC=10. Further details regarding data pre-processing and data splits are provided in Appendix B.2.For the organoid drug-screen experiments, we consider an ICNN architecture in addition to the Flow-matching models. The ICNN model is based on CellOT (Bunne et al., 2023a) ; a method for learning cell specific response to treatments. The ICNN (and likewise CellOT) counterparts our FM Table 2: Experimental results on the organoid drug-screen dataset for population prediction of treatment response across replicate populations averaged over co-culture conditions. Results are reported for models trained on data embedded into 10 principle components. We report the the 1-Wasserstein (W1), 2-Wasserstein (W2), and the maximum-mean-discrepancy (MMD) distributional distances. We consider two settings for MFM with varying nearest-neighbours parameter. For extended results in Table 4 .Test 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.W 1 W 2 MMD (×10 -3 ) W 1 W 2 MMD (×10 -3 ) FM 1.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. MFM (k = 10) performs best on generalization to unseen replicates. We include results reported for the separate cell cultures in Table 4 in Appendix F.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. MFM (k = 10) performs best on generalization to unseen replicates. We include results reported for the separate cell cultures in Table 5 in Appendix F.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. 
conclusion and future work: 
 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 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. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes]Justification: The paper discuss the compute resources and reproducibility in Appendix C.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.• The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.• The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: The research does conform with the NeurIPS Code of Ethics. The study presented involves only public or synthetic data, which is freely available online. The considered models do not impose risks of misuse or dual-use.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.• The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: The paper discusses the broader impact in Appendix E.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.• The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.• If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] .Justification: The models considered in the paper do not carry the risks of misuse or dual-use.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.• Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.• We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL.• The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.• For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.• If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA] .Justification: The paper does not release new assets.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.• The paper should discuss whether and how consent was obtained from people whose asset is used.• At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] .Justification: The empirical study presented in the paper is conducted on the synthetic or publicly available data.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. Justification: The empirical study presented in the paper is conducted on the synthetic or publicly available data.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.• For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: Theorem 1. Consider a dataset of populations D = {(π(x 0 , x 1 | i))} i generated from some unknown conditional model π(x 0 , x 1 | c)p(c). Then the following objectiveL(ω, θ) = E p(c) 1 0 dt E pt(xt | c) ∥v * t (x t | c) -v t (x t | φ(p 0 , θ), ω)∥ 2 (16)is equivalent to the Meta Flow Matching objectiveL MFM (ω, θ) = E i∼D E π(x0,x1 | i) 1 0 dt ∂ ∂t f t (x 0 , x 1 ) -v t (f t (x 0 , x 1 ) | φ(p 0 , θ); ω) 2 (17)up to an additive constant.Proof. The loss functionEQUATION)The last term does not depend on θ, the second term we can estimate, for the first term, we use the formula for the (from Eq. ( 8))EQUATIONThus, the loss is equivalent (up to a constant) toEQUATION)Note that in the final expression we do not need access to the probabilistic model of p(c) if the joints π(x 0 , x 1 | c) are already sampled in the data D. Thus, we haveEQUATION)B Experimental DetailsThe synthetic letters dataset contains 242 train populations a 10 test populations. Each population contains roughly between 750 and 2700 samples. In this dataset.The organoid drug-screen dataset contains 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.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=1e -4, β 1 =0.5, β 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 4 linear layers with 512 hidden units and SELU activation functions. 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 3 graph convolution layers with 512 hidden units. The final GCN model layer outputs an embedding representation e ∈ R d . For the Synthetic experiment, we found that d = 256 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 ω and the GCN model parameters θ. We alternate between updating the respective model parameters every epoch. 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.We implement all our experiments using PyTorch and PyTorch Geometric. We submitted our code as supplementary material with our submission.All experiments were conducted on a HPC cluster primarily on NVIDIA Tesla T4 16GB GPUs. Each individual seed experiment run required only 1 GPU. Each experiment ran between 3-11 hours and all experiments took approximately 500 GPU hours.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., 2022 (Huguet et al., , 2023)) , unbalanced optimal transport (Benamou, 2003; Yang and Uhler, 2019; Chizat et al., 2018) , aligned (Somnath et al., 2023; Liu et al., 2023) , or stochastic settings (Bunne et al., 2023a; Koshizuka and Sato, 2023 ) in this work.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.Table 4 : Experimental results on the organoid drug-screen dataset for population prediction of treatment response across replicate populations. Results are reported for models trained on data embedded into 10 principle components. We report the the 1-Wasserstein (W1), 2-Wasserstein (W2), and the maximum-mean-discrepancy (MMD) distributional distances. We consider 2 settings for MFM with varying nearest-neighbours parameter. Table 5 : Experimental results on the organoid drug-screen dataset for population prediction of treatment response across patient populations. Results are reported for models trained on data embedded into 10 principle components. We report the the 1-Wasserstein (W1), 2-Wasserstein (W2), and the maximum-mean-discrepancy (MMD) distributional distances. We consider 2 settings for MFM with varying nearest-neighbours parameter. Guidelines:Fibroblasts Train Test W 1 W 2 MMD (×10 -3 ) W 1 W 2 MMD (×10 -3 ) FM 1.584 ± 0.W 2 MMD (×10 -3 ) W 1 W 2 MMD (×10 -3 ) FM 2.002 ± 0.Fibroblasts Train Test W 1 W 2 MMD (×10 -3 ) W 1 W 2 MMD (×10 -3 ) FM 1.599 ± 0.• The answer NA means that the abstract and introduction do not include the claims made in the paper.• The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.• The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.• It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.We consider only the highest dosage and leave exploration of dose-dependent response to future work. 


Decision: This paper proposes an innovative approach to embedding the initial distribution in Flow Matching problems, using a GNN to learn embeddings that generalize to unseen initial distributions. While the idea shows promise, the overall enthusiasm among reviewers was moderate, with concerns about the significance of the results and scalability. Although the authors addressed many of these concerns in their response, the reviewers remained cautious. Reviewer BiEg provided a strong score of 7 and could have championed the paper, but other reviewers maintained scores ranging from 5 to 6. Given the mixed reviews, I recommend rejecting the paper at this time.

Reviews: review: paper_summary: paper_summary: This paper introduces Meta Flow Matching (MFM), a flow matching framework modeling interacting samples evolving over time by integrating vector fields on the Wasserstein manifold. The authors leverage a Graph Neural Network to embed populations of samples and thus generalize the method over different initial distributions. The authors demonstrate the method on individual treatment responses predictions on a large multi-patient single-cell drug screen dataset.
strength_weakness: strength_weakness: In Table 1 of the synthetic experiment, the authors compared the performance of FM, CGFM and MFM of k=0,1,10,50. MFM doesn't seem to beat existing methods on the metrics and from the visualizations, it's hard to tell MFM is actually doing better than FM. Also, for the various values of k, some explanations on how performance correlates with values of k and why might be necessary for readers to understand this table.

In both experiments, the authors only compared FM, CGFM, and in Table 2 also ICNN. Probably more methods, like diffusion, should also be taken into comparison. Also, in experiment 2, only W1, W2 and MMD are computed as metrics. While these are useful when modeling distributions, more metrics, especially those specific to this application may be applied.
questions: questions: In Table 1, the authors compared k=0,1,5,10,50 and in Table2, only k = 0 and 10 are listed. Some explanations on how this decision is made is probably helpful.
limitations: limitations: The authors have not addressed limitations.
paper_summary: paper_summary: The paper discussed the novel problem setup of generative modeling of the dynamics of probability distributions. The paper proposed Meta Flow Matching (MFM), an extension of the flow matching framework for implicitly learning the vector fields on the Wasserstein manifold of probability distributions. The paper demonstrated better transferability of the proposed framework on unseen distributions on both synthetic datasets and real-world drug-screen datasets.
strength_weakness: strength_weakness: - The proposed method seems to be a special case of a conditionally trained flow matching model where the conditions are continuous learnable embeddings. Such an idea has already been applied in various diffusion or flow matching models including image generation (conditioned on text embedding in the latent space), protein co-design [1] (conditioned on sequence, generate protein structure, or vice versa), and peptide design [2] (conditioned on receptor proteins, generate peptides).

- The idea of population embedding in the paper is similar to task embedding, which has been well-explored in the meta learning (e.g. [3]). Although the authors claimed their framework to be *meta* flow matching, related work in meta learning seems to lack.


[1] Campbell, Andrew, et al. "Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design." arXiv preprint arXiv:2402.04997 (2024).

[2] Li, Jiahan, et al. "Full-Atom Peptide Design based on Multi-modal Flow Matching." arXiv preprint arXiv:2406.00735 (2024).

[3] Achille, Alessandro, et al. "Task2vec: Task embedding for meta-learning." Proceedings of the IEEE/CVF international conference on computer vision. 2019.
questions: questions: 1. What is the major difference between the proposed method and a conditionally trained flow matching model? See Weakness 1 for some examples of conditional generative models in other fields. These models also rely on continuous learnable embeddings as conditions during training and sampling.

2. If MFM can recovers the conditional generation via flow matching (CGFM) (Proposition 1), what is the benefit of using the proposed scheme over the latter? Can you provide more details regarding training CGFM? For example, what are fed into the model as conditions for CGFM?

3. There is a (minor) mismatch between Figure 1 and the MFM objective in Eq.15. In Figure 1 (or Eq.14), the flow matching model operates on the probability distributions to output a vector field $v\_t(p\_t)$ in the tangent space (an affine subspace as the probabilities needs to be normalized). However, in Eq.15, the flow matching model still operates on the data space $v\_t(x\_t)$. This is probably because the probability distributions are only implicitly described via data samples. Nonetheless, the authors should avoid saying in the caption of Figure 1 that the model learns a vector field on the Wasserstein manifold. During both training and sampling, the learned vector field is always defined on the data space in this work.

4. GNNs is a reason choice for data with geometric properties. However, the single-cell data does not seem to a exhibit a simple Euclidean geometry for GNN to work. Can you provide more justifications? Is there a better choice for the population embedding?
limitations: limitations: The authors have adequately and properly discussed the limitations and potential societal impact of the paper in the Appendix.
paper_summary: paper_summary: This paper proposed an extension of the Conditional Generative Modeling via Flow Matching (CGFM) framework. Taken inspiration from the theory of Wasserstein Gradient Flow, this new framework, Meta Flow Matching, proposed to learn the push-forward mapping of multiple measures in the same measures-space. This is motivated by the realistic problem of modeling single-cell perturbation data where we want to see the response of populations of cells of patients when receiving different treatments. A novelty of Meta Flow Matching is that by combining amortized optimization and CGFM, the trained MFM velocity network can model newly observed populations _without_ knowing their labels/conditions. Two empirical benchmarks were performed to showcase the effectiveness of MFM compared to the Flow Matching (FM) and CGFM.
strength_weakness: strength_weakness: - The first part of the methodology section seems to be phrased as a new methodological contribution, but if I'm not mistaken this is just more or less restating the already established theory of W2 gradient flow and continuity equation (eq 14). I think the authors should put Section 3.1 into the background section (2nd Section).

- There is a lack of discussion on whether the 3 crucial assumptions (line 145-152) are satisfied in a realistic biological setting. For example, in theory, Assumption (iii) on the unique existence of the Cauchy problem stands when the velocity field satisfies some regularity conditions -- I'm not sure this can be extended to a parameterized neural network that takes input from another (graph) neural network as an embedding function, which is hardly Lipschitz smooth in most of the case.

- Algorithm boxes at the end of section 3 is highly welcome. Or if the authors cannot allocate the space, I highly recommend putting two (one for training and one for sampling) into the Appendix. It is quite hard to follow how the velocity is trained in reality. For example, what function $f_t(x_0, x_1)$ did the authors take for this work? Is it still linear interpolation as vanilla flow matching? Or does it involve adding some form of stochasticity as in stochastic interpolant or VP-SDE as in diffusion model? Is the coupling $(x_0, x_1)$ sampled to match randomly, or they are sampled to some form of alignment as in the multisample flow matching paper (Pooladian et al. 2023)?

- This might not be the original purpose of this work, but I would love to see how MFM perform on conditional image generation task. One can pick a simple small dataset such as CIFAR10 that already includes class labels, or better yet ImageNet dataset. The performance in this takse will be much more convincing than the synthetic experiment, where I would argue would target the same type of task.
questions: questions: See weaknesses.
limitations: limitations: See weaknesses.
paper_summary: paper_summary: This paper addresses the issue of the flow matching method's lack of dependence on the data population. It proposes incorporating the initial population density into the vector field through amortization—using a Graph Neural Network (GNN) to embed the populations and adding this embedding to the input of the vector field network. The paper argues that this dependence would better model the data due to sample interactions, demonstrating improved generalization on unseen initial distributions. The method's application is showcased in perturbation drug screening.
strength_weakness: strength_weakness: 1. The paper could explain more about the meta-learning aspect of this method.
2. It could include explanations and/or ablation studies on the role of meta-learning and the GNN, especially in the synthetic experiment.
3. More detail is needed on what properties of the Wasserstein manifold of probabilities are used in the model. It is unclear how the model proposed in section 3.2 depends on the properties of the Wasserstein manifold described in section 3.1.
questions: questions: 1. How is this model different from using a GNN encoder to add the embedding to the input of the vector field, trained end-to-end? Specifically, what makes it meta-learning? Considering that $\theta$ and $\omega$ are jointly optimized to minimize $\mathcal{L}_{MFM}$, is there anything unique about the training algorithm?
2. How would this meta-learning framework compare with using the embedding from a pretrained encoder, such as a VGAE, or a classification model?
3. For the synthetic experiment, how much does the orientation-invariant embedding from the GNN help the model? Is the observed "generalization" due to the unique inductive bias of the GNN? Would any other encoder, without orientation invariance, such as an MLP, still yield good performance?
4. How well does it scale to high dimensions, considering that only 10 dimensions are used in the experiments? How well does it scale to large datasets where you need to build large KNN graphs? Would it scale effectively to high-dimensional tasks such as image generation?
5. In line 261, it states, “We use CGFM to assess if other models are fitting the data.” Is this used as an evaluation metric in the experiments?
limitations: limitations: The authors have adequately addressed the limitations.

Comments: title: Thank you for the clarification
comment: I appreciate the authors' efforts in conducting ablation experiments and addressing the clarification questions. 
The additional experiments have resolved my concerns in scalability and the necessity of a GNN. 
I agree with the authors' clarifications (esp. with the algorithm box) on the novelty of joint training.
Therefore, I have raised my score.
I thank the authors for their response and raise my score from 4 to 5.
I thank the authors for providing more explanations and clarifications regarding the difference between the proposed method of MFM versus existing methods of CGFM and meta-learning. After clarifying the difference from existing work of conditional flow matching and task embedding in meta-learning, I believe this work indeed offers an alternative perspective that can be interesting to the community. In this regard, I raise my score from 5 to 6. I would suggest the authors highlight these distinctions in the revised manuscript.
Thank you for your time and effort in reading and responding to our rebuttal. Following your suggestion, we have moved the content of Section 3.1 to the background section to avoid any possible confusion and to make our contribution clearer.

We are happy to consider any additional suggestions and to make further changes to improve the overall quality of our work.
I have read the rebuttal of the authors. I remain in my opinion that the background part in Section 3.1 should be moved back to the earlier section to not be confused as a contribution. I also understand that it is quite hard to find theoretical analysis for this type of work (that leans more on methodological and algorithmic contributions). While the authors' rebuttal clarifies plenty of my concerns, I think the work still requires some reorganization and additional modification. Therefore, I keep my Borderline Acceptance score.