Мета Flow Matching: Интеграция векторных полей на многообразии Вассерштейна
Meta Flow Matching: Integrating Vector Fields on the Wasserstein Manifold
August 26, 2024
Авторы: Lazar Atanackovic, Xi Zhang, Brandon Amos, Mathieu Blanchette, Leo J. Lee, Yoshua Bengio, Alexander Tong, Kirill Neklyudov
cs.AI
Аннотация
Многочисленные биологические и физические процессы могут быть смоделированы как системы взаимодействующих сущностей, эволюционирующих непрерывно во времени, например, динамика взаимодействующих клеток или физических частиц. Изучение динамики таких систем крайне важно для прогнозирования временной эволюции популяций в новых образцах и неизученных условиях. Модели, основанные на потоках, позволяют изучать эту динамику на уровне популяции — они моделируют эволюцию всего распределения образцов. Однако современные модели, основанные на потоках, ограничены одной начальной популяцией и набором предопределенных условий, которые описывают различные динамики. Мы утверждаем, что множество процессов в естественных науках должны быть представлены как векторные поля на многообразии Вассерштейна вероятностных плотностей. То есть изменение популяции в любой момент времени зависит от самой популяции из-за взаимодействий между образцами. Это особенно важно для персонализированной медицины, где развитие заболеваний и их ответ на лечение зависят от микроокружения клеток, специфичного для каждого пациента. Мы предлагаем Meta Flow Matching (MFM) — практический подход к интегрированию вдоль этих векторных полей на многообразии Вассерштейна путем амортизации модели потока по начальным популяциям. А именно, мы используем графовые нейронные сети (GNN) для встраивания популяции образцов и применяем эти вложения для обучения модели Flow Matching. Это позволяет MFM обобщать начальные распределения, в отличие от ранее предложенных методов. Мы демонстрируем способность MFM улучшать прогнозирование индивидуальных ответов на лечение на крупномасштабном наборе данных одноклеточного скрининга лекарств для множества пациентов.
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.