ICML 2025

Wasserstein Flow Matching: Generative Modeling Over Families of Distributions
Doron Haviv * 1 2 Aram-Alexandre Pooladian * 3 Dana Pe’er 2 4 Brandon Amos 5
Abstract
Generative modeling typically concerns transport-
ing a single source distribution to a target distri-
bution via simple probability flows. However, in
fields like computer graphics and single-cell ge-
nomics, samples themselves can be viewed as dis-
tributions, where standard flow matching ignores
their inherent geometry. We propose Wasserstein
flow matching (WFM), which lifts flow matching
onto families of distributions using the Wasser-
stein geometry. Notably, WFM is the first algo-
rithm capable of generating distributions in high
dimensions, whether represented analytically (as
Gaussians) or empirically (as point-clouds). Our
theoretical analysis establishes that Wasserstein
geodesics constitute proper conditional flows over
the space of distributions, making for a valid FM
objective. Our algorithm leverages optimal trans-
port theory and the attention mechanism, demon-
strating versatility across computational regimes:
exploiting closed-form optimal transport paths for
Gaussian families, while using entropic estimates
on point-clouds for general distributions. WFM
successfully generates both 2D & 3D shapes
and high-dimensional cellular microenvironments
from spatial transcriptomics data. Code is avail-
able at Wasserstein Flow Matching.
1. Introduction
Today’s abundance of data and scalability of training mas-
sive neural networks has made it possible to generate hyper-
realistic images on the basis of training examples (OpenAI,
2022), as well as video and audio clips (Vyas et al., 2023;
Xing et al., 2023), and, of course, text (Bubeck et al., 2023).
All of these are instances of generative modeling: given
*Equal contribution 1Weill Cornell 2Memorial Sloan–Kettering
Cancer Center 3Center for Data Science, New York University
4Howard Hughes Medical Institute 5Meta AI. Correspondence to:
Doron Haviv ...