Riemannian Generative Decoder¶

会议: ICML 2026
arXiv: 2506.19133
代码: https://github.com/yhsure/riemannian-generative-decoder (有)
领域: 表示学习 / 几何深度学习 / 可解释生成模型
关键词: 黎曼流形、生成式解码器、几何正则化、MAP 估计、计算生物学

一句话总结¶

本文针对 Riemannian VAE 必须为每种流形手工设计复杂概率密度的痛点，提出 Riemannian Generative Decoder (RGD)——彻底丢掉 encoder，把每个样本的 latent 当作自由参数用黎曼优化器 (RiemannianAdam) 直接训，同时引入"按局部度量逆缩放的输入噪声"作为几何正则，在合成分支扩散树、人类线粒体 DNA、细胞周期 scRNA-seq 三个真实生物数据上恢复出更忠实的几何，且在高维下数值稳定胜过 VAE 基线。

研究背景与动机¶

领域现状：真实数据（进化树、社交网络、周期信号）常具有非欧几何结构，但主流表示学习几乎都假设 \(\mathbb{R}^d\) 隐空间，导致几何信息被强行压扁。Riemannian VAE 系列（\(\mathcal{S}\)-VAE 用 von Mises–Fisher、\(\mathcal{P}\)-VAE 用 Poincaré、\(\Delta\)VAE 用 Brownian motion 采样）尝试把 VAE 搬到球面、双曲、torus 等流形上，是当前 geometry-aware representation learning 的主流。

现有痛点：每种流形都需要专门设计概率密度，难处包括 (i) von Mises–Fisher 归一化常数含贝塞尔函数；(ii) Poincaré 上 Riemannian 正态分布的归一化常数 + volume correction 都得 Monte Carlo；(iii) wrapped normal 需要截断求和；(iv) \(\Delta\)VAE 走 Brownian motion 模拟。所有这些近似既数值不稳定（高维下尤其爆炸），又把流形选择牢牢绑在"是否能写出 tractable density"上，导致无法直接处理 ProductManifold 等异构组合几何。

核心矛盾：VAE 框架要求 encoder \(q_\phi(z|x)\) 与 prior \(p(z)\) 都在流形上有可算的密度，但绝大多数 Riemannian 流形上根本没有 closed-form density。强行近似的代价是 bias、unstable optimizer 与 limited manifold coverage——三者直接互锁。

本文目标：(i) 让表示学习能在任何 geoopt 支持的黎曼流形上做（包括 Product 组合）；(ii) 摆脱对 manifold density 的依赖；(iii) 给出一个几何感知的正则项，使 decoder 平滑度与流形局部度量对齐而不是固定欧式同向。

切入角度：Deep Generative Decoder (DGD) 思路——不要 encoder，把 latent \(z_i\) 当 free parameter 用 MAP 直接优化。把这一思路升级为 Riemannian DGD：用黎曼优化器替代欧式 Adam，每个 latent 始终留在流形上，从而绕过 density 近似。

核心 idea：抛掉 encoder + 用 RiemannianAdam 把 latent 当流形上的自由参数 + 在训练中加入协方差为 \(\sigma^2 G^{-1}(z)\) 的几何噪声作为正则，让 decoder 在度量大的方向自然变光滑。

方法详解¶

整体框架¶

给定数据 \(X=\{x_i\}_{i=1}^N \in \mathbb{R}^D\)，选一个 \(d\) 维黎曼流形 \((\mathcal{M},g)\) 作为 latent space。把 \(Z=\{z_i\}_{i=1}^N\) 视为每样本一个的自由参数，与 decoder \(f_\theta:\mathcal{M}\to\mathcal{X}\) 一起做 MAP 估计： \(\mathcal{L}(\theta, Z) = \sum_i \big(-\log p_\theta(x_i|z_i) - \log p(z_i)\big) - \log p(\theta)\) 其中 likelihood 通常取 isotropic Gaussian（重构损失即 MSE），prior 在 compact 流形上取 uniform、在 non-compact 流形上取 wrapped 或 Riemannian normal。训练交替：\(\theta\) 用 Adam 走欧式步、\(Z\) 用 RiemannianAdam 走黎曼步，每步用 retraction \(R_z(\cdot)\) 把 tangent 向量映回流形保证 \(z^{(t+1)}\in\mathcal{M}\)。借助 geoopt 库这套实现非常简洁。

关键设计¶

