AdaGO融合自适应步长与正交更新,提升大模型训练效率
AdaGrad Meets Muon: Adaptive Stepsizes for Orthogonal Updates
- 结合AdaGrad步长与正交更新方向,保持更新路径正交性
- 在CIFAR-10和函数回归任务中优于Muon与Adam
- 仅需新增一个标量变量,计算内存开销极低
最近提出的Muon优化器通过正交化动量更新权重矩阵,在大规模语言模型训练中表现出色。然而,如何为这类正交更新设定学习率仍不明确。相比之下,AdaGrad是一种广泛应用的自适应方法,通过累积历史梯度对随机梯度进行缩放。本文提出新算法AdaGO,将基于范数的AdaGrad型步长与正交更新方向结合,兼具两者优势。与其它改进版Muon不同,AdaGO保持更新方向的正交性,可视为谱下降方向,同时通过累积历史梯度范数来适应优化景观。实现上仅需对Muon做微小修改,额外增加一个标量变量(累积平方梯度范数),计算与内存开销极低。在标准光滑性和无偏有界噪声假设下,建立了非凸函数在随机与确定性设置下的最优理论收敛速率。在CIFAR-10分类和函数回归任务上的实验表明,AdaGO性能优于Muon与Adam。
原文摘要 · Abstract (English)
The recently proposed Muon optimizer updates weight matrices via orthogonalized momentum and has demonstrated strong empirical success in large language model training. However, it remains unclear how to determine the learning rates for such orthogonalized updates. AdaGrad, by contrast, is a widely used adaptive method that scales stochastic gradients by accumulated past gradients. We propose a new algorithm, AdaGO, which combines a norm-based AdaGrad-type stepsize with an orthogonalized update direction, bringing together the benefits of both approaches. Unlike other adaptive variants of Muon, AdaGO preserves the orthogonality of the update direction, which can be interpreted as a spectral descent direction, while adapting the stepsizes to the optimization landscape by scaling the direction with accumulated past gradient norms. The implementation of AdaGO requires only minimal modification to Muon, with a single additional scalar variable, the accumulated squared gradient norms, to be computed, making it computationally and memory efficient. Optimal theoretical convergence rates are established for nonconvex functions in both stochastic and deterministic settings under standard smoothness and unbiased bounded-variance noise assumptions. Empirical results on CIFAR-10 classification and function regression demonstrate that AdaGO outperforms Muon and Adam.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。