Meta Flow Matching: Integrating Vector Fields on the Wasserstein Manifold

We thank the reviewer for their time and effort in reviewing our paper. We are heartened to hear that the reviewer views that MFM is "clearly novel" and is "happy to see the conceptual novelty" of our method. We next answer all the important questions raised by the reviewer while refer to the general response to all reviewers for any additional experiments.

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?

The FP-equation is a special case of the PDE in equation (12) when, for instance, we take $v_t(x, p_t) = -\frac{1}{2}\nabla_x \log p_t(x)$, as we discuss in Example 2. In general, equation (12) is our assumption on the dynamics of the system (for letters we know it is diffusion, for cells it is unknown). To be precise, we assume that the vector field $v_t(x, p_t)$ can be represented as a function of density $p_t$, and try to learn this functional dependency via a GNN. However, clearly, we do not know the analytic or any other form of how $v_t(x, p_t)$ depends on $p_t$.

Can you method be extent to SDE case?

Yes. It is straightforward to extend MFM to the stochastic setting. We leave exploration of empirical performance of MFM in the stochastic setting to future work.

Analysis of population embeddings: To learn the population embeddings, I believe you should have enough example of populations otherwise the generalization is less promised.

We agree that generalization is likely to improve with increasing data. We refer the reviewer to Figure 5 for an ablation over the number of letter populations and its effect on generalization performance in the synthetic letters setting. We indeed observe that as we increase the number of training populations, the generalization performance of MFM also improves. In contrast, the performance of the baseline methods, which are not designed to operate in the regime of generalizing to unseen distributions, does not meaningfully improve.

I would like to see a detailed analysis of population embeddings for a sanity check in both experiments.

We thank the reviewer for their helpful suggestion. To illustrate the ability of MFM to learn meaningful population embeddings, we have added a detailed analysis of population embeddings both in the synthetic data setting and in the organoid drug-screen dataset in Appendix F.1. To do this, we first use Uniform Manifold Approximation and Projection (UMAP) to project embeddings in 2 dimensions. We then compute the pairwise distances between the distributions (letters silhouette populations for the synthetic data and control populations for organoid drug-screen data). Through this analysis, we found that the embeddings reflect data characteristics and clusters of similar groups. Please see Appendix F.1 and Figure 7 for details in the updated text. On the organoid drug-screen data, we found that cell populations from patients with similar chemotherapeutic drug responses also cluster together. Ramos Zapatero et al., (2023) identified patients who are chemosensitive (responsive to chemotherapy), chemorefractory (not responsive to chemotherapy). In turn, we observed PDO11 and PDO141 clustering together which are derived from different patients and both classified as chemorefractory.

We thank the reviewer for their time in effort in constructing this in-depth review of or paper, offering meaningful suggestions, and insightful questions. We are happy to see that the reviewer found our work well motivated, that our "idea to solve this problem is simple and elegant", and found our manuscript "overall well-written and organized". Below we address the clarifying questions and suggestions brought up by the reviewer.

First of all, the authors consider ICNN/CellOT as a baseline that can accommodate varying baseline distributions and conditioning. I think the introduction to ICNN presented 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.

Thank you for pointing this out. What we refer to as the "method" here is "CellOT", while the base architecture of CellOT is an ICNN. Throughout the paper and results, we use ICNN to denote the CellOT baseline (see lines 480-483). We also clarify that the CellOT/ICNN model can generalize to new cells, but is not designed to generalize to unseen distributions. We have updated lines 348-350 to clarify this in the text. Further, we reinforce this through our empirical results, where akin to FM, CellOT (ICNN) struggles to generalize across unseen distribution/populations relative to MFM.

We wish to clarify that the novelty of MFM is 2 fold: (1) we can train a generative model to condition on entire distributions through the use of a population embedding model which learns embeddings of entire distributions/populations, and (2) MFM can take into account additional interactions of cells/particles. This differs from CellOT which does not condition on entire distributions, does not learn embeddings for entire populations, and does not take into account interactions of cells. We also refer to the paragraph starting on line 57 where we discuss the difference between MFM and existing methods (including CellOT). On line 63-66, we state one of the limitations of these existing methods in that they are restricted to operating on a single measure.

