Muon优化器

参考：

1. https://kexue.fm/archives/10592

2. https://kexue.fm/archives/10739

3. https://kexue.fm/archives/10922

4. https://kexue.fm/archives/10996

核心摘要

• 核心思想：将参数更新从逐元素提升到矩阵层面，通过正交化梯度矩阵实现更高效、更稳定的参数更新。

• 关键创新：使用Newton-Schulz迭代法近似求解梯度矩阵的正交化形式，在保持计算效率的同时提升收敛速度。

• 实际效益：训练效率提升约2倍，显著降低大模型研发成本，并可直接复用AdamW的超参数，降低工程成本。

MUON优化器：矩阵正交化与高效训练的数学革命

在深度学习优化算法领域，随着模型规模的指数级增长，传统优化器如AdamW在处理超大规模模型时逐渐显现出局限性。2024年，Keller Jordan提出了一种革命性的新优化器——MUON（Matrix Update with Orthogonalization），其核心思想是将参数更新从逐元素处理提升到矩阵层面，通过正交化梯度矩阵实现更高效、更稳定的参数更新。MUON的创新性在于它不再简单地沿梯度方向更新，而是基于矩阵的几何特性，通过Newton-Schulz迭代法近似求解梯度矩阵的正交化形式，从而在保持计算效率的同时显著提升收敛速度和稳定性。

一、MUON的数学基础：矩阵正交化与Newton-Schulz迭代

1. 矩阵正交化目标

MUON的核心在于矩阵的正交化处理。对于一个矩阵参数 $W \in \mathbb{R}^{d_{in} \times d_{out}}$ ，其梯度为 $G \in \mathbb{R}^{d_{in} \times d_{out}}$ ，MUON的目标是找到一个正交矩阵$O$，使得 $O$ 尽可能接近梯度 $G$ ，同时满足正交矩阵的约束条件：

$$O = \arg\min_{O} \{\|O - G\|_F \quad \text{s.t.} \quad O^T O = I \text{ 或 } O O^T = I\}$$

这里的 $\| \cdot \|_F$ 表示Frobenius范数，$I$ 为单位矩阵。MUON通过矩阵的SVD分解来实现这一目标：

$$G = U \Sigma V^T$$

Frobenius范数：通常表示为 $\| A \|_F$ ，其定义为矩阵所有元素的平方和的平方根。

$$\| A \|_F = \sqrt{\sum_{i=1}^m\sum_{j=1}^n |a_{ij}|^2}$$

其中， $a_{ij}$ 表示矩阵 $A$ 的第 $i$ 行第 $j$ 列的元素，$m$ 表示矩阵的行数，$n$ 表示矩阵的列数。

MUON的更新方向为$UV^T$，即保留梯度矩阵的奇异向量，而丢弃奇异值。这种"对偶化"处理使得参数更新方向保持正交性，避免了梯度方向之间的相互干扰，从而能够更均匀地调整参数矩阵的各个维度。

核心洞察 (Core Insight)

MUON的更新方向为 $UV^T$，即保留梯度矩阵的奇异向量，而丢弃奇异值。这种“对偶化”处理使得参数更新方向保持正交性，避免了梯度方向之间的相互干扰，从而能够更均匀地调整参数矩阵的各个维度。

2. Newton-Schulz迭代方法

直接计算矩阵的正交化形式$UV^T$需要进行SVD分解，计算复杂度较高。MUON采用了一种高效的替代方法——Newton-Schulz迭代，通过五阶迭代多项式近似求解：

$$X_{k+1} = a X_k + b (X_k X_k^T) X_k + c (X_k X_k^T)^2 X_k$$

其中$X_k$是迭代过程中的矩阵，$M_k = X_k X_k^T$是$X_k$的格拉姆矩阵（Gram matrix），系数$a=3.4445$，$b=-4.7750$，$c=2.0315$。初始值$X_0$设置为：

$$X_0 = \frac{G}{\|G\|_F + \epsilon}$$

这里的$\epsilon$是一个极小的常数，防止除以零。Newton-Schulz迭代需要梯度矩阵$G$的谱范数小于1才能保证收敛，因此MUON会对梯度进行归一化处理。

