§1.3.28

Warmup + Cosine/Linear/InverseSqrt 学习率调度？

手写练习

—用 torch.optim.lr_scheduler 组合 Warmup + Cosine

核心概念

学习率调度（Learning Rate Scheduling）是在模型训练过程中动态调整优化器学习率（Learning Rate, LR）的策略。其核心思想是在训练的不同阶段采用不同的学习率，以平衡收敛速度和最终性能。Warmup（预热）是一种在训练初期使用一个非常小的学习率，然后逐渐线性增加到预设的基础学习率的策略。它能有效避免模型在训练初期因参数随机初始化而导致梯度过大、模型振荡甚至发散的问题，有助于稳定训练过程。Warmup 结束后，学习率会按照某种衰减策略（如余弦、线性、逆平方根等）逐渐降低，以帮助模型在接近最优点时更精细地搜索，从而收敛到一个更好的局部或全局最优解。

原理与推导

学习率调度可以分为两个阶段：Warmup 阶段和 Decay（衰减）阶段。我们定义以下关键变量：

$\eta_t$ : 在第 $t$ 个训练步（step）时的学习率。
$\eta_{base}$ : 预设的基础学习率，也是 Warmup 阶段结束时的目标学习率。
$t$ : 当前训练步，从 0 开始。
$T_{warmup}$ : Warmup 阶段的总步数。
$T_{total}$ : 训练的总步数。

1. Warmup 阶段 ( $0 \le t \le T_{warmup}$ )

在 Warmup 阶段，学习率从一个很小的值（通常为 0）线性增长到 $\eta_{base}$ 。

数学公式:

\eta_t = \eta_{base} \cdot \frac{t}{T_{warmup}}

推导与动机:
- 当 $t=0$ 时， $\eta_0 = 0$ 。
- 当 $t=T_{warmup}$ 时， $\eta_{T_{warmup}} = \eta_{base}$ 。
- 在 $0$ 到 $T_{warmup}$ 之间，学习率随 $t$ 线性增长。
直观解释:
- 训练初期，模型权重是随机的，损失函数的梯度可能非常大且方向不稳定。如果直接使用一个较大的学习率，权重更新的步子会迈得太大，可能导致模型“飞出”最优解所在的“山谷”，造成训练不稳定。Warmup 就像是给汽车引擎预热，先用低转速稳定运行，再逐步提高到正常转速，确保整个过程平稳。

2. Decay 阶段 ( $T_{warmup} < t \le T_{total}$ )

Warmup 结束后，学习率开始衰减。

(a) 余弦退火 (Cosine Annealing)

数学公式: 学习率从 $\eta_{base}$ 平滑地衰减到某个最小值 $\eta_{min}$ （通常为 0）。

\eta_t = \eta_{min} + \frac{1}{2}(\eta_{base} - \eta_{min}) \left(1 + \cos\left(\frac{t - T_{warmup}}{T_{total} - T_{warmup}}\pi\right)\right)

推导与动机:
- 该公式的核心是 $\cos(\theta)$ 函数。我们希望当 $t$ 从 $T_{warmup}$ 变化到 $T_{total}$ 时，学习率从 $\eta_{base}$ 变化到 $\eta_{min}$ 。
- 定义进度比例 $\alpha = \frac{t - T_{warmup}}{T_{total} - T_{warmup}}$ ， $\alpha$ 从 0 变化到 1。
- 我们让角度 $\theta = \alpha \cdot \pi$ ，于是 $\theta$ 从 0 变化到 $\pi$ 。
- $\cos(\theta)$ 从 $\cos(0)=1$ 变化到 $\cos(\pi)=-1$ 。
- 表达式 $\frac{1}{2}(1 + \cos(\theta))$ 将 $[1, -1]$ 的范围映射到 $[1, 0]$ 。
- 最后，通过线性变换 $y = y_{min} + (y_{max}-y_{min}) \cdot x$ （其中 $x \in [0,1]$ ），我们将 $[1, 0]$ 的变化范围映射到 $[\eta_{base}, \eta_{min}]$ ，得到最终公式。
几何解释: 学习率的变化曲线是一个余弦函数的半周期，下降过程先慢、再快、后慢。这种平滑的衰减方式被认为有助于模型在训练后期更好地探索损失函数的平坦区域，找到更鲁棒的最小值。

(b) 线性衰减 (Linear Decay)

数学公式: 学习率从 $\eta_{base}$ 线性衰减到 0。

\eta_t = \eta_{base} \cdot \left(1 - \frac{t - T_{warmup}}{T_{total} - T_{warmup}}\right)

推导与动机:
- 这是一个简单的线性插值。当 $t=T_{warmup}$ 时，分数项为 0， $\eta_t = \eta_{base}$ 。
- 当 $t=T_{total}$ 时，分数项为 1， $\eta_t = 0$ 。
几何解释: 学习率随时间呈一条直线下降。这是一种简单有效的衰减方式，但相比余弦退火，其在训练末期的衰减可能不够“温柔”。

