Wasserstein Flow Matching: Generative modeling over families of distributions
This addresses the problem of generating distributions as samples in fields like graphics and genomics, where standard methods ignore distributional geometry, representing a novel method for a known bottleneck.
The paper tackles generative modeling over families of distributions, such as in computer graphics and single-cell genomics, by proposing Wasserstein flow matching (WFM), which lifts flow matching onto distributions using Wasserstein geometry, enabling generation of distributions in high dimensions for shapes and cellular microenvironments.
Generative modeling typically concerns transporting a single source distribution to a target distribution via simple probability flows. However, in fields like computer graphics and single-cell genomics, samples themselves can be viewed as distributions, where standard flow matching ignores their inherent geometry. We propose Wasserstein flow matching (WFM), which lifts flow matching onto families of distributions using the Wasserstein geometry. Notably, WFM is the first algorithm 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 transport theory and the attention mechanism, demonstrating 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 available at https://github.com/DoronHav/WassersteinFlowMatching .