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

Ualibyek Nurgulan; Seungwoo Yoo; Prin Phunyaphibarn; Minhyuk Sung

cs.LGSep 30, 2026

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

Authors: Ualibyek Nurgulan, Seungwoo Yoo, Prin Phunyaphibarn, Minhyuk Sung

Organizations: KAIST

Abstract

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.

Figures & tables

Figure 1: One-step samples on general manifolds. Left: target eigenfunctions on the Bunny and Spot meshes for different eigenfunction indices $k$ . Right: samples generated by our method in a single function evaluation (NFE $=$ 1). Samples concentrate according to the geometry-induced target distributions on the underlying meshes.

Cost	$c(x,y)$	Gibbs kernel $e^{-c/\varepsilon}$	$\nabla_{\!g}^{(1)}c(x,y)$	Identifiable for
Squared Geodesic	$\tfrac{1}{2}d_{g}(x,y)^{2}$	$e^{-d_{g}(x,y)^{2}/2\varepsilon}$	$-\log_{x}(y)$	a.e. $\varepsilon$ on $\mathbb{S}^{d}$ , $\mathbb{T}^{d}$
Spectral	$-\varepsilon\log k_{\rho}(x,y)$	$\sum_{n}\rho(\lambda_{n})\,\phi_{n}(x)\phi_{n}(y)$	$-\varepsilon\,\nabla_{\!g}^{(1)}\!\log k_{\rho}(x,y)$	all $\varepsilon$ , any $\mathcal{M}$
Chordal	$\\|x-y\\|^{2}$	$e^{-\\|x-y\\|^{2}/\varepsilon}$	$2P_{x}(x-y)$	all $\varepsilon$ , any $\mathcal{M}\subset\mathbb{R}^{D}$
Geodesic	$d_{g}(x,y)$	$e^{-d_{g}(x,y)/\varepsilon}$	$-\log_{x}(y)/d_{g}(x,y)$	all $\varepsilon$ on $\mathbb{S}^{d}$

Table 1: Overview of cost functions considered.

P_{x}

denotes the orthogonal projection onto

T_{x}\mathcal{M}

for

\mathcal{M}\subset\mathbb{R}^{D}

. The last column states for which

\varepsilon

and on which manifolds the Gibbs kernel is nondegenerate, so that the velocity field is identifiable by Thm. 3.2 ; “a.e.” means all but countably many. On manifolds not listed, no guarantee is established.

Kernel	$\rho(\lambda_{n})$	Parameters
Matérn	$\sigma^{2}\bigl(\tfrac{2\nu}{\kappa^{2}}+\lambda_{n}\bigr)^{-\nu-d/2}$	$\nu>\tfrac{1}{2},\ \kappa,\sigma^{2}>0$
Heat	$e^{-t\lambda_{n}}$	$t>0$
Sub. Heat	$e^{-t\lambda_{n}^{\alpha}}$	$t>0,\ \alpha\in(0,1]$

Table 2: Spectral densities of nondegenerate spectral kernels. When used as costs in the Sinkhorn divergence Eqn. 4 , these kernels yield identifiable velocity fields. See App. D for details.

Figure 2: Generated samples on geospatial event data benchmarks at 1 NFE. Each column visualizes the dataset (top) and samples generated by GFM ( Davis et al., 2026 ) (middle) and RW-Flow -S (bottom). RW-Flow -S more accurately captures the underlying data distributions, whereas GFM produces diffuse samples that deviate substantially from the ground-truth distributions.

	Volcano			Earthquake			Flood			Fire
Method	kMMD $\downarrow$	COV $\uparrow$	1-NNA	kMMD $\downarrow$	COV $\uparrow$	1-NNA	kMMD $\downarrow$	COV $\uparrow$	1-NNA	kMMD $\downarrow$	COV $\uparrow$	1-NNA
Held-out	$0.153$	$0.451$	$0.555$	$0.026$	$0.570$	$0.471$	$0.025$	$0.548$	$0.509$	$0.024$	$0.543$	$0.504$
RFM	$0.292$	$0.431$	$0.835$	$0.291$	$0.278$	$0.845$	$0.307$	$0.345$	$0.790$	$0.357$	$0.176$	$0.913$
GFM-L	$\underline{0.099}$	$0.516$	$0.553$	$0.040$	$\mathbf{0.570}$	$0.562$	$0.062$	$\underline{0.569}$	$0.547$	$0.029$	$0.520$	$0.643$
GFM-E	$0.199$	$0.431$	$0.823$	$0.174$	$0.339$	$0.806$	$0.191$	$0.350$	$0.795$	$0.190$	$0.201$	$0.900$
GFM-S	$\mathbf{0.094}$	$0.512$	$0.677$	$0.042$	$0.478$	$0.693$	$0.060$	$0.498$	$0.651$	$0.031$	$0.392$	$0.768$