(c) 逆平方根衰减 (Inverse Square Root Decay)

这是 Transformer 论文《Attention Is All You Need》中提出的经典调度器，它将 Warmup 和 Decay 统一在一个公式中。

数学公式: 为了与 Transformer 论文保持一致，这里的 $\eta_{base}$ 被一个与模型维度相关的缩放因子代替。令基础学习率为 $lrate = d_{model}^{-0.5}$ 。

\eta_t = lrate \cdot \min(t^{-0.5}, t \cdot T_{warmup}^{-1.5})

推导与动机:
- Warmup 部分: $t \cdot T_{warmup}^{-1.5}$ 。这是一个线性增长项。当 $t=T_{warmup}$ 时，其值为 $T_{warmup} \cdot T_{warmup}^{-1.5} = T_{warmup}^{-0.5}$ 。
- Decay 部分: $t^{-0.5}$ 。这是一个逆平方根衰减项。当 $t=T_{warmup}$ 时，其值为 $T_{warmup}^{-0.5}$ 。
- $\min(...)$ 函数确保了在 $t \le T_{warmup}$ 时，选择线性增长部分（因为 $t \cdot T_{warmup}^{-1.5} < t^{-0.5}$ ）；在 $t > T_{warmup}$ 时，选择逆平方根衰减部分。在 $t = T_{warmup}$ 处，两者相等，保证了曲线的连续性。
信息论解释: 逆平方根衰减的衰减速度比线性和余弦都慢。在随机优化理论中，为了保证收敛，学习率衰减速度通常需要满足 $\sum \eta_t = \infty$ 和 $\sum \eta_t^2 < \infty$ 。 $t^{-0.5}$ 的衰减率处于收敛的边界上，这种缓慢的衰减可能给予模型更长的探索时间，对于复杂的 Transformer 模型被证明是有效的。
算法复杂度: 所有这些调度器在每个训练步的计算都是 $O(1)$ 的时间和空间复杂度，因为它们只依赖于当前的步数和预设的几个常数，对训练性能几乎没有影响。

代码实现

以下代码演示了如何使用 torch.optim.lr_scheduler 组合实现 Warmup + Cosine 学习率调度。我们将使用 SequentialLR 来串联一个线性 Warmup 调度器和一个余弦退火调度器。

python

 1import torch
 2import torch.nn as nn
 3import torch.optim as optim
 4from torch.optim.lr_scheduler import LambdaLR, CosineAnnealingLR, SequentialLR
 5import matplotlib.pyplot as plt
 6import numpy as np
 7
 8# --- 1. 设置参数 ---
 9# 定义一个虚拟模型和优化器
10model = nn.Linear(10, 2)
11# 将基础学习率（Warmup结束时的最大学习率）传递给优化器
12base_lr = 0.1
13optimizer = optim.SGD(model.parameters(), lr=base_lr)
14
15# 训练参数
16total_epochs = 100
17warmup_epochs = 10
18# 确保衰减阶段的 T_max 正确，它应该是衰减阶段的步数
19cosine_t_max = total_epochs - warmup_epochs
20
21# --- 2. 定义 Warmup 调度器 ---
22# LambdaLR 允许我们通过一个 lambda 函数自定义学习率的缩放因子
23# 在 warmup 阶段，我们希望学习率从 0 线性增加到 base_lr
24# 学习率公式: lr = base_lr * (epoch / warmup_epochs)
25# LambdaLR 的 lambda 函数接收当前的 epoch，返回一个乘法因子
26# 所以，因子应该是 epoch / warmup_epochs
27warmup_scheduler = LambdaLR(
28    optimizer,
29    lr_lambda=lambda epoch: epoch / warmup_epochs
30)
31
32# --- 3. 定义 Cosine 衰减调度器 ---
33# CosineAnnealingLR 在 T_max 个 epoch 内将学习率从上一个调度器的最终值（即 base_lr）
34# 退火到 eta_min
35cosine_scheduler = CosineAnnealingLR(
36    optimizer,
37    T_max=cosine_t_max, # T_max 是半个余弦周期的步数
38    eta_min=0.001       # 最小学习率
39)
40
41# --- 4. 使用 SequentialLR 组合两个调度器 ---
42# SequentialLR 按顺序链接多个调度器
43# milestones 参数指定了切换到下一个调度器的 epoch 点
44# 在第 warmup_epochs 个 epoch 结束时，从 warmup_scheduler 切换到 cosine_scheduler
45scheduler = SequentialLR(
46    optimizer,
47    schedulers=[warmup_scheduler, cosine_scheduler],
48    milestones=[warmup_epochs]
49)
50
51# --- 5. 模拟训练过程并记录学习率 ---
52print("开始模拟训练并记录学习率变化...")
53lrs = []
54for epoch in range(total_epochs):
55    # 获取当前学习率并记录
56    current_lr = optimizer.param_groups[0]['lr']
57    lrs.append(current_lr)
58    
59    # 模拟训练步骤
60    optimizer.step()  # 更新模型参数（在实际训练中会先计算梯度）
61    scheduler.step()  # 更新学习率
62
63print(f"在第 0 epoch, LR = {lrs[0]:.4f}")
64print(f"在第 {warmup_epochs-1} epoch (Warmup结束前), LR = {lrs[warmup_epochs-1]:.4f}")
65print(f"在第 {warmup_epochs} epoch (Cosine开始时), LR = {lrs[warmup_epochs]:.4f}")
66print(f"在第 {total_epochs-1} epoch (训练结束时), LR = {lrs[-1]:.4f}")
67
68# --- 6. 绘制学习率曲线 ---
69plt.figure(figsize=(10, 6))
70plt.plot(np.arange(total_epochs), lrs)
71plt.title("Warmup + Cosine Learning Rate Schedule")
72plt.xlabel("Epoch")
73plt.ylabel("Learning Rate")
74plt.grid(True)
75# 绘制Warmup结束的垂直线
76plt.axvline(x=warmup_epochs, color='r', linestyle='--', label=f'Warmup End (Epoch {warmup_epochs})')
77plt.legend()
78plt.show()