Since ICNN is conceptually the closest competitor to MFM, why did the authors not benchmark their method on the dataset consider in that manuscript? 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?

We thank the reviewer for bringing up this important related work. The CellOT datasets are not suitable for the MFM task/objective as there is no opportunity to generalize between initial distributions. The datasets in CellOT have a single control (initial) distribution with multiple treatment conditions. In contrast, the organoid drug-screen dataset considered in our work contains many distribution/population pairs $(p_0, p_1)$ -- specifically, after pre-processing, we have 927 control and treated distribution pairs $(p_0, p_1)$. The problem addressed by MFM is analogous to distributional regression, where each data point is an entire distribution, with the task to generalize across unseen distributions. This problem is ill-posed if there are not sufficient quantity of pairs $(p_0, p_1)$, hence rendering the datasets from CellOT unsuitable.

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)

This is a great suggestion. We have added the two additional baselines suggested by the reviewer. Namely, we denote $d(p_0, p_1)$ for comparison of the unperturbed source distribution $p_0$ with the base target distribution $p_1$. We use $d(p_0 - \mu_0 + \tilde{\mu}_1, p_1)$ to denote a comparison between the base target distribution with the source distribution shifted by the difference in means between the perturbed cells $\tilde{\mu_1}$ and the unperturbed cells $\mu_0$. We compute $\tilde{\mu_1}$ by taking the average of all $p_1$ population means in the training set. In Table 2 and Table 3, we see that all models perform better than these simple baselines.

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.

We thank the reviewer for pointing this. We agree this is perhaps too strong of a statement and have removed this statement from the text.

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? Could you plot some of the predicted distributions just as they did in Figure 3?

This is an excellent question. To draw insights towards downstream conclusions on the organoid drug-screen dataset, we compute and plot 2-D projections of embeddings outputted from the population embedding model $\varphi(p_0; \theta)$ (see Figure 7 in Appendix F.1). We observe that we recover grouping of patient populations which are generally consistent with the findings of (Zapatero et al. 2023). For instance, populations from patients 99 and 109 (both chemorefractory) are grouped together, and populations from patients 23 and 27 (both chemosensitive) show lower pairwise distances. We discuss these results on lines 1200-1205 in Appendix F.1.

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?

We thank the reviewer for bringing up this point. We agree and have updated the results for the patient split to include evaluation across left-out populations from 3 different patients. We report the mean and standard deviation of performance metrics across these 3 splits (please see updated Table 3). We observe that MFM outperforms baselines in this setting while also exhibiting robustness across 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.

For the replica split (Table 2) uncertainties over 3 independent model seeds (please refer to the last sentence of the paragraph on lines 365-373). Specifically, we use this to show that models are robust to changes in model initialization. Following the helpful suggestion of the reviewer, for the patients split (Table 3) we report mean and standard deviation over 3 different left-out patient splits to show performance robustness across different patients.

Independent matching may be suboptimal for all considered methods: Tong et al (2023) 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.

We thank the reviewer for their valuable point. We have added experiments that incorporate optimal transport (OT) couplings between source and target distributions for the synthetic letters dataset (Table 4) and for the organoid drug-screen dataset (Table 2).

We observe that using OT couplings generally improves performance across all methods, with MFM or MFM-OT still yielding the best performance on the left-out test data. We note that adding minibatch OT couplings alone does not provide a way of accounting for generalizing across source distributions.

We thank the reviewer again for this comment. We believe these additional experiments help clarify that it is not the use of optimal transport that allows generalization across initial populations, but the population embedding in MFM.

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.

In short, we report results for both. Namely, all the methods are conditioned on the treatment, then FM ignores the base distribution/patient identity (the former in your question), but CFM (as we indicate via CGFM) additionally conditions on the population identity conditions (the latter in your question). Each patient has numerous $(p_0, p_1)$ population pairs. We have updated lines 483-485 and Appendix D.2 to clarify this. For CGFM, we condition on these population identities as known conditions. Population identity conditions for CGFM are represented as one-hot vectors, hence the model cannot generalize to unseen population identities. We state this on lines 426-427.

