跳转至

Adaptive Riemannian Graph Neural Networks

会议: AAAI 2026
arXiv: 2508.02600
代码: 有(公开仓库,基于 PyG + Geoopt)
领域: 图神经网络 / 几何深度学习
关键词: 黎曼几何, 自适应度量张量, Ricci流正则化, 几何异质性, 消息传递

一句话总结

提出 ARGNN 框架,为图上每个节点学习一个连续的、各向异性的对角黎曼度量张量,从而自适应地捕获图中不同区域(层级结构 vs 密集社区)的局部几何特性,统一并超越了固定曲率和离散混合曲率的几何 GNN 方法。

背景与动机

真实图数据普遍存在几何异质性:同一网络中既有适合双曲空间表示的树状层级结构,也有适合球面空间的密集环状社区。现有几何 GNN 方法要么将整个图嵌入到单一固定曲率空间(Euclidean/Hyperbolic/Spherical),要么使用离散积空间(如 CUSP 的 ℍ×𝕊×𝔼),无法充分表达节点级别的连续几何变化。κ-GCN 为每个节点学习标量曲率,但仍是各向同性的,无法刻画方向性几何信息。

作者在 Wisconsin 网络上的可视化清晰展示了:不同区域的曲率(从平坦到强弯曲)差异显著,验证了固定几何空间必然在某些区域产生严重失真。

核心问题

如何在保持计算可行性的前提下,为图上每个节点学习一个对称正定(SPD)度量张量,使消息传递能够根据局部几何自适应调整,同时确保训练稳定性和理论可证明的表达能力?

方法详解

整体框架

ARGNN 包含三个核心组件:(1) 节点级对角度量张量的学习,(2) 基于学到的度量的几何消息传递,(3) Ricci 流启发的几何正则化。端到端联合学习几何场与节点表示。

关键设计

1. 对角度量张量参数化

每个节点 i 的度量张量被参数化为对角矩阵 \(\mathbf{G}_i = \text{diag}(\mathbf{g}_i)\),其中 \(\mathbf{g}_i \in \mathbb{R}^d_{++}\)。这不仅是计算上的简化(参数从 \(O(d^2)\) 降至 \(O(d)\)),而且对应着各向异性保角变换——每个特征维度有独立的局部缩放因子,形式上是全度量张量和标量曲率之间一个精巧的折中。

度量向量通过一个小型 MLP 生成:将节点特征 \(\mathbf{h}_i\) 和邻域聚合特征 \(\mathbf{a}_i\) 拼接后,经 softplus 激活保证严格正性:

\[\mathbf{g}_i = \text{softplus}\left(f_\theta^{(g)}([\mathbf{h}_i; \mathbf{a}_i])\right)\]

2. 几何消息传递

  • 测地距离:在度量 \(\mathbf{G}_i\) 下变为加权欧氏距离 \(d_{\mathbf{G}_i}(\mathbf{h}_i, \mathbf{h}_j) = \sqrt{\sum_k g_{i,k}(h_{i,k} - h_{j,k})^2}\)
  • 几何调制系数 \(\tau_{ij}\):将方向向量投影到主轴上,用 \(\tanh(-\log g_{i,k})\) 作为曲率开关——\(g_{i,k}\) 大时趋近 -1(空间收缩),小时趋近 +1(空间膨胀)
  • 几何注意力 \(\alpha_{ij}\):在各自度量下计算余弦相似度,每个节点的范数在自己的度量空间中测量

消息更新为 \(\mathbf{m}_{ij} = \tau_{ij} \cdot \sigma(\alpha_{ij}) \cdot \mathbf{W}_m \mathbf{h}_j\)

3. Ricci 流正则化

离散 Ricci 曲率在第 k 维近似为 \(\text{Ric}_{kk}^{(i)} = \frac{1}{2|\mathcal{N}(i)|} \sum_{j \in \mathcal{N}(i)} \frac{g_{i,k} - g_{j,k}}{d_{\text{graph}}(i,j)}\),两个正则项: - \(\mathcal{L}_{\text{Ricci}}\):惩罚 Ricci 曲率的平方和,鼓励 Ricci 平坦性 - \(\mathcal{L}_{\text{smooth}}\):惩罚相邻节点度量向量的差异,保证几何场平滑

损失函数 / 训练策略

\[\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{task}} + \alpha \mathcal{L}_{\text{Ricci}} + \beta \mathcal{L}_{\text{smooth}}\]

