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

会议: NeurIPS 2026
arXiv: 2609.25659
代码: https://github.com/doronhaviv/RWEFM
领域: 图像生成
关键词: 流匹配、黎曼流形、最优传输、Wasserstein空间、生成模型

一句话总结¶

提出了黎曼-Wasserstein熵流匹配（RWEFM），将流匹配提升至黎曼流形上的Wasserstein测度空间$\mathcal{P}_2(\mathcal{M})$，结合GPU高效的黎曼熵映射逼近最优传输，实现了非欧流形上测度分布的高保真生成。

研究背景与动机¶

现代众多前沿科学计算场景中的核心实体天然表现为非欧几里得几何域上的概率测度集合（即“分布的分布”或元分布）。例如，分子构象系综将单个分子刻画为三维空间旋转群与平移群上的构象概率分布；全球气候记录是二维球面上的气象变量分布；而在组织级单细胞转录组学（scRNA-seq）中，单个生物样本由处于超球面或双曲潜在流形上的数千个细胞状态构成的经验测度表达。对这类数据进行生成式建模，是实现计算机内抗体分子设计、假设检验、以及复杂物理与生物系统因果表征的关键使能路径。

然而，传统的生成扩散与流匹配模型主要面向欧几里得空间中的有限维单样本向量点。尽管近来 Wasserstein Flow Matching（WFM）与 Meta FM 开始尝试直接在概率测度的二次 Wasserstein 空间上学习生成流，但这些方法强依赖欧几里得平坦空间的全局线性结构，将底层基流形视作 $\mathcal{M} \equiv \mathbb{R}^d$。一旦数据本真几何为带曲率的球面、环面、双曲空间甚至通用三角网格表面，欧氏测度流匹配会因违背几何拓扑约束产生严重畸变；另一方面，现有的黎曼生成模型（如 RFM）仅能生成流形上的单点样本，完全丧失了捕获整个细胞样本或构象系综内粒子间多体关联与非独立同分布（non-i.i.d.）分布特性的能力。

面对“流形基几何约束”与“分布级高维测度空间”的双重理论与计算壁垒，直接在无穷维空间上构建测度流匹配的瓶颈在于：不仅需要确立流形测度空间上的连续性方程与拉格朗日流生成合法性，更要在有限粒子采样与 GPU 大规模训练下高效、严谨地求解流形间的 Monge 最优传输映射。本文的核心 idea 是：将条件流匹配提升至黎曼流形上的二次 Wasserstein 空间 $\mathcal{P}_2(\mathcal{M})$，以基于黎曼最优传输测地线的 McCann 位移插值为条件路径，并提出基于切空间重心投影的黎曼熵映射（Riemannian Entropic Map），实现了任意黎曼流形乃至无解析度量网格上概率测度分布的端到端可微流匹配。

方法详解¶

整体框架¶

RWEFM 的核心目标是学习一条连接先验测度分布 $\mathbb{P}_0 \in \mathcal{P}_2(\mathcal{P}_2(\mathcal{M}))$ 与目标真实测度分布 $\mathbb{P}_1 \in \mathcal{P}_2(\mathcal{P}_2(\mathcal{M}))$ 的速度场。输入端从先验与目标元分布中抽取配对经验测度点云，通过求解两点云间的黎曼最优传输获得粒子级别的传输位移；随后沿测地线 McCann 插值路径构建条件轨迹，并训练基于置换等变自注意力机制的神经网络拟合流形切向量场；推理时则从噪声点云出发沿学习到的黎曼向量场做几何投影数值积分生成目标经验分布。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["元分布采样<br/>配对经验测度点云 X0~μ, X1~ν"] --> B["黎曼熵映射<br/>切空间提升-加权平均-指数回弹投影"]
    B --> C["McCann 位移插值路径<br/>构造流形测地线粒子轨迹与条件速度场"]
    C --> D["自注意力等变向量场回归<br/>拟合切空间速度场并最小化黎曼切范数损失"]
    D --> E["流形欧拉积分生成<br/>沿切向量场推演并重投影至流形几何"]

关键设计¶

