课程/M11/反向传播中的矩阵与训练稳定性
M11 · 大模型专题核心约 40 分钟

11.12 反向传播中的矩阵与训练稳定性

本节位置:模块 11 · 大模型专题 | 前置:8.5 | 预计:40 分钟

一句话:反向传播是「Jacobian 的连乘」;深层网络稳不稳定,取决于这些矩阵的奇异值是否把梯度放大或压扁。

0. 场景引入

一个 2.1B 参数的语言模型(32 层、),在 2048 张加速卡上预训练,checkpoint 每 5,000 步存一次(约 44 分钟)。

第 10,000 步一切正常:loss 2.10,梯度范数 0.5。第 10,040 步,loss 跳到 8.7;第 10,055 步,日志里出现 nan。翻记录发现梯度范数从 0.5 涨到 300,attention 的 logits 变得极大、softmax 几乎 one-hot。

这不是玄学。反向传播要穿过几十层矩阵乘法,每一层 Jacobian 都在放大或缩小梯度;只要谱范数长期大于 1,梯度就指数增长。

值班的你只有三个开关:梯度裁剪阈值 、混合精度格式(fp16 + loss scaling,还是 bf16)、以及要不要整批跳过。三个选择的价码不一样:

  • 从最近的 checkpoint 回滚:2048 卡 × 44 分钟 × 5 元/卡时 ≈ 7,500 元,外加学习率调度错位;
  • 若 nan 已经写进 checkpoint 并继续被覆盖,只能退回三天前: 万元,进度全丢;
  • 若把 从 1.0 硬压到 0.3:梯度范数的中位数是 0.6,一半的步都会被裁,等于偷偷把学习率减半——不炸了,但收敛慢一大截。

真正要判断的是:这次爆炸是某一层 Jacobian 连乘放大出来的,还是某个 batch 的数值溢出? 前者得动结构(残差缩放、归一化、初始化),后者得动数值(裁剪、fp32 归约、bf16 的动态范围)。判错方向的代价,就是上面那张价目表。

读完本节,你能回答的三个问题是:

  1. 梯度穿过一层 到底被放大几倍?怎样用一个数(谱范数)估出 32 层后的总倍数?
  2. 残差、归一化、初始化各自把放大倍数拉回 1 的哪一段?梯度裁剪能救什么、救不了什么?
  3. bf16 与 fp16 的动态范围差多少?梯度什么时候下溢成 0、什么时候溢出成 inf,为什么 bf16 通常不用 loss scaling?

1. 直觉

把前向传播想成一条流水线,反向传播就是把「误差信号」沿原路送回去:

  • 每经过一个线性层 ,梯度要乘 ;
  • 每经过 softmax,梯度要乘它的 Jacobian;
  • 每经过归一化,梯度要投影掉某些方向。

如果每层都把信号放大 1.1 倍,100 层后就是 倍——爆炸。 如果每层缩小到 0.9,100 层后只剩 ——消失。

💡 一句话直觉:反向传播 = 梯度乘一串 Jacobian;稳定性 = 这串矩阵的奇异值不要长期偏离 1。

1.1 它在现实世界里对应什么

工程上,反向传播像一条回传的传送带,每一站都装着一组齿轮:前向时数据往前走,反向时误差信号往回走,每一站的放大倍数就是那一层 Jacobian 的谱范数。

量级:32 层,若每层平均放大 ,32 层后是 倍,还算温和;把 只抬到 ,就是 倍;若是 64 层配 ,则是 倍。反方向同样致命:、32 层,梯度只剩 ,也就是 3.4%。每层 10% 的偏差,64 层后就是 倍的差别——层数会把小偏差指数地兑现。

这个类比在哪里失效:齿轮比是刻死的常数,而 Jacobian 随输入变——同一个 在激活变大时谱范数不变,可 softmax 的 Jacobian 会从 缩到 ;而且梯度不是被「一个」矩阵放大,是被一串矩阵之积放大, 只是上界,实际倍数取决于这些矩阵的奇异向量是否互相对齐,可能比上界小好几个数量级。

2. 定义

2.1 梯度形状

对线性层 、损失 :

其中 ,形状与 相同。

批量和序列下,,:

2.2 链式法则(向量形式)

若 ,,则

是 的 Jacobian。深度网络就是这些 Jacobian 的乘积。

2.3 softmax 的 Jacobian

它把梯度中「沿概率向量方向」的分量删掉(因为 softmax 对整体平移不变)。

2.4 谱范数与 Lipschitz 常数

矩阵 的谱范数:

它等于最大拉伸倍数。对复合:

2.5 梯度裁剪

只改长度,不改方向。

3. 为什么这样定义

为什么要看谱范数而不是元素大小? 元素大不代表放大能力强;一个矩阵可能有很多大元素但相互抵消,也可能元素不大却沿某方向强拉伸。谱范数精确刻画「最大放大倍数」。

