Meta Flow Matching: Integrating Vector Fields on the Wasserstein Manifold

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

Track: Extended abstract

Keywords: Flow matching, Diffusion, Dynamics, Cell dynamics

TL;DR: We propose a new framework, Meta Flow Matching, for integrating vector fields on the Wasserstein manifold of probability densities and show the use case to improve prediction of patient specific treatment responses.

Abstract:

Numerous biological and physical processes can be modeled as systems of interacting samples 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, standard flow-based models are limited to a single initial population and a set of predefined conditions which describe different dynamics.
We propose _Meta Flow Matching_ (MFM), a practical approach to integrating along vector fields on the Wasserstein manifold by amortizing the flow model over the initial populations.
We demonstrate that MFM improves the prediction of individual treatment responses on a large scale multi-patient single-cell drug screen dataset.

Submission Number: 73