When Riemann flows with Wasserstein: Generative Modeling of Probability Distributions on Manifolds

arXiv:2609.25659 · cs.LG · Submitted 2026-09-22 · Read on arXiv

cs.LG

Submitted: 2026-09-22

Updated: 2026-09-22

Code: https://github.com/rois-codh/kmnist

License: http://creativecommons.org/licenses/by/4.0/

The gist: Many scientific datasets, such as molecular conformational ensembles or single-cell tissue measurements, are naturally modeled as meta-distributions: distributions over probability measures on

Terminology

Abstract

Many scientific datasets, such as molecular conformational ensembles or single-cell tissue measurements, are naturally modeled as meta-distributions: distributions over probability measures on non-Euclidean domains. Existing generative methods largely assume Euclidean geometry and fail to capture this structure. We introduce Riemannian Wasserstein Entropic Flow Matching (RWEFM), a generative framework on the Wasserstein space P 2(M) of a Riemannian manifold (M,g). RWEFM is trained by regressing a neural vector field onto Riemannian optimal transport velocities, using McCann displacement interpolations as conditional paths. We confirm theoretically that this construction leads to a valid flow matching approach on P 2(M) and introduce the Riemannian Entropic Map, a GPU-efficient approximation of the optimal transport map on manifolds. Our experiments show that by respecting the intrinsic geometry of the data, RWEFM can generate whole single-cell samples in hyperspherical latent spaces and protein conformational ensembles on the torus. As RWEFM requires only a geodesic distance and a projection operator, it is not restricted to manifolds with closed-form geometry, which we demonstrate by generating distributions on a general triangulated mesh.

Sources

Related papers