3.3 Learning Rate Scheduling

Learning rate 是 optimizer 中最像“温度”的超参数：太大则训练发散，太小则收敛缓慢或卡在糟糕区域。Scheduler 的作用不是装饰曲线，而是控制训练动力学的阶段转换。

Why Schedule Learning Rate

在二次近似

\[ L(\theta) \approx L(\theta^\star) +\frac{1}{2}(\theta-\theta^\star)^\top H(\theta-\theta^\star) \]

下，gradient descent 稳定需要

\[ 0<\eta<\frac{2}{\lambda_{\max}(H)}. \]

但深度网络的 effective curvature 会随训练变化。早期需要较大 step 探索，中后期需要较小 step 收敛，因此固定 learning rate 往往不是最优。

Theorem: Stability Range for Quadratic GD

For the quadratic objective \[ f(x)=\frac12x^\top Hx \] with symmetric positive definite \(H\), gradient descent \[ x_{t+1}=x_t-\eta Hx_t \] converges for all initial \(x_0\) if \[ 0<\eta<\frac{2}{\lambda_{\max}(H)}. \]

Proof

因为 \(H\) 对称正定，可以正交分解

\[ H=Q\Lambda Q^\top, \]

其中 \(\Lambda=\operatorname{diag}(\lambda_1,\ldots,\lambda_d)\)。令 \(z_t=Q^\top x_t\)，则

\[ z_{t+1} = Q^\top(I-\eta H)x_t = (I-\eta\Lambda)z_t. \]

每个特征方向独立更新：

\[ z_{t+1,i}=(1-\eta\lambda_i)z_{t,i}. \]

收敛需要所有方向满足

\[ |1-\eta\lambda_i|<1. \]

这等价于

\[ 0<\eta<\frac{2}{\lambda_i} \]

对所有 \(i\) 成立，因此最严格约束来自最大特征值：

\[ 0<\eta<\frac{2}{\lambda_{\max}(H)}. \]

这个定理给 scheduler 一个清楚的数学意义：learning rate 控制的是“每个曲率方向的收缩系数”。高曲率方向需要小步长，低曲率方向允许大步长。深度学习里 \(H\) 不固定也不总是正定，但这个局部二次图景仍然解释了为什么训练后期常要降 LR。

Descent Lemma View

Quadratic stability 给了一个线性系统视角；更一般的 smooth optimization 会用 Lipschitz gradient。若 \(f\) 是 \(L\)-smooth，即

\[ \|\nabla f(x)-\nabla f(y)\|\le L\|x-y\|, \]

则有 descent lemma：

\[ f(y) \le f(x)+\nabla f(x)^\top(y-x)+\frac{L}{2}\|y-x\|^2. \]

把 \(y=x-\eta\nabla f(x)\) 代入：

\[ f(x-\eta\nabla f(x)) \le f(x) - \eta\|\nabla f(x)\|^2 + \frac{L\eta^2}{2}\|\nabla f(x)\|^2. \]

因此

\[ f(x_{t+1}) \le f(x_t) - \eta\left(1-\frac{L\eta}{2}\right)\|\nabla f(x_t)\|^2. \]

只要 \(0<\eta<2/L\)，右侧就有下降项。这和二次情形的 \(2/\lambda_{\max}\) 一致，因为二次函数的 smoothness 常数就是最大特征值。

Proof Sketch

令 \(\phi(\tau)=f(x+\tau(y-x))\)。由微积分基本定理，

\[ f(y)-f(x) = \int_0^1 \nabla f(x+\tau(y-x))^\top(y-x)\,d\tau. \]

加减 \(\nabla f(x)\)：

\[ = \nabla f(x)^\top(y-x) + \int_0^1 [\nabla f(x+\tau(y-x))-\nabla f(x)]^\top(y-x) d\tau. \]

由 Cauchy-Schwarz 和 Lipschitz gradient，

\[ \left| [\nabla f(x+\tau(y-x))-\nabla f(x)]^\top(y-x) \right| \le L\tau\|y-x\|^2. \]

积分得到 \(\frac{L}{2}\|y-x\|^2\)。

