Aram-Alexandre Pooladian
Affiliation
New York University
Talk Title
Wasserstein Flow Matching: Generative modeling over families of distributions
Abstract
TThe task of generative modeling typically concerns the transport of a single source distribution to a single target distribution on the basis of samples. Transport can take the form of an optimal transport map, but recent work has shown that it is more tractable to learn simple probability flows through what is known as "flow matching", a framework which suggests to perform regression onto simple probability flows. In this work, we study the task of learning flows between families of distributions (i.e., many source measures to many target measures). Examples include families of means-covariances (e.g., Gaussians), or families of point-clouds --- in both of these examples, the geometry of the data is relevant to the generative task, and we wish to preserve this along the flows. We introduce Wasserstein flow matching (WFM), which appropriately lifts flow matching onto families of distributions by appealing to the Riemannian nature of the Wasserstein geometry. Our algorithm leverages theoretical and computational advances in entropic optimal transport, as well as the attention mechanism in our neural network architecture. As applications, we demonstrate how to generate representations of granular cell states from single-cell genomics data (via Bures--Wasserstein FM) and synthesize cellular microenvironments from spatial transcriptomics datasets. Code is available at [WassersteinFlowMatching]. This is joint work with Doron Haviv, Brandon Amos, and Dana Pe'er.
Bio
Aram-Alexandre Pooladian is a PhD candidate at NYU advised by Jonathan Niles-Weed. He is interested in the development of scalable algorithms for probabilistic inference, as well as theoretical and statistical questions stemming from the theory of optimal transport. Starting Fall 2025, he will be a Foundations of Data Science postdoc at Yale University.