Encoder-less MAP + RiemannianAdam 直接在流形上优化 latent：
- 功能：把"每条样本的隐变量"做成可学习的自由参数，让任意黎曼流形（Sphere、PoincareBall、Lorentz、SPD、UpperHalf、Stiefel、ProductManifold...）都立即可用。
- 核心思路：放弃 amortized inference。每步训练对 latent 用黎曼梯度 \(\nabla_z^{\mathcal{R}}\mathcal{L}=G(z)^{-1}\nabla_z^E\mathcal{L}\)，配合 retraction (通常是 exponential map) 更新：\(z^{(t+1)}=R_{z^{(t)}}(-\eta\,\nabla_{z^{(t)}}^{\mathcal{R}}\mathcal{L})\)。RiemannianAdam 在 tangent space 维护自适应方向，保证收敛速度与 Adam 类似，但每步严格在流形上。Compact 流形 prior 直接取 \(1/\text{Vol}(\mathcal{M})\)（常数，不影响梯度），non-compact 用 wrapped/Riemannian normal。
- 设计动机：encoder 是 manifold density 近似的根源问题，因为 \(q_\phi(z|x)\) 必须在曲面上是 tractable 概率分布。Goldberg-DGD 已经证明丢掉 encoder 的 MAP 范式在欧式空间能 work；本文把它直接 lift 到流形，绕过了所有 density 近似。一个意外收益是 ProductManifold 这种异构积流形也立即可用，因为 RGD 不需要为每个流形给 prior。
几何感知的输入噪声正则：
- 功能：让 decoder 的局部 Jacobian 与流形度量自动对齐，鼓励几何上"应该相似的点"被解到相似的输出。
- 核心思路：训练时给 latent 注入噪声 \(\epsilon\sim\mathcal{N}(0, \sigma^2 G^{-1}(z))\)（协方差用流形度量逆，使噪声在度量大的方向上更弱、在小的方向更强），用 exponential map \(z'=\text{Exp}_z(\epsilon)\approx z+\epsilon\) 注入。作者跟随 Bishop (1995) 的二阶 Taylor 展开推导，得到等价正则项：\(\mathbb{E}_\epsilon[L(z')]\approx L(z)+\sigma^2\,\text{Tr}(J(z)^\top G^{-1}(z) J(z))\)。其中 \(J(z)=\partial_z f\)，加号后的项就是被流形度量加权的 Jacobian 范数惩罚。
- 设计动机：欧式高斯噪声在曲面上会让模型在 metric 大的方向被过度惩罚、metric 小的方向欠惩罚；用 \(G^{-1}(z)\) 缩放的噪声等于把 isotropic 正则对齐到流形几何，在 homogeneous 流形（球面）上退化为近似 isotropic，在曲率非均匀的双曲流形上则按位置自适应。相比 Lee & Park (2023) 的二阶曲率正则（涉及 Hessian-vector product，公式整页），RGD 的方案只用一次 Jacobian 计算，scalability 强很多。
统一框架支持任意黎曼流形 + Product 组合：
- 功能：让用户只需把 prior 知识表达为流形选择，剩下交给框架，无需手写 ELBO 或密度近似。
- 核心思路：直接复用 geoopt 实现的所有流形（Euclidean、Sphere、Stereographic、PoincareBall、Lorentz、SPD、Stiefel、UpperHalf、ProductManifold...），只要这些流形给出 exponential map / retraction / 度量即可。Product 流形写成 \(\mathcal{M}=\mathcal{M}_1\times\cdots\times\mathcal{M}_K\)，度量直接取 direct sum，自动覆盖"一部分维度球面 + 一部分维度双曲"这种异构需求。
- 设计动机：以往工作每写一个新流形就要重新推 prior 与近似，研究者无法快速对比"哪种几何更适合我的数据"。RGD 把流形从一个算法假设变成一个配置项，让 hypothesis-based exploration 真正可行；细胞周期数据上 torus、sphere、Euclidean 都能一键切换比较。

损失函数 / 训练策略¶

目标为公式 (10) 的负后验：\(\mathcal{L}=\sum_i(-\log p_\theta(x_i|z_i)-\log p(z_i))-\log p(\theta)\)。\(\theta\) 用 Adam、\(Z\) 用 RiemannianAdam 交替更新。几何正则通过对 latent 加 \(\mathcal{N}(0,\sigma^2 G^{-1})\) 噪声实现（一次 retraction 把噪声映回流形）。重构 loss 按数据性质选（连续 → Gaussian/MSE；离散 → categorical）。无 KL 项、无 ELBO，无 Monte Carlo 估计归一化常数。

实验关键数据¶

主实验¶

三个真实/合成生物数据集：(a) 合成分支扩散树（7 层、\(d=50\)、6350 样本）→ 双曲流形天然适配；(b) 人类线粒体 DNA 67k 条序列 + haplogroup 标签 → 双曲适配 phylogeny；(c) 细胞周期 scRNA-seq 5367 细胞 × 189 基因 → 周期性，torus 适配。