Table 3: Results on geospatial event data benchmarks. kMMD , COV , and 1-NNA (closest to 0.5 is best) on the volcano, earthquake, flood, and fire datasets (NFE

=

1; bold = best, underline = second best; the held-out row is the val-vs-test reference and is excluded from ranking.) Suffixes denote variants: CT (Consistency Training) sCT (simplified Consistency Training) L (Lagrangian), E (Eulerian), S (semigroup), MF (mean flow); GG (geodesic Gaussian), GL (geodesic Laplacian), M (Matérn);

v

denotes

v

-prediction; RW-Flow costs: SG (squared geodesic), C (chordal), G (geodesic), H (heat), M (Matérn), S (subordinated). The metrics are averaged over 3 runs. Full results with standard deviations are in Tab. 9(b) .

	General			Glycine			Proline			Prepro			RNA
Method	kMMD $\downarrow$	COV $\uparrow$	1-NNA	kMMD $\downarrow$	COV $\uparrow$	1-NNA	kMMD $\downarrow$	COV $\uparrow$	1-NNA	kMMD $\downarrow$	COV $\uparrow$	1-NNA	kMMD $\downarrow$	COV $\uparrow$	1-NNA
Held-out	$0.007$	$0.585$	$0.499$	$0.023$	$0.590$	$0.491$	$0.020$	$0.602$	$0.489$	$0.027$	$0.572$	$0.514$	$0.036$	$0.588$	$0.504$
RFM	$0.469$	$0.191$	$0.831$	$0.281$	$0.314$	$0.741$	$0.478$	$0.160$	$0.880$	$0.486$	$0.225$	$0.820$	$0.559$	$0.205$	$0.942$
GFM-L	$\mathbf{0.018}$	$0.515$	$0.561$	$\underline{0.030}$	$0.562$	$\underline{0.523}$	$0.026$	$0.547$	$0.545$	$0.043$	$0.551$	$0.531$	$0.102$	$\underline{0.571}$	$0.572$
GFM-E	$0.327$	$0.257$	$0.777$	$0.134$	$0.389$	$0.681$	$0.388$	$0.162$	$0.878$	$0.324$	$0.288$	$0.794$	$0.540$	$0.313$	$0.888$
GFM-S	$0.067$	$0.362$	$0.690$	$0.056$	$0.450$	$0.633$	$0.059$	$0.378$	$0.699$	$0.080$	$0.423$	$0.672$	$0.183$	$0.473$	$0.756$

Table 4: Results on protein and RNA torsion-angle datasets. kMMD , COV , and 1-NNA (closest to 0.5 is best) on four protein and RNA torsion-angle datasets (NFE

=

1; bold = best, underline = second best; the held-out row is the val-vs-test reference and is excluded from ranking.) Variant suffixes as in Tab. 3 . The metrics are averaged over 3 runs. Full results with standard deviations are in Tab. 10(b) .

Bunny (

k=10

)

Bunny (

k=50

)

Spot (

k=10

)

Spot (

k=50

)

Method

kMMD

\downarrow

COV

\uparrow

1-NNA

Time (s)

\downarrow

kMMD

\downarrow

COV

\uparrow

1-NNA

Time (s)

\downarrow

kMMD

\downarrow

COV

\uparrow

1-NNA

Time (s)

\downarrow

kMMD

\downarrow

COV

\uparrow

1-NNA

Time (s)

\downarrow

Held-out

0.020

0.583

0.493

—

0.025

0.583

0.498

—

0.012

0.587

0.500

—

0.014

0.584

0.515

—

RFM 1000

\underline{0.048}

\mathbf{0.569}

\mathbf{0.514}

\underline{49.5}

\mathbf{0.028}

\mathbf{0.565}

\mathbf{0.516}

\underline{48.7}

