One step generation improves data models on curved spaces
RW-Flow: One-Step Generation on Compact Manifolds via Wasserstein Gradient Flows
Machine Learning
Summary
Many types of important data naturally live on curved shapes called manifolds, like Earth’s surface or molecule shapes. Existing methods to create new examples of such data need many steps, which can be slow. This paper introduces RW-Flow, a method that can generate new data samples in just one step on these curved spaces. The authors explain when this one-step approach can perfectly match the true data, and show their method works well on real examples like earthquake locations and protein shapes.
What this means in practice
- •For geospatial analysts: Generate accurate single-step models of event locations on Earth to improve forecasting and simulation of geospatial phenomena.
- •For biomolecular modelers: Create one-step generative models of protein and RNA torsion angles to aid structural biology and drug design workflows.
Authors
Ualibyek Nurgulan, Seungwoo Yoo, Prin Phunyaphibarn, Minhyuk Sung
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.