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

Sep 22, 2026cs.LG

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

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.

Doron Haviv, Edward De Brouwer, Rishabh Anand +4

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

Ask PDF

May 12, 2026cs.LG

One-Step Generative Modeling via Wasserstein Gradient Flows

Diffusion models and flow-based methods have shown impressive generative capability, especially for images, but their sampling is expensive because it requires many iterative updates. We introduce W-Flow, a framework for training a generator that transforms samples from a simple reference distribution into samples from a target data distribution in a single step. This is achieved in two steps: we first define an evolution from the reference distribution to the target distribution through a Wasserstein gradient flow that minimizes an energy functional; second, we train a static neural generator to compress this evolution into one-step generation. We instantiate the energy functional with the Sinkhorn divergence, which yields an efficient optimal-transport-based update rule that captures global distributional discrepancy and improves coverage of the target distribution. We further prove that the finite-sample training dynamics converge to the continuous-time distributional dynamics under suitable assumptions. Empirically, W-Flow sets a new state of the art for one-step ImageNet 256

\times

256 generation, achieving 1.29 FID, with improved mode coverage and domain transfer. Compared to multi-step diffusion models with similar FID scores, our method yields approximately 100

\times

faster sampling. These results show that Wasserstein gradient flows provide a principled and effective foundation for fast and high-fidelity generative modeling.

Jiaqi Han, Puheng Li, Qiushan Guo +3

Stanford University · ByteDance

Diffusion Models Generative Models Wasserstein Distance Imagenet Gradient

Ask PDF

Oct 1, 2026cs.LG

Discrete Wasserstein Flows for One-Step Generative Modeling

We introduce a new framework for one-step generative modelling on finite state spaces. To extend drifting beyond continuous domains, we use discrete Wasserstein geometry to define a target-relative KL gradient flow over the transitions of a reversible Markov kernel. We realize this probability flow at the particle level through Markov jumps and amortize the resulting transport updates into a latent-conditioned generator, so that the iterative dynamics are required only during training while inference remains one-step. In a controlled setting where the underlying distributions and transport dynamics can be computed exactly, we verify KL dissipation, consistency between the particle dynamics and the probability flow, and the predicted numerical scaling. We further show that a finite-capacity neural generator can track these exact transport targets while retaining one-step generation. These results validate the basic construction and provide a foundation for scaling Discrete Drifting to structured discrete data.

Alessandro Micheli, Andrea Zerio, Samir Bhatt

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

Wasserstein Gradient Flows Wasserstein Distance Generative Models Markov Particle Dynamics

Ask PDF