预训练收敛慢怎么办？Warmup Steps与Cosine Decay学习率调度策略详解

在预训练语言模型、视觉Transformer或大规模推荐模型时，训练曲线的常见表现是：前几千步损失几乎不下降，随后快速下降，但后期又出现明显震荡甚至验证指标反弹。很多团队把问题归因于模型容量或数据质量，却忽略了优化器中的学习率调度。固定学习率或简单阶梯衰减很难同时满足训练初期的稳定性与训练后期的收敛精度。Warmup Steps与Cosine Decay的组合，正是在不同训练阶段对参数更新步长进行精细控制的一种策略。

为什么固定学习率会让预训练收敛变慢

预训练模型通常使用大批量数据和较大的初始学习率来加速优化。但在训练最初的几百到几千步，网络权重还是随机初始化状态，梯度方向和尺度都极不稳定。此时如果直接把学习率拉到目标值，某些参数可能会在一次更新中跨越过大的距离，导致激活分布漂移、归一化层统计量失真，甚至让损失出现尖峰。这种早期的参数破坏不一定立刻表现为损失上升，但会延长模型进入稳定下降通道的时间。

从优化角度看，Adam和AdamW等自适应优化器虽然对每个参数做了梯度二阶矩归一化，但在梯度方差极大时，自适应步长仍然可能偏大。尤其是在多头注意力和前馈网络叠加的深层结构中，初始梯度范数差异明显，统一使用较大学习率会让不同层以非常不一致的节奏更新。学习率预热通过线性增加学习率，相当于在梯度方向尚不可靠时先用小步长探索，等梯度协方差结构稳定后再进入全量更新。

训练后期的固定学习率或阶梯衰减也有明显问题。当模型接近损失函数的平坦区域时，较大的学习率会让参数在极小值附近来回反弹，难以真正收敛到低损失区域。Cosine Decay使用余弦函数使学习率平滑下降，并且在训练结束时接近零，这正好匹配了模型从快速探索到精细收敛的节奏。

Warmup Steps的机制与实现

Warmup Steps指的是在训练开始阶段，让学习率从很小的值线性增加到预设峰值所需的步数。其核心目的不是提升收敛速度，而是保护训练初期的优化轨迹。具体实现通常将学习率乘以一个从0到1的线性系数，或者在极小初始学习率与峰值学习率之间进行线性插值。例如，若峰值学习率为5e-4，warmup步数为2000，则第0步学习率接近0，第1000步为2.5e-4，第2000步达到5e-4。

Warmup步数并不是越长越好。过长的预热会浪费大量算力在过低学习率上，导致前期损失下降非常缓慢；过短的预热则无法有效抑制初始梯度噪声。实际配置时，常见做法是把warmup步数设为总训练步数的1%到5%。对于总步数100,000的任务，1000到5000步都是合理范围。如果批大小很大，例如超过8192，梯度噪声相对较小，可以适当缩短warmup；如果使用小批量或高学习率，则需要更长的warmup来防止早期震荡。

下面是用PyTorch实现线性warmup与余弦衰减的代码。它使用LambdaLR将当前步数映射为学习率系数。

import math
import torch

def create_scheduler(optimizer, num_warmup_steps, num_training_steps, min_lr_ratio=0.0):
    def lr_lambda(current_step):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
        return min_lr_ratio + (1.0 - min_lr_ratio) * 0.5 * (1.0 + math.cos(math.pi * progress))
    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

optimizer = torch.optim.AdamW(model.parameters(), lr=5e-4)
scheduler = create_scheduler(optimizer, num_warmup_steps=2000, num_training_steps=100000, min_lr_ratio=0.0)

Hugging Face Transformers也提供了现成接口，适合不想手动编写调度逻辑的场景。该接口会自动处理warmup之后的余弦衰减，并且可以与Trainer的total_steps直接配合。

from transformers import get_cosine_schedule_with_warmup

optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)
scheduler = get_cosine_schedule_with_warmup(
    optimizer,
    num_warmup_steps=500,
    num_training_steps=20000
)

Cosine Decay的衰减逻辑与常见变体

Cosine Decay的核心公式是：在warmup结束后，学习率按照余弦函数从峰值平滑下降到最小学习率。当前进度为线性比例，余弦值从1变化到-1，经过缩放后使学习率从峰值逐渐接近最小值。公式可以写成：lr = min_lr + 0.5 * (max_lr - min_lr) * (1 + cos(pi * progress))。其中progress在0到1之间。该曲线的特点是开始下降较慢，中段加速下降，末段又变慢，这使得模型在训练中后期仍能保持较大探索能力，而在最后阶段逐渐进入精细收敛。