通过5次迭代，$X_5$将收敛到一个近似正交矩阵，可作为梯度矩阵$G$的正交化替代。这种迭代方法仅需5%的额外计算开销，却实现了矩阵层面的优化，显著提升了MUON的实用性。

二、MUON优化器的完整更新公式

MUON优化器的完整更新过程可分为三个主要步骤：动量累积、Newton-Schulz迭代正交化和参数更新。

更新流程 (Update Pipeline):

• 动量累积: 通过指数加权平均累积历史梯度，保留历史信息，加速收敛。公式: $M_t = \beta M_{t-1} + G_t$

• 正交化: 核心创新！对动量缓冲矩阵 $M_t$ 进行Newton-Schulz迭代，得到近似正交矩阵**$O_t$**，保证更新方向的正交性。

• 参数更新: 根据不同的MUON变体，使用正交矩阵**$O_t$**和相应的缩放因子来更新模型参数。

1. 动量累积

MUON的第一步与动量优化器类似，通过指数加权平均累积历史梯度：

$$M_t = \beta M_{t-1} + G_t$$

其中$M_t$是动量缓冲矩阵，$\beta$ 是SGD动量系数（通常设为0.9），$G_t$ 是当前时间步的梯度矩阵。这一步骤保留了历史梯度信息，有助于加速收敛和抑制震荡。

2. Newton-Schulz正交化

MUON的第二步是其创新所在，即通过Newton-Schulz迭代对动量缓冲矩阵$M_t$进行正交化处理：

$$X_{k+1} = a X_k + b (X_k X_k^T) X_k + c (X_k X_k^T)^2 X_k$$

初始值 $X_0 = \frac{M_t}{\|M_t\|_F + \epsilon}$ ，经过5次迭代后得到近似正交矩阵 $O_t = X_5$ 。这一步骤的关键在于：

• 保留了梯度的方向信息
• 丢弃了梯度的幅度信息（奇异值）
• 保证了更新方向的正交性
• 计算复杂度可控，仅增加约5%的额外开销

3. 参数更新

MUON的第三步是参数更新，根据不同的版本，更新公式有所不同：

朴素版MUON：

$$W_t = W_{t-1} - \eta \cdot \text{sign}(\|M_t\|_F) \cdot O_t$$

Keller Jordan原始版MUON：

$$W_t = W_{t-1} - \eta \cdot \sqrt{\max(1, \frac{d_{out}}{d_{in}})} \cdot O_t$$

MuP版MUON：

$$W_t = W_{t-1} - \eta \cdot \sqrt{\frac{d_{out}}{d_{in}}} \cdot O_t$$

Moonlight变体MUON（月之暗面改进版）：

$$W_t = W_{t-1} - \eta \cdot 0.2 \cdot \max(d_{out}, d_{in}) \cdot O_t$$

其中$\eta$是学习率，$d_{in}$和$d_{out}$是参数矩阵的输入和输出维度。Moonlight变体引入了形状感知学习率调整机制，通过$0.2 \cdot \max(d_{out}, d_{in})$的缩放因子，确保不同形状矩阵的更新RMS值与AdamW一致，从而解决了MUON在大规模训练中可能面临的学习率调整问题。

Moonlight 变体的关键改进

Moonlight变体引入了形状感知学习率调整机制，通过 $0.2 \cdot \max(d_{out}, d_{in})$ 的缩放因子，确保不同形状矩阵的更新RMS值与AdamW一致，从而解决了MUON在大规模训练中可能面临的学习率调整问题。

三、Moonlight变体：大规模训练的改进策略

月之暗面团队在MUON基础上进行了两项关键改进，形成了Moonlight变体，使其更适合大规模语言模型训练：

1. 权重衰减机制

MUON在大规模训练中存在一个潜在问题：权重和层输出的RMS值可能持续增长，超出bf16的高精度范围，导致模型性能下降。为解决这一问题，Moonlight变体引入了AdamW风格的权重衰减：

$$W_t = W_{t-1} - \eta \cdot (O_t \cdot 0.2 \cdot \max(d_{out}, d_{in}) + \lambda W_{t-1})$$