1. 测度空间 $\mathcal{P}_2(\mathcal{M})$ 上的弱连续性方程与叠加原理：保障无穷维流匹配的理论相容性
将流匹配从有限维欧氏空间推广至非欧流形上的测度空间，首要痛点在于 $\mathcal{P}_2(\mathcal{M})$ 本质是一个无穷维伪黎曼空间，必须严谨证明向量场沿弱连续性方程推演能唯一诱导出拉格朗日意义下的测度流动。本文通过定义光滑紧支柱状泛函（Cylinder Functionals）$\varphi_t(\mu) = F(t, \int V_1 d\mu, \dots, \int V_k d\mu)$，构建了 $\mathcal{P}_2(\mathcal{M})$ 上的弱连续性方程： $$\int_0^T \int_{\mathcal{P}_2(\mathcal{M})} \left( \partial_t \varphi_t(\mu) + \langle \nabla_{\mathcal{W}} \varphi_t(\mu), V_t(\mu) \rangle_{L^2(\mu)} \right) d\mathbb{P}_t(\mu) dt = 0$$ 结合测度空间叠加原理（Superposition Principle for Random Measures），定理证明只要边际向量场 $V_t$ 满足二阶动能积分有限，则对 $\mathbb{P}_0$-几乎每个测度 $\mu$，均存在唯一的 $W_2$-绝对连续曲线 $\dot{\gamma}_\mu(t) = V_t(\gamma_\mu(t))$，且边缘流满足 $\mathbb{P}_t = (\Phi_t)_\# \mathbb{P}_0$。进一步证明了边缘向量场与条件向量场在条件流匹配损失下的梯度等价性 $\nabla_\theta \mathcal{L}_{CFM} = \nabla_\theta \mathcal{L}_{FM}$，使在经验点云上做条件回归具备了严密的测度动力学支撑。

2. 黎曼熵映射（Riemannian Entropic Map）：GPU 上快速可扩展的最优传输逼近
流匹配条件路径构建依赖于点云对之间真实的 Monge 最优传输映射 $T^{\mu \to \nu}$，然而在流形上精确求解非正则化 Monge 问题的计算复杂度高达立方级，且缺乏封闭解。针对这一瓶颈，本文将欧氏熵正则化最优传输拓展至黎曼流形，构建了“提升-平均-回弹”（lift-average-retract）三步黎曼熵映射估计器： $$T_\varepsilon(x) := \exp_x \left( \int_\mathcal{M} \log_x(y) d\pi_\varepsilon(y|x) \right)$$ 其中 $\pi_\varepsilon$ 由基于测地代价矩阵 $C_{ij} = \frac{1}{2} d_g^2(x_i, y_j)$ 的 Sinkhorn 算法在 GPU 上极速求解得到。理论上证明了该估计器在对偶势函数下严格满足 $T_\varepsilon(x) = \exp_x(-\nabla f_\varepsilon(x))$，并在几何与密度正则假设下给出了非渐近统计收敛界：在 $d$ 维流形上其实际估计误差满足 $\mathcal{O}(\varepsilon \log(e/\varepsilon) + \varepsilon^{-d/2} n^{-1/2} \log n)$，当选取正则化强度 $\varepsilon \asymp n^{-1/(d+2)}$ 时达到最佳收敛率 $\widetilde{\mathcal{O}}(n^{-1/(d+2)})$。

3. McCann 位移测地插值与切空间回归速度场：内蕴几何约束下的流形流
为确保经验测度点云在任意时间步 $t \in [0, 1]$ 均严格落在底层流形 $\mathcal{M}$ 的几何约束上，传统欧氏线性内插 $(1-t)X_0 + t X_1$ 必须被舍弃。本文将条件概率路径定义为流形上的 McCann 测地线位移插值： $$z_i = \exp_{x_i}(t \log_{x_i}(T_\varepsilon(x_i)))$$ 与之切相容的目标速度场为该测地线的时间导数 $v_i = \frac{1}{1-t} \log_{z_i}(T_\varepsilon(x_i)) \in T_{z_i}\mathcal{M}$。神经网络 $v_\theta(X_t, t)$ 采用 6 层多头自注意力机制（MHA，4 头）构建以保持粒子置换等变性，网络输出由非嵌入层直接投影至粒子所在的外在切空间，最终通过最小化点云切空间均方黎曼范数误差进行端到端优化： $$\mathcal{L} = \frac{1}{N} \sum_{i=1}^N \| \hat{v}_i - v_i \|_{g(z_i)}^2$$