与阶梯衰减相比，Cosine Decay不需要手动设置衰减节点和衰减倍率，减少了超参数数量。与指数衰减相比，余弦衰减在训练末段不会过早把学习率压到极小，而是保留一定探索空间，随后平滑收束。因此它在视觉Transformer、BERT、GPT等预训练任务中表现稳定。

实际使用时可以有多种变体。一是带重启的余弦退火，它在多个周期中反复执行余弦曲线，适合需要从局部最优跳出的任务；二是设置最小学习率不为零，例如峰值学习率的5%到10%，这可以避免最后阶段更新量过小；三是组合线性warmup与余弦衰减，这也是当前预训练中最常见的方案。下面代码展示了如何计算每一步的学习率，并打印前几步验证曲线。

import math

max_lr = 5e-4
min_lr = 5e-5
warmup_steps = 2000
total_steps = 100000

for step in range(0, 100000, 5000):
    if step < warmup_steps:
        lr = max_lr * step / warmup_steps
    else:
        progress = (step - warmup_steps) / (total_steps - warmup_steps)
        lr = min_lr + 0.5 * (max_lr - min_lr) * (1.0 + math.cos(math.pi * progress))
    print(step, lr)

Warmup与Cosine Decay协同配置的关键经验

要把这套调度策略用好，需要根据总训练步数、批大小、模型规模和优化器类型来设置参数。首先是峰值学习率。对于AdamW优化器，常见峰值学习率在1e-4到1e-3之间，小模型可以取更高，大模型则倾向于更低。过高的峰值学习率即使有warmup保护，也可能在中后期造成训练不稳定；过低则会让整个训练过程偏慢。可以先从小规模代理任务中搜索峰值学习率，再迁移到大规模预训练。

其次是warmup步数。总步数在50,000以下的任务，warmup步数可以设为总步数的2%到5%；总步数在100,000以上的任务，1%通常足够。如果使用非常大的批大小，例如32,768，梯度估计更准确，warmup步数可以减少到几百步。但如果同时提高了峰值学习率，则需要同步增加warmup步数，让模型有足够时间适应大学习率。

最后是最小学习率和训练周期。若训练1个epoch，最小学习率可以设为零或非常小；若训练多个epoch，建议保留峰值学习率的5%到10%作为最小学习率，避免模型在epoch之间进入完全停滞状态。对于带重启的余弦退火，周期长度要尽量与epoch边界对齐，否则会在epoch中间突然抬升学习率，破坏当前batch的统计特性。

一个常见误区是只调整warmup步数而忽略总步数变化。修改数据集规模或微调epoch后，总训练步数会变化，如果warmup步数不变，预热占比就会偏移。更稳妥的做法是把warmup表示为总步数的比例，并在代码中动态计算，而不是写死一个整数。

实验对比与典型训练曲线分析

假设在相同的Transformer预训练任务中，固定学习率5e-4的模型在20,000步后损失仍在2.8附近震荡，而使用2,000步warmup加余弦衰减的模型在同样步数下损失降到2.4，验证困惑度也低约0.3。差异主要来自训练初期的稳定性和末期的精细收敛。固定学习率在后期由于步长过大，参数在最优区域周围持续震荡，无法进一步降低损失；而余弦衰减让学习率在最后5%的训练步中已经降到极低水平，参数更新幅度很小，更容易停在损失曲面的低谷。

观察学习率曲线还可以发现，warmup阶段损失下降幅度不一定很大，但梯度范数会逐渐稳定。这意味着预热的主要作用不是立刻降低损失，而是为后续快速下降建立良好的参数状态。很多训练日志中，warmup结束后的几千步才会出现损失快速下降，这是正常现象，不要误判为预热无效。

此外，不同层对学习率的敏感度不同。在深层预训练模型中，注意力层的输出投影矩阵通常对学习率更敏感，前馈网络的中间层则相对鲁棒。如果出现训练早期某个层的梯度突然增大，可以适当延长warmup步数，而不是直接降低全局峰值学习率。这样既能保持整体收敛速度，又能避免局部梯度爆炸。

总结

Warmup Steps与Cosine Decay组合解决预训练收敛慢的本质，是让优化器在不同阶段采用不同的更新尺度。早期小步长保护随机初始化状态，中期大步长快速探索，后期余弦衰减平滑收束。相比固定学习率和阶梯衰减，这套策略只需要调整warmup占比、峰值学习率和最小学习率三个关键参数，就能在多数预训练任务中获得更稳定的收敛和更好的泛化。实际工程中应把warmup与总步数绑定，根据批大小和优化器调整峰值学习率，并通过记录学习率与梯度范数曲线来验证调度是否合理。

学习率调度WarmupCosine Decay修改时间：2026-08-22 02:09:48