Meta Flow Matching: Integración de Campos Vectoriales en la Variedad 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

Resumen

Numerosos procesos biológicos y físicos pueden modelarse como sistemas de entidades interactuantes que evolucionan continuamente en el tiempo, por ejemplo, la dinámica de células en comunicación o partículas físicas. Aprender la dinámica de tales sistemas es esencial para predecir la evolución temporal de poblaciones en muestras novedosas y entornos no vistos. Los modelos basados en flujo permiten aprender estas dinámicas a nivel poblacional: modelan la evolución de la distribución completa de muestras. Sin embargo, los modelos basados en flujo actuales están limitados a una única población inicial y un conjunto de condiciones predefinidas que describen diferentes dinámicas. Argumentamos que múltiples procesos en las ciencias naturales deben representarse como campos vectoriales en la variedad de Wasserstein de densidades de probabilidad. Es decir, el cambio de la población en cualquier momento depende de la población misma debido a las interacciones entre las muestras. En particular, esto es crucial para la medicina personalizada, donde el desarrollo de enfermedades y su respectiva respuesta al tratamiento dependen del microentorno de células específico de cada paciente. Proponemos Meta Flow Matching (MFM), un enfoque práctico para integrar a lo largo de estos campos vectoriales en la variedad de Wasserstein mediante la amortización del modelo de flujo sobre las poblaciones iniciales. Específicamente, incrustamos la población de muestras utilizando una Red Neuronal de Grafos (GNN) y usamos estas incrustaciones para entrenar un modelo de Flow Matching. Esto otorga a MFM la capacidad de generalizar sobre las distribuciones iniciales, a diferencia de métodos propuestos previamente. Demostramos la capacidad de MFM para mejorar la predicción de respuestas individuales al tratamiento en un conjunto de datos a gran escala de cribado de fármacos a nivel de célula única en múltiples pacientes.

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.