# RW-Flow: One-Step Generation on Compact Manifolds via Wasserstein Gradient Flows

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 28 of its 52 references that are in the map.
- First posted on arXiv 30 Sep 2026.

## Builds on (most cited first, 28 of 28)

- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [One-step Diffusion with Distribution Matching Distillation](https://mlsanity.com/paper/2311.18828/) (2023)
- [Mean Flows for One-step Generative Modeling](https://mlsanity.com/paper/2505.13447/) (2025)
- [Improved Techniques for Training Consistency Models](https://mlsanity.com/paper/2310.14189/) (2023)
- [Learning Generative Models with Sinkhorn Divergences](https://mlsanity.com/paper/1706.00292/) (2017)
- [Flow Matching on General Geometries](https://mlsanity.com/paper/2302.03660/) (2023)
- [Riemannian Score-Based Generative Modelling](https://mlsanity.com/paper/2202.02763/) (2022)
- [Score identity Distillation: Exponentially Fast Distillation of Pretrained Diffusion Models for One-Step Generation](https://mlsanity.com/paper/2404.04057/) (2024)
- [On Wasserstein Two Sample Testing and Related Families of Nonparametric Tests](https://mlsanity.com/paper/1509.02237/) (2015)
- [Riemannian Continuous Normalizing Flows](https://mlsanity.com/paper/2006.10605/) (2020)
- [Improving GANs Using Optimal Transport](https://mlsanity.com/paper/1803.05573/) (2018)
- [Riemannian Diffusion Models](https://mlsanity.com/paper/2208.07949/) (2022)
- [Normalizing Flows on Tori and Spheres](https://mlsanity.com/paper/2002.02428/) (2020)
- [Generative Modeling via Drifting](https://mlsanity.com/paper/2602.04770/) (2026)
- [How to build a consistency model: Learning flow maps via self-distillation](https://mlsanity.com/paper/2505.18825/) (2025)
- [Mat'ern Gaussian processes on Riemannian manifolds](https://mlsanity.com/paper/2006.10160/) (2020)
- [Latent Variable Modelling with Hyperbolic Normalizing Flows](https://mlsanity.com/paper/2002.06336/) (2020)
- [Geodesic Exponential Kernels: When Curvature and Linearity Conflict](https://mlsanity.com/paper/1411.0296/) (2014)
- [Stationary Kernels and Gaussian Processes on Lie Groups and their Homogeneous Spaces I: the compact case](https://mlsanity.com/paper/2208.14960/) (2022)
- [One-Step Generative Modeling via Wasserstein Gradient Flows](https://mlsanity.com/paper/2605.11755/) (2026)
- [Generalised Flow Maps for Few-Step Generative Modelling on Riemannian Manifolds](https://mlsanity.com/paper/2510.21608/) (2025)
- [Riemannian Consistency Model](https://mlsanity.com/paper/2510.00983/) (2025)
- [Riemannian MeanFlow](https://mlsanity.com/paper/2602.07744/) (2026)
- [Kernel-Gradient Drifting Models](https://mlsanity.com/paper/2605.10727/) (2026)
- [Invariant kernels on Riemannian symmetric spaces: a harmonic-analytic approach](https://mlsanity.com/paper/2310.19270/) (2023)
- [The Gaussian kernel on the circle and spaces that admit isometric embeddings of the circle](https://mlsanity.com/paper/2302.10623/) (2023)
- [Gaussian kernels on non-simply-connected closed Riemannian manifolds are never positive definite](https://mlsanity.com/paper/2303.06558/) (2023)

## Abstract (arXiv)

Manifold-valued data, and consequently the distributions they induce, are prevalent across many domains, ranging from the locations of geospatial events, such as earthquakes, to biomolecular torsion angles that encode information about three-dimensional structure. While diffusion and flow-based generative models have been successfully extended to compact manifolds, sampling typically requires tens or hundreds of sequential network evaluations. We introduce RW-Flow, a theoretically grounded framework for learning one-step generative models on compact manifolds via Wasserstein gradient flows. The main challenge is identifiability: driving the velocity field to zero should guarantee that the model distribution matches the target distribution. We establish a necessary and sufficient condition for identifiability on compact, connected Riemannian manifolds. We specifically show that, for a symmetric, Lipschitz-continuous cost function, the velocity field induced by the Sinkhorn divergence is identifiable if and only if the associated Gibbs kernel is nondegenerate. This characterization provides a general principle for designing identifiable costs on compact manifolds. It also reveals that the squared geodesic distance, the natural manifold analogue of the squared Euclidean distance, does not always guarantee identifiability. Across benchmarks involving geospatial events, protein side chain torsion angles, RNA backbone torsion angles, and general manifolds discretized as triangular meshes, RW-Flow outperforms existing one-step methods in nearly all settings under fair comparison conditions.