Transformer训练中的Warm-up（预热）策略

知识点图片

在Transformer模型的训练中，Warm-up（预热）是一种非常关键的学习率调度策略。几乎所有主流的大语言模型（如BERT、GPT系列、LLaMA等）在训练时都会使用这一策略。

简单来说，Warm-up就是在训练刚开始时，不直接使用设定好的最大学习率，而是让学习率从零（或一个极小的值）开始，经过一定步数（或比例）的训练后，逐渐线性增加到最大值，然后再按照特定的衰减策略（如余弦衰减、线性衰减）逐渐降低。

以下是对Transformer训练中Warm-up策略的全面解析：

一、为什么Transformer需要Warm-up？

在CNN或RNN时代，Warm-up并不是必须的，但在Transformer中它却至关重要。主要原因有以下几点：

1. 缓解Adam优化器的“冷启动”问题（核心原因）

Transformer通常使用Adam或AdamW优化器。Adam优化器会根据梯度的二阶矩（即梯度的平方的滑动平均）来动态调整每个参数的学习率。

在训练的最初阶段，模型参数是随机初始化的，梯度往往非常大且极不稳定。
此时梯度的二阶矩估计（方差）还未收敛，如果此时给一个较大的学习率，Adam计算出的自适应步长会异常巨大，导致模型参数发生剧烈且盲目的更新，直接把模型“训崩”（甚至出现NaN或Loss不再下降）。
Warm-up的作用：以极小的学习率让模型“慢走”几步，给Adam优化器积累足够多、足够稳定的梯度二阶矩信息，使优化器状态趋于平稳。

2. 防止早期梯度爆炸和破坏初始化

Transformer模型往往很深，且包含大量的自注意力机制。如果一开始学习率很大，巨大的梯度会迅速破坏掉精心设计的随机初始化分布（如Xavier或Kaiming初始化），导致模型陷入极差的局部最优解。预热可以让参数在早期平滑过渡。

3. Post-LN 架构的固有缺陷（历史原因）

在最早的《Attention Is All You Need》论文中，Transformer使用的是Post-LN（后置层归一化）结构。

研究（如《On Layer Normalization in the Transformer Architecture》）表明，在Post-LN结构下，靠近输出层的梯度非常大，而靠近输入层的梯度非常小。这种极端的梯度不平衡在训练初期极易导致不收敛。
Warm-up强行压制了前期的更新步长，是Post-LN能够成功训练的必要条件。
注：现代大模型（如GPT系列、LLaMA）大多改用了Pre-LN（前置层归一化），Pre-LN的梯度更稳定，对Warm-up的依赖有所降低，但为了获得更好的泛化能力和最终性能，Warm-up依然被作为标配保留。

二、常见的Warm-up调度策略

Warm-up通常只占据整个训练过程的一小部分（例如前1%~10%的步数），真正的区别往往在于Warm-up结束后的衰减（Decay）阶段。

1. 原始Transformer策略（Inverse Square Root）

《Attention Is All You Need》中提出的经典公式：学习率首先线性增加，达到峰值后，按步数的倒数平方根衰减。

公式： $lr = d_{model}^{-0.5} \cdot \min(step\_num^{-0.5}, step\_num \cdot warmup\_steps^{-1.5})$
特点：前期成正比例上升，后期缓慢下降。

2. Linear Warm-up + Linear Decay (BERT常用)

学习率线性上升到最高点，然后再线性下降到0。

特点：简单直观，在微调（Fine-tuning）阶段非常常用。

3. Linear Warm-up + Cosine Decay (现代大模型标配)

学习率线性上升到最高点，然后按照余弦函数（Cosine）的曲线缓慢下降到最高学习率的10%左右。

应用：GPT-3、LLaMA、Chinchilla等预训练几乎全部采用此策略。
优点：余弦衰减在训练中期能保持相对较高的学习率，有利于跳出局部最优；在训练末期学习率下降非常平缓，有利于模型在最优解附近精细收敛。

三、如何设置Warm-up的参数？

在实际训练中，Warm-up阶段的长短通常有两种设置方式：

按绝对步数设置：例如设置 warmup_steps = 4000 或 10000。在超大规模预训练中经常使用。
按总步数的比例设置：通常设置为总训练步数的 1% 到 10%（如 warmup_ratio = 0.05）。
- 预训练（Pre-training）：数据量极大，通常只需较小的比例（如1%-5%）。
- 微调（Fine-tuning）：总步数较少，通常设置 5% - 10%。

四、代码实现示例 (基于 Hugging Face `transformers`)

在PyTorch和Hugging Face生态中，实现Warm-up非常简单。以下是一个使用“线性预热+余弦衰减”的例子：

python

from transformers import get_cosine_schedule_with_warmup
from torch.optim import AdamW

# 假设模型和数据加载器已经准备好
optimizer = AdamW(model.parameters(), lr=5e-5)

total_steps = 10000        # 总训练步数
warmup_steps = 1000        # 预热步数 (这里占10%)

# 创建调度器
scheduler = get_cosine_schedule_with_warmup(
    optimizer=optimizer,
    num_warmup_steps=warmup_steps,
    num_training_steps=total_steps
)

# 在训练循环中调用
for step in range(total_steps):
    # 1. 前向传播、计算Loss
    # 2. loss.backward()
    
    optimizer.step()     # 更新参数
    scheduler.step()     # 更新学习率 (每个step调用一次)
    optimizer.zero_grad()

总结

在Transformer的训练中，Warm-up不是可选项，而是必须项。它像是在汽车高速行驶前先挂低挡起步，避免了“冷机轰大油门”导致的发动机（优化器状态和参数初始化）损坏，从而保障了深度模型能够平稳地驶入最佳的优化路径。

右滑查看面试常问

播面