We thank the reviewer for their detailed feedback and constructive comments, which gave us an opportunity to improve the clarity of our work significantly. We are pleased to see that the reviewer found our work intriguing and intuitively meaningful. We now address key clarification questions raised by the reviewer.

First and foremost, the article is written in a very confusing way.

Thank you for bringing this up, we incorporated the changes you suggested into the manuscript.

It is unclear what a Wasserstein manifold is to start with.

Thank you for the suggestion, we updated the manuscript correspondingly (see Appendix A, lines 157-161). Note that the Wasserstein manifold is "the Riemannian interpretation of the Wasserstein distance developed by Otto" (Ambrosio et al., page 168) and is a well-defined and a well-studied abstract formalism. Indeed, the metric of the Wasserstein manifold is the Wasserstein metric, and the tangent space is defined by the gradient flows as we discuss in lines 157-161. Figure 2 is just an illustration rather than a precise depiction of the Wasserstein manifold. Note that we clearly state in the caption and describe in the text that this is the space of distributions $\mathcal{P}_2(X)$. We are confident that under no circumstances the reader might think that a sphere serves as a precise depiction of the infinite-dimensional space of distributions.

At least an appendix with all definitions should be added.

Thank you for your suggestion, we added Appendix A including the definitions of the concepts we are using in the paper. We believe this increases the clarity of our method.

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.

Thank you for raising this concern. We respectfully disagree that the appendix shows that the model is not particularly outperforming. We believe this misunderstanding maybe due to the inclusion of evaluation on the train data for reference in the appendix. We note that we primarily care about model performance on the test data and not the train data in this work. To clarify this, we have added bolding to the best performing method in the appendix tables to further illustrate that our model outperforms baseline methods. We hope this clarifies that MFM outperforms baselines in all settings on the test data.

We thank the reviewer for their valuable feedback and great questions. We believe through incorporating the feedback provided by the reviewer, we have improved the clarity and quality of our manuscript. We hope that our rebuttal fully addresses all the salient points raised by the reviewer and we kindly ask the reviewer to potentially upgrade their score if the reviewer is satisfied with our responses and updated manuscript. We are also more than happy to answer any further questions that arise.

We thank the reviewer for their time and positive appraisal of our work. We are thrilled that the reviewer viewed our work to be "very well written and easy to understand" and "provides a solid theoretical basis for the proposed method". We now provide responses to the main questions raised by the reviewer.

Figure 1.c is not that helpful in my opinion.

We are happy to clarify and amend the visual illustrations in our manuscript. In regards to Figure 1, we have updated the caption to try and clarify the advantages of MFM and how it differs from current approaches. We are open to suggestions on how to improve the visual depiction in this figure to further reinforce the ideas in our work. In regards to Figure 2-c, in the general response, we provide a brief clarification regarding our illustration of the Wasserstein manifold and how it is considered in MFM. We are happy to incorporate any suggestions the reviewer may have to further improve and simplify visual depictions of our framework and method.

In Eq.17, shouldn't the condition $c$ be part of $\phi$ or $v_t$ as well?

We thank the reviewer for pointing to this important detail. In Theorem 1, $c$ corresponds to the ideal condition on the population that we’re trying to learn. Eq. 17 defines an objective function used to optimize a $v(\cdot; \varphi(p_0; \theta) \omega)$, where $\varphi(p_0; \theta)$ is a population embedding model that learns to represent the entire initial population, approximating the ideal condition. Hence, Eq. 17 defines, in a sense, an unconditional MFM objective in so far that it does not contain any known conditions (such as treatments in our biological setting). From here it is trivial to extend Eq. 17 to the conditional setting with any amount of known conditions. For instance, we show this in Algorithm 1 and Algorithm 2, where we use $c^i$ to denote a known treatment condition for population $i$. The condition $c^i$ and the embedding of the population play different roles, e.g. the condition $c^i$ carries the information about the treatment, while the embedding has the information about the patient through the population of cells.

