One-Step Generation via Riemannian Wasserstein Gradient Flows

2026-10-03

AI Summary

This study addresses the issue of ineffective updates in existing generative models within constrained spaces caused by neglecting geometric structure, proposing a general one-step generation framework adapted to Riemannian manifolds. Methodologically, multiple geometry-aware fields are derived and compared, integrating Riemannian Wasserstein gradient flows with drifting models to construct a unified training paradigm. The research reveals the critical impact of objective function selection on generative performance. Experiments demonstrate that the proposed framework achieves competitive one-step generation results across data with diverse structures, effectively validating the superiority and generalization capability of incorporating geometry-aware mechanisms.

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.

Problem

• One-Step Generation
• Wasserstein Gradient Flows
• Constrained Spaces
• Riemannian Geometry
• Drifting Models

Innovation

• Riemannian Wasserstein Gradient Flows
• One-Step Generation
• Geometry-aware Fields
• Constrained Spaces
• Drifting Models