Theorem: Constant LR Rate for Strongly Convex Smooth GD

If \(f\) is \(\mu\)-strongly convex and \(L\)-smooth, gradient descent with \(0<\eta\le 1/L\) satisfies \[ f(x_t)-f^\star \le (1-\eta\mu)^t\,[f(x_0)-f^\star]. \]

Proof Sketch

由 descent lemma，取 \(y=x-\eta\nabla f(x)\)：

\[ f(x^+) \le f(x) - \eta\left(1-\frac{L\eta}{2}\right)\|\nabla f(x)\|^2. \]

当 \(\eta\le 1/L\) 时，

\[ 1-\frac{L\eta}{2}\ge \frac12, \]

所以

\[ f(x^+) \le f(x)-\frac{\eta}{2}\|\nabla f(x)\|^2. \]

对 \(\mu\)-strongly convex function，有 Polyak-Lojasiewicz 型不等式

\[ \|\nabla f(x)\|^2 \ge 2\mu(f(x)-f^\star). \]

代入得到

\[ f(x^+)-f^\star \le (1-\eta\mu)(f(x)-f^\star). \]

递推 \(t\) 次即得结论。

这个式子解释了两个 scheduler 现象：

早期如果 curvature estimate 不稳定，直接用接近 \(1/L\) 的 LR 可能炸；
后期若想继续减少 suboptimality，降低 LR 可以减少噪声震荡，但也会让线性收敛因子接近 \(1\)。

在强凸二次函数里，最佳固定 LR 还可以写得更精确。若特征值在 \([\mu,L]\)，最坏方向收缩率为

\[ \rho(\eta)=\max_{\lambda\in[\mu,L]}|1-\eta\lambda|. \]

令两端误差相等：

\[ 1-\eta\mu=-(1-\eta L), \]

得到

\[ \eta^\star=\frac{2}{L+\mu}, \qquad \rho^\star=\frac{L-\mu}{L+\mu} = \frac{\kappa-1}{\kappa+1}, \]

其中 \(\kappa=L/\mu\) 是 condition number。深度网络当然不是全局强凸，但这个推导非常有用：LR schedule 本质是在不断追踪一个变化中的 effective curvature 和 noise scale。

LR as Noise Temperature

SGD 的 mini-batch gradient 可以写成

\[ g_B(\theta)=\nabla L(\theta)+\xi_B, \]

其中 \(\xi_B\) 是由采样 batch 引入的噪声。若 batch size 是 \(B\)，在粗略独立假设下，

\[ \operatorname{Var}(\xi_B)\propto\frac1B. \]

参数更新为

\[ \theta_{t+1} = \theta_t-\eta\nabla L(\theta_t)-\eta\xi_B. \]

所以随机噪声项的尺度近似随

\[ \frac{\eta}{\sqrt{B}} \]

变化，而扩散强度常用

\[ \frac{\eta}{B} \]

作为粗略 proxy。于是 batch size、learning rate、warmup、weight decay 不能孤立调：增大 batch 会降低噪声，增大 LR 会提高更新尺度和噪声注入。

这也是为什么大 batch 训练不是“吞吐更高所以一定更好”。它改变了优化过程的随机性，scheduler 需要一起改变。

更细一点，可以看 update noise 的协方差。设单样本梯度为 \(g_i\)，总体梯度为 \(g=\mathbb{E}[g_i]\)，单样本梯度噪声协方差为 \(\Sigma\)。batch 平均梯度

\[ g_B=\frac{1}{B}\sum_{i=1}^{B}g_i \]

满足

\[ \mathbb{E}[g_B]=g, \qquad \mathrm{Cov}(g_B)=\frac{1}{B}\Sigma. \]

SGD update 是

\[ \Delta\theta=-\eta g_B =-\eta g-\eta(g_B-g). \]

所以 noise covariance 约为

\[ \mathrm{Cov}(\Delta\theta_{\text{noise}}) = \frac{\eta^2}{B}\Sigma. \]

这解释了为什么有时讨论 \(\eta/\sqrt{B}\)，有时讨论 \(\eta/B\)：前者对应 update noise 的标准差尺度，后者常出现在连续时间 SDE 近似里的温度 proxy。它们都在说同一件事：LR 和 batch size 共同决定随机优化的温度。

