技术摘要：量化超参数迁移与嵌入层学习率的重要性

问题陈述

训练大型语言模型（LLM）需要大量的超参数调优，随着模型规模达到万亿参数级别，这种调优在计算上变得难以承受。超参数迁移通过在小规模模型上寻找最优超参数，并利用缩放定律或结构参数化将其外推至更大规模，从而解决这一问题。

存在两种主要方法：

1. 功能缩放定律：拟合曲线以预测最优超参数（如学习率 $\eta$）如何随模型规模缩放。虽然有效，但这些方法通常缺乏超出其拟合域的泛化能力，且需要昂贵的搜索来推导稳健的函数形式。
2. 结构参数化：设计模型使得训练动态在不同规模下保持不变。最大更新参数化（$\mu$P）是一个突出的例子，其基于激活及其更新应与宽度无关的假设推导而来。经验表明，$\mu$P 能够实现高质量的学习率迁移，允许从业者将小型代理模型的最优学习率用于大型模型。

然而，在理解为何 $\mu$P 在实践中表现如此出色方面仍存在显著差距。$\mu$P 的理论推导依赖于现实中往往不成立的假设（例如无限宽度极限、权重与激活的完全对齐、固定数据集大小和固定训练步数）。尽管存在这些理论上的不匹配，$\mu$P 始终优于标准参数化（SP）。本文旨在量化超参数迁移的质量，并隔离驱动 $\mu$P 成功的具体机制。

方法论

1. 超参数迁移的量化框架

作者提出了一个框架，利用三个不同的指标严格评估迁移质量，超越了简单的性能比较：

• 损失可预测误差（$E$）：衡量损失景观拟合所提出的缩放假设的程度。低 $E$ 值表明损失是可预测的，外推是可靠的。高 $E$ 值则表明存在复杂的景观、不稳定性或相变，使得外推不可靠。
• 迁移鲁棒性指数（$\kappa$）：量化随着模型规模扩大，损失对外推学习率误差的敏感度。定义为 $\kappa = \alpha - 2\beta + \gamma$（源自损失、学习率和海森曲率的缩放定律），负的 $\kappa$ 意味着景观随宽度变平，从而抑制外推误差；正的 $\kappa$ 意味着景观变陡，放大误差。
• 渐近损失退化（$R(\infty)$）：衡量特定参数化的不可约损失与无限规模下所有参数化所能达到的最佳损失之间的差距。这反映了迁移友好的参数化是否牺牲了绝对性能。

2. 实验设置

• 模型：在 FineWeb-Edu 数据集上训练的仅解码器 GPT 风格 Transformer。
• 优化器：带有 Warmup-Stable-Decay (WSD) 调度策略的 AdamW。
• 缩放：宽度（嵌入维度）从 $n=128$ 缩放至 $n=2048$。
• 模式：实验在固定步数（10,000 步）和计算最优（固定每参数 token 数）设置下进行。
• 消融策略：作者将 SP 和 $\mu$P 都表示为统一的数学形式（表 1），并识别出四个关键差异：
  1. 嵌入层学习率缩放（SP 中为 $\Theta(1/n)$，$\mu$P 中为 $\Theta(1)$）。
  2. 最后一层初始化方差（SP 中为 $\Theta(1/n)$，$\mu$P 中为 $\Theta(1/n^2)$）。
  3. LayerNorm 学习率缩放（SP 中为 $\Theta(1/n)$，$\mu$P 中为 $\Theta(1)$）。
  4. 注意力缩放（SP 中为 $1/\sqrt{d}$，$\mu$P 中为 $1/d$）。

作者进行了系统的消融实验，逐个修改 SP 以匹配 $\mu$P（反之亦然），以隔离迁移质量的主要驱动因素。

主要结果

1. 嵌入层学习率的首要地位

核心发现是，嵌入层学习率是区分 $\mu$P 与 SP 的决定性因素。

• SP + 嵌入层：当修改 SP 使其嵌入层使用 $\Theta(1)$ 学习率（同时保持其他层为 SP 缩放）时，其迁移指标（$E$、$\kappa$、$R(\infty)$）变得与 $\mu$P 几乎相同。
• $\mu$P - 嵌入层：反之，将 $\mu$P 中嵌入层的学习率降低至 $\Theta(1/n)$（匹配 SP），会重新引入 SP 特有的训练不稳定性和高可预测误差（$E$）。
• 机制：在 SP 中，嵌入层学习率充当瓶颈。由于嵌入层执行的是与宽度无关的查找操作（不同于对宽度求和的隐藏层），$\Theta(1/n)$ 的学习率过小，导致学习缓慢和训练不稳定，特别是在训练早期。将该率最大化至 $\Theta(1)$ 可稳定动态。

2. 次要因素

• 注意力缩放：如果不与其他变化配对，改变注意力缩放（$1/\sqrt{d} \to 1/d$）会对鲁棒性（$\kappa$）产生微妙但负面的影响，使迁移变得脆弱。
• LayerNorm：在 SP 中增加 LayerNorm 学习率会破坏训练稳定性，而在 $\mu$P 中降低它则影响微乎其微，这表明 LayerNorm 参数可以缓慢训练而不损害迁移效果。
• 最后一层初始化：最后一层初始化方差的差异对迁移指标的影响微乎其微，这可能是由学习率预热所缓解的。

3. 权重衰减与计算最优训练

• 固定步数设置：权重衰减通过正则化景观改善了损失可预测性（$E$），但一致地降低了渐近性能（$R(\infty)$）。
• 计算最优设置：在固定每参数 token 数的模式下（训练步数随 $n^2$ 缩放），标准的权重衰减惯例（$\eta \cdot \lambda = \Theta(1)$）会导致鲁棒性（$\kappa$）下降。作者建议，在此模式下权重衰减的适当缩放可能需要为 $\eta \cdot \lambda = \Theta(1/n^2)$，但这仍是未来工作有待解决的问题。

4. 泛化性

在 CIFAR-100 上训练的 CNN 实验证实，输入层（类似于嵌入层）的重要性超出了 Transformer 的范围。在 SP 中以 $\Theta(1)$ 学习率训练第一层，可恢复与 $\mu$P 相似的迁移质量。

意义与主张

本文声称提供了首个用于评估超参数迁移的量化框架，区分了可预测性、鲁棒性和绝对性能。

该工作的主要意义在于去神秘化 $\mu$P。作者证明，实现高质量学习率迁移并不需要 $\mu$P 的全部条件集合。相反，在 AdamW 训练的语境下，$\mu$P 的“魔力”主要归因于一个简单直观的调整：最大化嵌入层的学习率。

这一发现具有即时的实际意义：

• 使用标准参数化（SP）的从业者只需修正嵌入层的学习率，即可实现与 $\mu$P 相当的迁移性能，而无需采用完整的 $\mu$P 框架。
• 它识别了一个此前被忽视的训练不稳定性来源：将嵌入层学习率设为瓶颈。
• 它强调了 $\mu$P 的理论推导虽然数学上优雅，但可能依赖于并非观察到的经验效益所严格必需的假设（如完全对齐）。

作者对其研究范围保持谦逊，指出其实验仅限于使用 AdamW 和特定数据集的仅解码器 Transformer。他们承认，计算最优模式下的适当权重衰减缩放以及其他优化器（如 SGD 或 Muon）下的行为需要进一步研究。