为什么残差有用? 残差的 Jacobian 是 ,即使 很小,梯度也能沿 传回去;不会因为连乘而指数衰减。

为什么归一化有用? LayerNorm/RMSNorm 把激活尺度拉回受控范围,使后续线性层的输入不会爆炸;同时它的 Jacobian 会抑制径向分量。

3.1 反事实:把残差缩放和裁剪拿掉,具体坏在哪

反例一:残差缩放系数写错一个数。

Pre-LN 的残差分支通常乘一个系数 ,GPT-2 式取法是 ; 时 。设残差分支 的 Jacobian 谱范数 ,那么每层在最坏方向上的增益约是 :

  • (正确):每层 ,32 层后 ;
  • (忘了缩放):每层 ,32 层后 。

同一个网络,只差一个常数,总放大倍数差 倍。

更隐蔽的是它的死法: 本身还没超出 bf16 的上界 ,所以日志里看不到 inf;但 Adam 存的是 , 已经溢出, 变成 inf,更新量 变成 。训练没有崩,它安静地停了——loss 从此一条直线,没有任何报错。

反例二:不做梯度裁剪,一个 spike 会在二阶矩里留下 1000 步的疤。

正常步的梯度范数是 0.5,尖峰步是 300,差 600 倍,平方差 倍。Adam 的二阶矩是滑动平均 ,取 :

  • 不裁剪:尖峰前 ;尖峰后 ,,而正常时 。由于更新量是 ,接下来约 步的更新幅度都被压小约 19 倍;
  • 裁剪到 再更新 :,——疤几乎不存在。

这就是「裁剪要在 optimizer.step() 之前、作用在 更新之前」的原因:对 Adam,裁剪的主要作用不是限制步长,而是防止一次尖峰污染二阶矩。

反例三:只裁剪、不跳过,坏数据连续来时照样把权重推飞。

假设坏 batch 连续 40 步,每步都被裁到 、方向几乎相同。Adam 每元素的步长 ,40 步累计 ;而这个模型的权重元素尺度只有 ——77% 的权重被同一个坏方向改写。裁剪保证「一步不迈大」,保证不了「四十步不朝一个方向走」。所以工程上要的是 if not isfinite(g_norm) or g_norm > 3.5: skip,而不是把 调到很小硬扛。

4. 性质与推导

4.1 线性层连乘

若网络是 层线性层(无非线性),则

梯度大小受 控制。

4.2 残差的梯度

对 :

特征值约为 。只要 不把特征值推到 或很大的正数,梯度就能稳定传播。

4.3 softmax 饱和

当 logits 很大时, 接近 one-hot, 接近零矩阵,梯度消失。这是「attention 太自信 → 学不动」的原因之一。

4.4 梯度裁剪的几何

裁剪把梯度投影到半径 的球内(若超出)。它保持方向,只缩短步长;因此不会改变「往哪走」,只防止「一步迈太大」。

4.5 权重衰减与谱

L2 正则倾向于减小 ,间接减小谱范数;但 Frobenius 范数与谱范数不等价(第 10.2 节)。

4.6 数值下溢、溢出与 bf16 的动态范围

浮点数 = 符号位 + 指数位 + 尾数位。指数位决定「能表示多大/多小」,尾数位决定「一位有效数字有多细」。三兄弟的差别全在这两栏:

格式指数位尾数位(显式)最大正规数最小正规数相对精度
fp32823
fp16510
bf1687

bf16 说穿了就是「fp32 的指数 + 截短的尾数」:动态范围和 fp32 一样宽(约 ),但只有 8 位有效数字。

下溢(梯度悄悄变 0)。 梯度里从 30 层深处传回来的分量可能只有 量级。fp16 的最小正规数是 , 直接掉进次正规区甚至被 flush 成 0——这个参数的更新从此消失,而且没有任何报错。解决办法是 loss scaling:先把 loss 乘 再反传,梯度整体抬到 ,更新前再除回去。bf16 不需要这一步: 在 bf16 里正常表示。

溢出(梯度直接变 inf)。 反向的另一端: 到 就超出 bf16 的上界。注意 存的是 , 已经溢出, 变 inf,Adam 的 变 0——这就是 §3.1 反例一里那种「不崩但停」的死法。同一件事在 fp16 上要早得多:梯度范数到 ,它的平方就刚好越过 fp16 的上界,而 bf16 要撑到 才溢出——在梯度范数上差了约 17 个数量级。

归约精度(bf16 的 被长求和放大)。 对 个随机符号的项求和,绝对误差随 增长,相对误差大致是 ——这是量级估算,意思是「可能到 25% 左右」。所以工程上的默认组合是:bf16 存数据,fp32 做累加;跨卡 all-reduce 也常用 fp32 桶,而不是直接拿 bf16 求和。

