メタフローマッチング：ワッサースタイン多様体上のベクトル場の統合

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

要旨

数多くの生物学的・物理的プロセスは、時間とともに連続的に進化する相互作用するエンティティのシステムとしてモデル化できます。例えば、細胞間のコミュニケーションや物理粒子のダイナミクスなどです。このようなシステムのダイナミクスを学習することは、新しいサンプルや未見の環境における個体群の時間的進化を予測するために不可欠です。フローベースモデルは、これらのダイナミクスを個体群レベルで学習することを可能にします。つまり、サンプルの分布全体の進化をモデル化します。しかし、現在のフローベースモデルは、単一の初期個体群と異なるダイナミクスを記述する一連の事前定義された条件に限定されています。我々は、自然科学における複数のプロセスが、確率密度のWasserstein多様体上のベクトル場として表現される必要があると主張します。つまり、任意の時点での個体群の変化は、サンプル間の相互作用により、個体群そのものに依存します。特に、これは個別化医療において重要です。なぜなら、疾患の進行やそれに対する治療反応は、各患者に特有の細胞の微小環境に依存するからです。我々は、Meta Flow Matching（MFM）を提案します。これは、初期個体群にわたってフローモデルを償却することで、Wasserstein多様体上のこれらのベクトル場に沿って統合する実用的なアプローチです。具体的には、Graph Neural Network（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.