Hyperparameter Transfer Enables Consistent Gains of Matrix-Preconditioned Optimizers Across Scales¶

会议: NEURIPS2025
arXiv: 2512.05620
代码: 待确认
领域: LLM效率 / 优化
关键词: optimizer scaling, μP, Shampoo, SOAP, Muon, hyperparameter transfer, matrix preconditioning

一句话总结¶

研究矩阵预条件优化器（Shampoo/SOAP/Muon）的超参数随模型宽度和深度的缩放规则（基于 μP），发现正确的超参缩放是实现一致加速的关键：使用 μP + 1/width weight decay，三者在 190M 到 1.4B 参数的 Llama 模型上一致实现约 1.4× 加速。

研究背景与动机¶

领域现状：多个矩阵预条件优化器（Shampoo、SOAP、Muon）在小规模实验中展示了相对 AdamW 的显著加速（1.5-2×）。Muon 在 OpenAI 万亿参数训练中已被使用。
现有痛点：复现工作报告的加速效果不一致——有的团队报告 2× 加速，有的只报告 1.1×，有的发现加速随规模快速消失。根本原因是缺乏可靠的超参数缩放规则。
核心矛盾：语言建模的 loss-compute 缩放指数很小（~0.05），意味着 2% 的 loss 差异对应 40% 的 compute 差异。在大模型上不可能做网格搜索，必须有可靠的超参迁移。
本文要解决什么？ 推导 Shampoo/SOAP/Muon 等的 μP 学习率缩放规则，并验证正确缩放对跨规模一致加速的关键作用。
切入角度：基于 μP（Maximal Update Parameterization）理论推导每个优化器的学习率、weight decay 随宽度和深度的缩放。
核心idea一句话：矩阵预条件优化器的加速是真实的，但需要正确的 μP 缩放才能在大规模上稳定展现。

方法详解¶

整体框架¶

对每种矩阵预条件优化器，推导其 μP 缩放规则（学习率、weight decay、正则化参数随模型宽度/深度如何缩放），然后在 190M→1.4B 语言模型上验证。

关键设计¶

μP 推导（通用方法）：
对每层权重矩阵 \(W \in \mathbb{R}^{n \times m}\)，确保更新量 \(\Delta W\) 的 RMS 缩放在宽度→∞时一致
不同优化器因为预条件方式不同，学习率缩放也不同：Adam 是 \(1/d_{\text{in}}\)，Muon 是 \(\sqrt{d_{\text{out}}/d_{\text{in}}}\)，Shampoo 取决于指数 \(e_L, e_R\)
有限宽度偏差的缓解：
μP 保证无限宽度极限下学习率迁移，但实际有限宽度下最优学习率可能偏移
Blocking（将大权重矩阵分块处理）和 Spectral Normalization 可有效缓解偏移
Grafting（用另一优化器的范数归一化更新方向）也影响缩放
Weight Decay 缩放：
发现独立 weight decay（不随学习率缩放）按 \(1/\text{width}\) 缩放在所有优化器上近乎最优
这与 Xiao (2024) 的建议一致

损失函数 / 训练策略¶

在 FineWeb 数据集上训练 Llama 架构语言模型（190M, 470M, 1B, 1.4B）
对每个优化器在小模型上调参，然后按 μP 缩放迁移到大模型

实验关键数据¶

主实验¶

在 190M→1.4B Llama 模型上的 compute-equivalent 比较：

优化器	μP 缩放	190M 加速	470M 加速	1B 加速	1.4B 加速
Muon	✅	~1.4×	~1.4×	~1.4×	~1.4×
SOAP	✅	~1.4×	~1.4×	~1.4×	~1.4×
Shampoo	✅	~1.4×	~1.4×	~1.3×	~1.3×
Muon	❌ (SP)	~1.4×	~1.2×	~1.0×	消失

消融实验¶

配置	结果	说明
μP vs 标准参数化	μP 下加速一致，SP 下随规模消失	超参迁移是关键
Blocking (128)	改善有限宽度偏差	实用且有效
Spectral norm	降低学习率敏感性	与 μP 互补
WD=1/width vs 固定 WD	1/width 近乎最优	跨优化器一致
Compute-optimal 模型大小	矩阵预条件优化器偏好更大模型	与 AdamW 不同的 scaling law

关键发现¶

超参缩放错误是之前复现失败的主要原因：标准参数化下，Muon/SOAP 的加速在 1B+ 模型上几乎消失
1.4× 加速是一致且可靠的：在正确 μP 下，三种优化器在所有测试规模上一致实现约 1.4× compute savings
μP 对不同优化器的缩放不同：不能直接把 Adam 的 μP 用在 Muon 上（常见错误）

亮点与洞察¶

系统性地解决了"矩阵预条件优化器到底有没有用"的争议——答案是有用，但需要正确缩放。之前的混乱结果都可以用超参错误解释
μP 推导方法的通用性：给出了一个简单通用的流程，可以对任何新优化器快速推导缩放规则
实用指导：给出了每种优化器的完整缩放公式（Table 1），可以直接使用

局限性 / 可改进方向¶

最大模型只到 1.4B：10B+ 模型上的验证缺失
仅测试语言建模：视觉/多模态等其他任务未涉及
加速相对保守（1.4×）：相比一些工作报告的 2× 偏保守，可能是因为 AdamW baseline 更充分调参
改进方向：更大规模验证；与学习率 warmup/cooldown schedule 的交互分析；自动超参迁移工具

评分¶

新颖性: ⭐⭐⭐⭐ μP 推广到新优化器的方法简洁，但核心框架基于已有 μP
实验充分度: ⭐⭐⭐⭐⭐ 多优化器×多规模×消融，非常彻底
写作质量: ⭐⭐⭐⭐⭐ 清晰系统，Table 1 总结缩放规则一目了然
价值: ⭐⭐⭐⭐⭐ 直接解决了社区争议，提供了可操作的实用指导