Meta Flow Matching: Integratie van Vectorvelden op de Wasserstein-variëteit

Meta Flow Matching: Integrating Vector Fields on the Wasserstein Manifold

August 26, 2024

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

cs.AI

Samenvatting

Talloze biologische en fysische processen kunnen worden gemodelleerd als systemen van interagerende entiteiten die continu in de tijd evolueren, zoals de dynamiek van communicerende cellen of fysische deeltjes. Het leren van de dynamiek van dergelijke systemen is essentieel voor het voorspellen van de temporele evolutie van populaties in nieuwe samples en onbekende omgevingen. Op stromingen gebaseerde modellen maken het mogelijk deze dynamiek op populatieniveau te leren - zij modelleren de evolutie van de volledige verdeling van samples. Huidige op stromingen gebaseerde modellen zijn echter beperkt tot een enkele initiële populatie en een set vooraf gedefinieerde condities die verschillende dynamieken beschrijven. Wij stellen dat meerdere processen in de natuurwetenschappen moeten worden voorgesteld als vectorvelden op de Wasserstein-variëteit van kansdichtheden. Dat wil zeggen, de verandering van de populatie op elk moment hangt af van de populatie zelf vanwege de interacties tussen samples. Dit is met name cruciaal voor gepersonaliseerde geneeskunde, waar de ontwikkeling van ziekten en hun respectievelijke behandelrespons afhangt van de micro-omgeving van cellen die specifiek is voor elke patiënt. Wij stellen Meta Flow Matching (MFM) voor, een praktische benadering om langs deze vectorvelden op de Wasserstein-variëteit te integreren door het stroommodel te amortiseren over de initiële populaties. Concreet embedden we de populatie van samples met behulp van een Graph Neural Network (GNN) en gebruiken we deze embeddings om een Flow Matching-model te trainen. Hierdoor kan MFM, in tegenstelling tot eerder voorgestelde methoden, generaliseren over de initiële verdelingen. Wij demonstreren het vermogen van MFM om de voorspelling van individuele behandelresponsen te verbeteren op een grootschalige multi-patiënt single-cell drug screen dataset.

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.