# Wasserstein Flow Matching: Generative modeling over families of distributions

arXiv: https://arxiv.org/abs/2411.00698 · HTML: https://mlsanity.com/paper/2411.00698/

## Key facts

- Cited by 28 papers in the map (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 46 of its 80 references that are in the map.
- First posted on arXiv 1 Nov 2024.

## Placed in

- Diffusion models / [Flow matching, time series and science](https://mlsanity.com/diffusion/flow-matching-time-series-and-science/)

## Builds on (most cited first, 30 of 46)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [ShapeNet: An Information-Rich 3D Model Repository](https://mlsanity.com/paper/1512.03012/) (2015)
- [Sparks of Artificial General Intelligence: Early experiments with GPT-4](https://mlsanity.com/paper/2303.12712/) (2023)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [Computational Optimal Transport](https://mlsanity.com/paper/1803.00567/) (2018)
- [3D Shape Generation and Completion through Point-Voxel Diffusion](https://mlsanity.com/paper/2104.03670/) (2021)
- [PointFlow: 3D Point Cloud Generation with Continuous Normalizing Flows](https://mlsanity.com/paper/1906.12320/) (2019)
- [Point-Voxel CNN for Efficient 3D Deep Learning](https://mlsanity.com/paper/1907.03739/) (2019)
- [Multisample Flow Matching: Straightening Flows with Minibatch Couplings](https://mlsanity.com/paper/2304.14772/) (2023)
- [Learning Gradient Fields for Shape Generation](https://mlsanity.com/paper/2008.06520/) (2020)
- [Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration](https://mlsanity.com/paper/1705.09634/) (2017)
- [A Survey on Video Diffusion Models](https://mlsanity.com/paper/2310.10647/) (2023)
- [Audiobox: Unified Audio Generation with Natural Language Prompts](https://mlsanity.com/paper/2312.15821/) (2023)
- [Dirichlet Flow Matching with Applications to DNA Sequence Design](https://mlsanity.com/paper/2402.05841/) (2024)
- [OSMnx: New Methods for Acquiring, Constructing, Analyzing, and Visualizing Complex Street Networks](https://mlsanity.com/paper/1611.01890/) (2016)
- [Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein](https://mlsanity.com/paper/2201.12324/) (2022)
- [SoftFlow: Probabilistic Framework for Normalizing Flow on Manifolds](https://mlsanity.com/paper/2006.04604/) (2020)
- [Multiresolution Tree Networks for 3D Point Cloud Processing](https://mlsanity.com/paper/1807.03520/) (2018)
- [Entropic estimation of optimal transport maps](https://mlsanity.com/paper/2109.12004/) (2021)
- [Discrete Point Flow Networks for Efficient Point Cloud Generation](https://mlsanity.com/paper/2007.10170/) (2020)
- [Statistical bounds for entropic optimal transport: sample complexity and the central limit theorem](https://mlsanity.com/paper/1905.11882/) (2019)
- [Variational inference via Wasserstein gradient flows](https://mlsanity.com/paper/2205.15902/) (2022)
- [An elementary introduction to information geometry](https://mlsanity.com/paper/1808.08271/) (2018)

## Built on by (most cited first, 28 of 28)

- [Curly Flow Matching for Learning Non-gradient Field Dynamics](https://mlsanity.com/paper/2510.26645/) (2025)
- [Flow Matching Meets Biology and Life Science: A Survey](https://mlsanity.com/paper/2507.17731/) (2025)
- [Learning of Population Dynamics: Inverse Optimization Meets JKO Scheme](https://mlsanity.com/paper/2506.01502/) (2025)
- [Bures-Wasserstein Flow Matching for Graph Generation](https://mlsanity.com/paper/2506.14020/) (2025)
- [Better Source, Better Flow: Learning Condition-Dependent Source Distribution for Flow Matching](https://mlsanity.com/paper/2602.05951/) (2026)
- [Modeling Microenvironment Trajectories on Spatial Transcriptomics with NicheFlow](https://mlsanity.com/paper/2511.00977/) (2025)
- [Towards High-Order Mean Flow Generative Models: Feasibility, Expressivity, and Provably Efficient Criteria](https://mlsanity.com/paper/2508.07102/) (2025)
- [Slicing Wasserstein Over Wasserstein Via Functional Optimal Transport](https://mlsanity.com/paper/2509.22138/) (2025)
- [Generative Distribution Embeddings: Lifting autoencoders to the space of distributions for multiscale representation learning](https://mlsanity.com/paper/2505.18150/) (2025)
- [Modeling Stochastic Conditional Dynamics from Sparse Observations via Kernel-Stabilized Flow Matching](https://mlsanity.com/paper/2411.08314/) (2024)
- [MFM-point: Multi-scale Flow Matching for Point Cloud Generation](https://mlsanity.com/paper/2511.20041/) (2025)
- [Generalized Wasserstein Flow Matching: Transport Plans, Everywhere, All at Once](https://mlsanity.com/paper/2605.08424/) (2026)
- [Dense associative memory for Gaussian distributions](https://mlsanity.com/paper/2509.23162/) (2025)
- [Learning Discrete Diffusion of Graphs via Free-Energy Gradient Flows](https://mlsanity.com/paper/2604.11311/) (2026)
- [Measure-to-measure Regression with Transformers](https://mlsanity.com/paper/2605.28075/) (2026)
- [Busemann Functions in the Wasserstein Space: Existence, Closed-Forms, and Applications to Slicing](https://mlsanity.com/paper/2510.04579/) (2025)
- [Entropy Across the Bridge: Conditional-Marginal Discretization for Flow and Schr"odinger Samplers](https://mlsanity.com/paper/2605.16126/) (2026)
- [Trajectory inference via Acceleration Matching](https://mlsanity.com/paper/2608.03916/) (2026)
- [MODE: Learning compositional representations of complex systems with Mixtures Of Dynamical Experts](https://mlsanity.com/paper/2510.09594/) (2025)
- [Tensor Gauge Flow Models](https://mlsanity.com/paper/2511.17616/) (2025)
- [Inverse Problems Conditioned on Observation Ensembles: Applications and Methods](https://mlsanity.com/paper/2601.22029/) (2026)
- [EvoEGF-Mol: Evolving Exponential Geodesic Flow for Structure-based Drug Design](https://mlsanity.com/paper/2601.22466/) (2026)
- [An Explicit Surrogate for Gaussian Mixture Flow Matching with Wasserstein Gap Bounds](https://mlsanity.com/paper/2603.28992/) (2026)
- [ITSPACE: Monotone Gaussian Optimal Transport Updates](https://mlsanity.com/paper/2606.30523/) (2026)
- [Gromov-Monge Flow Matching for Equivariant Graph Generation](https://mlsanity.com/paper/2608.26961/) (2026)
- [When Riemann flows with Wasserstein: Generative Modeling of Probability Distributions on Manifolds](https://mlsanity.com/paper/2609.25659/) (2026)
- [GLASS: Global Latent Aggregation with Slot-based Set Decoding for Scalable All-Atom Crystal Generation](https://mlsanity.com/paper/2609.37158/) (2026)
- [Finsler Flow Matching: Dynamics-Aware Geodesic Interpolation for Single-Snapshot Trajectory Inference](https://mlsanity.com/paper/2610.11318/) (2026)

## Abstract (arXiv)

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 .