图神经网络优化器换几何:GraphMuon按传播后特征定步长,三个OGB基准比AdamW快1.32至2.42倍
图神经网络的权重矩阵作用在经图传播的特征上,因此一个在参数空间里归一化良好的更新,仍可能被图结构放大成很大的表示层扰动。GraphMuon把Muon的矩阵正交化改到由图传播输入协方差定义的度量下进行,并按更新对层表示的实际影响缩放步长。在三个大规模OGB基准上,它达到AdamW验证损失目标所需的实际耗时缩短为原来的1/2.32至1/1.32,即1.32至2.42倍加速。
图神经网络的权重矩阵作用在经图传播的特征上,因此一个在参数空间里归一化良好的更新,仍可能被图结构放大成很大的表示层扰动。GraphMuon把Muon的矩阵正交化改到由图传播输入协方差定义的度量下进行,并按更新对层表示的实际影响缩放步长。在三个大规模OGB基准上,它达到AdamW验证损失目标所需的实际耗时缩短为原来的1/2.32至1/1.32,即1.32至2.42倍加速。
该工作针对 Muon 优化器在动量累积与调优有限步 Newton–Schulz 正交化下的收敛问题,证明在足够宽的两层 ReLU 网络全批量训练中,只要学习率与 (1-μ)√ε 成正比,Muon 即可以高概率达到任意目标经验平方损失 ε,命中时间上界为 O((1-μ)-1ε-1/2),且所需宽度与目标精度和动量无关。
已读完全部报道