Research map / Probabilistic and scientific modeling / Variational and generative sampling
Wasserstein Flow Matching: Generative Modeling Over Families of Distributions
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Lifts Riemannian flow matching to the Wasserstein space of distributions, learning flows between distributions over Gaussians (closed-form Bures-Wasserstein geodesics) and distributions over point-clouds (entropic OT with transformers).
Problem
Generative modeling when each sample is itself a probability distribution (e.g., shapes as point-clouds, cellular niches as Gaussians or point-clouds), where standard flow matching ignores the geometry of distributions and cannot handle variable-sized or high-dimensional point-cloud realizations.
Approach
Apply Riemannian flow matching on the 2-Wasserstein space: regress a neural field onto Wasserstein geodesic (McCann interpolation) velocities between uncoupled source and target distributions of distributions; use closed-form OT/BW geometry for Gaussians and entropic OT maps estimated via Sinkhorn with a transformer for point-clouds.
Key results
- On Gaussian families, BW-FM (Riemannian) achieves lowest average min W2^2 across all 8 datasets, e.g., 1.88e-4 on Cities versus 1.75e-3 for Frobenius FM and 7.26e-3 for BW-FM (Euclidean).
- On ShapeNet 1-NN with Chamfer distance, WFM reports the best CD on Airplane-type/Car comparisons (69.88 airplane, 53.41 car) among compared methods, while EMD is competitive rather than best.
- WFM generates variable-sized point-cloud distributions (MNIST, Letters) and 254-gene MERFISH niches, settings where prior voxel or fixed-size methods were reported as not applicable (NA).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Riemannian BW-FM instead of BW-FM (E) | Gaussian generation · Spiral - 16 (2D) | minimum W2^2 distance: 4.00 \cdot 10^{-4} -> 2.98 \cdot 10^{-4} | Helped |
| Use Riemannian BW-FM instead of BW-FM (E) | Gaussian generation · Spirals - 128 (2D) | minimum W2^2 distance: 1.70 \cdot 10^{-3} -> 1.28 \cdot 10^{-3} | Helped |
| Use Riemannian BW-FM instead of BW-FM (E) | Gaussian generation · Two Moons (2D) | minimum W2^2 distance: 8.96 \cdot 10^{-4} -> 1.84 \cdot 10^{-4} | Helped |
| Use Riemannian BW-FM instead of BW-FM (E) | Gaussian generation · Sphere (3D) | minimum W2^2 distance: 2.14 \cdot 10^{-3} -> 6.65 \cdot 10^{-4} | Helped |
| Use Riemannian BW-FM instead of Frobenius FM | Gaussian generation · Cities (2D) | minimum W2^2 distance: 1.75 \cdot 10^{-3} -> 1.88 \cdot 10^{-4} | Helped |
| Use Riemannian BW-FM instead of BW-FM (E) | Gaussian generation · ECG (15D) | minimum W2^2 distance: 3.26 \cdot 10^{-1} -> 9.24 \cdot 10^{-2} | Helped |
| Use Riemannian BW-FM instead of BW-FM (E) | Gaussian generation · seqFISH (16D) | minimum W2^2 distance: 5.31\cdot10^{-1} -> 3.11\cdot10^{-1} | Helped |
| Use Riemannian BW-FM instead of BW-FM (E) | Gaussian generation · scRNA-seq (32D) | minimum W2^2 distance: 2.74 -> 1.31 | Helped |
| Use WFM instead of PSF | ShapeNet generation · Airplane | CD: 71.11 -> 69.88 | Helped |
| Use WFM instead of PSF | ShapeNet generation · Airplane | EMD: 61.09 -> 64.44 | Hurt |
| Use WFM instead of PVD | ShapeNet generation · Chair | CD: 56.26 -> 57.62 | Hurt |
| Use WFM instead of PVD | ShapeNet generation · Chair | EMD: 53.32 -> 57.93 | Hurt |
Surprising findings
- Frobenius FM and Euclidean BW-FM generated covariances do not respect the BW manifold and required projection onto PSD matrices.
Limitations
- Computational complexity is quadratic in point-cloud size, limiting scalability.
- Theoretical guarantees rely on exact OT maps while practice uses entropic approximations; general-distribution computations are non-rigorous.
- Entropic OT map estimation error scales poorly with dimension (rate n^{-2/(d+4)}).
Experimental setup
Datasets: ShapeNet ModelNet40 MNIST Letters seqFISH mouse embryogenesis MERFISH motor cortex 254-gene atlas COVID-19 scRNA-seq atlas Cities ECG XENIUM
Benchmarks: ShapeNet MNIST Letters MERFISH XENIUM seqFISH scRNA-seq COVID Cities ECG
Metrics: 1-NN accuracy with Chamfer distance 1-NN accuracy with earth mover's distance average minimum squared 2-Wasserstein distance between gener
Baselines: Frobenius FM BW-FM (Euclidean) PointFlow SoftFlow DPF-Net Shape-GF PVD PSF
Builds on: Riemannian flow matching flow matching entropic optimal transport McCann interpolation multisample flow matching
Closest papers by meaning
- Generalized Wasserstein Flow Matching: Transport Plans, Everywhere, All at Once (NeurIPS 2026 · similarity 0.91)
- Flowing Datasets with Wasserstein over Wasserstein Gradient Flows (ICML 2025 · similarity 0.88)
- GENOT: Entropic (Gromov) Wasserstein Flow Matching with Applications to Single-Cell Genomics (NeurIPS 2024 · similarity 0.88)
- Scalable Wasserstein Gradient Flow for Generative Modeling through Unbalanced Optimal Transport (ICML 2024 · similarity 0.88)
- Wasserstein Wormhole: Scalable Optimal Transport Distance with Transformer (ICML 2024 · similarity 0.87)
- Flow matching achieves almost minimax optimal convergence (ICLR 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.