从零手写大模型 · 训练篇 05 - Learning Rate Scheduling（学习率调度）

复习：Adam solve了什么，没solve什么

上一篇我们从"固定学习率的两个老大难问题"——震荡和各方向步长不匹配——推导出了Momentum、RMSProp，最后拼成了Adam。Adam确实聪明：它给每个参数配了一个"自适应步长"，不再是一刀切。

但这里有一个容易被忽略的点：Adam的自适应，指的是参数之间步长不同（这个方向陡就走小步，那个方向缓就走大步），而不是训练阶段之间步长的变化。也就是说，哪怕用了Adam，你设的那个初始学习率lr=0.001，在训练第1步和训练第10000步，本质上还是同一个量级在起作用。

这就引出今天的问题：训练的"早期"和"后期"，需要的学习率其实是不一样的。而且不一样到，如果你不管它，会出两种典型的问题。

为什么固定学习率会两头不讨好

想象一下损失函数是一个山谷（哪怕Adam帮你把各个方向的坡度磨平了，整体轮廓还在）。

训练早期：参数还是随机初始化的样子，离山谷底部很远，loss很大，梯度往往也比较大、比较"混乱"（各个batch算出来的梯度方向可能差异很大）。这时候如果学习率一上来就设得比较大，很容易一步迈太猛，直接冲过头，甚至在山谷两壁之间来回震荡，训练不稳定，loss曲线一开始就上蹿下跳。

训练后期：参数已经接近山谷底部了，这时候你反而希望走小碎步，慢慢逼近最优点。可如果学习率还保持一开始那个不大不小的值，你会发现loss降到某个程度后就开始在最优点附近"打转"——每一步都跨得太大，刚好又跨过了那个最低点，落到对面去，来回踏步但不再下降。

用我们熟悉的"猫吃鱼"的说法：这就像猫扑鱼的时候，一开始鱼离得远，猫应该大步冲过去（大学习率，快速接近目标）；但等到猫已经扑到鱼跟前了，还继续按刚才冲刺的步子扑，那大概率是直接扑空、鱼都被拍飞了（学习率太大导致跳过最优点）。这时候猫需要的是收着爪子，一点点挪近乎（学习率衰减）。

所以我们需要的不是一个"数"，而是一条随训练进程变化的曲线——这就是学习率调度（Learning Rate Scheduling）要解决的问题。

两种最常用的策略：Warmup 和 Cosine Decay

学习率调度的花样很多，但今天我们只讲工程上最常用、也最值得先理解透的两个：Warmup（预热）和Cosine Annealing（余弦退火），以及它们组合起来用的方式。

Warmup：训练一开始，先别急

上面说了，训练早期梯度不稳定，容易一步迈太猛。Warmup的思路很直接：开局的时候，学习率不要一上来就是设定值，而是从一个很小的值，用几百到几千步的时间，线性爬升到目标学习率。

比如目标学习率是lr_max = 3e-4，warmup步数T_warmup = 100，那么第t步（t ≤ T_warmup）的学习率就是：

lr(t) = lr_max × (t / T_warmup)

在t=1的时候，学习率几乎是0；到t=100的时候，正好爬到lr_max。这段时间参数在做的事情，更像是"先小步试探，把最离谱的梯度方向和参数尺度先调顺了"，等到网络的状态稳定下来一些，再放开手脚用完整的学习率去学。

Cosine Annealing：训练后期，慢慢收

Warmup爬到顶之后呢？如果之后一直保持lr_max不变，我们前面说的"后期打转"问题还是会出现。这时候需要在剩下的训练步数里，把学习率平滑地降下来。

为什么是"余弦"而不是直接线性往下降？因为余弦曲线有一个很好的形状特点：一开始降得慢、中间降得快、快结束时又降得慢。也就是说，在学习率还比较大的那段（离目标不远但还有优化空间），它下降得温和一些，给模型足够时间继续有效学习；到了训练中段，果断快速下降；快训练完的时候，学习率已经很小了，此时下降速度又放缓，让参数能在很小的步长下做最后的精细收敛。

公式长这样（假设总训练步数T_total，从T_warmup步之后开始算余弦部分的进度p）：

p = (t - T_warmup) / (T_total - T_warmup)
lr(t) = lr_min + 0.5 × (lr_max - lr_min) × (1 + cos(π × p))

p从0走到1，cos(π×p)就从1走到-1，整个lr(t)就从lr_max平滑地降到lr_min（lr_min通常设成一个很小的值，比如0，或者lr_max的1%）。

两段拼起来，完整的学习率曲线就是：先线性爬坡（Warmup），到达峰值后，再走一段余弦下坡（Cosine Decay）。这也是目前训练Transformer类模型最常见的默认配置。

