图神经网络优化器换几何：GraphMuon按传播后特征定步长，三个OGB基准比AdamW快1.32至2.42倍
核心概要
图神经网络的权重矩阵作用在经图传播的特征上，因此一个在参数空间里归一化良好的更新，仍可能被图结构放大成很大的表示层扰动。GraphMuon把Muon的矩阵正交化改到由图传播输入协方差定义的度量下进行，并按更新对层表示的实际影响缩放步长。在三个大规模OGB基准上，它达到AdamW验证损失目标所需的实际耗时缩短为原来的1/2.32至1/1.32，即1.32至2.42倍加速。
深度剖析
图神经网络的优化器需要按图传播后的表示来度量更新，而不只是按参数矩阵本身。 AdamW按坐标逐维归一化梯度矩，Muon按权重矩阵做近似正交化，两者都不区分被图传播放大的方向与被衰减的方向。 在OGBN-ARXIV、OGBN-MAG、OGBN-PRODUCTS三个大规模OGB基准的GCN骨干上，GraphMuon达到AdamW验证损失目标的实际耗时加速分别为2.32倍、2.42倍和1.32倍，对应95%自助法置信区间为[2.01, 2.69]、[2.29, 2.53]和[1.19, 1.37]。
做法是把Muon的极分解式更新搬到图诱导的度量里，再按表示层扰动定步长。 GraphMuon先由每层图传播输入算出协方差，用其幂次把动量映射进该度量、做Newton–Schulz正交化再映射回来，并用同一协方差把更新缩放到其引起的表示层扰动大小。 在覆盖节点分类、归纳式图分类和图回归的15个数据集–架构组合上，GraphMuon在全部15个组合中都比AdamW更少轮次达到目标，在13个组合中实际耗时最短，平均加速4.62倍，同时保持有竞争力或更好的预测性能。
启示与展望
评估覆盖九个图学习基准上的15个数据集–架构组合，包括直推与异配节点分类、归纳式分子图分类和ZINC图回归，骨干为GCN、GAT、GINE、GraphGPS和GraphGCN，全部在单块NVIDIA RTX 6000上运行。
当图很小、节点数相对隐藏维度偏少时，图诱导协方差可能低秩且噪声大，此时控制表示层扰动不一定带来更快收敛或更好泛化，AdamW仍可能保持竞争力。在Citeseer配GCN和ZINC配GraphGCN两个配置上，GraphMuon轮次更少但累计实际耗时略慢于AdamW。图感知强度τ的最优取值依数据集而异。