4. 基于度量与投影的通用网格扩展机制：摆脱封闭解析测地线的限制
针对三角化网格表面、隐式曲面等缺乏解析 $\exp$ 与 $\log$ 映射的复杂科学几何域，本文提出了无需闭式导出的通用黎曼熵映射机制。利用测地距离函数的平方梯度性质，将对数映射改写为负外在梯度经投影算子雅可比矩阵在切平面的投影：$\log_x(y) = - J_\Pi(x) (d_g(x, y) \nabla_x d_g(x, y))$，利用 Sinkhorn 反向传播即可免额外计算开销获取。在回弹阶段，采用基于外在欧拉步、表面就近投影 $\Pi_\mathcal{M}$ 与速度重归一化的切空间投影欧拉积分（Projected Euler Integration），成功使 RWEFM 具备了在斯坦福兔子等非解析三角离散曲面上生成复杂经验分布的能力。

损失函数 / 训练策略¶

训练时，先验噪声分布采用双层分层高斯采样后投影至流形，以保持点云均值与协方差的跨样本变异性。每个 mini-batch 内部首先计算点云之间的内蕴几何 Chamfer 距离矩阵，并通过批内 Sinkhorn 匹配实现 Wasserstein 空间上的批次最优配对。模型使用 Adam 优化器在单个 NVIDIA B200 GPU 上训练 500,000 步，初始学习率 $3 \times 10^{-4}$，每 5,000 步衰减 0.99；熵正则化参数设为 $\varepsilon = 0.002$。在条件生成中引入分类器无关引导（CFG，权重 $w=2$，空条件丢弃率 $p=0.1$）。

实验关键数据¶

主实验¶

论文在双曲空间 $\mathbb{H}^2$（EMNIST 字符）、球面 $\mathbb{S}^2$（MNIST 数字）、环面 $\mathbb{T}^2$（KMNIST 假名）以及通用三角网格 $\mathcal{M}_{\text{mesh}}$（斯坦福兔子）上评估点云分布生成质量，采用 Chamfer Distance（CD）和 Earth Mover's Distance（EMD）评估类别级最近邻偏差（1-NN-D，范围 $[0, 0.5]$，越低越好，0 对应 50% 完美分类准确率）。

流形空间	任务类别	指标 (1-NN-D)	RWEFM (本文)	WFM (欧氏基线)	SetRFM (无OT流形)	RFM (点级流形)
双曲空间 $\mathbb{H}^2$	EMNIST 'h'	CD / EMD	0.15 / 0.09	0.30 / 0.28	0.30 / 0.31	0.48 / 0.48
双曲空间 $\mathbb{H}^2$	EMNIST 'w'	CD / EMD	0.23 / 0.10	0.41 / 0.39	0.34 / 0.33	0.44 / 0.45
球面 $\mathbb{S}^2$	MNIST '3'	CD / EMD	0.18 / 0.18	0.29 / 0.28	0.38 / 0.38	0.30 / 0.38
球面 $\mathbb{S}^2$	MNIST '4'	CD / EMD	0.19 / 0.18	0.29 / 0.27	0.38 / 0.38	0.42 / 0.45
平坦环面 $\mathbb{T}^2$	KMNIST 'ki'	CD / EMD	0.17 / 0.18	0.22 / 0.17	0.26 / 0.28	0.47 / 0.46
兔子网格 $\mathcal{M}_{\text{mesh}}$	MNIST '0'	CD / EMD	0.30 / 0.28	0.50 / 0.50	0.30 / 0.26	–
兔子网格 $\mathcal{M}_{\text{mesh}}$	MNIST '2'	CD / EMD	0.26 / 0.21	0.40 / 0.42	0.29 / 0.23	–

在真实科学计算场景中，RWEFM 在单细胞超球面隐空间 $\mathbb{S}^{127}$ 全组织样本生成及 mdCATH 蛋白质主链二面角分布（$\mathbb{T}^2$）生成上全面超越基准：

评估领域与基流形	评估模型	1-NN-D (CD) ↓	MMD (EMD) ↓	$W_2$ ↓	$W_1$ ↓	MMD-CD ↓
单细胞 scRNA-seq ($\mathbb{S}^{127}$)	SetFM	0.4460	0.1026	–	–	0.1174
	WFM	0.2945	0.0324	–	–	0.0355
	SetRFM	0.3754	0.0300	–	–	0.0341
	RWEFM	0.1889	0.0140	–	–	0.0137
