One-Step Generation via Riemannian Wasserstein Gradient Flows

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. Existing formulations, however, have been developed predominantly for continuous Euclidean domains, such as image spaces, where particles admit unconstrained additive updates. We develop a unified formulation of Wasserstein Gradient Flow on general smooth state spaces equipped with a chosen Riemannian geometry. Given any sufficiently regular distributional functional, our framework converts its first variation into an intrinsic particle velocity, defines the corresponding flow of probability laws, and provides geometry-compatible finite updates. Under standard assumptions, we prove that the flow decreases the objective and that its discrete update is consistent. The resulting updates can be estimated from finite samples and amortized by a neural generator. We empirically demonstrate one-step generation across diverse domains: geospatial data, MNIST-Binary, and LM1B. Together, our theory provides a general framework for fast objective-driven generation on Riemannian state spaces.