Learning Rate Warmup/zh
| Article | |
|---|---|
| Topic area | Optimization |
| Prerequisites | Stochastic Gradient Descent, Learning Rate Schedule, Adam |
概述
学习率热身是一类学习率调度,其中学习率最初设为较小的值,并在训练的前几百至数千次迭代中逐步增加到目标峰值,然后才进入任何标准的衰减阶段。该技术于 2010 年代中期非正式地出现在深度学习实践中,2017 年由 Goyal 等人为大批量 ImageNet 训练给出了正式方案,并在最初的transformer论文和 BERT 论文中使用之后,成为 Transformer 模型训练流水线中几乎通用的组成部分。[1][2]
其动机在于训练最初的若干次更新具有异常高的风险:参数处于随机初始化状态,梯度幅值很大且与真实下降方向相关性较差,而像 Adam 这样的自适应优化器尚未积累可靠的二阶矩估计。在此期间使用较小的学习率可以防止这些早期更新将迭代点推入损失曲面的不良区域,之后便可在训练的主体阶段安全地使用更高的速率。
线性热身
最常见的形式是线性热身。给定目标峰值学习率 $ \eta_{\max} $ 以及以迭代(有时为轮次)计的热身长度 $ T_w $,第 $ t $ 次迭代时的速率为
$ {\displaystyle \eta_t = \eta_{\max} \cdot \min\!\left(1, \frac{t}{T_w}\right).} $
在第 $ T_w $ 步之后,调度会继续按所规划的衰减进行:余弦退火、逆平方根衰减、阶梯衰减或保持常数速率。典型的热身长度介于总训练量的 1% 到 10% 之间,语言模型预训练中常见的取值为 500 到 4000 次迭代。
一个近亲,常数热身,在前 $ T_w $ 步使用一个固定的小速率 $ \eta_0 \ll \eta_{\max} $,然后跳跃到 $ \eta_{\max} $。Goyal 等人观察到,当 $ \eta_0 $ 与 $ \eta_{\max} $ 之间差距较大时,常数热身较为脆弱,因而推荐渐进式热身(即上述线性形式)作为大批量训练更可靠的默认方案。
Transformer 中的逆平方根热身
最初的transformer论文将线性热身与随后的逆平方根衰减结合起来,产生了一种具有单个平滑可微峰值的调度:
$ {\displaystyle \eta_t = d_{\text{model}}^{-1/2} \cdot \min\!\left( t^{-1/2},\, t \cdot T_w^{-3/2} \right),} $
其中 $ d_{\text{model}} $ 是 transformer 模型维度,$ T_w $ 是热身长度(通常为 4000 步)。最小值的两条分支在 $ t = T_w $ 处相交,调度从线性增长的热身阶段过渡到逆平方根衰减。这一方案以微小的改动被许多后续的序列到序列系统所继承,直到线性热身加余弦衰减在自回归 语言模型预训练中取代了它。
热身为什么有效
已有若干互补的解释,但都尚未完全定论。Liu 等人给出了最具影响力的分析:在 Adam 及其变体中,自适应二阶矩估计 $ v_t $ 在最初几次更新中方差极高,这使得每参数的有效学习率 $ \eta / \sqrt{v_t + \epsilon} $ 不可靠。[3]短暂的热身阶段使 $ v_t $ 在全局步长变大之前有时间趋于稳定。同一篇论文还提出了 RAdam,这是一种 Adam 的变体,包含方差修正项,旨在使显式热身变得不必要;在实际中,RAdam 与显式热身仍然都在使用。
第二条理由涉及初始化附近的损失曲面。随机权重产生的激活值和梯度远离最终的训练轨迹;早期的大步长可能放大层间的不平衡(例如transformer中注意力子层与前馈子层之间的不平衡),并引发数值不稳定性,如softmax前 logits 的爆炸。热身可在此过渡期内限制每一步造成的损害。
第三种解释专门针对大批量训练,援引梯度噪声尺度。当批量非常大时,随机梯度接近真实梯度,因而优化器的行为几乎是确定性的;此时较高的初始学习率可能导致迭代点持续过冲。Goyal 等人将热身与线性缩放规则(峰值速率与批量大小成正比)结合起来,把带动量的SGD扩展到了 ImageNet 上 8192 张图像的批量,且未损失准确率。
变体与组合
除了线性形式和逆平方根形式外,实践中还使用其他若干形状。
- 余弦热身:速率从 $ 0 $ 沿半个余弦曲线上升到 $ \eta_{\max} $,比线性斜坡更为平滑。该形状随后由余弦衰减镜像,形成单一的钟形调度。
- 指数热身:$ \eta_t = \eta_{\max} \cdot (1 - e^{-t / \tau}) $,其中 $ \tau $ 为某一时间常数。在深度学习中较少见,但在强化学习和元学习中有所出现。
- 逐层热身:在 LAMB 及相关的逐层自适应优化器中,信任比机制产生了一种隐式热身,其长度因层而异。通常仍会在其之上保留一个显式的全局热身。
热身可以自然地与大多数衰减调度组合。在现代 LLM 预训练中占主导地位的模式是线性热身后接 余弦衰减,衰减到峰值速率的某一比例(通常为 10%)。热身也常在微调开始时再次应用,理论相同:新初始化的输出头以及由框架重置的优化器状态,都会从短暂的低速率阶段中获益。
实践考量
原则上热身会引入两个超参数($ T_w $ 和起始速率,常为零),但大多数方案会将起始速率固定为零或一个小常数(如 $ 10^{-7} $),仅调整 $ T_w $。经验性指引:
- 对于 Transformer 预训练,将 $ T_w $ 设置为总训练步数的 1% 到 5%。小于 0.5% 的取值经常引发损失尖刺,尤其是在模型宽度较大时。
- 对于有监督微调,典型的热身步数为 50 到 500 步;最优长度随数据集规模而变化。
- 扩大批量大小或模型规模时增大 $ T_w $;这两种调整都会放大热身所要应对的早期训练不稳定性。
- 当使用 混合精度或 FSDP / 流水线并行时,更长的热身有助于容忍损失缩放因子和层归一化的初始失配。
热身与 AdamW 中的权重衰减存在相互作用:由于解耦的权重衰减与梯度无关,热身阶段的有效衰减强度并未受到影响,但因为梯度更新较小,朝零方向的相对拉力更大。因此,一些方案会按照热身阶段学习率所用的相同系数来缩放权重衰减,不过这并非通用做法。
局限性与替代方案
热身在经验上稳健,但在理论上解释尚不充分。其有效性取决于尚未完全刻画的细节:优化器类型、模型架构、运算精度以及初始化方案。RAdam 试图将热身阶段吸收进优化器本身;Adafactor 与 Lion 同样在某些情形下减少了对热身的需求,但并未将其完全消除。架构层面的修正,如 Pre-Norm transformer和精心设计的初始化(Fixup、DeepNorm、ReZero),可降低实现稳定所需的热身长度,但在大规模训练中很少能让省略热身成为安全选项。
在缺乏单一封闭形式理论的情况下,多个研究团队达成的共识是:在任何使用自适应优化器、大批量或 transformer 类架构的训练流水线中,都应保留显式的热身阶段,并在首次训练某个新模型类别时,将其长度视为值得做小规模定向扫描的超参数。