蛋白质构象系综 ($\mathbb{T}^2$)	SetFM	–	0.0114	0.4472	0.4097	–
	WFM	–	0.0120	0.4059	0.4023	–
	SetRFM	–	0.0119	0.4608	0.4126	–
	RWEFM	–	0.0097	0.3970	0.3850	–

消融实验与分析¶

针对传输映射形式（熵映射重心投影 vs 采样单粒子映射）、正则化强度 $\varepsilon$ 与训练粒子数 $N$ 的消融研究如下：

配置 / 变体	目标流形 & 数据集	1-NN-D (CD) ↓	1-NN-D (EMD) ↓	训练耗时 / 说明
RWEFM (黎曼熵映射)	$\mathbb{H}^2$ (EMNIST 'h')	0.149	0.085	4h 06m（连续重心投影，保真度最高）
RWEFM (随机采样映射)	$\mathbb{H}^2$ (EMNIST 'h')	0.192	0.157	~3h 40m（单粒子匹配，易受高维噪声干扰）
RWEFM ($\varepsilon=0.0002, N=256$)	$\mathbb{S}^2$ (MNIST '3')	0.172	0.165	>4h（OT精度极高，收敛最慢）
RWEFM ($\varepsilon=0.002, N=128$)	$\mathbb{S}^2$ (MNIST '3')	0.182	0.179	2h 45m（计算效率与生成质量最佳平衡）
RWEFM ($\varepsilon=0.2, N=32$)	$\mathbb{S}^2$ (MNIST '3')	0.385	0.392	<1h（熵过度平滑导致传输路径偏离测地线）

关键发现¶

内蕴流形几何的作用极其显著：在双曲空间 $\mathbb{H}^2$ 和通用网格上，忽略几何的欧氏 WFM 的 1-NN-D 普遍飙升至 0.40-0.50（接近完全不可分的上限），而 RWEFM 能够保持在 0.10-0.28；而在曲率为 0 的平坦环面 $\mathbb{T}^2$ 上，WFM 与 RWEFM 的差距明显缩小，从反面印证了流形曲率是决定非欧流匹配优势的关键因素。
高维空间的“采样映射”退化选择：在单细胞 $\mathbb{S}^{127}$ 极高维隐空间中，由于测度集中现象与维度灾难，熵映射的重心投影容易收缩至不真实的“平均粒子”，此时退化采用从输运计划中采样的离散对应点映射反而显著提升了生物样本异质性表达。
超参数单调权衡规律：训练速度与生成质量在 $(\varepsilon, N)$ 构成的网格中表现出严格的单调趋势，用户可根据显存与算力预算平滑调节精度与效率。

亮点与洞察¶

无穷维 Wasserstein 空间的微分几何与流匹配严密闭环：不仅在理论上完成了非欧流形上测度空间连续性方程的拉格朗日解的存在唯一性证明，而且给出非欧测度流回归的损失等价定理，填补了流形上元分布生成的数学空白。
切空间内的黎曼熵映射设计极具工程巧思：巧妙利用“对数映射提升到欧氏切空间求加权平均，再通过指数映射拉回流形”的代数结构，使复杂非欧流形上的最优点映射运算可以借助标准矩阵乘法在 GPU 上高效矢量化。
打通通用非解析网格与基础模型潜空间的实用接口：通过仅依赖测地距离与流形最近点投影算子，摆脱了传统微分几何对黎曼度量解析张量的强依赖，直接释放了其在单细胞基础模型与蛋白质动力学系综中的开箱即用价值。

局限与展望¶

截断轨迹与割迹（Cut Locus）的奇异性处理：理论分析假设对数映射全局光滑且单值，但在包含割迹的紧致流形上（如球面背对点、环面绕圈），测地线不唯一，当前工程实现通过数值强制选单值规避，缺乏更严谨的割迹边界测度处理。
超高维流形下熵平滑的重心塌缩：当底层流形维度极高（如细胞隐空间 $d=128$）时，熵正则化加权平均倾向于生成均值模长缩小的收缩点，未来需探索基于切空间稀疏注意力或测地极值投影的抗平滑估计器。
网格距离的拉普拉斯谱前度量依赖：在通用网格上仍需依赖双调和谱前度量（Biharmonic premetric）近似测地距离，对于数十万面元的高精度曲面，特征分解的显存开销构成了新的预处理瓶颈。

