谱分配视角下的Muon优化器:从矩阵奇异值到超越Adam的更新策略
优化器是深度学习训练中的“隐形引擎”。刚开始接触 PyTorch 时,我们通常会直接使用 AdamW 训练 Transformer,很少有人深究为什么默认选择 Adam、而不是 SGD;直到最近 Muon 这类优化器出现,社区开始重新审视“逐元素归一化到底丢失了什么”。本文从谱分配(Spectral Allocation)的角度解释 Muon 为什么能在部分场景中超越 Adam,并给出可以落地的改进思路。
本文适合正在训练语言模型、Transformer 或 CNN 的工程师阅读。读完你会理解:权重矩阵的奇异值谱对优化器行为有什么影响;Adam 的逐元素归一化在矩阵参数上存在哪些缺陷;Muon 的正交化更新做了什么;以及如何在自己项目中安全地引入和改进 Muon。
1. 为什么优化器决定了大模型训练的成败
1.1 优化器不只是“梯度下降的按钮”
在深度学习中,优化器的核心任务不是“让 loss 下降”,而是在有限的计算预算内,找到泛化能力足够好的参数。对于大模型训练,优化器的影响往往比想象中更大,同一个模型架构、同一份数据,使用 AdamW 和 SGD 训练,最终效果可能相差很大。优化器还决定了训练稳定性、收敛速度、对学习率的敏感程度以及显存占用。
过去几年,Adam 系优化器几乎成为 Transformer 训练的默认选项。Adam 的优点是自适应学习率、对稀疏梯度友好、超参数相对鲁棒。但 Adam 并非没有代价:它对每个参数维度独立归一化,忽略了参数之间的结构关系,尤其是矩阵参数。这个问题在参数规模小的时候不明显,但在大模型训练中会被放大。
1.2 从 Adam 到 Muon 的转变
Muon 是 Keller Jordan 等人提出的优化器,核心思路很直接:对于 2D 矩阵参数,不再对梯度做逐元素归一化,而是先将梯度正交化,再用正交化后的方向更新参数。这里的“正交化”并非像 SVD 那样精确计算,而是通过 Newton-Schulz 迭代近似完成。
Muon 在训练 Transformer、Mamba 等模型时,表现出比 AdamW 更高的 token 效率。它的出现让研究者重新关注优化器背后的“谱结构”,也就是本文标题中的 Spectral Allocation 概念。简单理解,Adam 把每个梯度坐标当作独立变量,Muon 则把整个梯度矩阵当作一个线性变换来归一化。
2. 谱分配:从权重矩阵的奇异值说起
2.1 什么是权重矩阵的谱
任意一个实矩阵 W 都可以做奇异值分解:W = UΣV^T。其中 U 和 V 是正交矩阵,Σ 是对角矩阵,对角线上的数值就是奇异值。这些奇异值从大到小排列,就构成了矩阵的“谱”。
谱描述了 W 在不同方向上的放大能力。最大奇异值对应的方向,是 W 作用后能量被放大最多的方向;较小奇异值对应的方向,则几乎被 W 压缩。这一组奇异值如果集中在前几个,称矩阵“谱集中”;如果分布比较均匀,称矩阵“谱平坦”。
2.2 谱集中与谱平坦:两类典型矩阵
在深层网络中,不同层、不同模块的权重矩阵,谱形态差异很大。例如:
- 自注意力中的 QKV 投影矩阵,经过训练后通常呈谱集中形态,少数大奇异值主导信息流动。
- MLP 中某些扩展层的权重,奇异值分布相对平坦。
- Embedding 矩阵则又不同,它的行之间往往具有稀疏相关性。
如果优化器对所有矩阵都采用同一种更新策略,等于默认所有矩阵的谱结构相同,这是不合理的。理想情况下,更新方向应该适配每个矩阵的谱形态,这正是“谱分配”想表达的思想:根据参数矩阵的谱结构,分配不同的更新方式和更新强度。
2.3 从谱的角度看优化器
从谱分配视角,可以重新理解经典优化器:
- SGD 使用原始梯度方向更新,相当于对所有奇异方向做等比例缩放,不做任何谱自适应。
- Adam 对梯度逐元素归一化,相当于把每个坐标统一缩放到同一尺度,但这并不等于把“方向”放到同一尺度。
- 自适应二阶方法(如 K-FAC)尝试显式建模矩阵的谱结构,但计算代价高。
- Muon 则是一种折中:通过正交化,将梯度矩阵的奇异值整体向 1 靠近,从而让所有主要方向获得相近的更新尺度。
这样看来,Muon 并不是唯一可行的谱分配策略,但它用很小的计算代价,换来了接近二阶方法的谱感知能力。
3. Adam 为什么在矩阵参数上“不够聪明”
3.1 Adam 的更新公式回顾
Adam 的更新过程可以简化为:
- 计算梯度 g。
- 更新一阶动量 m = β1·m + (1-β1)·g。
- 更新二阶动量 v = β2·v + (1-β2)·g²。
- 用 m / (√v + ε) 作为更新方向。
这里的“逐元素归一化”体现在第 4 步:梯度的每个分量,都被自身历史平方根缩放。对于稀疏特征,这个设计很好;但对于密集矩阵参数,问题就出现了。
3.2 AdamW 与 Adam 的区别
AdamW 是 Adam 的修正版本,主要差别在于权重衰减的实现位置。
- Adam 在梯度上先加上 weight decay,再进入一阶、二阶矩估计,这会导致权重衰减被自适应学习率放大。
- AdamW 则先计算自适应更新方向,再单独从参数中减去 weight_decay × param,解耦了权重衰减和自适应步长。
所以在现代大模型训练中,AdamW 几乎完全取代了 Adam。本文后续提到 Adam,均默认指 Adam/AdamW 系优化器。
3.3 Adam 的谱缺陷
Adam 的逐元素归一化,本质上是把每个坐标的方差调整为 1。但矩阵参数的方向性,或者说谱结构,没有被考虑。举例来说,假设某个权重矩阵的梯度中,90% 的能量集中在一个大奇异值方向,其余方向是噪声。Adam 会把每个坐标都缩放到类似尺度,结果大方向的相对优势被削弱,噪声方向反而被放大。
Adam 在非平稳、稀疏场景下收敛快,但在矩阵参数上,这种“无差别缩放”会导致更新方向偏离真实下降方向。这也是很多实验观察到 Adam 在深层网络中虽然收敛快,但泛化不如带着 momentum、经过平衡方向处理的更新方式的原因。
4. Muon 的核心思路:正交化更新
4.1 Muon 的基本流程
Muon 的关键思想是:将梯度矩阵先做某种“谱归一化”,再将归一化后的矩阵作为更新方向。
给定一个 2D 梯度矩阵 G,Muon 的更新流程是:
- 更新动量缓冲 M = μ·M + G,其中 μ 通常取 0.95。
- 应用 Nesterov 修正:N = G + μ·M。
- 对 N 做正交化,得到 O。
- 参数更新:W ← W − lr·O。
这里的正交化并不是让更新矩阵严格正交,而是让它的奇异值尽量靠近 1。这样更新方向中各主要方向的尺度趋于一致,同时又不会像 Adam 那样破坏矩阵的整体方向结构。
4.2 Newton-Schulz 迭代
正交化操作需要通过迭代近似计算。最常用的是 Newton-Schulz 迭代,它的目标是从一个矩阵 A 出发,逐步逼近 A(A^T A)^{-1/2},也就是对 A 做极分解中的正交因子近似。
一个常用的迭代公式是:
其中 a、b、c 是固定系数,例如 a=3.4445、b=-4.7750、c=2.0315。这个系数来自一个多项式近似,目的是让迭代快速收敛到正交矩阵。通常迭代 4~6 次就能获得较好的近似。
4.3 为什么正交化有效
正交化更新的本质,是对梯度矩阵的谱做“均匀化”:所有奇异值被拉向 1。这等于告诉优化器:我们不知道哪个方向更重要,所以先让所有方向获得均匀的更新尺度。
为什么这样的策略在深层网络中有效?一种解释是,深度神经网络中信息流动依赖矩阵的“有效秩”。如果更新方向被少数大奇异值主导,训练容易陷入某些方向的过拟合;如果更新方向包含太多噪声方向,又会影响收敛稳定性。正交化更新正好介于两者之间,在保留方向信息的同时抑制极端谱。
Muon 还结合了 momentum 和 Nesterov,使得更新方向更平滑,进一步减少噪声。可以说,Muon 在谱分配上比 Adam 更“公平”,比 SGD 更“聪明”。
5. Muon 的完整代码实现与运行验证
5.1 最小项目结构
为了便于实践,我们搭建一个最小实验工程:
示例中使用 PyTorch 2.0 以上版本,Python 3.10 以上即可。重点演示核心逻辑,不引入分布式依赖。
5.2 核心代码:正交化函数
先实现 Newton-Schulz 正交化函数。
这里需要解释几点:
- 归一化是为了避免迭代过程中数值过大。
- 转置操作是为了减小矩阵乘法计算量。若 m > n,则先转置再迭代,最后转置回来,得到的结果仍然满足正交化需求。
- 迭代步数并非越多越好。过多步会让矩阵更接近正交,但也会增加计算开销,一般取 5 即可。
5.3 完整 Muon 优化器类
接下来实现一个简化版 Muon 优化器。为了方便演示,我们将 2D 参数交给 Muon,1D 参数交给内部维护的 AdamW。
这个类是教学演示用的,未实现分布式、混合精度等复杂特性。生产环境建议参考官方仓库或社区实现。
5.4 训练循环示例
写一个简单的训练循环,验证 Muon 能够正常完成参数更新。
运行结果示例:
这里 loss 下降不代表模型已经训练好,只是验证优化器可以正常 work。读者在自己的任务中,需要结合验证集指标评估优化器优劣。
6. Muon 与 Adam/SGD/AdamW 的行为对比
6.1 收敛速度与 token 效率
Muon 在 Transformer 语言模型上的主要优势是 token 效率。论文与社区复现中,Muon 在相同 token 预算下,通常能达到比 AdamW 更低的 loss 或更高的下游指标。原因是其正交化更新在深层网络中的有效信号更强,减少了噪声方向的干扰。
需要提醒的是,具体提升幅度与模型规模、数据分布、超参数关系很大。不能拿别人的 benchmark 直接当效果保证,必须在自己的任务上做 A/B 实验。
6.2 泛化能力与稳定性
SGD 通常被认为泛化好但收敛慢;Adam 收敛快但在某些平坦极小值附近泛化不如 SGD。Muon 介于两者之间:它保留了自适应优化器的快速收敛特性,又通过正交化减少了逐元素归一化带来的方向失真。
在稳定性方面,Muon 对梯度噪声更敏感,需要配合梯度裁剪。批量较小时,建议增大 batch 或使用梯度累积。
6.3 计算代价对比
从计算量看,Muon 比 SGD 多出动量更新和 Newton-Schulz 迭代,比 Adam 多出矩阵乘法。矩阵维度越大,Newton-Schulz 迭代的 $O(n^3)$ 成本越明显。实际中,5 次迭代的额外开销通常在 10%~20% 左右,对于训练大模型来说可以接受。
如果矩阵维度很大,可以先转置让迭代在较小维度上进行,或使用 bfloat16 计算正交化方向。这些都属于工程优化,后面会展开。
7. 如何改进 Muon:算法与工程实践
7.1 参数分区与混合优化策略
一个最重要的改进经验是:不要对所有参数无脑使用 Muon。常见做法是:
- Embedding 矩阵、输出头(lm_head)保留 AdamW,因为这些位置的梯度稀疏性强,正交化反而可能引入噪声。
- LayerNorm、bias 等 1D 参数使用 AdamW。
- 只有隐藏层中的 2D 权重矩阵使用 Muon。
在我的实验中,这个分区策略比全参数 Muon 更稳定,收敛曲线也更平滑。
7.2 谱感知自适应正交化
标准的 Muon 对所有矩阵使用相同的正交化强度。但根据谱分配的思想,不同矩阵的谱结构差异很大,可以动态调整迭代步数或缩放强度。
一个示例思路:先通过少量迭代估计矩阵的谱平坦度,如果谱越集中,就对更新方向做更强的正交化;如果谱越平坦,就减少迭代步数。
这种方法需要更多实验验证,但方向是合理的:把“谱分配”从理论概念变成可计算的自适应策略。
7.3 降低正交化计算开销
在实际大模型训练中,Newton-Schulz 迭代带来的矩阵乘法会占用额外算力。几个经验:
- 迭代步数从 5 减到 4,效果差异很小,但计算量降低 20%。
- 在 bfloat16 下执行正交化,主权重保持 fp32,方向质量损失不大。
- 对非方阵,先比较行列数,让迭代始终在小维度上进行。
- 在张量并行场景中,可以把梯度 AllReduce 与正交化合并,减少一次通信等待。
这些优化不会改变优化器数学性质,但能显著提升训练吞吐。
7.4 学习率与调度调整
Muon 的默认学习率通常比 AdamW 高一个数量级。例如:
| 优化器 | 常用学习率范围 | 备注 |
|---|---|---|
| AdamW | 3e-4 ~ 3e-3 | 稳定,但收敛速度偏慢 |
| Muon | 1e-2 ~ 3e-2 | 需要配合 warmup 和梯度裁剪 |
Muon 对 warmup 更敏感。建议先设置 1000~2000 步 warmup,让 momentum 缓冲稳定后再逐步提高学习率,避免前期震荡。
8. 使用 Muon 的常见问题与排查清单
在实际使用中,最常见的问题集中在以下几个方向。
| 问题现象 | 可能原因 | 解决思路 |
|---|---|---|
| 训练早期 loss 震荡 | 学习率过高、warmup 不足 | 降低 lr,增加 warmup 步数,检查梯度裁剪 |
| 更新方向不稳定 | 批量过小、梯度噪声大 | 增大 batch 或梯度累积,提高 momentum 权重 |
| 验证集指标不如 AdamW | 参数分区不合理 | 将 Embedding、head 切回 AdamW |
| 显存或算力开销高 | Newton-Schulz 迭代次数过多 | 减少迭代步数,使用 bfloat16 计算方向 |
| 一维参数发散 | 错误地全部走 Muon | 只对 2D 参数使用 Muon |
| loss 下降但泛化差 | 权重衰减过大或过小 | 单独搜索 weight_decay,不要直接沿用 AdamW 的值 |
如果训练崩溃,建议按以下顺序排查:
- 先确认梯度裁剪是否生效,建议 max_norm 从 1.0 开始。
- 再确认学习率,Muon 使用 0.01 时是否出现 NaN。
- 检查参数分区,是否把 LayerNorm 权重误传给 Muon。
- 最后检查 Newton-Schulz 迭代的数值稳定性,确保输入不是极端大值。
9. 最佳实践与工程建议
9.1 先小规模复现,再大规模切换
不要直接把生产模型的优化器从 AdamW 切换到 Muon。建议先在参数量较小的模型上做完整对比实验,固定数据顺序、随机种子、token 预算,只改变优化器。比较训练 loss 曲线、验证集指标和训练吞吐三个维度。
9.2 保持超参数独立性
AdamW 的一组优秀超参数,不经过调整就迁移到 Muon 是不现实的。Muon 学习率、warmup、momentum、weight_decay 都需要重新搜索。建议先固定学习率和 momentum,其他参数按经验默认,再逐步调整。
9.3 记录并监控梯度范数
Muon 对梯度尺度敏感。建议在训练循环中定期打印梯度范数,观察是否出现尖峰。如果梯度范数突然上升,及时调低学习率或减小 max_norm。
9.4 与现有训练体系兼容
Muon 可以作为一个独立优化器接入现有 PyTorch、Lightning 或 Hugging Face Trainer。需要注意:
- 自定义优化器可能无法直接使用 Trainer 自带的学习率调度器,需要单独设置。
- 混合精度训练时,确认正交化计算的精度策略。
- 分布式训练时,确保梯度同步完成后才执行 step。
9.5 将谱分配作为通用调参视角
即使不切换到 Muon,谱分配思想也有指导意义:当你发现 AdamW 收敛慢时,可以先分析权重矩阵的奇异值分布,判断是否存在过度的谱集中现象。如果某些矩阵奇异值差距过大,可以考虑对这些矩阵手动放大学习率或使用正交化更新。
10. 下一步:谱感知优化器还有什么可以探索
Muon 的提出让优化器设计重新回到“矩阵结构”本身。谱分配是一个比“正交化”更广的概念,未来还有不少方向值得尝试。
- 自适应迭代步数:根据矩阵谱的平坦度动态决定 Newton-Schulz 迭代次数。
- 谱缩放混合策略:不完全正交化,而是将奇异值向目标分布压缩。
- 与低秩适配技术结合:在 LoRA 微调中,将正交化应用于低秩增量矩阵。
- 分布式通信优化:把正交化前的梯度压缩与通信合并,降低大模型训练开销。
如果读者想动手实践,我建议先在小规模语言模型上跑一组 Muon vs AdamW 的对比实验,记录两条 loss 曲线,再尝试把第 7 节中的参数分区和谱感知自适应方案加进去。理解优化器背后的谱结构,比单纯换一个优化器更有价值。