메타 플로우 매칭: 와서스테인 매니폴드 상의 벡터 필드 통합

초록

수많은 생물학적 및 물리적 과정은 시간에 따라 지속적으로 진화하는 상호작용 개체들의 시스템으로 모델링될 수 있습니다. 예를 들어, 세포 간의 통신이나 물리적 입자의 역학이 이에 해당합니다. 이러한 시스템의 역학을 학습하는 것은 새로운 샘플과 보이지 않는 환경에서의 개체군의 시간적 진화를 예측하는 데 필수적입니다. 플로우 기반 모델은 이러한 역학을 개체군 수준에서 학습할 수 있게 해줍니다. 이 모델은 샘플 전체의 분포 진화를 모델링합니다. 그러나 현재의 플로우 기반 모델은 단일 초기 개체군과 서로 다른 역학을 설명하는 미리 정의된 조건들로 제한되어 있습니다. 우리는 자연과학의 여러 과정이 확률 밀도의 와서스테인 매니폴드(Wasserstein manifold) 상의 벡터 필드로 표현되어야 한다고 주장합니다. 즉, 시간의 어느 순간에서든 개체군의 변화는 샘플 간의 상호작용으로 인해 개체군 자체에 의존합니다. 특히, 이는 개인 맞춤형 의학에서 질병의 발달과 그에 따른 치료 반응이 각 환자에게 특정한 세포의 미세 환경에 의존하는 경우에 중요합니다. 우리는 메타 플로우 매칭(Meta Flow Matching, MFM)을 제안합니다. 이는 초기 개체군에 대해 플로우 모델을 분산시켜 와서스테인 매니폴드 상의 이러한 벡터 필드를 통합하는 실용적인 접근법입니다. 구체적으로, 우리는 그래프 신경망(Graph Neural Network, GNN)을 사용하여 샘플 개체군을 임베딩하고, 이러한 임베딩을 사용하여 플로우 매칭 모델을 학습시킵니다. 이는 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.

메타 플로우 매칭: 와서스테인 매니폴드 상의 벡터 필드 통합

Meta Flow Matching: Integrating Vector Fields on the Wasserstein Manifold

초록

Support