量化超参数迁移及嵌入层学习率的重要性

来源：https://arxiv.org/html/2605.21486

摘要

超参数迁移能够将最优优化超参数从小规模外推至大规模，这对训练大型语言模型（LLMs）至关重要。实现方式有两种：一是为超参数拟合标度律，二是通过巧妙选择参数化方案（如最大更新参数化 μP），使最优超参数近似尺度不变。本文首先提出了一个量化超参数迁移的框架，包含三个指标：（1）标度律拟合质量，（2）对外推误差的鲁棒性，（3）因参数化选择导致的渐近损失惩罚。随后，我们通过一系列全面的消融实验，探讨了 μP 为何相对于标准参数化（SP）能提供高质量的学习率迁移——现有理论尚不足以解释这一点。我们发现，使用 AdamW 训练时，μP 相对于 SP 的巨大优势仅源于最大化嵌入层学习率。在 SP 中，嵌入层学习率是导致训练不稳定的瓶颈；将其按宽度因子增大以匹配 μP，可显著平滑训练过程并改善超参数迁移。我们还发现，权重衰减能改善标度律拟合质量，但在固定 token-parameter 比例设置下，会损害外推的鲁棒性。

11footnotetext: 马里兰大学帕克分校物理系 22footnotetext: 马里兰大学帕克分校计算机科学系 33footnotetext: 马里兰大学帕克分校联合量子研究所 44footnotetext: Meta 超级智能实验室，基础人工智能研究

1 引言

