arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2602.17080cs.LGmath.OC
收藏

Adam改进Muon:基于正交动量的自适应矩估计

Adam Improves Muon: Adaptive Moment Estimation with Orthogonalized Momentum

Minxin Zhang, Yuxuan Liu, Hayden Schaeffer

AI总结:

本文提出NAMO和NAMO-D优化器,通过正交化动量与基于范数的Adam型噪声适应相结合,提升大语言模型训练性能。

摘要:

高效的随机优化通常将在确定性领域表现良好的更新方向与适应随机扰动的机制相结合。虽然Adam使用自适应矩估计来促进稳定性,但Muon通过利用权重层的矩阵结构,通过正交化动量实现,展现出在大语言模型训练中的优越性能。我们提出了一种新的优化器和对角扩展,NAMO和NAMO-D,实现了正交化动量与基于范数的Adam型噪声适应的首次系统整合。NAMO通过单个自适应步长缩放正交化动量,保留正交性的同时在可忽略的额外成本上优于Muon。NAMO-D则通过右乘一个具有钳制条目的对角矩阵来正交化动量。这种设计实现了神经元层面的噪声适应,并与常见的近块对角Hessian结构相一致。在标准假设下,我们为这两种算法在确定性设置中建立了最优收敛率,并表明在随机设置中,其收敛保证能够适应随机梯度的噪声水平。在预训练GPT-2模型的实验中,NAMO和NAMO-D均优于AdamW和Muon基线,其中NAMO-D通过一个额外的钳制超参数在维持良好条件更新方向和利用细粒度噪声适应之间取得进一步收益。

英文摘要:

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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