工程实践

使用场景:
- 大型 Transformer 模型: 对于 BERT、GPT、ViT 等大型模型，Warmup + Cosine/Linear 衰减几乎是标配。这些模型对初始化和学习率非常敏感，Warmup 提供了必要的训练稳定性。
- 从零开始训练: 在没有预训练权重，从头开始训练深度神经网络时，Warmup 可以有效防止早期梯度爆炸。
- 迁移学习/微调: 即使在微调阶段，一个短暂的 Warmup（例如，总步数的 1-5%）加上一个温和的衰减也能帮助模型在新数据上更好地适应，同时避免破坏预训练学到的知识。
超参数选择经验:
- warmup_steps 或 warmup_ratio: 通常设置为总训练步数的 5% 到 10%。例如，如果总共训练 100,000 步，可以选择 5,000 到 10,000 步作为 Warmup。过长的 Warmup 会浪费计算资源在低效的学习率上；过短则可能无法起到稳定作用。
- base_lr: 这是最关键的超参数之一，通常需要通过实验（如 LR Range Test）来确定。调度器的作用是在这个 base_lr 周围进行调整，而不是帮你找到它。
- min_lr (for Cosine): 通常设置为 base_lr 的 10% 或直接设为 0。在非常长的训练中，保持一个小的 min_lr 可以让模型在最后阶段仍有能力进行微调。
性能 / 显存 / 吞吐的权衡:
- 学习率调度器本身的计算开销极小，对显存和吞吐量（Throughput）几乎没有影响。
- 其核心价值在于通过优化训练动态过程，用更少的总训练步数（或在相同步数下）达到更高的模型性能，从而节省大量的计算时间和成本。一个好的调度策略是提升训练效率的关键杠杆。
常见坑和调试技巧:
- 第一准则：可视化！ 在开始正式、昂贵的训练前，务必编写一个简单的脚本（如上文代码）来绘制你配置的学习率曲线。90% 的调度器配置错误（如总步数算错、Warmup 比例不当）都可以通过一张图直观地发现。
- Step-based vs. Epoch-based: PyTorch 的调度器默认是基于 epoch 的。但在处理大型数据集时，通常采用基于 step 的更新。这意味着 scheduler.step() 应该在每个 batch 的 optimizer.step() 之后被调用，而不是在每个 epoch 结束时。这是初学者最常犯的错误。
- 检查当前学习率: 在训练日志中打印 optimizer.param_groups[0]['lr']，以确保学习率确实在按预期变化。如果它一直不变，很可能是你忘记调用 scheduler.step() 或者调用位置错误。

常见误区与边界情况

误区1：优化器中的 lr 参数: 传递给优化器构造函数（如 optim.SGD(..., lr=0.1)) 的 lr 应该是你的基础学习率（base_lr），即 Warmup 想要达到的峰值。调度器会在此基础上进行缩放，而不是从一个你手动设置的初始小学习率开始。
误区2：scheduler.step() 的调用时机: 自 PyTorch 1.1.0 之后，官方推荐的调用顺序是 optimizer.step() 之后再调用 scheduler.step()。在旧版本中，顺序相反。遵循最新实践可以避免一些警告和潜在问题。

边界情况1：恢复训练 (Resuming Training): 这是一个非常重要的面试考点。当你从一个 checkpoint 恢复训练时，必须同时保存和加载优化器和调度器的状态。