Meta Flow Matching: Integration von Vektorfeldern auf der Wasserstein-Mannigfaltigkeit

Meta Flow Matching: Integrating Vector Fields on the Wasserstein Manifold

August 26, 2024

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

cs.AI

papers.abstract

Zahlreiche biologische und physikalische Prozesse können als Systeme interagierender Entitäten modelliert werden, die sich kontinuierlich über die Zeit entwickeln, z. B. die Dynamik kommunizierender Zellen oder physikalischer Teilchen. Das Erlernen der Dynamik solcher Systeme ist entscheidend, um die zeitliche Entwicklung von Populationen über neue Proben und unbekannte Umgebungen hinweg vorherzusagen. Flow-basierte Modelle ermöglichen es, diese Dynamik auf Populationsebene zu lernen – sie modellieren die Entwicklung der gesamten Verteilung der Proben. Allerdings sind aktuelle Flow-basierte Modelle auf eine einzige Ausgangspopulation und eine Reihe vordefinierter Bedingungen beschränkt, die unterschiedliche Dynamiken beschreiben. Wir argumentieren, dass multiple Prozesse in den Naturwissenschaften als Vektorfelder auf der Wasserstein-Mannigfaltigkeit von Wahrscheinlichkeitsdichten dargestellt werden müssen. Das bedeutet, dass die Veränderung der Population zu jedem Zeitpunkt von der Population selbst abhängt, bedingt durch die Interaktionen zwischen den Proben. Dies ist insbesondere für die personalisierte Medizin entscheidend, wo die Entwicklung von Krankheiten und deren jeweilige Behandlungsreaktion von der zellspezifischen Mikroumgebung jedes Patienten abhängt. Wir schlagen Meta Flow Matching (MFM) vor, einen praktischen Ansatz zur Integration entlang dieser Vektorfelder auf der Wasserstein-Mannigfaltigkeit, indem das Flow-Modell über die Ausgangspopulationen amortisiert wird. Konkret betten wir die Population der Proben mit einem Graph Neural Network (GNN) ein und nutzen diese Einbettungen, um ein Flow-Matching-Modell zu trainieren. Dies verleiht MFM die Fähigkeit, über die Ausgangsverteilungen zu generalisieren, im Gegensatz zu bisher vorgeschlagenen Methoden. Wir demonstrieren die Fähigkeit von MFM, die Vorhersage individueller Behandlungsreaktionen auf einem groß angelegten Multi-Patienten-Einzellzell-Drug-Screen-Datensatz zu verbessern.

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.