Muon 优化器

Muon(MomentUm Orthogonalized by Newton-schulz)是一种专门用于优化神经网络二维权重矩阵的优化器,由 Keller Jordan 等人在 2024 年提出。它的特点是先计算带动量的梯度,再通过牛顿-舒尔茨(Newton-Schulz)迭代把该更新矩阵近似正交化,使各方向的更新幅度更加均衡。

概述

传统优化器如 AdamW 对参数逐元素地自适应缩放学习率,而 Muon 把权重视为矩阵整体处理。研究者发现,深度网络中许多梯度更新在少数奇异方向上能量集中,导致更新方向不平衡;Muon 通过正交化把奇异值近似压平,让更新更均匀地分布到各方向,从而提高训练效率。它通常只作用于线性层与注意力投影等二维参数,而嵌入层、输出头和标量参数仍交给 AdamW。

核心原理

• 动量累积:先像常规做法一样维护梯度的一阶动量。
• 正交化近似:用几次牛顿-舒尔茨迭代把动量矩阵近似映射为正交矩阵,避免了昂贵的奇异值分解。
• 矩阵级更新:更新以整个权重矩阵为单位施加,而非逐元素缩放。
• 混合使用:一维参数与嵌入等仍用 AdamW,二者配合。

应用场景

Muon 主要用于大语言模型和视觉模型的预训练阶段,在相同算力预算下往往能更快达到目标损失,或用更少 token 达到同等效果。一些开源模型训练报告显示,采用 Muon 可显著降低训练成本,因此它在追求训练效率的团队中受到关注。

常见问题

问:Muon 能完全取代 AdamW 吗?答:通常不能。Muon 面向二维权重矩阵,而嵌入层、输出层和归一化参数一般仍使用 AdamW,二者是配合关系而非替代关系。

问:牛顿-舒尔茨迭代会不会很慢?答:它只需少数几次矩阵乘法即可得到足够好的近似,开销远小于精确奇异值分解,因此额外成本可控。