理论推导了最优超参数与同质率 \(\mathcal{H}\) 的关系:\(\alpha^* \propto \mathcal{H}/L\)\(\beta^* \propto d/|\mathcal{V}|\)。常数 \(c_1 = (1 - \mathcal{H}) + 0.1\)\(c_2 = 0.1(1 + \mathcal{H})\),实验验证与网格搜索最优值差距在 0.5% 以内。使用 Adam + Riemannian Adam (Geoopt) 优化器。

实验关键数据

在 9 个基准数据集上进行节点分类和链接预测:

指标 Cora Actor Wisconsin 对比对象
F1 (分类) 86.83 42.18 90.65 CUSP: 83.45/41.91/88.30
AUROC (链接) 91.03 76.40 77.48 CUSP: 89.85/74.20/74.50
  • 节点分类:9 个数据集全部最优,Cora 比 CUSP 高 3.38%,Wisconsin 比最佳基线高 2.35%
  • 链接预测:9 个数据集全部最优,Actor 76.40% vs GNRF 73.50%
  • 效率:比 CUSP 快约 35%,内存比全张量方法低约 40%,与 HGCN 接近

消融实验要点

  • 去掉 Ricci 正则化:异质图损失更大(Actor -1.3%, Wisconsin -1.8%)
  • 去掉平滑正则化:Wisconsin 下降 3.5%(混合同质结构需要平滑过渡)
  • 固定几何 vs 自适应:ARGNN 在异质图上比固定几何高约 5%(Actor)
  • 理论指导的超参数与网格搜索差距仅 0.3-0.5%
  • 最优层数 L=3,嵌入维度 d=128

亮点

  1. 参数化简洁优雅:对角度量张量在计算效率、几何表达力和可解释性之间取得精巧平衡——每个 \(g_{i,k}\) 直接量化了第 k 个特征维度对节点 i 的几何重要性
  2. 理论完备性强:证明了收敛保证、通用近似性(统一 Euclidean/Hyperbolic/Spherical/Product 空间为特例)、泛化界和鲁棒性
  3. 理论指导实践:基于同质率的超参数公式在实验中高度有效,减少 100× 调参开销
  4. 可解释的学习几何:学到的曲率分布与图的同质率高度一致——异质图曲率更大、度量离散度更高,可视化结果提供了对图结构的直观几何理解

局限性 / 可改进方向

  1. 对角约束的表达力瓶颈:对角度量无法捕获特征维度间的相关性和旋转几何,对于高度纠缠的特征空间可能受限。作者也承认低秩分解 \(\mathbf{G}_i = \mathbf{L}_i \mathbf{L}_i^T\) 是值得探索的中间方案
  2. 可扩展性:尽管复杂度与标准 GNN 同阶,但为每个节点维护 \(O(d)\) 的度量向量在超大规模图上仍有负担(作者提到可通过聚类共享度量缓解)
  3. 深层网络退化:L>3 后性能下降,几何正则化未能完全解决过平滑问题
  4. 实验场景有限:仅测试了节点分类和链接预测,未涉及图级任务和异构图

与相关工作的对比

方法 几何类型 节点自适应 各向异性 端到端学习
HGCN 固定双曲
κ-GCN 标量曲率
CUSP 离散积空间 部分
GNRF 固定几何+Ricci演化
ARGNN 连续对角度量场

ARGNN 是首个学习连续且各向异性度量张量场的图学习框架,在表达层次结构和理论统一性上均超越先前方法。

启发与关联

  • 对角度量张量的思路可推广到点云处理、分子图等其他几何数据——用 per-point 的各向异性度量替代全局度量
  • Ricci 流正则化的离散化方式值得在其他流形学习场景借鉴(如表示学习中的曲率控制)
  • 同质率感知的超参数选择策略可迁移到其他需要图结构先验知识的方法中

评分

  • 新颖性: ⭐⭐⭐⭐⭐ — 首次将连续各向异性黎曼度量场引入 GNN,概念上是对几何 GNN 范式的本质推进
  • 理论深度: ⭐⭐⭐⭐⭐ — 收敛性、通用性、泛化界、鲁棒性证明齐全,理论-实验闭环
  • 实验质量: ⭐⭐⭐⭐ — 全面的消融和多指标评估,但缺少图级任务和更大规模数据集
  • 实用性: ⭐⭐⭐⭐ — 计算高效,理论指导调参,但需测试工业场景适用性
  • 写作质量: ⭐⭐⭐⭐⭐ — 结构清晰,动机阐述充分,理论与实验衔接自然