Definition: Critical Batch Size

The critical batch size is the rough scale beyond which increasing batch size gives diminishing optimization-speed returns because gradient noise is already small relative to useful descent signal.

在临界 batch size 以下，增大 batch 往往能提高吞吐并允许更大 LR；超过以后，更多 batch 主要减少噪声，却不一定减少达到同等 loss 所需的 token 数。对 LLM 预训练来说，更常用的横轴不是 epoch，而是 consumed tokens：

\[ \text{tokens} = \text{optimizer steps}\times B_{\text{global}}\times T_{\text{seq}}. \]

因此 scheduler 最好能用 optimizer step 和 token budget 两种方式互相换算。若改变 global batch 而不改 total tokens，total_steps 也要随之改变：

\[ S_{\text{total}} = \frac{N_{\text{tokens,total}}} {B_{\text{global}}T_{\text{seq}}}. \]

这就是为什么同一条 cosine curve 在不同 batch size 下不能简单复用 step 数。

Warmup

Definition: Learning Rate Warmup

Warmup increases the learning rate from a small value to the target value during the first \(T_w\) steps: \[ \eta_t=\eta_{\max}\frac{t}{T_w}, \qquad 0\leq t\leq T_w. \]

Transformer 训练通常需要 warmup，因为初始参数、LayerNorm statistics、Adam moment estimates 都还不稳定。直接用大 LR 可能让 attention logits、residual stream 或 embedding 层出现过大更新。

Warmup and Adam Bias Correction

Adam 的一阶、二阶动量为

\[ m_t=\beta_1m_{t-1}+(1-\beta_1)g_t, \qquad v_t=\beta_2v_{t-1}+(1-\beta_2)g_t^2. \]

早期 \(m_t,v_t\) 还在从零启动。虽然 bias correction 会除以 \(1-\beta^t\)，但网络激活、loss scale、gradient statistics 本身也在剧烈变化。Warmup 的作用是让有效更新

\[ \Delta\theta_t = -\eta_t\frac{\hat{m}_t}{\sqrt{\hat{v}_t}+\epsilon} \]

逐步放大，而不是在模型还没稳定时直接使用 peak LR。

可以把 Adam 的 per-coordinate effective learning rate 写成

\[ \eta_{t,j}^{\text{eff}} = \eta_t\frac{1}{\sqrt{\hat{v}_{t,j}}+\epsilon}. \]

训练早期 \(\hat{v}_{t,j}\) 对梯度尺度估计很粗，某些 coordinate 的 denominator 可能过小，导致局部 effective LR 极大。Warmup 不只是“让 LR 曲线好看”，而是在 Adam 的预条件器还没稳定时限制更新半径。

一个有用的诊断量是 update ratio：

\[ r_t = \frac{\|\Delta\theta_t\|_2}{\|\theta_t\|_2+\epsilon}. \]

如果 warmup 期间 \(r_t\) 突然尖峰，说明表面 LR 可能正常，但 actual update 已经异常。大模型训练日志里，除了 lr，最好同时记录：

metric	meaning
`grad_norm`	raw gradient scale
`update_norm / param_norm`	actual step size relative to weights
`adam_v_mean` or percentile	second-moment scale
`loss_scale`	AMP dynamic loss scale
`skipped_steps`	overflow-induced skipped updates

Warmup length 也要和模型/数据规模相配。太短会在 embedding、attention logits、LayerNorm/RMSNorm 附近产生早期冲击；太长则浪费高 LR 的探索阶段。经验上 LLM pretraining 常按 total steps 的小比例或固定 token 数 warmup；SFT/LoRA 微调则经常用更短 warmup，因为 base model 已经稳定。

Pitfall: Warmup Is Counted in Optimizer Steps

With gradient accumulation, the scheduler should usually step once per optimizer update, not once per micro-batch. Otherwise warmup ends too early.

Step Decay and Exponential Decay

Step decay:

\[ \eta_t = \eta_0\gamma^{\lfloor t/T\rfloor}. \]

