Meta Flow Matching: Integrating Vector Fields on the Wasserstein Manifold
Authors: Lazar Atanackovic, Xi (Nicole) Zhang, Brandon Amos, Mathieu Blanchette, Leo J Lee, Yoshua Bengio, Alexander Tong, Kirill Neklyudov
ICLR 2025 | Venue PDF | Archive PDF | Plain Text | LLM Run Details
| Reproducibility Variable | Result | LLM Response |
|---|---|---|
| Research Type | Experimental | 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. ... Next, we explore how MFM can be applied to model single-cell perturbation data... 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... 5 EXPERIMENTS To show the effectiveness of MFM to generalize under previously unseen populations for the population prediction task, we consider two experimental settings. (i) A synthetic experiment with well defined coupled populations, and (ii) experiments on a publicly available single-cell dataset... |
| Researcher Affiliation | Collaboration | 1University of Toronto 2Vector Institute 3Mila Quebec AI Institute 4Mc Gill University 5Meta 6Université de Montréal 7CIFAR Fellow |
| Pseudocode | Yes | Algorithm 1: Meta Flow Matching (training) ... Algorithm 2: Meta Flow Matching (sampling) ... We show pseudo-code for training in Algorithm 1 and sampling in Algorithm 2. |
| Open Source Code | Yes | Our code is available at: https://github.com/lazaratan/meta-flow-matching ... To ensure reproducibility of our findings and results we submitted our source code as supplementary materials. |
| Open Datasets | Yes | We evaluate MFM on predicting the response of patient-derived cells to chemotherapy treatments in a recently published large scale single-cell drug screen dataset where there are known to be patient-specific responses (Ramos Zapatero et al., 2023). ... Organoid drug-screen data. For experiments on biological data, we use the organoid drug-screen dataset from Ramos Zapatero et al. (2023). |
| Dataset Splits | Yes | We construct two data splits for the organoid drugscreen dataset (see Fig. 4-right). Replicate split; here we consider 2 settings, (i) leaveout replicates evenly across all patients for testing (replica-1, depicted in Fig. 4 see Table 5 for results on this split), and (ii) leaving out a batch of replicates from a single patient while including the remaining replicas from the same patient for training (replica-2, results in Table 2). The second setting is the Patients split (results in Table 3); here we leave-out replicates fully in one patients in this setting, we are testing the ability of the model to generalize population prediction of treatment response for unseen patients. We do this for 3 different patients and report results across these independent splits. Further details regarding the organoid drug-screen dataset, data pre-processing, and data splits are provided in Appendix D.2. Baselines. ... In the replica-1 split, we use 713 populations for training, 111 left-out population for validation, and 103 left-out populations for testing. In the replica-2 split, we use 861 populations for training, 33 left-out populations for validation, and 33 left-out populations for testing. ... For the patients split, we consider 3 different patient splits where we independently leave out all populations from either patient 21, patient 27, and patients 75. ... In the patients splits, we only split data into train and test sets, resulting in 839 population pairs for training and 88 population pairs for test evaluation for PDO-21 and PDO-27 splits. For PDO-75, the train split contain 830 train population pairs and 97 test population pairs. |
| Hardware Specification | Yes | All experiments were conducted on a HPC cluster primarily on NVIDIA Tesla T4 16GB and A40 48GB GPUs. |
| Software Dependencies | No | We implement all our experiments using Py Torch and Py Torch Geometric. The paper mentions software names (Py Torch, Py Torch Geometric) but does not provide specific version numbers for these dependencies. |
| Experiment Setup | Yes | All vector field models vt 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. ... 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. ... 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. ... For the syntehtic experiment, FM and CGFM model were trained for 12000 epochs, while MFM models were trained for 24000 epochs, with a population batch size of 10 and a sample batch size of 700. |