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

2026-09-30

AI Summary

This study addresses the inefficiency of multi-step iterative sampling in generative models on manifolds by proposing a one-step generation framework based on Wasserstein gradient flows. Methodologically, the work integrates Riemannian geometry to establish necessary and sufficient conditions for the identifiability of Sinkhorn divergence velocity fields, while revealing the modeling limitations of conventional geodesic distances to enable efficient representations of compact manifolds. Experiments demonstrate that the proposed framework significantly outperforms existing one-step methods on geographic and biomolecular data benchmarks. By combining rigorous theoretical guarantees with computational efficiency, this work presents a novel paradigm for generative modeling on manifolds.

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.

Problem

• compact manifolds
• one-step generation
• identifiability
• Wasserstein gradient flows
• manifold-valued data

Innovation

• Wasserstein gradient flows
• compact manifolds
• one-step generation
• identifiability
• Sinkhorn divergence