Exponential decay:

\[ \eta_t = \eta_0\gamma^{t}. \]

这类 scheduler 的风格是“训练到某些阶段后突然降温”。它在传统 CNN 训练中很常见，但在大模型预训练里不如 cosine decay 平滑。

step decay 的问题是 LR 不连续。对 momentum/Adam 来说，优化器 state 仍然携带过去大 LR 阶段的速度和二阶统计，但外层 \(\eta_t\) 突然缩小：

\[ \Delta\theta_t = -\eta_t u_t. \]

若 \(\eta_t\) 在某一步乘上 \(\gamma\)，update norm 也会突然乘上 \(\gamma\)，而 \(u_t\) 的方向和统计没有同步“冷却”。这在传统训练中常常可接受，甚至能作为手动 phase transition；但在长程预训练里，平滑 schedule 更容易让 loss 曲线和 optimizer state 连续。

Cosine Decay

Definition: Cosine Learning Rate Decay

After warmup, cosine decay sets \[ \eta_t = \eta_{\min} +\frac{1}{2}(\eta_{\max}-\eta_{\min}) \left(1+\cos\frac{\pi(t-T_w)}{T-T_w}\right). \]

cosine decay 的优点是平滑、无突变、结尾自然接近小 LR。LLM pretraining 中常见配置是 linear warmup + cosine decay。

cosine schedule 的平滑性可以从导数看出。令

\[ u=\frac{t-T_w}{T-T_w}. \]

则

\[ \eta(u) = \eta_{\min} + \frac12(\eta_{\max}-\eta_{\min})(1+\cos\pi u). \]

导数为

\[ \frac{d\eta}{du} = -\frac{\pi}{2}(\eta_{\max}-\eta_{\min})\sin\pi u. \]

在 \(u=0\) 和 \(u=1\) 处导数都为 \(0\)，所以从 peak LR 进入 decay 和最终接近 floor 都没有尖锐拐点。这里的 floor \(\eta_{\min}\) 不是随便加的：若 floor 太低，训练末期几乎不再学习；若 floor 太高，loss 会在噪声带附近震荡，尤其是小数据 SFT 或 preference optimization。

常见配置可以写成：

schedule	best for	caveat
step decay	short supervised training, manual phases	discontinuous update scale
linear decay	simple token-budget allocation	end phase may be too abrupt
cosine decay	long pretraining	needs correct total_steps
inverse-sqrt	long tail training	early/late scale harder to tune
constant with warmup	continued training / some fine-tuning	may not cool enough near end

Polynomial and Inverse-Sqrt Schedules

Polynomial decay:

\[ \eta_t = (\eta_{\max}-\eta_{\min}) \left(1-\frac{t-T_w}{T-T_w}\right)^p +\eta_{\min}. \]

当 \(p=1\) 时是 linear decay；\(p=2\) 时后期降得更快。Inverse-square-root schedule 常见于早期 Transformer：

\[ \eta_t = d_{\text{model}}^{-1/2} \min(t^{-1/2},tT_w^{-3/2}). \]

它先线性 warmup，再按 \(t^{-1/2}\) 下降。这个 schedule 的动机是训练越久，更新幅度越保守，同时保留比指数衰减更长的尾巴。

One-Cycle Policy

One-cycle policy 先升 LR 再降 LR，同时常配合 momentum 反向变化。直觉是：

早期用较大 LR 找到宽 basin；
中期保持足够噪声跳出 sharp region；
后期快速降 LR 做 fine convergence。

实践中可以用 LR range test 粗略找 peak LR：从很小 LR 开始指数增加，每隔若干 step 记录 loss。当 loss 开始明显发散前的 LR 往往是上界，peak LR 取其若干分之一。概念代码：

for step, batch in enumerate(loader):
    lr = min_lr * (max_lr / min_lr) ** (step / total_probe_steps)
    set_lr(optimizer, lr)
    loss = train_one_step(batch)
    log(step=step, lr=lr, loss=loss)

这个 test 只能给数量级，不能替代正式训练。原因是 probe 期间 optimizer state、data order、warmup 和 loss scale 都与真实 run 不完全一致。