训练大型神经网络需要仔细调整众多超参数，包括学习率、权重衰减和批大小等 [28 (https://arxiv.org/html/2605.21486#bib.bib88),7 (https://arxiv.org/html/2605.21486#bib.bib73),38 (https://arxiv.org/html/2605.21486#bib.bib128),30 (https://arxiv.org/html/2605.21486#bib.bib129),34 (https://arxiv.org/html/2605.21486#bib.bib61),2 (https://arxiv.org/html/2605.21486#bib.bib130)]。随着模型规模扩展到万亿参数，在大规模下进行超参数调优的成本变得难以承受。解决方法是超参数迁移：先在小规模下找到最优超参数以及相应的标度律，然后利用该标度律外推大规模下的最优超参数。实践中主要有两种方法。第一种方法拟合函数形式，以预测最优学习率 η* 随模型和/或数据规模变化的规律 [7 (https://arxiv.org/html/2605.21486#bib.bib73),36 (https://arxiv.org/html/2605.21486#bib.bib75),3 (https://arxiv.org/html/2605.21486#bib.bib74)]。这种方法虽然有效，但拟合的关系可能无法推广到拟合域之外，且本身需要昂贵的超参数搜索才能获得稳健的函数形式。

第二种方法更具结构性——通过参数化模型，使激活值及其更新与宽度无关 [48 (https://arxiv.org/html/2605.21486#bib.bib90)]，从而在不同规模下保持训练动力学不变 [6 (https://arxiv.org/html/2605.21486#bib.bib120),46 (https://arxiv.org/html/2605.21486#bib.bib53)]。Yang 和 Hu [48 (https://arxiv.org/html/2605.21486#bib.bib90)] 在此期望下推导出最大更新参数化 (μP)，最初用于研究无限宽度极限下的特征学习。实验表明，μP 具有一个理想的实用性质——在不同宽度下保持相当一致的最优学习率。这一性质称为学习率迁移，使 practitioners 能够将从小型代理模型中获取的最优学习率用于训练更大规模的模型 [47 (https://arxiv.org/html/2605.21486#bib.bib51)]，从而减少大规模超参数搜索的需求。后续有多项工作将此框架扩展到深度缩放 [5 (https://arxiv.org/html/2605.21486#bib.bib52),49 (https://arxiv.org/html/2605.21486#bib.bib76),9 (https://arxiv.org/html/2605.21486#bib.bib94)]、替代架构 [4 (https://arxiv.org/html/2605.21486#bib.bib107),8 (https://arxiv.org/html/2605.21486#bib.bib96),43 (https://arxiv.org/html/2605.21486#bib.bib91),małaśnicki2025muparametrizationmixtureexperts,19 (https://arxiv.org/html/2605.21486#bib.bib121)] 和优化器 [29 (https://arxiv.org/html/2605.21486#bib.bib77),13 (https://arxiv.org/html/2605.21486#bib.bib92),37 (https://arxiv.org/html/2605.21486#bib.bib110)]，表明其底层原理可推广至现代网络的整个设计空间。

当前实践引发了几个问题。首先，在为超参数拟合标度律后，总可以将其重新参数化为（近似）尺度不变的。这种情况下，μP 及其变体是否有本质区别？这表明迫切需要量化超参数迁移的质量，以便在不同方案之间进行比较。

其次，μP 及其变体的理论推导做了几个实际中不成立的假设。它假设在无限宽度极限下训练步数有限，而实际训练则处于相反极限——训练步数远大于宽度。它还假设权重更新与激活值完全对齐 [48 (https://arxiv.org/html/2605.21486#bib.bib90)]，但实际中这种对齐在训练初期较低，即使训练后期也从未达到完全对齐 [10 (https://arxiv.org/html/2605.21486#bib.bib54)]。它没有考虑学习率调度（如学习率预热的方面），而如我们在附录 F (https://arxiv.org/html/2605.21486#A6) 所示，这些方面对观察到迁移至关重要。最后，该理论是在固定数据集大小的设定下推导的，而实际训练往往同时缩放数据和参数量，通常采用每参数 token 数 (TPP) 固定的设置 [17 (https://arxiv.org/html/2605.21486#bib.bib48)]。令人惊讶的是，尽管这些假设不成立，μP 仍表现出高质量的学习率迁移，这引出了一个问题：为什么？

本文首先开发了一个量化超参数迁移的框架，提出了三个关键指标。第一个是标度律拟合本身的质量。第二个是鲁棒性指标：随着模型规模扩大，最优超参数对扰动的敏感性如何变化。第三个涉及损失的渐近性能：某些参数化可能以牺牲损失绝对值为代价实现更高质量的迁移，这也需要予以考虑。我们将看到，这三个指标可以相互权衡。

有了量化超参数迁移质量的框架，我们继而更详细地探讨为何 μP 表现出高质量迁移。我们证明，在使用 AdamW [31 (https://arxiv.org/html/2605.21486#bib.bib123)] 训练的标准仅解码器 Transformer [44 (https://arxiv.org/html/2605.21486#bib.bib139)] 中，μP 相对于 SP 有四个关键变化。通过检查所有 16 种消融，我们隔离出嵌入层学习率是关键因素。μP 的嵌入层学习率远大于 SP。只需对 SP 做此一处改动，SP 的超参数迁移质量就基本与 μP 相当。我们进一步研究了训练动力学，发现嵌入层学习率在整个训练过程中都很重要，尤其是在训练早期；如果其不够大，就会导致训练不稳定。因此，训练不稳定性的一个有趣来源是学习过程可能受到嵌入层学习率的瓶颈限制。

这些结果凸显了嵌入层学习率的根本重要性。已知嵌入层从数据中学习了大量结构 [32 (https://arxiv.org/html/2605.21486#bib.bib132),16 (https://arxiv.org/html/2605.21486#bib.bib134),42 (https://arxiv.org/html/2605.21486#bib.bib102),24 (https://arxiv.org/html/2605.21486#bib.bib133)]，因此，将其保持过低会严重拖累训练，而最大化它有助于稳定学习过程，这或许并不令人意外。

本文贡献：本文包含以下贡献：

• • 一个用于衡量超参数迁移质量的量化框架。该框架包含三个指标：(1) 用误差 E 表示的标度律拟合质量，(2) 迁移鲁棒性指数 κ，用于衡量超参数从小规模到大尺度外推时对误差的敏感性，以及 (3) 渐近损失退化 R(∞)，用于衡量在目标规模下相对于最优参数化的性能。
• • 聚焦于 GPT 预训练及 AdamW 优化器，我们将 SP 和 μP 表示为统一形式（表 1 (https://arxiv.org/html/2605.21486#S2.T1)）。两种参数化在四个关键方面存在差异 (SP → μP)：(1) 嵌入层学习率 Θ(1/n) → Θ(1)，(2) 最后一层初始化方差 Θ(1/n) → Θ(1/n²)，(3) LayerNorm 学习率 Θ(1/n) → Θ(1)，(4) 注意力缩放 1/√d → 1/d。
• • 我们检查了区分 μP 与 SP 的全部 16 种消融，并隔离出嵌入层学习率是主要因素：适当缩放嵌入层学习率的 SP 在超参数迁移质量上基本与 μP 持平。
• • 我们发现，如果嵌入层学习率不够大，可能会引起训练不稳定。
• • 我们研究了权重衰减的效果，表明：(1) 它可以改善标度律拟合质量（更小的 E），(2) 在固定每参数 token 数设置下，它会损害外推的鲁棒性。

2 预备知识：神经网络参数化

表 1：针对 AdamW 的标准参数化 (SP) 与最大更新参数化 (μP) 对比。权重衰减随学习率缩放，使得其对权重更新的贡献 η·λ·θ 在不同宽度下保持 Θ(1)。

我们考虑一个神经网络 f_θ，其参数为 θ，通过 AdamW 优化器 [31 (https://arxiv.org/html/2605.21486#bib.bib123)] 最小化损失 L(θ) 进行训练，学习率为 η，权重衰减为 λ。模型的性能可通过沿不同维度（如宽度、深度、上下文长度和数据）的缩放得到提升。参数化 则是一组规则，指定随着模型规模扩大，关键网络和优化器超参数应如何调整，以维持稳定的训练动力学。本文中，我们关注宽度 n 作为缩放维度，学习率 η 作为待缩放超参数。

遵循 Yang 和 Hu [48 (https://arxiv.org/html/2605.21486#bib.bib90)]，我们使用每层 l 的四个标量指数 {a_l, b_l, c_l, d_l} 来参数化网络和优化器，这些指数控制宽度缩放的不同方面。指数 a_l 控制前向传递缩放：h^(l+1) = n^{-a_l} W^(l+1) φ(h^(l))，其中 h^(l) 和 W^(l) 分别表示第 l 层的预激活值和权重，φ(·) 是激活函数。指数 b_l 缩放初始化方差：W^(l) ~ N(0, n^{-2b_l})。指数 c_l 缩放逐层学习率：η^(l) = η · n^{-c_l}，其中 η 是全局学习率。最后，指数 d_l 缩放权重衰减强度：λ^l = λ · n^{-d_l}，使得乘积 η^l·λ^l 与宽度无关。对于 Transformer，通常的缩放点积注意力包含一个按头维度 d 的缩放，标准选择为 1/√d [45 (https://arxiv.org/html/2605.21486#bib.bib131)]。

在此通用设定下，通过对训练动力学施加稳定性条件，可得到特定的参数化方案。abcd 参数化的两个典型例子是 SP [41 (https://arxiv.org/html/2605.21486#bib.bib122)] 和 μP [48 (https://arxiv.org/html/2605.21486#bib.bib90)]。SP 要求初始化时激活值不会随宽度增长而爆炸或消失，这对每层施加了一个约束；剩余的自由度用于做出最简单的选择，例如所有层使用统一的学习率。μP 施加了更强的条件：激活值及其更新都必须与宽度无关，这导致每层有两个约束。为便于解析处理，这些额外约束是在三个附加假设下推导的：(1) 在宽度 n→∞ 时训练步数有限，(2) 权重更新与它们作用的激活值完全对齐 [48 (https://arxiv.org/html/2605.21486#bib.bib90)]，以及 (3) 当 n 缩放时数据集固定。我们在附录 H (https://arxiv.org/html/2605.21486#A8) 中提供了针对 SGD 和 Adam 优化器的 μP 自包含推导。

由于我们的目标之一是理解 μP 的哪些方面对迁移至关重要，我们需要将 SP 和 μP 表示为统一形式。为此，SP 定义时使用权重初始化方差为 Θ(1/fan-in)，以及全局学习率 η_l = η · n^{-1} 应用于所有层。注意，我们选择了一种约定，将学习率中的 1/n 因子剥离出来，这与其他约定不同。为了使 SP 与 μP 的比较更透明，我们利用 abcd 参数化的对称性 [48 (https://arxiv.org/html/2605.21486#bib.bib90)]，将我们的 μP 描述中的乘子 a 设为 0。如表 1 (https://arxiv.org/html/2605.21486#S2.T1) 所示，两者在四个关键方面存在差异：(1) 嵌入层学习率，(2) 最后一层初始化方差，(3) LayerNorm 学习率，以及 (4) 注意力缩放。权重衰减强度 λ 随学习率缩放，使得其贡献 ηλθ 在两种参数化中均为 Θ(1)。

这些变化的直观原因如下。SP 对所有层应用统一的 Θ(1/n) 学习率。这种缩放对于隐藏层是必需的，因为这些层将其输出计算为 n 个输入项的和：h_i^(l) = ∑_{j=1}^n W_ij^{l+1} h_j^(l)；如果没有 Θ(1/n) 的学习率，激活更新 Δh_i^(l+1) 将以 Θ(n) 比例缩放。相比之下，嵌入层（每个 token 的查找表）…