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

1 INTRODUCTION

• Existing few-step manifold methods reduce sampling cost but are harder to optimize because they learn maps across timestep pairs rather than the inference-time one-step map.Manifold diffusion and flow models typically require many network evaluations for sampling.Source: §1 · p. 1
• The paper asks which costs induce identifiable velocity fields on compact manifolds, where the Euclidean squared-distance guarantee does not directly settle the question.W-Flow provides the one-step Wasserstein-gradient-flow framework motivating this extension.Source: §1 · pp. 1–2
• For compact, connected Riemannian manifolds, identifiability holds exactly when the associated Gibbs kernel is nondegenerate for symmetric, Lipschitz costs.This characterization supports designing identifiable costs; squared geodesic cost does not always guarantee identifiability.Source: §1 · p. 2
• RW-Flow surpasses existing methods in nearly all standard benchmark settings under matched model capacity and training time.The benchmarks include Earth and climate data on spheres and protein and RNA torsion angles on tori.Source: §1 · p. 2

2 BACKGROUND: W-FLOW

• A Wasserstein gradient flow evolves a distribution along the steepest-descent direction of an energy functional measuring discrepancy from the target.The flow can be expressed as a continuity equation with a velocity field.Source: §2 · p. 2
• Choosing the energy as a divergence from the target induces a velocity field that transports the model distribution toward that target.W-Flow transports generated samples along this field and regresses generator outputs toward the transported samples.Source: §2 · p. 2
• The one-step training objective is equivalent to a scaled expected squared velocity, so it vanishes exactly when the velocity field does.A zero field guarantees distribution matching only if the field is identifiable.Source: §2 · p. 2
• In Euclidean spaces, the Sinkhorn divergence with squared Euclidean-distance cost is identifiable, but corresponding guarantees on compact manifolds remain unclear.The unresolved question is which costs guarantee identifiability on compact manifolds and how to characterize them.Source: §2 · p. 3

3 RW-FLOW

• RW-Flow extends W-Flow to compact, connected Riemannian manifolds and derives the manifold Wasserstein-gradient-flow velocity for the Sinkhorn divergence.The manifold objective replaces Euclidean updates with the exponential map and squared geodesic distance; the velocity is tangent-valued.Source: §3 · p. 3, §3.1 · p. 3

3.1 VELOCITY FIELDS ON COMPACT MANIFOLDS

• The Sinkhorn-gradient velocity combines attraction toward coupled data points with repulsion among generated samples that cancels entropic bias at q = p.The cost function determines the velocity field and therefore its identifiability.Source: §3.1 · p. 4

3.2 IDENTIFIABILITY VIA NONDEGENERACY

• For symmetric, Lipschitz costs on compact, connected manifolds, the Sinkhorn-induced velocity field is identifiable if and only if its Gibbs kernel is nondegenerate.Nondegeneracy can be checked spectrally or, for positive-definite kernels, through universality; spectral nondegeneracy requires every coefficient to be nonzero, even if some are negative.Source: §3.2 · p. 4

3.3 COST FUNCTIONS ON COMPACT MANIFOLDS

• Squared geodesic cost gives identifiability on spheres and tori for almost every ε, but this argument does not extend to general compact manifolds.The squared geodesic cost’s associated spectral coefficient vanishes for at most countably many ε values on these domains.Source: §3.3 · p. 5

3.3.1 SPECTRAL COSTS

• A spectral kernel is nondegenerate when its spectral density is nonzero at every Laplace-Beltrami eigenvalue, yielding identifiable velocity fields under the stated positivity and regularity conditions.The corresponding cost is −ε log kρ when kρ is positive and log kρ is Lipschitz.Source: §3.3.1 · pp. 5–6

3.3.2 DISTANCE-BASED COSTS

• On the sphere, geodesic distance itself induces an identifiable velocity field, unlike squared geodesic distance.The associated Laplacian kernel is universal for every λ > 0, and λ = 1/ε gives the Gibbs kernel for geodesic distance.Source: §3.3.2 · p. 6

Generative Modeling on Manifolds

• Manifold diffusion and flow models still require multiple sequential SDE or ODE integration steps, making sampling computationally expensive.Earlier work extended normalizing flows to hyperbolic spaces, spheres, and tori, followed by diffusion and flow-based models on Riemannian manifolds.Source: §4 · p. 6

Few-Step Generative Models

• Few-step generators either distill pretrained multi-step models or train from scratch using the underlying SDE or ODE dynamics.Distillation adds a stage after teacher pretraining; drifting models learn one-step generators from scratch using a discrepancy-related velocity field, and W-Flow uses an identifiable field.Source: §4 · pp. 6–7