手工数值演示：同样20步，固定lr vs Warmup+Cosine

光说曲线形状比较抽象，我们还是拿"猫吃鱼"的两层小网络，做一次简化的手工对比，直观感受一下"同样的训练步数，学习率曲线不同，结果差多少"。

设定：d_model=4，用EP3里那套梯度下降的更新规则，训练20步，分别用两种学习率策略：

策略A（固定学习率）：全程lr=0.1

策略B（Warmup+Cosine）：T_warmup=4，T_total=20，lr_max=0.15，lr_min=0.01

策略B在前4步的学习率：

t=1: lr = 0.15 × (1/4) = 0.0375
t=2: lr = 0.15 × (2/4) = 0.075
t=3: lr = 0.15 × (3/4) = 0.1125
t=4: lr = 0.15 × (4/4) = 0.15

第4步之后进入余弦下降段，比如到第12步（p = (12-4)/(20-4) = 0.5）：

lr(12) = 0.01 + 0.5×(0.15-0.01)×(1+cos(π×0.5))
       = 0.01 + 0.5×0.14×(1+0) = 0.08

到第20步（p=1）：

lr(20) = 0.01 + 0.5×0.14×(1+cos(π)) = 0.01 + 0.5×0.14×0 = 0.01

把这两条学习率曲线分别代入EP3那套手工梯度下降的迭代，跑完20步之后你会观察到一个典型现象：策略A在前几步loss下降很快（因为一开始lr就是0.1，比策略B前几步的lr都大），但训练到12步左右开始，loss不再稳定下降，而是在某个值附近小幅震荡——这正是"步子太大，跨过了最优点又跨回来"的表现。策略B因为前4步lr较小，起步稍慢，但从第5步开始lr超过0.1、逐步走高又逐步回落，中段推进力度更足，同时后段lr降到0.01这种很小的值，最后几步loss还能继续往下挤一点点，最终收敛到的loss值比策略A更低、曲线也更平滑，没有明显震荡。

这也印证了前面"猫扑鱼"的比喻：不是学习率越大冲得越快就越好，关键是"什么时候该冲、什么时候该收爪子"。

PyTorch实现

手写这套调度逻辑不难，但PyTorch的torch.optim.lr_scheduler已经把常见策略封装好了，工程上直接调用即可。

先看最基础的用法骨架：

import torch
from torch.optim.lr_scheduler import LambdaLR
import math

optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)

T_warmup = 100
T_total = 2000

def lr_lambda(step):
    if step < T_warmup:
        # warmup阶段：线性爬升，返回的是相对于lr_max的比例
        return step / max(1, T_warmup)
    else:
        # cosine decay阶段
        progress = (step - T_warmup) / max(1, T_total - T_warmup)
        return 0.5 * (1.0 + math.cos(math.pi * progress))

scheduler = LambdaLR(optimizer, lr_lambda=lr_lambda)

for step in range(T_total):
    # ... 正常的前向、loss、backward ...
    optimizer.step()
    scheduler.step()      # 每个optimizer step之后，更新一次学习率
    optimizer.zero_grad()

这里有一个容易踩的点：LambdaLR的lr_lambda返回的不是学习率本身，而是相对于optimizer里设置的初始lr的一个乘法系数。所以上面代码里optimizer初始化时设的lr=3e-4，其实就是我们公式里的lr_max，lr_lambda返回1.0的时候，实际学习率就是3e-4；返回0.5的时候，实际学习率是1.5e-4。

如果不想自己手写这个函数，transformers库（HuggingFace）里有现成的：

from transformers import get_cosine_schedule_with_warmup

scheduler = get_cosine_schedule_with_warmup(
    optimizer,
    num_warmup_steps=100,
    num_training_steps=2000
)

效果和上面手写的lr_lambda是一样的，只是把warmup+cosine这套组合直接封装成了一个函数，这也是目前训练GPT类模型时最常用的调用方式之一。

小结

这一篇的核心就一句话：学习率不该是一个常数，而应该是一条先爬升、后衰减的曲线——Warmup让训练早期别冲太猛，Cosine Decay让训练后期能精细收敛。Adam解决的是"不同参数该走多大步子"，学习率调度解决的是"同一个参数在不同训练阶段该走多大步子"，两者是互补的，而不是互相替代的关系。

至此，我们的"猫吃鱼"两层小网络，已经集齐了梯度下降（EP3）、Adam优化器（EP4）、学习率调度（EP5）——一套相对完整的现代训练配置。训练篇到这里正式完结，共 5 篇：从损失函数、反向传播，到梯度下降、Adam 优化器与学习率调度，串起模型训练的基础流程。

← 返回训练篇目录