Meta Flow Matching: Integração de Campos Vetoriais na Variedade de Wasserstein

Meta Flow Matching: Integrating Vector Fields on the Wasserstein Manifold

August 26, 2024

Autores: Lazar Atanackovic, Xi Zhang, Brandon Amos, Mathieu Blanchette, Leo J. Lee, Yoshua Bengio, Alexander Tong, Kirill Neklyudov

cs.AI

Resumo

Numerosos processos biológicos e físicos podem ser modelados como sistemas de entidades interagentes que evoluem continuamente ao longo do tempo, por exemplo, a dinâmica de células em comunicação ou partículas físicas. Aprender a dinâmica de tais sistemas é essencial para prever a evolução temporal de populações em novas amostras e ambientes não observados. Modelos baseados em fluxo permitem aprender essas dinâmicas ao nível populacional - eles modelam a evolução de toda a distribuição de amostras. No entanto, os modelos atuais baseados em fluxo estão limitados a uma única população inicial e a um conjunto de condições predefinidas que descrevem diferentes dinâmicas. Argumentamos que múltiplos processos nas ciências naturais devem ser representados como campos vetoriais na variedade de Wasserstein de densidades de probabilidade. Isto é, a mudança da população em qualquer momento no tempo depende da própria população devido às interações entre as amostras. Em particular, isso é crucial para a medicina personalizada, onde o desenvolvimento de doenças e a sua resposta respectiva ao tratamento dependem do microambiente celular específico de cada paciente. Propomos o *Meta Flow Matching* (MFM), uma abordagem prática para integrar ao longo desses campos vetoriais na variedade de Wasserstein, amortizando o modelo de fluxo sobre as populações iniciais. Nomeadamente, incorporamos a população de amostras usando uma Rede Neural de Grafos (GNN) e usamos essas incorporações para treinar um modelo de *Flow Matching*. Isso confere ao MFM a capacidade de generalizar sobre as distribuições iniciais, ao contrário de métodos propostos anteriormente. Demonstramos a capacidade do MFM para melhorar a previsão de respostas individuais ao tratamento num conjunto de dados de rastreio de fármacos em células únicas de múltiplos pacientes em larga escala.

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.