当黎曼流与Wasserstein相遇：流形上概率分布的生成建模
When Riemann flows with Wasserstein: Generative Modeling of Probability Distributions on Manifolds
- Genentech Inc.（基因泰克公司）
- Yale University（耶鲁大学）
机构由 AI 辅助整理，请以论文原文为准。
AI总结：
针对非欧几里得域上的元分布生成问题，提出RWEFM框架，在黎曼流形的Wasserstein空间上利用最优传输速度回归训练，可生成单细胞和蛋白质构象，且适用于一般三角网格。
AI中文摘要：
许多科学数据集，例如分子构象系综或单细胞组织测量，天然地被建模为元分布：非欧几里得域上概率测度的分布。现有的生成方法大多假设欧几里得几何，无法捕捉这种结构。我们引入了黎曼Wasserstein熵流匹配（RWEFM），一种在黎曼流形$(\mathcal{M},g)$的Wasserstein空间$\mathcal{P}_2(\mathcal{M})$上的生成框架。RWEFM通过将神经向量场回归到黎曼最优传输速度上进行训练，使用McCann位移插值作为条件路径。我们在理论上证实了这种构造在$\mathcal{P}_2(\mathcal{M})$上产生有效的流匹配方法，并引入了黎曼熵映射，一种在流形上最优传输映射的GPU高效近似。我们的实验表明，通过尊重数据的内在几何，RWEFM能够在超球面潜空间中生成完整的单细胞样本，并在环面上生成蛋白质构象系综。由于RWEFM仅需要测地距离和投影算子，它不限于具有闭式几何的流形，我们通过在一般三角网格上生成分布来证明这一点。
英文摘要：
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.
