Benchmarking Optimizers for Large Language Model Pretraining¶
会议: NeurIPS 2026
arXiv: 2509.01440
代码: https://github.com/epfml/llm-optimizer-benchmark
领域: 优化算法
关键词: 大语言模型预训练、优化器基准测试、AdEMAMix、D-Muon、超参数优化
一句话总结¶
通过在统一设置下对 11 种前沿优化器进行涵盖 4 种稠密模型规模、MoE 架构、多种批大小及不同训练时长的系统性评估（耗费约 30,000 GPU 时，训练 2,900 个模型），全面厘清了大模型预训练中新型优化器的真实性能边界与关键超参数实践准则。
研究背景与动机¶
大语言模型预训练的成功高度依赖于数据质量、模型架构与可扩展优化算法三大支柱。尽管过去五年中自回归语言模型在规模与下游能力上取得了飞跃，底层优化器的选择却显得异乎寻常地保守——自 2014 年提出以来，Adam 及引入解耦权重衰减的 AdamW 几乎统治了深度学习与大模型预训练长达十年。近年来，优化理论与系统社区涌现出大量宣称在收敛速度、显存占用、通信开销或最终性能上超越 AdamW 的新型优化器，包括双动量机制（AdEMAMix）、符号梯度方法（Lion、Signum）、近似二阶与正交投影算法（Sophia、SOAP、Muon / D-Muon）、无调度器算法（Schedule-Free AdamW、Prodigy）以及方差缩减架构（MARS）。
然而，当前学术界与工业界对这些新型优化器的评估存在严重的协议割裂与局部视角。部分工作仅在极短的 speedrun 设定（如 nanoGPT 竞赛）下对比训练损失，往往将短程激进下降误当作长期最优；部分测试忽略了权重衰减实现（如误用耦合 \(\ell_2\) 正则化导致符号方法失效）、学习率衰减幅度（传统仅衰减至 10% 极大限制了后期收敛）或辅助正则项（如误加无益且干扰排序的 z-loss）对基准的干扰；更关键的是，许多宣称“无需调参”或“超越 AdamW”的结论往往建立在对 AdamW 超参数欠调、而对新方法过度调优的不公平对比之上。
这种评估缺失直接导致从业者难以判断是否应该在生产级规模切换优化器，研究人员也难以识别具备真正长期扩展潜力的优化方向。本文的核心 idea 是：构建一个受控、标准化、精确核算算力成本的大模型预训练优化器基准测试套件，在固定数据（FineWeb）与 Llama 架构下系统性横跨 124M 至 720M 稠密模型与 520M MoE 模型，严格消融批大小、训练步数（从不足 Chinchilla 最优到超长训练）、学习率调度与权重衰减，提供切实可靠的优化器选型与调优指南。
方法详解¶
整体框架¶
本文构建的评估体系围绕标准化大语言模型预训练管线展开。模型架构统一采用具有词表嵌入权重绑定（weight tying）的现代 Llama 变体，包含 SwiGLU 激活函数、RMSNorm 归一化与 RoPE 旋转位置编码，涵盖 124M、210M、583M、720M 稠密模型以及 520M 稀疏混合专家（MoE，8 专家、Top-2 路由）模型；预训练数据统一采用 FineWeb 数据集的 100B token 高质量去重子集，分词器采用 GPT-2 词表。
在优化算法库的设计上，作者将当前主流的 11 种新型优化器归纳为五大技术流派，并在统一的形式化框架下重新审视了它们的更新机制与参数解耦策略。评测流水线严格遵循“在近 Chinchilla 计算最优区间进行系统网格超参数调优，随后外推并针对长程训练验证”的原则。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入：FineWeb 文本数据与 Llama 架构<br/>(124M-720M 稠密 / 520M MoE)"] --> B["1. 统一优化器算法族与张量维度解耦<br/>(1D/标量 AdamW + 2D 矩阵专用更新)"]
B --> C["2. 正确解耦权重衰减机制<br/>(修复符号方法与 Muon 矩阵衰减漏洞)"]
C --> D["3. 学习率退火与调度器调优<br/>(Cosine 衰减至 0.01×γ_max + 适配预热)"]
D --> E["4. 跨规模与训练时域系统评测<br/>(小/大批次 × 短程/Chinchilla/长程扩展)"]
E --> F["输出：预训练验证损失、帕累托前沿与选型指南"]关键设计¶
1. 优化器算法族的形式化统一与张量维度分流：规范五大流派的参数更新边界
大模型参数在不同层级具有异质的几何结构与数值敏感性。本文将评测方法统一为五大类别，并系统规范了高维矩阵与低维张量的处理方式： - 自适应一阶矩扩展（AdamW、ADOPT、AdEMAMix）：ADOPT 通过调整二阶矩归一化与动量累积的先后顺序，在步进前仅依赖 \(\boldsymbol{v}_{t-1}\) 并对梯度除法施加自适应截断 \(\mathtt{clamp}(\boldsymbol{g}_t / \max\{\sqrt{\boldsymbol{v}_{t-1}}, \varepsilon\}, t^{1/4})\)；AdEMAMix 则在传统快速动量 \(\boldsymbol{m}_t\)（\(\beta_1\)）基础上引入一条超慢动量缓冲 \(\boldsymbol{m}_t^{\mathrm{slow}}\)（\(\beta_3 \approx 0.999 \sim 0.9999\)），通过权重系数 \(\alpha\) 混合历史长周期梯度。 - 符号梯度族（Signum、Lion）：Lion 通过符号操作 \(\mathtt{sign}(\beta_1 \boldsymbol{m}_{t-1} + (1-\beta_1) \boldsymbol{g}_t)\) 实现动量插值更新并更新二阶缓冲区；Signum 则使用动量累积梯度的纯符号方向。 - 近似二阶与正交投影族（Muon / D-Muon、SOAP、Sophia）：针对 2D 权重矩阵，Muon 利用 5 次 Newton-Schulz 迭代计算动量梯度的近似正交极分解（SVD 极基）；SOAP 利用梯度协方差特征基定期做 QR 分解，在旋转坐标系中维护二阶矩统计量；Sophia 则利用 Gauss-Newton-Bartlett 估计器在 softmax 对数几率采样下每隔 \(\phi\) 步估计 Hessian 对角元。 - 无调度器自适应族（SF-AdamW、Prodigy）：SF-AdamW 通过迭代点加权平均取代学习率衰减；Prodigy 利用双重指数移动平均自适应估计距离初始点的有效步长 \(d_t\)。 - 方差缩减族（MARS）：在批次间引入方差缩减校正量 \(\boldsymbol{c}_t = \boldsymbol{g}_t + \eta \frac{\beta_1}{1-\beta_1}(\boldsymbol{g}_t - \boldsymbol{g}_{t-1})\) 代替原始梯度。
针对 SOAP、Muon、MARS 等结构化优化器，本文统一采用张量维度分流策略：Embedding 层、RMSNorm 标量参数和最终分类头 \(\mathtt{lm\_head}\) 等一维或特异性参数统一交由 AdamW 更新，仅将二维注意力与前馈隐藏层权重送入专用更新核，保证训练数值稳定与低显存开销。
2. 解耦权重衰减的纠偏与规范：纠正传统实现对符号方法与矩阵投影的隐性损害
权重衰减在防止 Transformer 权重范数无节制膨胀中扮演关键平衡角色。然而，广泛存在的开源实现误区严重失真了优化器评测： 首先，在 PyTorch 的标准 SGD 动量实现中，权重衰减常被耦合进动量项内部。若直接用于 Signum，更新项变为 \(\mathtt{sign}(\beta \boldsymbol{m}_{t-1} + (1-\beta) \boldsymbol{g}_t - \lambda(1-\beta)\boldsymbol{g}_t)\)，权重衰减会直接污染符号运算的判别方向，导致优化失效。本文重写了带解耦权重衰减与 Nesterov 动量的 Signum 算法： $\(\boldsymbol{x}_{t+1} \leftarrow \boldsymbol{x}_t - \gamma_t \left( \mathtt{sign}(\beta \boldsymbol{m}_t + \boldsymbol{g}_t) + \lambda \boldsymbol{x}_t \right)\)$ 其次，原版 Muon 仅在 AdamW 负责的 1D 参数上施加权重衰减，2D 核心矩阵参数完全缺少衰减，导致在大模型与长程训练中表现显著退化。本文验证并推广了 D-Muon 的全参数解耦衰减与 RMS 动态缩放原则，使矩阵优化器重获正则化平衡。
3. 学习率退火深度与调度准则：打破 10% 终值教条与预热适配
大模型训练文献普遍沿用 Chinchilla 与 Llama 的惯例，将学习率按余弦衰减至最大值的 10%（\(\gamma_{\mathrm{end}} = 0.1 \times \gamma_{\max}\)）。本文系统消融证明，这一停滞点使优化器在后期过早丧失精细搜索能力。将终值学习率进一步衰减至 \(0.01 \times \gamma_{\max}\) 甚至更低，能够无额外计算代价地带来稳定增益。同时，评测厘清了预热（Warmup）与调度器的匹配特性：虽然多数自适应方法在常规 2,000 步线性预热下表现稳健，但符号方法（Signum、Lion）、Sophia 与 SF-AdamW 极度依赖长达总步数 25% 的长预热以平抑早期的梯度噪声；在调度器选择上，精调的 Cosine 整体全面优于 WSD（Warmup-Stable-Decay）与线性衰减（仅 Muon 显示出对 WSD 的轻微偏好）。
损失函数 / 训练策略¶
训练采用标准自回归交叉熵损失，输入序列长度固定为 512 tokens。消融实验明确排除了在近期 OLMo 等代码库中流行的辅助 z-loss 正则项（\(\mathcal{L}_z = 10^{-4} \log^2 \sum \exp(z_i)\)），证明其对最终损失与训练稳定性无正面贡献，反而与较小的权重衰减结合时会扭曲优化器排序。 训练实施基于扩展版 nanoGPT 与 FlashAttention-2，采用 bfloat16 混合精度（归一化层与 Softmax 保持 float32），优化器状态使用 float32 存储。在分布式数据并行（DDP）环境下，硬件覆盖 A100-80GB、H100-80GB 及 GH200-120GB 超算节点。
实验关键数据¶
主实验¶
在 720M 参数规模、1M tokens 批大小（1984 序列 × 512 长度）的标准预训练设置下，各优化器在不同训练长度（8B、16B 以及超越 Chinchilla 最优的 48B tokens）下的最终验证损失对比见下表（摘自论文图 1 与图 17 数据集）：
| 优化器类别 | 优化器名称 | 8B Tokens (短程) | 16B Tokens (Chinchilla ~14.4B) | 48B Tokens (长程扩展) | 相对 AdamW 表现 |
|---|---|---|---|---|---|
| 基准 | AdamW | 3.238 | 3.125 | 2.998 | 基准锚点 |
| 双动量自适应 | AdEMAMix | 3.208 | 3.088 | 2.965 | 全周期显著领先 |
| 方差缩减架构 | MARS | 3.212 | 3.092 | 2.968 | 紧随 AdEMAMix，大批次极强 |
| 矩阵正交投影 | D-Muon | 3.220 | 3.102 | 2.981 | 显著优于基础 Muon，全程优于 AdamW |
| 特征空间旋转 | SOAP | 3.225 | 3.112 | 3.004 | 短程优于 AdamW，长程被 AdamW 反超 |
| 自适应无调参 | Prodigy | 3.242 | 3.129 | 3.010 | 大模型大批次表现接近 AdamW |
| 动量重构自适应 | ADOPT | 3.235 | 3.122 | 2.995 | 略优于或持平 AdamW |
| 符号动量优化 | Lion | 3.285 | 3.180 | 3.065 | 大模型大批次扩展性变差 |
| 符号梯度方法 | Signum | 3.290 | 3.195 | 3.080 | 大模型大批次扩展性变差 |
| 无调度器自适应 | SF-AdamW | 3.265 | 3.162 | 3.050 | 大模型长程明显落后 |
| 二阶估计 | Sophia | 3.230 | 3.160 | 3.095 | 短程收敛快，后期发散或退化 |
消融实验¶
1. 权重衰减在 124M 模型短程与长程预训练中的敏感性消融（摘自图 7 与图 23）
| 优化器 | 权重衰减 \(\lambda\) | 2.1B Tokens (短程) 验证损失 | 16.8B Tokens (长程) 验证损失 | 结论与机制分析 |
|---|---|---|---|---|
| AdamW | \(\lambda = 0\) (无衰减) | 3.652 | 3.295 | 严重欠正则化，参数范数无节制膨胀 |
| AdamW | \(\lambda = 0.1\) (标准) | 3.630 | 3.245 | 长程训练最佳设定 |
| AdamW | \(\lambda = 0.5\) (高衰减) | 3.595 | 3.268 | 短程训练加速收敛显著，长程轻微欠拟合 |
| Lion | \(\lambda = 0.1\) (标准) | 3.710 | 3.310 | 符号梯度需配合高衰减平衡步长 |
| Lion | \(\lambda = 0.5\) (高衰减) | 3.612 | 3.325 | 高衰减下在短程反超标准 AdamW (\(\lambda=0.1\)) |
| Signum | \(\lambda = 0.1\) (标准) | 3.725 | 3.330 | 配合解耦衰减保证更新方向纯净 |
| Signum | \(\lambda = 0.5\) (高衰减) | 3.620 | 3.340 | 短程大幅缩小与自适应方法的差距 |
| Muon | \(\lambda_{\mathrm{1D}} = 0.1\) (原版) | 3.625 | 3.280 | 2D 矩阵无衰减，长程性能急剧受损 |
| D-Muon | \(\lambda = 0.1\) (全参数) | 3.590 | 3.220 | 引入 2D 权重衰减后全面压制基础 Muon |
2. 学习率终值衰减深度消融（210M 模型，Cosine 调度，摘自图 15）
| 衰减终值比例 \(\gamma_{\mathrm{end}} / \gamma_{\max}\) | 最终验证损失 | 困惑度 (PPL) | 说明 |
|---|---|---|---|
| \(0.1\) (传统默认惯例 10%) | 3.185 | 24.17 | 后期步长过大，无法收敛至狭窄极小值平坦区 |
| \(0.01\) (本文推荐准则 1%) | 3.142 | 23.15 | 显著超越传统惯例，成为所有规模的最佳默认值 |
| \(0.001\) (0.1%) | 3.146 | 23.24 | 性能依然明显优于 10%，但在 Cosine 下略逊于 1% |
| \(10^{-6}\) (趋近于零) | 3.155 | 23.45 | 过早停滞，后期更新动力不足 |
关键发现¶
- AdEMAMix 与 MARS 的全面领先：AdEMAMix 凭借快慢双动量设计在所有模型尺度（124M 至 720M）、所有批大小以及所有训练时长中均展现出极其稳健的统治力；MARS 则在大批次（1M、2M tokens）下释放巨大潜力，与 AdEMAMix 并列第一梯队，远抛其他优化器。
- SOAP 的长程衰减与算力开销：SOAP 尽管在短程训练中表现出色彩斑斓的收敛曲线，但随着训练跨越 Chinchilla 最优步数，AdamW 稳步缩小与其差距并在长程反超；更严重的是，单卡墙上时钟测试表明 SOAP 是唯一耗时随模型参数量非线性大幅攀升的优化器（特征矩阵分解开销巨大）。
- Sophia 的长程发散危机：Sophia 虽然在部分短程与 MoE 任务上展现出激进收敛，但在小批次长步数以及大批次超长步数（超过 130k 步）时均会出现灾难性的数值发散（梯度范数爆炸），其实用可靠性亟待改进。
- 符号方法的批大小依赖性：Lion 与 Signum 在小批次（16k tokens）下由于单样本梯度噪声导致符号量剧烈抖动，表现极差；增大批次至 262k tokens 后性能大幅逼近 AdamW；但在模型进一步扩大至 720M 时，符号方法的尺度扩展性再次遇到瓶颈。
亮点与洞察¶
- 高权重衰减是短程训练的强效加速器：在固定学习率前提下，短周期预训练采用 \(\lambda=0.5\)（而非通常的 0.1）能压制参数范数并显著提升泛化损失，甚至使符号方法反超未调优的 AdamW。
- 修复 PyTorch 开源实现的隐形暗坑：纠正了 Signum 误用耦合衰减导致符号翻转的问题，并为 Muon 的核心 2D 矩阵补充解耦权重衰减（D-Muon），直接逆转了算法的性能表现。
- 大模型长程训练需调大 \(\beta_2\) 与 \(\beta_3\)：训练时域扩展至数万步以上时，将 AdamW 的 \(\beta_2\) 从常用的 0.95 提升至 0.999 乃至 0.9999，并将 AdEMAMix 的慢动量 \(\beta_3\) 从 0.999 提至 0.9999，能有效重现长程梯度平滑优势。
局限与展望¶
- 下游评测与更大参数规模缺口：作者受算力约束最大仅评测至 720M 参数稠密模型与约 50B tokens，尚未检验在 7B+ 模型及万亿级 token 下的损失演化；同时评测完全以验证损失为目标，未直接评测下游常识推理或少样本评估基准（下游能力与预训练损失偶尔存在非单调偏离）。
- 分布式显存与 Sharding 框架兼容度：评测侧重单卡/DDP 下的数学行为与时间开销，尚未评估 FSDP、Megatron-LM 或 ZeRO-3 下各类优化器状态（如 SOAP 协方差基、AdEMAMix 额外慢动量缓冲区）的通信量与切分兼容性。
相关工作与启发¶
- vs. Vyas et al. (SOAP, 2024): SOAP 原始论文声称大批次下大幅领先 AdamW。本文复现指出其基线采用了劣质配置（未充分衰减学习率、加入有害 z-loss、权重衰减过低 0.0001）；在修复超参数后，AdamW 在长程训练中即可追平并超越 SOAP，且避免了巨大的特征值分解墙上时钟延迟。
- vs. Keller Jordan et al. (Muon, 2024): Muon 专为短程 nanoGPT speedrun 设计，缺少 2D 矩阵的权重衰减；本文通过评测表明 D-Muon 补齐权重衰减并结合 RMS 缩放后，方能成为全程稳定超越 AdamW 的高效矩阵优化器。
评分¶
- 新颖性: ⭐⭐⭐⭐☆ [首次针对大模型预训练构建大规模、高精度、涵盖 11 种新型优化器的受控统一基准]
- 实验充分度: ⭐⭐⭐⭐⭐ [耗费 30,000 GPU 时，系统消融 2,900 个模型，覆盖多规模、批大小、时长、调度器与衰减参数]
- 写作质量: ⭐⭐⭐⭐⭐ [实验脉络清晰严谨，深入解剖开源实现的各类缺陷，提供扎实透明的配置指南]
- 价值: ⭐⭐⭐⭐⭐ [为学术界与工业界大模型预训练的优化器选型与调优提供了极高实用价值的事实基准]