训练

优化器

SGD、动量、Adam、AdamW 与学习率调度

图中数值
玩具曲面 2 个权重
AdamW
β₁ 0.9, β₂ 0.95 λ 0.1 · LLaMA 2
调度
预热 + 余弦

全部步骤

  1. 逆着梯度走一步

    一个关于两个权重的玩具损失，山谷横向比纵向陡 100 倍。朴素梯度下降逆着梯度走。一个在横向上安全的步长，在纵向上就小得可怜，所以它来回折返、缓慢爬行：20 步之后损失为 0.39。

    w ← w − η · g

  2. 动量

    动量保留过去梯度的累计和。横跨山谷的折返互相抵消，沿山谷的稳定推力不断累加，所以每一步走得更远。

    v ← β v + g · w ← w − η v

  3. Adam：每个权重各有步长

    Adam 把每个权重的步长除以它自身梯度的滑动大小。于是陡峭和平缓的方向不论尺度如何，都以由 η 决定的相近速度移动。它是 Transformer 的默认选择。

    w ← w − η · m̂ / (√v̂ + ε)

  4. AdamW：单独的权重衰减

    权重衰减每一步都把每个权重往零拉一点，防止它们无谓地变大。AdamW 直接把它作用在权重上，而不是通过梯度，所以 Adam 的缩放不会削弱它。

    w ← w − η (m̂/(√v̂+ε) + λ w)

  5. 先预热，再衰减

    步长在训练中不断变化。LLaMA 2 先线性预热 2,000 步（Adam 的滑动平均一开始噪声很大），然后在大约五十万步里沿余弦曲线降到峰值 3 × 10⁻⁴ 的 10%。

    预热 2,000 · 余弦降到 10%

  6. 优化器保存了什么

    Adam 为每个权重保存两个数 m 和 v，通常是 32 位，此外还有一份 32 位的权重主副本。训练一个模型所需的内存是运行它的好几倍。

    每个权重约 16 字节，推理只需 2

代码

opt = torch.optim.AdamW(model.parameters(), lr=3e-4, betas=(0.9, 0.95), weight_decay=0.1)
sched = get_cosine_schedule_with_warmup(opt, num_warmup_steps=2000, num_training_steps=500_000)
for batch in data:
    loss = model(batch).loss; loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step(); sched.step(); opt.zero_grad()
# inside opt.step(), per weight: m = β1·m + (1−β1)·g;  v = β2·v + (1−β2)·g²
#   w -= lr · (m̂ / (√v̂ + ε) + weight_decay · w)

延伸阅读