\mathbf{0.016}

\mathbf{0.561}

\mathbf{0.522}

\underline{54.8}

\underline{0.023}

\underline{0.540}

\mathbf{0.549}

\underline{54.7}

RW-Flow -H

\mathbf{0.030}

\underline{0.558}

\underline{0.537}

\mathbf{0.16}

\underline{0.034}

\underline{0.544}

\underline{0.550}

\mathbf{0.13}

\underline{0.018}

\underline{0.493}

\underline{0.601}

\mathbf{0.24}

\mathbf{0.015}

\mathbf{0.549}

\underline{0.559}

\mathbf{0.24}

Table 5: Results on mesh datasets. kMMD , COV , 1-NNA (closest to 0.5 is best), and sampling time on the Bunny and Spot datasets ( bold = best, underline = second best; the held-out row is the val-vs-test reference and is excluded from ranking). RFM is sampled with 1,000 steps of its projected Euler integrator (RFM 1000 ); RW-Flow -H uses the heat cost and is sampled with NFE

=

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol	Meaning
Manifold and geometry
$(\mathcal{M},g)$	compact, connected, smooth Riemannian manifold with metric $g$ ; $\mathbb{S}^{d}$ and $\mathbb{T}^{d}$ denote the unit sphere and the flat torus of dimension $d$
$\mathbb{R}^{D}$	ambient Euclidean space in which $\mathcal{M}$ is embedded for the network parametrization
$T_{x}\mathcal{M}$	tangent space of $\mathcal{M}$ at $x$ ; $\langle\cdot,\cdot\rangle_{g}$ and $\\|\cdot\\|_{g}$ are the metric and its norm on $T_{x}\mathcal{M}$
$\exp_{x}$ , $\log_{x}$	exponential map $T_{x}\mathcal{M}\to\mathcal{M}$ at $x$ and its inverse on the domain where it is defined
$d_{g}(x,y)$	geodesic distance on $(\mathcal{M},g)$

Appendix

Table 6: Notations used in the paper .

Energy Functional	$\bm{T_{q,p}(x)}$	Identifiability Condition
Sinkhorn	$\displaystyle-\int\nabla_{\!g}^{(1)}c(x,y)\,\pi^{qp}(\,\mathrm{d}y\mid x)$	$k$ nondegenerate, $c$ symmetric Lipschitz
MMD	$\displaystyle\int\nabla_{\!g}^{(1)}k(x,y)\,\,\mathrm{d}p(y)$	$k$ characteristic
Smoothed KL (KGD)	$\displaystyle\frac{\int\nabla_{\!g}^{(1)}k(x,y)\,\,\mathrm{d}p(y)}{\int k(x,y)\,\,\mathrm{d}p(y)}$	$k$ characteristic, normalized

Appendix

Table 7: Summary of velocity terms and identifiability conditions for different energy functionals. The velocity field of an energy functional is

V_{q,p}=T_{q,p}-T_{q,q}

Data
Datasets	$\mathbb{S}^{2}$ : volcano, earthquake, fire, flood ( Mathieu & Nickel, 2020 ) ; $\mathbb{T}^{2}$ : Top500 backbone torsion angles $(\phi,\psi)$ ( Lovell et al., 2003 ) by amino-acid type (General, Glycine, Proline, Pre-Pro); $\mathbb{T}^{7}$ : RNA torsion angles ( Murray et al., 2003 ) ; all following the conventions of Chen & Lipman (2024) .
Coordinates	ambient unit vector on $\mathbb{S}^{2}$ ; angles in $[0,2\pi)$ on $\mathbb{T}^{d}$
Split	80/10/10 train/validation/test, computed once with a split seed independent of the training seed and read by every implementation; validation for model selection, test once for the reported number
Network
Architecture	MLP, $\mathrm{Linear}(d_{\mathrm{in}},H)\to\mathrm{SiLU}\to[\mathrm{Linear}(H,H)\to\mathrm{SiLU}]^{\times 3}\to\mathrm{Linear}(H,d)$ ; no normalisation, residual connections or time embedding; PyTorch default initialisation
Width $H$	1024 on $\mathbb{S}^{2}$ ; 512 on $\mathbb{T}^{2}$ and $\mathbb{T}^{7}$

Appendix

Table 8: Experimental setup, shared by every method.

Table 9: Results on earth datasets (full). Mean

