# When Riemann flows with Wasserstein: Generative Modeling of Probability Distributions on Manifolds

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 18 of its 50 references that are in the map.
- First posted on arXiv 22 Sep 2026.

## 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, 18 of 18)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [Flow Matching Guide and Code](https://mlsanity.com/paper/2412.06264/) (2024)
- [Multisample Flow Matching: Straightening Flows with Minibatch Couplings](https://mlsanity.com/paper/2304.14772/) (2023)
- [SDXL-Lightning: Progressive Adversarial Diffusion Distillation](https://mlsanity.com/paper/2402.13929/) (2024)
- [Flow Matching on General Geometries](https://mlsanity.com/paper/2302.03660/) (2023)
- [SE(3)-Stochastic Flow Matching for Protein Backbone Generation](https://mlsanity.com/paper/2310.02391/) (2023)
- [Dirichlet Flow Matching with Applications to DNA Sequence Design](https://mlsanity.com/paper/2402.05841/) (2024)
- [Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein](https://mlsanity.com/paper/2201.12324/) (2022)
- [Entropic estimation of optimal transport maps](https://mlsanity.com/paper/2109.12004/) (2021)
- [Categorical Flow Matching on Statistical Manifolds](https://mlsanity.com/paper/2405.16441/) (2024)
- [Meta Flow Matching: Integrating Vector Fields on the Wasserstein Manifold](https://mlsanity.com/paper/2408.14608/) (2024)
- [Wasserstein Flow Matching: Generative modeling over families of distributions](https://mlsanity.com/paper/2411.00698/) (2024)
- [Generalised Flow Maps for Few-Step Generative Modelling on Riemannian Manifolds](https://mlsanity.com/paper/2510.21608/) (2025)
- [Flowing Datasets with Wasserstein over Wasserstein Gradient Flows](https://mlsanity.com/paper/2506.07534/) (2025)
- [Generalized Wasserstein Flow Matching: Transport Plans, Everywhere, All at Once](https://mlsanity.com/paper/2605.08424/) (2026)
- [Barycentric Projections of Optimal Transport Plans on Riemannian Manifolds](https://mlsanity.com/paper/2606.07926/) (2026)

## Abstract (arXiv)

Many scientific datasets, such as molecular conformational ensembles or single-cell tissue measurements, are naturally modeled as meta-distributions: distributions over probability measures on non-Euclidean domains. Existing generative methods largely assume Euclidean geometry and fail to capture this structure. We introduce Riemannian Wasserstein Entropic Flow Matching (RWEFM), a generative framework on the Wasserstein space P2(M) of a Riemannian manifold (M,g). RWEFM is trained by regressing a neural vector field onto Riemannian optimal transport velocities, using McCann displacement interpolations as conditional paths. We confirm theoretically that this construction leads to a valid flow matching approach on P2(M) and introduce the Riemannian Entropic Map, a GPU-efficient approximation of the optimal transport map on manifolds. Our experiments show that by respecting the intrinsic geometry of the data, RWEFM can generate whole single-cell samples in hyperspherical latent spaces and protein conformational ensembles on the torus. As RWEFM requires only a geodesic distance and a projection operator, it is not restricted to manifolds with closed-form geometry, which we demonstrate by generating distributions on a general triangulated mesh.