其中$\lambda$是权重衰减系数。这一改进机制类似于AdamW的权重衰减，但整合到了MUON的矩阵更新框架中。实验证明，虽然原始MUON在初期收敛速度更快，但引入权重衰减后，模型在长期训练中表现更稳定，过拟合现象减少。

潜在问题与解决方案

原始MUON的潜在问题：权重和层输出的RMS值可能持续增长，超出bf16的高精度范围，导致模型性能下降。Moonlight变体的解决方案：引入AdamW风格的权重衰减机制，使长期训练更稳定。

2. 更新缩放机制

MUON的更新缩放机制基于其数学特性。对于形状 $[A,B]$ 的全秩矩阵，MUON的更新RMS值为 $\frac{1}{\max(A,B)}$ 。为使不同形状矩阵的更新量级一致，Moonlight变体引入了更新缩放：

$$O_t^{\text{scaled}} = O_t \cdot \max(A,B) \cdot 0.2$$

其中$0.2$的系数是为了与AdamW的默认表现相匹配。这种缩放机制确保了MUON在大规模训练中能够直接复用AdamW的超参数，无需重新调整。

四、MUON与AdamW的比较分析

MUON与AdamW作为两种不同的优化策略，各有其优势和适用场景：

特性	MUON	AdamW
更新策略	矩阵正交化更新	逐元素自适应学习率
计算开销	低（仅需5%额外开销）	中等
内存占用	低（仅需存储动量缓冲）	中等
小规模模型收敛速度	快（比AdamW快约2倍）	中等
大规模模型收敛稳定性	需Moonlight变体改进	稳定
超参数敏感度	对形状敏感（需调整）	相对稳定
理论基础	基于矩阵几何与牛顿迭代	基于矩估计与动量法

MUON的核心优势在于其矩阵正交化处理，使得参数更新方向更加"均匀"和"有效"，尤其在小规模模型上收敛速度显著提升。实验表明，在800M参数模型上训练100B tokens时，MUON的验证损失明显低于AdamW，但原始MUON在大规模模型上表现不佳。

核心优势 (Core Advantage)

MUON的核心优势在于其矩阵正交化处理，使得参数更新方向更加“均匀”和“有效”，尤其在小规模模型上收敛速度显著提升。

五、MUON在实际应用中的表现与挑战

1. 实际应用表现

MUON在实际应用中展现出以下优势：

• 训练效率翻倍：在Moonlight-16B模型的训练中，MUON仅用5.7T tokens就达到了传统方法需要约11T tokens才能达到的学习效果，计算效率提升约2倍。
• MoE架构优化：MUON特别适合Mixture-of-Experts（MoE）架构，通过分组路由机制和Scaling Factor优化，降低了通信开销47%，使16B参数的Moonlight模型在单卡A10上即可实现INT4量化部署，显存占用仅8.7GB。
• 超参数复用：Moonlight变体通过更新缩放机制，使得MUON可以直接复用AdamW的超参数，大幅降低了工程成本。

2. 大规模训练挑战

MUON在大规模训练中面临的主要挑战包括：

• 权重RMS失控：原始MUON在大规模训练中可能出现权重RMS值持续增长的问题，导致数值不稳定和模型性能下降。Moonlight变体通过引入权重衰减机制有效缓解了这一问题。
• 梯度归一化要求：Newton-Schulz迭代要求梯度矩阵的谱范数小于1，MUON通过除以Frobenius范数来近似满足这一条件，但在实际训练中仍需精细控制。
• 通信开销：尽管MUON的计算开销低，但在分布式训练环境中，其正交化处理可能增加通信复杂度。Moonlight团队基于ZeRO-1优化策略实现了通信高效的MUON版本，将通信带宽需求降低至同类模型的65%。
• 实现复杂度：MUON的矩阵正交化处理需要特定的实现优化，如使用PyTorch的torch.addmm等融合操作提升计算速度和数值稳定性。

六、MUON的实现细节与最佳实践

1. 代码实现关键点

MUON的代码实现有几个关键点需要注意：

