Meta Flow Matching : Intégration de champs vectoriels sur la variété de Wasserstein

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

De nombreux processus biologiques et physiques peuvent être modélisés comme des systèmes d'entités en interaction évoluant de manière continue dans le temps, par exemple la dynamique de cellules en communication ou de particules physiques. Apprendre la dynamique de tels systèmes est essentiel pour prédire l'évolution temporelle des populations dans de nouveaux échantillons et environnements non observés. Les modèles à base de flux permettent d'apprendre ces dynamiques au niveau de la population - ils modélisent l'évolution de la distribution entière des échantillons. Cependant, les modèles actuels à base de flux se limitent à une population initiale unique et à un ensemble de conditions prédéfinies qui décrivent différentes dynamiques. Nous soutenons que de multiples processus dans les sciences naturelles doivent être représentés comme des champs de vecteurs sur la variété de Wasserstein des densités de probabilité. Autrement dit, le changement de la population à tout instant dépend de la population elle-même en raison des interactions entre les échantillons. Ceci est particulièrement crucial pour la médecine personnalisée où le développement des maladies et leur réponse respective au traitement dépendent du microenvironnement cellulaire spécifique à chaque patient. Nous proposons Meta Flow Matching (MFM), une approche pratique pour intégrer le long de ces champs de vecteurs sur la variété de Wasserstein en amortissant le modèle de flux sur les populations initiales. À savoir, nous intégrons la population d'échantillons à l'aide d'un Réseau de Neurones Graphiques (GNN) et utilisons ces intégrations pour entraîner un modèle Flow Matching. Cela confère à MFM la capacité de généraliser sur les distributions initiales, contrairement aux méthodes proposées précédemment. Nous démontrons la capacité de MFM à améliorer la prédiction des réponses individuelles aux traitements sur un jeu de données à grande échelle de criblage médicamenteux unicellulaire multi-patients.

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.