SOAP, Muon, and Beyond: Pushing LLM Pretraining Scales
论文地址：https://arxiv.org/abs/2607.20548
项目页面：https://github.com/NVIDIA-NeMo/Emerging-Optimizers
学术交流：922230617
目录
1. 引言
2. 优化器概览：AdamW、Muon 与 SOAP
2.1 AdamW：简单但受限
2.2 Shampoo：结构感知的 Kronecker 因子分解
2.3 SOAP：在预处理特征基中执行 Adam
2.4 Muon：直接进行谱正交化
3. 大规模 Batch 训练的挑战与解决方案
3.1 MoE 模型中的有效 Batch Size
3.2 学习率缩放：平方根规则
3.3 分布式策略：全矩阵 vs. 分片
4. 背景与相关工作
5. 预训练实验：Muon vs. SOAP vs. AdamW
5.1 模型与数据
5.2 公平对比：update-RMS 匹配
5.3 Muon 在大 Batch 下显著优于 AdamW
5.4 解决 SOAP 的训练不稳定性
5.4.1 “弹弓” 失稳现象
5.4.2 修复：每步 QR + 包含当前梯度
5.4.3 引入 KL-散度协方差估计
5.5 Muon vs. SOAP 正面比较
6. 系统实现：层间分布式优化器
7. 结论与未来方向
1. 引言
在大模型训练中，优化器不仅仅是 “调参工具”——它决定了分布式系统的内存布局（优化器状态常比模型参数更占显存）、通信开销（batch size 极限）以及最终的数据效率和泛化能力。长期以来，AdamW 凭借其元素级自适应的简单性和可扩展性，成为事实标准。但这种 “无视矩阵结构” 的做法，忽略了梯度间的相关性，也放弃了利用曲率信息加速收敛的机会。
高阶优化器（如 Shampoo、SOAP、Muon）通过Tensor 级别的预处理（Kronecker 因子或正交化）来逼近二阶信息，理论上能走更大步长、更快收敛。然而，它们的计算开销、数值稳定性、分布式实现复杂性，一直阻碍其在千亿级模型上落地。
NVIDIA 团队在最新技术报告中，系统性地解决了这些问题，并在数万亿 token 的预训练任务中验证了 SOAP 和 Muon 的优越性。核心贡献包括：
大 batch 规模下的稳定性：MoE 模型 batch 达 100M token 时，Muon/SOAP 仍优于 AdamW；
修复 SOAP 的 “失稳”：通过每步 QR 正交化和 KL 散度协方差估计，消除 loss 尖峰；
公平对比协议：基于 update-RMS 匹配，实现优化器间学习率公平迁移；
可扩展的分布式实现：在 Megatron-LM 中引入层间分布式优化器，隐藏通信开销。
2. 优化器概览：AdamW、Muon 与 SOAP
现代 LLM 预训练优化器遵循一条从 “逐元素自适应” 到 “结构感知” 的演进路线。为了理解 Muon 和 SOAP 的优势，首先需要理清它们与奠基性方法 Shampoo 以及最通用的 AdamW 之间的血缘关系。
2.1 AdamW：简单但受限
AdamW 是目前最通用的基准优化器。它通过维护梯度的一阶矩（动量 m_t）和二阶矩（方差 v_t）的指数移动平均（EMA），为每个参数元素独立调整更新步长：
一阶矩：
二阶矩：
其预处理更新方向为：
这种对角预处理使得 AdamW 内存友好且易于分片，但它完全忽略了权重矩阵的行列结构及相关性。
2.2 Shampoo：结构感知的 Kronecker 因子分解
为了捕捉梯度的结构性关联，Shampoo 应运而生。它不将梯度展平为一维向量，而是保留其矩阵形状 G_t ∈ R^{m×n}，并分别维护行和列的协方差矩阵：
其更新方向通过 Kronecker 积形成的全矩阵预处理生成：
Shampoo 通过 Kronecker 分解近似了二阶曲率信息，但计算矩阵逆平方根（R^{−1/4}）的代价较高，且内存占用显著大于 AdamW。
2.3 SOAP：在预处理特征基中执行 Adam
SOAP 直接建立在 Shampoo 的数学框架之上，但做了一个巧妙的 “嫁接”：它利用 Shampoo 维护的特征向量矩阵（Q_L, Q_R）将梯度旋转到预处理器的对角化特征基中，然后在该基中应用 Adam 的逐元素自适应更新，最后再旋转回原空间：
这种设计使得 SOAP 既保留了 Shampoo 的结构感知能力（大曲率方向被抑制），又继承了 AdamW 平滑的自适应特性，理论上更易于超参数的迁移。但其代价是需要维护完整的 Kronecker 因子、特征基，且对数值稳定性极为敏感。
2.4 Muon：直接进行谱正交化
与 SOAP 的 “复杂嫁接” 不同，Muon 走出了另一条截然不同的高效路径。它完全不计算协方差矩阵或预处理器，而是先像 AdamW 一样计算动量矩阵 M_t，随后直接通过牛顿-舒尔茨（Newton-Schulz）迭代近似计算该动量的极分解（Polar Decomposition），即最近的正交矩阵：
Muon 的更新本质上是谱更新：它不调整每个元素的大小，而是改变更新方向，使特征值分布更加均衡。相比 SOAP，Muon 省去了二阶矩内存，但要求矩阵运算在完整的 2D 张量上进行，无法像 AdamW 那样按元素任意分片。
小结：Shampoo 通过全矩阵预处理奠定了结构感知的基础；SOAP 在 Shampoo 的基础上嫁接了 Adam 的自适应更新，表达能力最强但开销最大；Muon 则通过极分解简化了预处理，以最小的额外内存开销实现了高效的谱更新。
3. 大规模 Batch 训练的挑战与解决方案
3.1 MoE 模型中的有效 Batch Size
对于混合专家（Mixture-of-Experts，MoE）模型，由于稀疏路由机制，每个专家实际处理的 token 数远小于全局 batch size。
假设路由器负载均衡理想，任何单个专家所看到的有效 batch size 由以下公式给出：
其中，k 为 Top-K 路由数量，N 为专家总数。这意味着增加全局 batch size 主要对稠密参数（如注意力层）构成压力，而稀疏专家参数仍处于较易优化的低 batch regime。
3.2 学习率缩放：平方根规则
当将批量大小从初始大小 B 调整为目标大小 B' 时，为保证不同 batch size 下更新方差的一致性，
求解得采用平方根缩放规则（Square Root Scaling Rule）：
这使得在扩大 batch size 时，参数更新的随机波动保持相似水平。
3.3 分布式策略：全矩阵 vs. 分片
AdamW 的 element-wise 状态可任意分片；但 Muon/SOAP 需要完整 2D 矩阵才能计算正交化或 Kronecker 因子。
为此，文中设计了层间分布式优化器（见第 6 节），而非传统的 ZeRO/FSDP 按元素切分。
4. 背景与相关工作
Shampoo 变体：Eigen-Shampoo、KL-Shampoo、DASH 等改善预处理估计。
谱优化：Muon、Scion 及后续 Muon+、AdaMuon、NorMuon 等。
系统层面：veScale-FSDP 支持非对称切分，Canzona 实现异步矩阵收集。
5. 预训练实验：Muon vs. SOAP vs. AdamW
5.1 模型与数据
模型：8B 稠密 GPT，3B/30B 纯 Transformer MoE，8B/72B 混合 Mamba-Transformer MoE。
训练数据：1T~3T token 的 Nemotron-3 数据集。
配置：序列长度 8192，全局 batch 25M token（基线），学习率采用 WSD（Warmup-Stable-Decay）计划。
5.2 公平对比：update-RMS 匹配
由于旋转矩阵不改变 Frobenius 范数，SOAP 与 AdamW 的更新 RMS 自然一致。而 Muon 需要乘以因子
来匹配更新尺度。文中明确采用了 Kimi-Moonshot 团队的 RMS 匹配协议。
5.3 Muon 在大 Batch 下显著优于 AdamW
实验设计：batch size 从 25M 逐步升至 50M、75M（带 ramp-up），学习率按平方根同步调整。
结果（图1、图2，表5）：
AdamW 在 500M 时 loss 明显退化，而 Muon 保持降低 loss。
在 30B-A3B 和 72B-A8B MoE 上，Muon 在编程和常识推理任务上 gains 最显著。
注意：Mamba 的 Conv1D 滤波器不适合 Muon，回退至 AdamW 效果更好。
5.4 解决 SOAP 的训练不稳定性
5.4.1 “弹弓” 失稳现象
原 SOAP 实现为省计算，每 10 步才更新特征基，且不包含当前步梯度。在大 batch 早期，loss 变化剧烈，陈旧预处理导致梯度范数振荡，进而引发 loss 尖峰（图4）。在 8B 模型上直接导致训练发散（图5）。
5.4.2 修复：每步 QR + 包含当前梯度
修改为每步重新计算特征基（QR 分解），并纳入当前梯度，即可完全消除尖峰（图7）。QR 比 eig 更快且效果无差别（附录 E 图13）。
5.4.3 引入 KL-散度协方差估计
令在时间步 𝑡 处梯度的奇异值分解为 G𝑡 = UΣVᵀ。假设 Kronecker 因子 S𝑎 和 S𝑏 与梯度的特征基大致对齐，使得：
其中 Λ𝑎 和 Λ𝑏 分别表示 Kronecker 因子所捕获的特征值。
将原 Shampoo 的 G·G^T 累加替换为 KL-Shampoo 的耦合更新：
这使 Kronecker 因子的条件数从 (σ_max / σ_min)^2 降为 σ_max / σ_min（平方根改善），极大增强数值稳定性（图6）（推导见论文附录）。
文中推荐KL-SOAP（即加了 KL 更新的 SOAP）为最强版本。
5.5 Muon vs. SOAP 正面比较
控制变量：统一不拆分 QKV，统一 blocking 策略。
结果（图7、图8）：
在大 batch（1x、2x、4x）下，KL-SOAP 的 CE loss 始终略低于 Muon（差距稳定但微小）。
使用精确 SVD 极分解的 MOP（Muon 变体）也略优于原 Muon，说明正交化精度有提升空间。
ϵ 超参数差异：
- SOAP 的 ϵ 类似 Adam，作为二阶估计分母的软下界。
- Muon 的 ϵ 用于归一化矩阵 Frobenius 范数下限。
- 因此，将 𝜖 视为特定优化器的超参数：对于SOAP，它控制预处理自适应矩和因子反演中的数值下限；而对于Muon，则控制牛顿-舒尔茨迭代中进入的最小归一化尺度。
局限：文中未系统调优，将作为后续研究。
6. 系统实现：层间分布式优化器
为实现 Muon/SOAP 的高效分布式训练，文中在 Megatron-LM 中构建了专门方案：
负载均衡：将整个参数矩阵（不分片）按大小排序后 round-robin 分配到各 DP rank，确保每卡内存均匀。
独立更新：每卡只更新自己负责的完整矩阵，计算完成后 flatten 到 buffer。
异步 All-Gather 重叠：利用 Megatron 的 DDP bucket，按层顺序发出可变长 All-Gather-V 集体通信，在计算前向的同时拉取下一 bucket 的更新参数，有效隐藏通信延迟。
该实现不牺牲数学等价性（全矩阵预处理）。
7. 结论与未来方向
结论。
Muon 和 SOAP 在大 batch（≥50M token）下明显优于 AdamW，尤其适合 MoE 的密集层。
KL-SOAP综合表现最佳，若显存充裕，优先推荐 KL-SOAP。
未来工作。
SOAP 的 TP 支持与 QKV 分裂：进一步适配 Tensor Parallelism。
优化器驱动的内存布局：让 DDP buffer 按矩阵形状分配，减少碎片和通信开销。
架构协同设计：研究 MLA、LoRA 等结构与全秩预处理的匹配性，设计混合策略。
更精确的正交化：大模型中小奇异值易受浮点噪声影响，需探索截断/正则化正交化。
Batch size 标度律：将系统约束（最大化 GPU 利用率）与理论标度律结合，动态优化 batch 计划。