• 梯度预处理：将梯度转换为bfloat16格式，提高数值稳定性。
• 梯度归一化：通过除以L2范数（norm().clamp(min=eps)）确保梯度矩阵的谱范数近似小于1，满足Newton-Schulz迭代的收敛条件。
• Newton-Schulz迭代：使用五阶 polynomial 迭代公式，通过两次torch.addmm操作实现：
  • 第一次计算$b M + c M^2$ （gram_update）
  • 第二次计算 $a Y + \text{gram\_update} ⋅ Y$

2. 超参数调整建议

MUON的超参数调整有几个关键点：

• 学习率调整：由于MUON的更新方向是正交化的，其有效学习率与AdamW不同，通常需要重新调整。
• 动量系数$\beta$：建议设置为0.9，与AdamW的动量系数一致。
• 权重衰减系数$\lambda$：建议与AdamW使用相同的权重衰减系数，Moonlight变体已确保兼容性。
• 归一化系数$\epsilon$：通常设置为$10^{-8}$，防止除以零。

MUON的形状感知学习率调整是其区别于传统优化器的关键特性，不同的参数矩阵形状（如$[A,B]$）可能需要不同的学习率缩放因子，但Moonlight变体通过$0.2 \cdot \max(A,B)$的缩放因子，使得不同形状的矩阵可以使用相同的学习率，大幅简化了超参数调整过程。

关键特性: 形状感知学习率

MUON的形状感知学习率调整是其区别于传统优化器的关键特性。Moonlight变体通过 $0.2 \cdot \max(A,B)$ 的缩放因子，使得不同形状的矩阵可以使用相同的学习率，大幅简化了超参数调整过程。

七、MUON的应用前景与研究方向

1. 应用前景

MUON优化器的应用前景广阔：

• 大模型训练效率革命：MUON的训练效率提升2倍，可显著降低大模型研发成本，使16B参数模型的训练成本从数百万美元级降至百万美元以内。
• MoE架构优化：MUON特别适合MoE架构，通过分组路由机制和Scaling Factor优化，可有效解决专家模型中的通信开销和参数利用率问题。
• 绿色AI发展：MUON的效率提升直接减少了能源消耗和碳排放，据估算，可使大模型训练过程的碳足迹减少约40-50%，符合AI可持续发展的行业趋势。

2. 研究方向

MUON的未来发展有几个重要研究方向：

• 二阶矩估计整合：探索如何将二阶矩估计整合到MUON框架中，可能进一步提升其性能。
• 混合精度优化：研究MUON在不同混合精度训练策略下的表现，如bf16+fp32或int8量化。
• 超大规模模型适配：探索MUON在万亿参数模型上的表现，以及如何进一步优化其通信效率和内存占用。
• 理论基础深化：深入研究MUON的理论基础，特别是其矩阵正交化处理与模型泛化能力之间的关系。

八、结论

MUON优化器代表了深度学习优化算法的一个重要方向——从逐元素处理上升到矩阵层面的优化。其核心思想是通过矩阵正交化处理梯度，保留梯度方向信息的同时丢弃幅度信息，从而实现更均匀、更有效的参数更新。Moonlight变体通过引入权重衰减和更新缩放机制，解决了MUON在大规模训练中的稳定性问题，使其成为训练大规模语言模型的有力工具。

MUON的革命性在于它不再简单地沿梯度方向更新，而是基于矩阵的几何特性，通过Newton-Schulz迭代法近似求解梯度矩阵的正交化形式。这种创新方法在保持计算效率的同时，显著提升了收敛速度和稳定性，为大模型训练带来了新的可能性。

随着MUON的开源和应用，我们有理由相信：千亿级参数模型的训练成本将大幅降低，定制化大模型的开发门槛将进一步下放，最终推动AI技术在更多行业场景的深度应用。MUON的出现标志着大语言模型发展从"参数竞赛"转向"效率竞赛"的关键转折点，为AI的长期健康发展奠定了基础。

未来展望 (Future Outlook)

MUON的出现标志着大语言模型发展从“参数竞赛”转向“效率竞赛”的关键转折点，为AI的长期健康发展奠定了基础。