Meta Flow Matching: Integrazione di Campi Vettoriali sulla Varietà di Wasserstein

Meta Flow Matching: Integrating Vector Fields on the Wasserstein Manifold

August 26, 2024

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

cs.AI

Abstract

Numerosi processi biologici e fisici possono essere modellati come sistemi di entità interagenti che evolvono continuamente nel tempo, ad esempio le dinamiche di cellule che comunicano o particelle fisiche. Apprendere le dinamiche di tali sistemi è essenziale per prevedere l'evoluzione temporale di popolazioni in nuovi campioni e ambienti non osservati. I modelli basati sul flusso (flow-based) consentono di apprendere queste dinamiche a livello di popolazione: essi modellano l'evoluzione dell'intera distribuzione dei campioni. Tuttavia, gli attuali modelli basati sul flusso sono limitati a una singola popolazione iniziale e a un insieme di condizioni predefinite che descrivono dinamiche diverse. Sosteniamo che molteplici processi nelle scienze naturali debbano essere rappresentati come campi vettoriali sulla varietà di Wasserstein delle densità di probabilità. Cioè, la variazione della popolazione in ogni momento dipende dalla popolazione stessa a causa delle interazioni tra i campioni. Ciò è cruciale in particolare per la medicina personalizzata, dove lo sviluppo delle malattie e la risposta rispettiva al trattamento dipendono dal microambiente cellulare specifico di ogni paziente. Proponiamo il Meta Flow Matching (MFM), un approccio pratico per integrare lungo questi campi vettoriali sulla varietà di Wasserstein ammortizzando il modello di flusso sulle popolazioni iniziali. Nello specifico, incorporiamo la popolazione di campioni utilizzando una Graph Neural Network (GNN) e utilizziamo questi incorporamenti per addestrare un modello Flow Matching. Ciò conferisce a MFM la capacità di generalizzare sulle distribuzioni iniziali, a differenza dei metodi proposti in precedenza. Dimostriamo la capacità di MFM di migliorare la previsione delle risposte individuali ai trattamenti su un ampio dataset di screening farmacologico su singola cellula multi-paziente.

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.