Parameter Groups and Layerwise LR

真实优化器经常有多个 parameter groups：

optimizer = AdamW(
    [
        {"params": decay_params, "lr": base_lr, "weight_decay": 0.1},
        {"params": norm_bias_params, "lr": base_lr, "weight_decay": 0.0},
        {"params": adapter_params, "lr": adapter_lr, "weight_decay": 0.0},
    ],
)

scheduler 需要明确是：

给所有 group 乘同一个 scalar；
还是每个 group 有独立 peak/floor；
或者冻结某些 group 后期再解冻。

对 fine-tuning 来说，常见做法是 base model 用小 LR，adapter/head 用较大 LR。若所有参数共用一个 LR，可能出现两种失败：

failure	cause
base model forgetting	LR 对 pretrained weights 太大
adapter underfitting	LR 为了保护 base 而太小

Layerwise LR decay 进一步让靠近输入层的参数 LR 更小：

\[ \eta_\ell=\eta_{\text{top}}\alpha^{L-\ell}, \qquad 0<\alpha<1. \]

它的直觉是底层表示更通用，顶层更 task-specific。但这也增加了配置复杂度：日志里只看一个 lr 不够，至少要记录每个 parameter group 的 LR。

Batch Size and Learning Rate Scaling

当 batch size 从 \(B\) 增加到 \(kB\)，梯度噪声大约降低。经验上有 linear scaling rule：

\[ \eta(kB)\approx k\eta(B). \]

但这个规则只在一定范围内成立。过大 batch 会降低 gradient noise，可能导致泛化变差或训练卡在 sharp minima。

Pitfall: Larger Batch Is Not Free Speed

Increasing batch size reduces steps per epoch but also changes the stochastic dynamics. If the learning rate, warmup, and regularization are not adjusted, large-batch training can become unstable or generalize worse.

Scheduler State in Real Training

真实训练里 scheduler 的 step 计数必须和 optimizer update 对齐。设：

\[ B_{\text{global}} = B_{\text{micro}}\times K_{\text{accum}}\times N_{\text{workers}}. \]

若数据集 token 数为 \(D_{\text{epoch}}\)，每个 optimizer step 消耗约 \(B_{\text{global}}\) 个样本或 token，则每个 epoch 的 optimizer steps 是

\[ S_{\text{epoch}} = \left\lceil \frac{D_{\text{epoch}}}{B_{\text{global}}} \right\rceil. \]

Scheduler 的 total_steps 应该基于 optimizer steps，而不是 dataloader iterations 或 raw micro-batches。

for micro_step, batch in enumerate(loader):
    loss = model(batch).loss / grad_accum
    loss.backward()

    if (micro_step + 1) % grad_accum == 0:
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad(set_to_none=True)

Resume 时还要恢复：

optimizer state；
scheduler state；
global optimizer step；
dataloader/sampler epoch and seed。

只恢复模型权重而不恢复 scheduler，训练曲线会出现“学习率时间穿越”。

Schedulers for LLM Training

Stage	LR behavior	Reason
Warmup	linear rise	stabilize Adam moments and activations
Main pretraining	cosine decay	smooth compute allocation
Continued pretraining	lower peak LR	avoid overwriting pretrained knowledge
SFT	much lower LR	align behavior without damaging base model
RL/DPO	often tiny LR	reward/preference gradients are noisy

对于大模型，scheduler 也是“不要把已有能力冲坏”的保护机制。预训练可以粗暴一些，post-training 必须保守很多。

Minimal Scheduler Pattern

import math


def warmup_cosine_lr(step: int, total: int, warmup: int, peak: float, floor: float):
    if step < warmup:
        return peak * step / max(1, warmup)
    progress = (step - warmup) / max(1, total - warmup)
    coeff = 0.5 * (1.0 + math.cos(math.pi * progress))
    return floor + (peak - floor) * coeff

这个函数是很多训练系统的最小骨架。真正工程里还需要 resume step、gradient accumulation 后的 global step、以及 distributed workers 间的一致计数。

更真实的实现会把 step 作为 checkpoint state 的一部分：

