纯 SGD 无法训练 ANDREA
随机梯度下降,起点
反向传播为每个参数计算梯度 g。纯随机梯度下降 (SGD) 使用 p -= lr * g 更新每个参数。一个学习率,每步一个方向,没有过去梯度的记忆。
纯 SGD 在大规模时崩溃有两个原因:
1. 梯度在参数间的幅度差异极大。一个罕见标记的嵌入大多数步骤接收微小的梯度;一个层归一化缩放接收很大的梯度。一个学习率无法适合两者。
2. 梯度震荡。来自16源语料库的噪声小批量将参数推向左边,然后右边,然后左边。普通SGD浪费步骤自相矛盾。
Adam (Kingma & Ba, 2015) 通过每个参数的两个运行平均值修复了这两个问题。
一阶矩 & 二阶矩
m: 平滑方向
一阶矩 m 对最近的梯度进行指数平均:
m = beta1 m + (1 - beta1) g
使用 beta1 = 0.9。经过几步后,m 承载了一个平滑的方向;一个坏的批次几乎不会改变它。
v:平滑幅度
二阶矩 v 对最近的平方梯度进行平均:
v = beta2 v + (1 - beta2) g^2
使用 beta2 = 0.999。v 跟踪每个参数的梯度通常有多大。梯度大的参数得到大的 v;梯度很小的参数得到小的 v。
每个参数的自适应学习率
将平滑方向除以平滑幅度的平方根,将每个参数重新缩放到可比的水平:
adam_step = m / sqrt(v + eps)
微梯度嵌入被放大;大梯度层归一化被缩小。现在一个全局 lr 适合所有参数。
解读 Moments
为什么早期步骤需要偏差校正
冷启动偏差
m 和 v 从零开始。第1步后,m = 0.1 g_1 & v = 0.001 g_1^2。两个估计值都远远低于长期平均值。没有校正,优化器起步谨慎,缓慢加速,浪费宝贵的早期步骤,而这些步骤正是表示形成的关键时期。
校正方法
Adam 通过 1 / (1 - beta^t) 来缩放每个估计值,其中 t 是步骤编号:
m_hat = m / (1 - beta1^t)
v_hat = v / (1 - beta2^t)
在第 1 步时,beta1 = 0.9,除数 (1 - 0.9) = 0.1,因此 m_hat = m / 0.1 = 10 * m。偏差校正估计与长期平均值所预测的结果相匹配。随着 t 增大,beta^t 趋近于 0,校正因子趋近于 1,校正值与未校正值趋于收敛。
解耦权重衰减(AdamW 的创新)
L2 正则化 vs 权重衰减
经典 L2 正则化向损失添加惩罚项:L_total = L_data + (lambda / 2) sum(p^2)。反向传播将该惩罚视为梯度的一部分:g_total = g_data + lambda p。L2 项会流经 Adam 的 m 和 v 更新,被按参数幅度的平滑和重缩放。
Loshchilov & Hutter (2019) 证明了通过 Adam 平滑正则化器会破坏两者。Adam 的自适应缩放会缩小大梯度参数上的权重衰减(衰减应在此处最努力对抗过拟合)& 放大小梯度参数上的权重衰减。
AdamW:直接应用衰减
AdamW 将权重衰减与梯度解耦。在参数更新期间,衰减直接应用于每个参数,绝不触及 m 或 v:
p -= lr (m_hat / (sqrt(v_hat) + eps) + weight_decay p)
现在有两个项驱动每一步:
1. Adam 项: m_hat / (sqrt(v_hat) + eps) 根据每个参数的幅度历史重新缩放梯度方向。
2. 衰减项: weight_decay * p 将每个参数均匀地向零收缩,而不经过 Adam 的平滑处理。
ANDREA-120M v2 设置 weight_decay = 0.01。每一步,每个参数向零收缩 1%,除了 Adam 项所做的之外。
为什么解耦重要
实证证据
v1 崩溃(无权重衰减)
ANDREA-120M v1 使用 vanilla Adam 训练了 165K 步。采样输出:
- 步骤 80K: region region region region region region region
- 步骤 110K: ''''' ''''' '' ''' '' ''' '''?' ''' ' '' '' '
- 步骤 140K: games, games, games, games, games, games, games
- 步骤 165K: Budy Budy Budy Budy Budy Budy Budy Budy Budy
损失数字保持在合理范围内(步骤 110K 时 EMA 最小值为 3.23,与随机几率 9.04 相比)。仅凭损失无法揭示重复崩溃:一个永远记忆一个标记的模型在该标记出现的每一步都能实现低交叉熵。
v2 稳定性 (weight_decay = 0.01)
v2 添加了 AdamW(加上梯度裁剪、LR 预热、样本监控)。在第 ~112K 步,生成的样本:
- 卡罗莱纳鹦鹉于 1939 年被宣布灭绝(事实正确)
- 傅里叶变换将信号分解为频率分量(教科书定义)
- 雨的节奏反复,窗上的小溪,生活之痛的喘息(俳句约束满足)
外部评分将 v2 样本评为 9.5/10,称其为“在此规模下令人印象深刻的连贯性和知识保留”。
12M 模型在没有 AdamW 的情况下存活了。为什么?
ANDREA-12M 在 vanilla Adam 上训练,没有崩溃。在 12M 参数规模下,权重矩阵保持足够小,Adam 的自适应缩放无法将单个权重推入导致重复的失控幅度。在 120M 规模下,每步权重幅度漂移更远并累积;均匀衰减施加一个向零的恒定恢复力。随着模型规模扩大,解耦权重衰减变得更重要。
选择 weight_decay = 0.01
相关活动
AdamW 与本课程中的三个兄弟活动互锁:
- 活动 11:LR 预热 + 余弦衰减。 仅靠 AdamW 无法从新鲜初始化的权重上的瞬时峰值学习率中拯救模型。预热在 2000 步内逐步增加 lr,从而让 AdamW 的偏差校正和权重衰减有时间稳定表示。
- 活动 12:梯度裁剪。 AdamW 假设梯度具有有界幅度。ANDREA 的 bandit 中源每 7 到 42 步转换一次,会产生偶尔的梯度峰值;裁剪在 AdamW 触及 m、v 或 p 之前将它们限制在 L2 范数 1.0。
- 活动 13:FP32 / FP16 / FP8 精度。 AdamW 为每个参数存储 m 和 v,使内存占用是权重本身的双倍。FP16 将其内存占用减半;FP8 再次减半。精度选择与优化器稳定性相互作用。
AdamW、warmup、clipping 和 precision 构成了一个四叶草。丢掉一片叶子,看着 ANDREA 崩溃。