面向图神经网络的模块感知优化

arXiv:2611.08379v1 公告类型：新论文 摘要：不同的优化器为参数更新赋予不同的几何结构。Adam 和 AdamW 等自适应方法采用逐坐标更新，而 Muon 等矩阵感知优化器则通过正交化更新利用权重矩阵的结构。然而，对于图神经网络（GNNs），其权重矩阵的有效几何结构不仅由学习到的通道混合权重决定，还受到其所作用的经图传播的表示的影响。因此，即使某个更新在矩阵层面得到了良好的归一化，当它应用于经图传播的特征时，对层输出的影响仍可能被放大。为弥补这一缺口，我们提出 GraphMuon，这是一种图感知优化器，可使 Muon 的矩阵更新适应经图传播的层输入的协方差结构。GraphMuon 在由这些传播后的表示定义的度量下执行正交化，并根据更新对层表示所产生的影响对步长进行归一化。我们从理论上证明，原始 Muon 可能受到图引起的更新放大效应的影响，而 GraphMuon 能够缓解这一效应。因此，GraphMuon 给出了 Muon 极分解更新步骤在图度量下的对应形式。在实证评估中，在三个大规模 OGB 基准上，GraphMuon 达到目标验证损失所需的时间相较于 AdamW 获得了 $1.32\times$-$2.42\times$ 的加速。在涵盖节点分类、归纳式图分类和图回归的 15 种数据集与架构配置中，它在 13 种配置下降低了达到目标所需的时间，同时保持了具有竞争力或更优的预测性能。代码可在 https://github.com/guy-hadad/GraphMuon 获取。