元流匹配：在Wasserstein流形上整合向量場

摘要

許多生物和物理過程可以被建模為隨時間持續演變的互動實體系統，例如通訊細胞或物理粒子的動態。學習這些系統的動態對於預測在新樣本和未見環境中人口的時間演變至關重要。基於流的模型允許在人口層面學習這些動態-它們模擬樣本整個分佈的演變。然而，目前的基於流的模型僅限於單個初始人口和一組預定義描述不同動態的條件。我們認為自然科學中的多個過程必須被表示為概率密度的Wasserstein流形上的向量場。也就是說，任何時間點人口的變化取決於人口本身，這是由於樣本之間的互動。特別是對於個性化醫學非常重要，疾病的發展及其相應的治療反應取決於每位患者特定的細胞微環境。我們提出了元流匹配（MFM），這是一種實際方法，通過攤銷初始人口上的流模型來整合Wasserstein流形上的這些向量場。換句話說，我們使用圖神經網絡（GNN）嵌入樣本人口，並使用這些嵌入來訓練流匹配模型。這使MFM能夠廣泛應用於初始分佈，不同於先前提出的方法。我們展示了MFM在大規模多患者單細胞藥物篩選數據集上改善個別治療反應預測的能力。

English

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.

元流匹配：在Wasserstein流形上整合向量場

Meta Flow Matching: Integrating Vector Fields on the Wasserstein Manifold

摘要

Support