\pm

s.d. over 3 seeds of kMMD , MMD , COV , and 1-NNA (closest to 0.5 is best) on the volcano, earthquake, flood, and fire datasets (NFE

=

x

v

denote

x

- and

v

-prediction; GG (geodesic Gaussian), GL (geodesic Laplacian), M (Matérn); MMDF = Wasserstein gradient flow of

\tfrac{1}{2}\,\mathrm{MMD}^{2}

with G (Gaussian), L (Laplacian), H (heat), M (Matérn) kernels; RW-Flow costs: SG (squared geodesic), C (chordal), G (geodesic), H (heat), M (Matérn), S (subordinated). Means only are summarised in Table 3 .

Table 10: Results on torus datasets (full). Mean

\pm

s.d. over 3 seeds of kMMD , MMD , COV , and 1-NNA (closest to 0.5 is best) on the four

\mathbb{T}^{2}

Ramachandran datasets and the

\mathbb{T}^{7}

RNA dataset (NFE

=

1; bold = best, underline = second best; the held-out row is the val-vs-test reference and is excluded from ranking.) Variant suffixes as in Table 3 . Means only are summarised in Table 4 .

	RNA
Method	kMMD $\downarrow$	MMD $\downarrow$	COV $\uparrow$	1-NNA
Held-out	$0.036$	$0.384$	$0.588$	$0.504$
RFM	$0.559\pm 0.000$	$1.665\pm 0.157$	$0.205\pm 0.002$	$0.942\pm 0.004$
GFM-L	$0.102\pm 0.013$	$0.433\pm 0.004$	$\underline{0.571}\pm 0.010$	$0.572\pm 0.017$
GFM-E	$0.540\pm 0.002$	$0.822\pm 0.074$	$0.313\pm 0.008$	$0.888\pm 0.006$
GFM-S	$0.183\pm 0.008$	$0.493\pm 0.002$	$0.473\pm 0.012$	$0.756\pm 0.006$

Appendix

(a) RNA

Table 12: Mesh results under RFM’s own training budget (full). We report kMMD , MMD , COV , 1-NNA (closest to 0.5 is best), and sampling time on the Bunny and Spot datasets ( bold = best, underline = second best; the held-out row is the val-vs-test reference and is excluded from ranking). RFM is trained for its shipped default of 100k iterations and sampled with 1,000 steps of its projected Euler integrator (RFM 1000 ); RW-Flow -H uses the heat cost, and is sampled with NFE

=

1. Time is the wall time to draw 2000 samples on one RTX 3090.

Explore similar work

When Riemann flows with Wasserstein: Generative Modeling of Probability Distributions on Manifolds
Sep 22, 2026Doron Haviv, Edward De Brouwer, Rishabh Anand +4Riemannian Flow Matching Data Manifold
Genentech Inc., South San Francisco, CA. · Yale University, New Haven, CT.
One-Step Generative Modeling via Wasserstein Gradient Flows
May 12, 2026Jiaqi Han, Puheng Li, Qiushan Guo +3Diffusion Models Generative Models
Stanford University · ByteDance
Discrete Wasserstein Flows for One-Step Generative Modeling
Oct 1, 2026Alessandro Micheli, Andrea Zerio, Samir BhattWasserstein Gradient Flows Wasserstein Distance
Imperial College London London, United Kingdom · Department of Computer Science, Aalborg University, Copenhagen, Denmark; Centre for Frontier AI Research (CFAR), Institute of Advanced Intelligence and Computing (IAIC), A*STAR, Singapore · University of Copenhagen Copenhagen, Denmark

Sep 22, 2026, cs.LGJ/K move · Enter open · S save

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

Doron Haviv, Edward De Brouwer, Rishabh Anand, Rex Ying+3

Genentech Inc., South San Francisco, CA. · Yale University, New Haven, CT.

Riemannian Flow Matching Data Manifold Wasserstein Distance Riemannian Manifolds Manifolds Spherical Latent Space+2

Open paper AskPDF 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

\mathcal{P}_2(\mathcal{M})

of a Riemannian manifold

(\mathcal{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

\mathcal{P}_2(\mathcal{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.

Abstract

Explore similar work

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

One-Step Generative Modeling via Wasserstein Gradient Flows

Discrete Wasserstein Flows for One-Step Generative Modeling