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

2026-09-22

AI Summary

为了解决非欧几何域上的概率分布生成问题，提出了一种基于黎曼流形的Wasserstein熵流匹配方法（RWEFM），通过回归神经向量场到黎曼最优传输速度实现。

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 $\mathcal{P}_2(\mathcal{M})$ of a Riemannian manifold $(\mathcal{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 $\mathcal{P}_2(\mathcal{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.

Problem

• non-Euclidean
• generative modeling
• probability distributions
• manifolds

Innovation

• Riemannian Wasserstein Entropic Flow Matching
• McCann displacement interpolations
• Riemannian Entropic Map
• non-Euclidean domains
• generative modeling