评分¶

新颖性: ⭐⭐⭐⭐⭐ [首次完整建立黎曼流形 Wasserstein 空间的流匹配理论并提出黎曼熵映射]
实验充分度: ⭐⭐⭐⭐⭐ [涵盖合成多曲率流形、通用离散兔子网格、真实生物单细胞与蛋白质动力学]
写作质量: ⭐⭐⭐⭐⭐ [理论推导严谨自洽，从附录测度叠加原理到工程实现细节完备透彻]
价值: ⭐⭐⭐⭐⭐ [为单细胞转录组学与计算生物物理构象系综提供了坚实的基础几何生成范式]

流形空间	任务类别	指标 (1-NN-D)	RWEFM (本文)	WFM (欧氏基线)	SetRFM (无OT流形)	RFM (点级流形)
双曲空间 \(\mathbb{H}^2\)	EMNIST 'h'	CD / EMD	0.15 / 0.09	0.30 / 0.28	0.30 / 0.31	0.48 / 0.48
双曲空间 \(\mathbb{H}^2\)	EMNIST 'w'	CD / EMD	0.23 / 0.10	0.41 / 0.39	0.34 / 0.33	0.44 / 0.45
球面 \(\mathbb{S}^2\)	MNIST '3'	CD / EMD	0.18 / 0.18	0.29 / 0.28	0.38 / 0.38	0.30 / 0.38
球面 \(\mathbb{S}^2\)	MNIST '4'	CD / EMD	0.19 / 0.18	0.29 / 0.27	0.38 / 0.38	0.42 / 0.45
平坦环面 \(\mathbb{T}^2\)	KMNIST 'ki'	CD / EMD	0.17 / 0.18	0.22 / 0.17	0.26 / 0.28	0.47 / 0.46
兔子网格 \(\mathcal{M}_{\text{mesh}}\)	MNIST '0'	CD / EMD	0.30 / 0.28	0.50 / 0.50	0.30 / 0.26	–
兔子网格 \(\mathcal{M}_{\text{mesh}}\)	MNIST '2'	CD / EMD	0.26 / 0.21	0.40 / 0.42	0.29 / 0.23	–

评估领域与基流形	评估模型	1-NN-D (CD) ↓	MMD (EMD) ↓	\(W_2\) ↓	\(W_1\) ↓	MMD-CD ↓
单细胞 scRNA-seq (\(\mathbb{S}^{127}\))	SetFM	0.4460	0.1026	–	–	0.1174
	WFM	0.2945	0.0324	–	–	0.0355
	SetRFM	0.3754	0.0300	–	–	0.0341
	RWEFM	0.1889	0.0140	–	–	0.0137
蛋白质构象系综 (\(\mathbb{T}^2\))	SetFM	–	0.0114	0.4472	0.4097	–
	WFM	–	0.0120	0.4059	0.4023	–
	SetRFM	–	0.0119	0.4608	0.4126	–
	RWEFM	–	0.0097	0.3970	0.3850	–

配置 / 变体	目标流形 & 数据集	1-NN-D (CD) ↓	1-NN-D (EMD) ↓	训练耗时 / 说明
RWEFM (黎曼熵映射)	\(\mathbb{H}^2\) (EMNIST 'h')	0.149	0.085	4h 06m（连续重心投影，保真度最高）
RWEFM (随机采样映射)	\(\mathbb{H}^2\) (EMNIST 'h')	0.192	0.157	~3h 40m（单粒子匹配，易受高维噪声干扰）
RWEFM (\(\varepsilon=0.0002, N=256\))	\(\mathbb{S}^2\) (MNIST '3')	0.172	0.165	>4h（OT精度极高，收敛最慢）
RWEFM (\(\varepsilon=0.002, N=128\))	\(\mathbb{S}^2\) (MNIST '3')	0.182	0.179	2h 45m（计算效率与生成质量最佳平衡）
RWEFM (\(\varepsilon=0.2, N=32\))	\(\mathbb{S}^2\) (MNIST '3')	0.385	0.392	<1h（熵过度平滑导致传输路径偏离测地线）