I am wondering in what cases I should not use MFM to model the change of a distribution using a flow?

Excellent question. MFM relies on there existing a learnable relationship between the source and target distributions. If there is no dependence between the initial and target distributions we would not expect MFM to work well. For instance, we do not expect MFM to be anyhow helpful for the classical setting of generative modelling because there is no any dependence between a given empirical distribution (e.g. natural images) and the Gaussian prior.

We thank the reviewer again for their valuable feedback and great questions. We have implemented the reviewer's suggestions and hope that our rebuttal addresses their questions. We are also more than happy to answer any further questions that arise.

Thanks for your rebuttal.

I increased my Presentation and Contribution Score after the rebuttal because the manuscript became better in that sense, and I strongly believe this is an important contribution to learning dynamics. I am sure many people in my research field will build upon this, and I hope the other reviewers will see it similarly.

In the remainder of this general response we address two points: (1) an overview of the new experiments we ran to address shared questions raised in the reviews; (2) an overview of changes to the manuscript to answer clarifying questions. Please refer to the updated manuscript for all changes mentioned in the general response and responses to individual reviewers. Changes are presented as blue text in the updated manuscript.

1. Experiments:
   • Reviewer Y1RT raised questions regarding the empirical experiments and provided meaningful suggestions to strengthen our results. We incorporated all suggestions from reviewer Y1RT, in turn improving the overall quality of our empirical results. Firstly, we added patient split experiments over 3 independent patient splits (Table 3 in the updated manuscript) to show MFM that is robust across different settings with differing left-out patient populations. Secondly, we added experiments with optimal transport (OT) couplings between samples of source and target distributions for a replica split (updated Table 2). This addition helped improve the overall performance of MFM while also increasing the strength of comparison to stronger baselines, i.e. FM-OT, CGFM-OT. Lastly, for experiments on the organoid drug-screen dataset, we added trivial baselines to demonstrate that the models learn meaningful population dynamics beyond trivial biological phenomena.
   • Reviewer kKjh asked about analyzing the population embeddings to verify that our population embedding model learns valuable representations. We added a section to the appendix to confirm this (Appendix F.1, Figure 7). Through this, we also addressed reviewer Y1RT's question regarding using our method to draw meaningful insights in the biological dataset.
2. Clarifications: Reviewers had clarifying questions which we addressed through the individual responses, and in some cases adding clarification to the manuscript. Below we outline and clarify some shared questions asked by reviewers.
   • Reviewer's Y1RT and kKjh asked clarifying questions regarding comparisons with Bunne et al. 2023 [1]. We note that in our work we indeed do compare with the method in [1]. In the individual responses, we clarified how MFM differs from the method in [1] and how datasets in [1] are not suitable for the setting which MFM is addressing. We made brief changes in the text to further clarify these points.
   • Reviewer Kj6P and jcAc asked clarifying questions pertaining the theoretical ideas presented in our work. Kj6P asked about our use of the Wasserstein manifold and the respective theory surrounding it. To help with understanding and improve clarity, we provided a clarifying explanation in the individual response, briefly clarified some notation in section 2.3, and added formal definitions in Appendix A of the concepts we are using in the paper. We note that our depiction of the Wasserstein manifold in Figure 2 is only an illustration. We use this figure to illustrate the notion that MFM learns to integrate vector fields on $\mathcal{P}_2(X)$, and because of this, can generalize to unseen populations.

We once again thank all the reviewers for their valuable time and insightful feedback. We believe that through the meaningful feedback and suggestions raised by the reviewers, we have improved the clarity, impact, and significance of our work. Through this, we believe that we have addressed all the concerns and questions posed by the reviewers. If the reviewers agree, we hope that the they will consider increasing their scores.

[1] Bunne, Charlotte, et al. "Learning single-cell perturbation responses using neural optimal transport." Nature methods 20.11 (2023): 1759-1768.