数据集	任务	最佳几何	关键指标	备注
Cell cycle scRNA-seq	相位距离 vs latent 距离相关	Sphere \(\mathbb{S}^2\)	Train Pearson 0.58、Test 0.60，重构 MAE 0.31	超过 \(\mathcal{S}\)-VAE / \(\Delta\)VAE
Branching diffusion	树距离 vs latent 测地距离相关	Lorentz \(\mathbb{H}^2\)（\(\sigma=1.0\)）	Train Pearson 0.81、Test 0.80	\(\mathcal{P}\)-VAE 仅 0.68
hmtDNA haplogroup 分类	下游 24/128 类 logistic 回归准确率	Hyperbolic \(\mathbb{H}^2_{\sigma=0.5}\)	24-way LR 0.70 / XGB 0.85；128-way LR 0.43	全面胜 Euclidean 和 \(\mathcal{P}\)-VAE

消融实验¶

配置	关键发现
几何噪声 \(\sigma\) 扫 0→2.6（双曲，分支扩散）	\(\sigma\approx 0.9\) 前相关性快速上升，之后噪声过大 decoder 无法保持局部精度 → 局部-全局 trade-off
双曲 \(\mathbb{H}^2\) vs Euclidean / Sphere（分支扩散）	\(\mathbb{H}^2\) 显著恢复树拓扑（Pearson 0.81 vs 0.53/0.56）
UMAP（分支扩散）	完全看不出树拓扑 → 验证 RGD 在解释性上的明显优势
与 \(\mathcal{P}\)-VAE/\(\mathcal{S}\)-VAE/\(\Delta\)VAE 比	高维下 VAE baseline 数值崩溃，RGD 仍可用
生成判别（XGBClassifier 区分真伪重构）	RGD Sphere 0.58 = \(\mathcal{S}\)-VAE Sphere 0.58 < \(\Delta\)VAE 0.62（越接近 0.5 越逼真）

关键发现¶

encoder-less 在高维数据下更稳：在 cell cycle 全基因（高维）+ ProductManifold 等设置下，\(\mathcal{S}\)/\(\Delta\)/\(\mathcal{P}\)-VAE 都因为 normalizing constant 或采样 Monte Carlo 数值崩溃，RGD 仍可训练并给合理重构。
几何噪声是真正学到几何的关键：移除它后双曲模型的距离-相关性接近 Euclidean baseline，说明流形选择本身只给了空间，真正把 decoder 推到与几何对齐还得靠正则。
流形选择 = 假设测试：分支扩散数据用双曲 latent 一秒还原树结构、用 UMAP 完全看不出，说明合理的几何先验比通用降维更可解释。
生成质量与 VAE 持平：discrimination test 上 RGD 的合成样本可分性与 \(\mathcal{S}\)-VAE 相同（0.58），说明丢掉 encoder 没损失生成保真度。

亮点与洞察¶

思路足够干净：把 DGD 从欧式空间一步推到任意黎曼流形，绕过了 Riemannian VAE 的整条 density approximation 链——这是少有的"减法贡献"，但带来的可扩展性极大。
几何噪声推导简单且 actionable：用 Bishop 噪声-正则等价的经典结论加一层度量逆缩放，就能在任何流形上得到 metric-aware Jacobian 惩罚；公式只有一行，工程实现也只是替换协方差。
使流形成为 hypothesis 而非算法负担：让科学家可以快速对同一份数据试 torus、sphere、双曲、SPD，pick 最符合先验的几何，对 computational biology / phylogenetics 有直接生产力提升。
可迁移到 LLM/VLM 隐表示：把 LLM 的 routing/expert 表示放到 Stiefel 或 Sphere、把 vision-language alignment 放到 ProductManifold，都能套用这套框架。

局限与展望¶

没有 amortized encoder 意味着新样本必须重新优化 latent（per-sample MAP 步）才能 encode，online inference 比 VAE 慢；可加 amortized warm-start。
几何正则中 \(\sigma\) 与流形曲率耦合，需要手工 sweep；自动调度是显然的下一步。
prior 在 non-compact 流形上仍用 wrapped/Riemannian normal，理论上引入轻微 bias，但作者明确把生成质量当作 sanity check。
实验全部在生物/合成数据上，未在 NLP/图像等大规模任务上验证（虽然原理可直接套用）。
ProductManifold 的维度分配方案需要人工设定，对复杂数据可能不易选；自动几何选择是未来方向。

评分¶

新颖性: ⭐⭐⭐⭐ "丢掉 encoder 把 DGD 升级到任意黎曼流形" + "metric-逆噪声" 都是干净独立的贡献，框架级简化效果明显。
实验充分度: ⭐⭐⭐⭐ 三个数据集覆盖周期/树/层次三种几何，多基线对比 + 噪声扫描 + downstream + 生成判别俱全；缺大规模 NLP/视觉验证。
写作质量: ⭐⭐⭐⭐⭐ 公式推导清晰，附录详尽，对前置工作介绍极完整，理论与工程的过渡非常顺。
价值: ⭐⭐⭐⭐ 在 geometric DL 中把 manifold 从"算法门槛"降为"配置项"，对 representation learning、computational biology 生产力提升显著。