class WarmupCosine:
    def __init__(self, optimizer, total, warmup, peak, floor):
        self.optimizer = optimizer
        self.total = total
        self.warmup = warmup
        self.peak = peak
        self.floor = floor
        self.step_id = 0

    def lr_at(self, step):
        if step < self.warmup:
            return self.peak * step / max(1, self.warmup)
        progress = (step - self.warmup) / max(1, self.total - self.warmup)
        progress = min(1.0, max(0.0, progress))
        coeff = 0.5 * (1.0 + math.cos(math.pi * progress))
        return self.floor + (self.peak - self.floor) * coeff

    def step(self):
        lr = self.lr_at(self.step_id)
        for group in self.optimizer.param_groups:
            scale = group.get("lr_scale", 1.0)
            group["lr"] = lr * scale
        self.step_id += 1
        return lr

    def state_dict(self):
        return {"step_id": self.step_id}

    def load_state_dict(self, state):
        self.step_id = state["step_id"]

训练循环里要把 scheduler step 绑定到 successful optimizer update：

for batch in loader:
    loss = forward_backward(batch)
    if ready_to_update:
        did_step = scaler_step_if_finite(optimizer, scaler)
        if did_step:
            scheduler.step()
            optimizer.zero_grad(set_to_none=True)

这里的 did_step 很重要：AMP overflow 时如果没有真正更新参数，却推进了 scheduler，LR 时间轴就会快于参数时间轴。长训练里这种偏差可能不明显，但 resume/debug 时会非常恼人。

在分布式训练中，所有 rank 必须看到同一个 scheduler state。通常只要每个 rank 的 optimizer step 次数一致即可；如果存在 uneven batches、elastic training、gradient overflow 只在部分 rank 触发，就要显式同步 did_step 或 global step。

Pitfall: Scheduler State Is Training State

Checkpointing only model weights is not enough. Optimizer state, scheduler step, scaler state, sampler position, and global token count together define where training actually is.

Practical Checks

训练开始前最好直接打印或画出 LR 曲线：

lrs = [
    warmup_cosine_lr(s, total=100_000, warmup=2_000, peak=3e-4, floor=3e-5)
    for s in range(100_000)
]

要检查：

step 0 的 LR 是否为预期值；
warmup 结束 step 是否正确；
peak LR 是否达到；
final LR 是否等于 floor；
resume 后 LR 是否连续；
多 worker 日志中 LR 是否一致；
scheduler 是否在 skipped optimizer step 后仍误更新。

Pitfall: Scheduler Step on Skipped AMP Updates

If AMP detects overflow and skips optimizer.step(), stepping the scheduler anyway changes LR without changing parameters. Some training stacks guard scheduler updates on successful optimizer steps.

Implementation Checklist

配置 LR schedule 前后，可以逐项检查：

total_steps 是 optimizer updates，不是 micro-batches；
warmup_steps 和 gradient accumulation 后的 update 计数一致；
改 global batch 后，total_steps 是否按 token budget 重新计算；
peak_lr、floor_lr 是否对所有 parameter groups 都有明确规则；
weight decay、norm/bias、adapter/head 是否使用合适的 group；
AMP skipped update 是否阻止 scheduler 前进；
resume 后第一步 LR 是否等于中断前下一步 LR；
多 rank 的 global step 和 LR 是否完全一致；
日志是否记录 lr、grad_norm、update_ratio、loss_scale；
LR 曲线是否在训练前画出来并人工检查端点；
SFT/RL/DPO 是否使用比 pretraining 更保守的 peak LR；
schedule 的单位是 steps、epochs 还是 tokens，配置文件里是否写清楚。

一个很小的 smoke test：

sched = WarmupCosine(optimizer, total=10, warmup=2, peak=1e-3, floor=1e-5)
lrs = [sched.step() for _ in range(10)]
assert lrs[0] == 0.0
assert max(lrs) <= 1e-3
assert abs(sched.lr_at(10) - 1e-5) < 1e-6
state = sched.state_dict()

更严格一点，还要测试 resume continuity：保存 state 后重新构造 scheduler，加载 state，确认下一步 LR 与原 scheduler 下一步完全一致。