arXiv:2602.17080cs.LGmath.OC2026-02被引 11

将正交动量与自适应步长结合,提升大模型训练效率

Adam Improves Muon: Adaptive Moment Estimation with Orthogonalized Momentum

  • 用自适应步长调节正交动量,保持更新方向稳定
  • 在GPT-2预训练中优于AdamW和Muon,NAMO-D性能更优
  • 适合追求高效训练的大模型研究者使用

高效随机优化通常结合确定性场景下表现良好的更新方向与适应随机扰动的机制。虽然Adam通过自适应矩估计提升稳定性,Muon则利用权重层的矩阵结构实现正交动量,在大语言模型训练中表现更优。本文提出新优化器NAMO及其对角扩展NAMO-D,首次将正交动量与基于范数的Adam型噪声适应有机结合。NAMO通过单一自适应步长缩放正交动量,保持正交性且开销极低;NAMO-D则通过右乘带截断项的对角矩阵,实现神经元级噪声适应,契合近似块对角海森结构。在标准假设下,两类算法在确定性设定中达到最优收敛率,并在随机设定中其收敛保证可随梯度噪声水平自适应调整。在GPT-2预训练实验中,两者均优于AdamW和Muon基准,其中NAMO-D通过额外的截断超参数进一步提升性能,平衡了更新方向条件良好与细粒度噪声适应之间的矛盾。

原文摘要 · Abstract (English)

Efficient stochastic optimization typically integrates an update direction that performs well in the deterministic regime with a mechanism adapting to stochastic perturbations. While Adam uses adaptive moment estimates to promote stability, Muon utilizes the weight layers' matrix structure via orthogonalized momentum, showing superior performance in large language model training. We propose a new optimizer and a diagonal extension, NAMO and NAMO-D, providing the first principled integration of orthogonalized momentum with norm-based Adam-type noise adaptation. NAMO scales orthogonalized momentum using a single adaptive stepsize, preserving orthogonality while improving upon Muon at negligible additional cost. NAMO-D instead right-multiplies orthogonalized momentum by a diagonal matrix with clamped entries. This design enables neuron-wise noise adaptation and aligns with the common near block-diagonal Hessian structure. Under standard assumptions, we establish optimal convergence rates for both algorithms in the deterministic setting and show that, in the stochastic setting, their convergence guarantees adapt to the noise level of stochastic gradients. Experiments on pretraining GPT-2 models demonstrate improved performance of both NAMO and NAMO-D compared to the AdamW and Muon baselines, with NAMO-D achieving further gains over NAMO via an additional clamping hyperparameter that balances the competing goals of maintaining a well-conditioned update direction and leveraging fine-grained noise adaptation.

优化器大模型训练正交动量自适应

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。