One-Step Generation via Riemannian Wasserstein Gradient Flows

arXiv:2610.04454v1 Announce Type: new Abstract: Recently, Drifting Models and Wasserstein Gradient Flows have attracted substantial attention because they move iterative distributional refinement to training and amortize it into a generator, enabling fast inference. However, existing formulations have been developed largely for continuous Euclidean domains, such as image spaces, where particles admit unconstrained additive updates. On constrained spaces, these updates can leave the valid domain or ignore its geometry, making them unsuitable targets for training. Recent work has adapted updates to these spaces, but has focused on particular fields or offered limited empirical comparison. We derive and compare several geometry-aware fields within a common training framework for one-step generators. We test the method on data with different structures and obtain competitive one-step results in each setting. The best-performing field varies by task, showing why the choice of objective matters in practice.

Related

• RW-Flow: One-Step Generation on Compact Manifolds via Wasserstein Gradient Flows
• On the Wasserstein Gradient Flow Interpretation of Drifting Models
• Riemannian MeanFlow for One-Step Generation on Manifolds
• A Unified Framework for Data-Free One-Step Sampling via Wasserstein Gradient Flows

Source: arXiv cs.LG | 2026-10-06