Wasserstein Flow Matching: Generative modeling over families of distributions

TL;DR

Wasserstein Flow Matching (WFM) introduces a novel approach to generative modeling over families of distributions using Wasserstein geometry, enabling high-dimensional distribution generation with theoretical guarantees and versatile computational methods.

Contribution

WFM is the first algorithm to generate distributions in high dimensions by lifting flow matching onto families of distributions using Wasserstein geometry.

Findings

Successfully generates 2D and 3D shapes.

Generates high-dimensional cellular microenvironments.

Combines optimal transport theory with attention mechanisms.

Abstract

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 .