在大型语言模型训练中,高效的随机优化算法至关重要。现有方法如Adam通过自适应矩估计来提升稳定性,而Muon则利用权重层的矩阵结构,通过正交化动量取得了优异表现。
本研究提出了两种新的优化器:NAMO及其对角扩展版本NAMO-D。它们首次从原理上实现了正交化动量与基于范数的Adam型噪声自适应机制的整合。NAMO采用单一自适应步长来缩放正交化动量,在保持正交性的同时,以可忽略的额外成本改进了Muon的性能。
NAMO-D则通过一个带截断项的对角矩阵右乘正交化动量。这种设计实现了神经元级的噪声自适应,并且与常见的近似块对角海森矩阵结构相吻合。在标准假设下,研究证明了两种算法在确定性设置下的最优收敛速率,并表明在随机设置中,其收敛保证能够自适应随机梯度的噪声水平。
在GPT-2模型的预训练实验中,NAMO和NAMO-D均表现出比AdamW和Muon基线更优的性能。其中,NAMO-D通过一个额外的截断超参数,在维持良好条件更新方向与利用细粒度噪声自适应这两个竞争目标之间取得平衡,从而实现了比NAMO进一步的性能提升。该研究为优化器设计提供了新的思路,特别是在结合结构感知动量与噪声自适应机制方面。