和 softmax 的联动。 第 11.11 节说 softmax 要先减 max,原因在这里更清楚:fp16 的指数上界只到 ,任何满足 的候选都会被 flush 成 0——fp16 下的 softmax 天然只有约 11 个 nat 的动态范围,长尾在采样之前就已经被砍掉了一次。这也是很多实现在 fp32 里算 softmax、再把结果转回低精度的原因。

5. 例子

5.1 手算:两层标量网络

设 ,,。

前向:。 梯度:

若 ,则 ,,。 若 10 层、每层权重 2,梯度会含 的因子——爆炸。

5.2 手算:梯度裁剪

梯度 ,,阈值 :

方向不变,长度 1。

5.3 场景:loss spike 的常见链条

  1. 某个 batch 的 attention logits 变大;
  2. softmax 变尖锐,熵下降,梯度变小但输出极端;
  3. 后续层收到极端输入,产生大梯度;
  4. 参数被推到一个坏区域,下一步 loss 飙升;
  5. 若没有裁剪/预热,可能直接 NaN。

对应手段:QK-norm、z-loss、梯度裁剪、学习率预热、降低 logits 尺度。

5.4 场景闭环:2.1B 参数、一次裁剪

输入:一个训练 step,全局 batch 是 2048 条序列 × 4096 token 个 token。模型 32 层、,参数量 。

计算:

  • FLOPs:训练约 FLOP/step;2048 张卡各按有效 TFLOP/s(合计 FLOP/s,约 MFU 一成多)算,单步 s。
  • 显存:每个参数要存 bf16 梯度(2 B)+ fp32 master 权重(4 B)+ Adam 的 (8 B) B,共 GB,再加 bf16 权重 4.3 GB,约 34 GB,80 GB 的卡放得下。注意其中 26 GB 是优化器状态(fp32 master + ),参数本身只占 4.3 GB——显存的大头从来不是权重。
  • 梯度范数:把 个数平方求和,读一遍 bf16 梯度 4.3 GB,2 TB/s 下约 2.1 ms;但每张卡要收发 GB 做 all-reduce,50 GB/s 网卡下约 0.17 s。算范数本身几乎免费,贵的是把它归约到全局。

输出:一个标量 与裁剪系数 ,随后走两个分支:(历史 p99)照常 optimizer.step(); 或非有限,则跳过这一步——不更新参数、不写 checkpoint、把 batch id 记进日志。

成本:单步 卡 卡时 元。跳一步只值 1.5 元和 0.53 秒,而从最近的 checkpoint 回滚一次要 7,500 元——差 5,000 倍,所以正确的做法是激进地 skip,而不是把阈值往下调硬扛。整轮训练 步、约 17.5 小时、 卡时 万元。

失败条件与兜底:

  • 只裁不跳:坏 batch 连来 40 步,同一方向累计改写 77% 权重(§3.1 反例三)。兜底:isfinite 检查 + 超 p99 直接 skip。
  • 阈值定得太低: 而梯度范数中位数 0.6,一半的步被裁,等效学习率减半,收敛变慢且难以归因。兜底:先跑 500 步只记录不裁剪,用实际 p50/p99 定阈值。
  • nan 让裁剪失效:min(1, c/NaN) 还是 NaN,NaN 与任何数比大小都为 false,裁剪「成功执行」却什么也没拦住。兜底:先判 isfinite 再判范数。
  • bf16 归约误差:8.6 GB 梯度用 bf16 直接求和,抵消会把相对误差放大到 量级。兜底:fp32 桶归约。
  • checkpoint 写坏:nan 权重被存盘后就再也回不去。兜底:存盘前校验 loss 与范数有限,存完异步上传并保留 3 个历史版本。

6. 常见误区

  1. 误区:梯度裁剪会改变优化方向。 纠正:它只按比例缩放整个梯度,方向不变(除非逐元素裁剪)。

  2. 误区:有了 LayerNorm 就不会梯度爆炸。 纠正:归一化只控制激活尺度;权重矩阵的谱范数、残差累积、softmax 饱和仍可能导致不稳定。

  3. 误区:loss spike 一定是数据问题。 纠正:可能是优化/数值问题(logits 爆炸、学习率过大、bf16 溢出),需要看梯度范数与激活统计。

  4. 误区:谱范数就是最大元素。 纠正:谱范数是最大奇异值,衡量算子放大能力;最大元素只是元素级指标。

7. 概念网络

8. 通往大模型:稳定性工具箱

手段作用
梯度裁剪限制单步更新幅度
学习率预热训练初期避免大更新
QK-norm控制 attention logits 尺度
z-loss惩罚过大 logits,稳定 softmax
残差 + Pre-Norm保梯度高速公路
权重衰减控制权重范数
混合精度 + loss scaling防 fp16 下溢