Few-step Generative Modeling on Manifolds

• Manifold few-step methods use shortcut mappings or consistency models, while concurrent KGD represents drifting fields as weighted averages of kernel gradients.KGD's field corresponds to the Wasserstein gradient flow of KL divergence between smoothed densities; this work instead uses Sinkhorn divergence and also investigates squared MMD.Source: §4 · p. 7

5.1 EXPERIMENT SETUP

• Benchmarks cover sphere, torus, and mesh distributions, using KMMD, MMD, 1-NNA, and COV against test samples.Baselines include RFM, GFM, RMF, RCM, and KGD; parameter counts, training time, and optimizer configurations are matched.Source: §5.1 · p. 7

5 EXPERIMENTS

• Across spheres and tori, RW-FLOW variants outperform competing methods in nearly all one-step settings; on meshes, RW-FLOW matches RFM using one step.The mesh comparison is against RFM sampled with 1,000 Euler steps, whose sampling time is two orders of magnitude longer.Source: §5.2 · p. 7, Table 3 · p. 8, §5.3 · p. 8

5.2 SPHERE – S2

• On sphere benchmarks, geodesic and subordinated heat costs achieve the strongest overall performance, and both induce identifiable velocity fields.The Wildfire and Flood visualizations show RW-FLOW-S more closely capturing reference distributions, while GFM samples are overly diffuse.Source: §5.2 · p. 7, Figure 2 · p. 8

5.3 FLAT TORUS – Td

• On protein and RNA torsion datasets, RW-FLOW variants outperform Lagrangian and other GFM self-distillation variants in nearly all settings.RW-FLOW-SG performs comparably to RW-FLOW-H; on T^d, squared geodesic cost is identifiable for all but countably many ε, although not for every ε > 0.Source: §5.3 · p. 8, Table 4 · p. 9

5.4 GENERAL MANIFOLDS

• On triangular meshes, RW-FLOW matches RFM with one sampling step, while RFM uses 1,000 Euler steps and takes two orders of magnitude longer.The reported mesh benchmarks are Bunny and Spot.Source: §5.4 · p. 9, Table 5 · p. 9

6 CONCLUSION

• RW-FLOW uses the Wasserstein gradient flow of Sinkhorn divergence and establishes a broad class of costs that guarantee identifiable velocity fields.The conclusion reports strong empirical performance alongside these theoretical guarantees.Source: §6 · p. 9
• The authors report using AI tools to assist with proofs, implementation, analysis, and feedback, and state that they reviewed and verified the work.They state that experiments use publicly available datasets and that code and experiment configurations will be released upon acceptance.Source: REPRODUCIBILITY STATEMENT · p. 10

APPENDIX

• The appendix collects deferred proofs, notation, alternative energy functionals, spectral-kernel analysis, and experimental setup and full results.It also discusses squared MMD and KL divergence over smoothed densities.Source: APPENDIX · p. 13

A NOTATION

• Table 6 collects the paper’s notation, including symbols for manifold geometry, measures and transport, and training.The listed assumptions include a compact, connected Riemannian manifold and a symmetric Lipschitz cost.Source: Table 6 · p. 13, A NOTATION · p. 13, B PROOFS · p. 14

B.1 PROOF OF THEOREM 3.1

• Theorem 3.1 expresses the velocity as the negative Riemannian gradient of the first variation of the Sinkhorn-divergence energy.The derivation is stated to differ little from the Euclidean computation despite the general cost and manifold setting.Source: §B.1 · p. 14
• The velocity formula uses the conditional distribution of the optimal entropy-regularized transport coupling and the gradient in the cost’s first argument.The coupling is disintegrated with respect to the model distribution’s point.Source: §B.1 · p. 14
• The proof differentiates the dual transport objective at its optimum while holding the dual potentials fixed, then obtains the first variation of the Sinkhorn divergence.The envelope theorem justifies this differentiation step.Source: §B.1 · p. 14

B.2 PROOF OF THEOREM 3.2

• On a compact, connected Riemannian manifold with symmetric Lipschitz cost, the induced velocity field is identifiable if and only if its Gibbs kernel is nondegenerate.The zero-velocity condition reduces to equality of the relevant dual potentials up to a constant.Source: §B.2 · p. 15
• Symmetry gives a kernel equation relating the two distributions when the velocity vanishes, which lets the proof equate the dual potentials up to a constant.The proof uses the zero-velocity lemma and the symmetric-kernel relation.Source: §B.2 · p. 15
• If the Gibbs kernel is degenerate, the proof constructs distinct probability measures whose induced velocity field is not identifiable.A nonzero signed measure in the kernel’s null space is used to construct the pair.Source: §B.2 · p. 16