可检验结论:若每层 Jacobian 的谱范数平均为 , 层后梯度放大约 。要让梯度稳定,需要 ;初始化(如 He/Xavier)和归一化就是在逼近这个目标。

# 观察梯度范数随层数的变化
import numpy as np
np.random.seed(0)
x = np.random.randn(64)
for scale in [0.9, 1.0, 1.1]:
    h = x.copy()
    for _ in range(50):
        W = np.random.randn(64, 64) * scale / np.sqrt(64)
        h = W @ h
    print(f"scale={scale}: ||h||={np.linalg.norm(h):.3e}")

8.1 算错时用户看到什么

这里没有用户直接面对梯度。你在屏幕上看的是 loss 曲线;真正落到用户身上的,是这条曲线背后已经被改坏的那几个参数。

  • 训练工程师看到的:loss 尖峰后「自行恢复」,但 p99 面板里梯度范数一直贴着 3 走——那是裁剪在替一次真实的发散打掩护。
  • 下游看到的:灰度模型开始复读同一句话、答非所问,或者长上下文下突然返回空字符串。
  • 最糟的一种:没有任何报错,loss 变成一条不再下降的直线(bf16 的 ,更新恒为 0),直到三天后有人发现「训练了什么也没学到」。
  • 回滚时看到的:checkpoint 里已经存了 nan 权重,重启后第一步 loss 就是 nan,再往回找,上一个能用的版本是三天前。
  • 看起来像别的问题的:用逐元素裁剪(clip_grad_value_)代替按范数裁剪,方向被改掉,收敛变慢却不报警——现场表现是「学习率好像没调好」。

共同点:这些症状都不会在训练日志里写成 inf。日志里能看到的 inf 是幸运的那一种。

9. 练习

  1. 概念:为什么残差连接能缓解梯度消失?用 Jacobian 解释。
  2. 计算:,,,求 ;若 20 层,放大多少倍?
  3. 计算:梯度 ,裁剪阈值 2,结果是什么?
  4. 场景:softmax 饱和为什么会让梯度消失?写出 Jacobian 并解释。
  5. 开放:如果只看 loss 曲线,能否判断是梯度爆炸还是过拟合?还需要哪些指标?
参考答案
  1. ,特征值含 1,梯度至少有一条「恒等通道」,不会被连乘压没。
  2. ;20 层 倍。
  3. 范数 3,缩放 ,结果 。
  4. 饱和时 接近 one-hot,,梯度被压为零。
  5. 不能;还需梯度范数、激活/logits 统计、验证集 loss、学习率、数据分布等。过拟合通常训练 loss 降、验证 loss 升;爆炸则是 loss 突跳 + 梯度范数飙升。

10. 一页速记

  • 反向传播 = Jacobian 连乘;稳定性看谱范数是否接近 1。
  • ,。
  • 残差 Jacobian 提供恒等通道。
  • softmax 饱和 ⇒ Jacobian 趋零 ⇒ 梯度消失。
  • 裁剪保方向、限步长;QK-norm/z-loss 控制 logits。
  • 物理意义 / 现实对应:谱范数是这台「矩阵机器」的最大放大倍数,梯度范数是这一批数据对参数的总推力,裁剪是给推力装上限位器,残差是一条不受阻的恒等回传高速路,bf16 的动态范围是量尺的量程—— 到 都能量,但只读得出两三有效位。

11. 参考文献与致谢

这一页的知识不是凭空来的。下面列出本页直接借用的原始工作与权威整理;正文中的人名、定理名已链接到对应条目。链接优先指向 DOI、arXiv 或可长期访问的原文存档。

  • Rumelhart 等(1986):Learning representations by back-propagating errors。反向传播的经典出处:把误差沿网络逐层回传写成可训练算法,也就是本页 §2.2 那串 连乘。Nature 323(DOI)
  • He 等(2016):Deep Residual Learning for Image Recognition。残差连接 的出处;本页 §4.2 的 恒等通道就是它的反向传播解释。CVPR 2016(DOI)
  • Ba 等(2016):Layer Normalization。LayerNorm 的原始论文:对每个 token 的隐藏维归一化、再用可学习的 恢复尺度,对应本页 §3 与 §8 的 Pre-Norm 稳定性。arXiv:1607.06450
  • Glorot 与 Bengio(2010):Understanding the difficulty of training deep feedforward neural networks。Xavier 初始化:按 fan-in/fan-out 平衡每层权重的方差,使前向与反向的尺度都不随层数漂移,正是本页 §8 结尾要逼近的 。PMLR v9

下一篇 → M12 进阶与全景

返回 M11 导览去练习场

随堂测 · 4 题

单选难度 2第 1 / 4 题
对 ,记 ,则 等于?

源文件:11-大模型专题/12-反向传播与稳定性.md

输入关键词开始搜索