B.3 PROOF OF LEMMA 3.3

• A continuous symmetric kernel diagonalized by a complete Laplace–Beltrami eigenbasis is nondegenerate exactly when every spectral coefficient is nonzero.The proof uses this equivalence to connect kernel nondegeneracy to spectral tests.Source: §B.3 · p. 16

B.4 PROOF OF PROPOSITION 3.3.2

• For squared geodesic costs, degenerate regularization strengths are at most countable on S^d for d ≥ 2 and countably infinite on T^d for d ≥ 1.The proof identifies torus degeneracy through vanishing Fourier-mode eigenvalues; each even mode has one zero, while odd modes have none.Source: §B.4 · pp. 17–18
• On spheres, degeneracy occurs only at a countable set of regularization strengths, as each spherical-harmonic eigenvalue has at most countably many zeros.The proof identifies degeneracy by the vanishing of an eigenvalue for some harmonic degree.Source: §B.4 · p. 17

C OTHER ENERGY FUNCTIONALS

• The appendix compares alternative energy functionals by their velocity fields and identifiability conditions.It includes squared MMD and KL divergence over smoothed densities among the alternatives discussed.Source: C OTHER ENERGY FUNCTIONALS · p. 18

C.1 SQUARED MAXIMUM MEAN DISCREPANCY (MMD)

• The velocity field induced by squared MMD is identifiable if and only if its kernel is characteristic.The proof uses injectivity of the kernel mean embedding to connect a zero velocity field with equality of the distributions.Source: §C.1 · pp. 18–19

C.2 SMOOTHED KL DIVERGENCE

• For a positive normalized kernel, the smoothed-KL flow is identifiable if and only if the kernel is characteristic.The guiding functional is KL divergence between the smoothed model and target densities.Source: §C.2 · p. 19
• MMD squared and smoothed KL can be identifiable when k is characteristic, while Sinkhorn requires a nondegenerate Gibbs kernel.For positive definite kernels, the latter condition is universality.Source: §C.2 · p. 20

D DISCUSSION ON SPECTRAL KERNELS

• Heat kernels are positive and have bounded score, so log K is Lipschitz and the heat kernel is valid under Thm. 3.2; Matérn kernels are valid for ν > 1/2.Source: D DISCUSSION ON SPECTRAL KERNELS · pp. 20–21
• Truncating the spectral series makes the cost unidentifiable because the resulting kernel has finite rank; the truncation error is controlled by nmax.The appendix adapts nmax to the heat-kernel decay rate.Source: D DISCUSSION ON SPECTRAL KERNELS · p. 21

E.1 TRAINING & INFERENCE ALGORITHMS

• The generator is trained with the velocity field as its target; in practice, predicted ambient-space velocities are projected onto the manifold’s tangent space.The paper updates the network by gradient descent on the stated loss.Source: §E.1 · pp. 21–22
• RW-Flow estimates the training velocity by solving entropy-regularized minibatch optimal-transport problems with Sinkhorn.The exact velocity would require integration over all points from the model and target distributions.Source: §E.1 · p. 22

• Every baseline and RW-Flow cost variant uses the same network, optimization settings, wall-clock budget, model-selection rule, and standalone scorer.The scorer imports nothing from method repositories and rejects non-finite or off-manifold samples rather than projecting them.Source: §E.2 · p. 23
• Evaluation reports KMMD, MMD, coverage, and 1-NNA, with 0.5 indicating indistinguishable samples for 1-NNA.Coverage detects mode collapse, while MMD can be blind to it; the metrics use geodesic distance.Source: §E.2 · p. 23
• The reported compute hardware is an NVIDIA RTX 3090 with 24 GB of memory.Source: §E.2 · p. 23

• Sinkhorn-divergence flows perform well on all datasets, whereas MMD-flow performance depends strongly on kernel and dataset.MMDF-L fails on General, and MMDF-G and MMDF-H are unstable on RNA for some seeds.Source: §E.3 · p. 24
• The full results cover Earth datasets, four T2 Ramachandran datasets, and an RNA dataset, with mean and standard deviation over three seeds.The Earth and torus tables report KMMD, MMD, coverage, and 1-NNA.Source: Table 9 · p. 25, Table 10 · p. 26
• The reported MMD-flow variants use Gaussian, Laplacian, heat, or Matérn kernels; RW-Flow variants include squared-geodesic, chordal, geodesic, heat, Matérn, and subordinated costs.Source: §E.3 · p. 25
• The mesh results compare Bunny and Spot runs, including RW-FLOW-H and RFM under RFM’s stated training budget.The table caption specifies KMMD, MMD, coverage, 1-NNA, and sampling time.Source: Table 11 · p. 27