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 的动态范围)。判错方向的代价,就是上面那张价目表。
读完本节,你能回答的三个问题是:
- 梯度穿过一层 到底被放大几倍?怎样用一个数(谱范数)估出 32 层后的总倍数?
- 残差、归一化、初始化各自把放大倍数拉回 1 的哪一段?梯度裁剪能救什么、救不了什么?
- 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 的动态范围
浮点数 = 符号位 + 指数位 + 尾数位。指数位决定「能表示多大/多小」,尾数位决定「一位有效数字有多细」。三兄弟的差别全在这两栏:
| 格式 | 指数位 | 尾数位(显式) | 最大正规数 | 最小正规数 | 相对精度 |
|---|---|---|---|---|---|
| fp32 | 8 | 23 | |||
| fp16 | 5 | 10 | |||
| bf16 | 8 | 7 |
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 的常见链条
- 某个 batch 的 attention logits 变大;
- softmax 变尖锐,熵下降,梯度变小但输出极端;
- 后续层收到极端输入,产生大梯度;
- 参数被推到一个坏区域,下一步 loss 飙升;
- 若没有裁剪/预热,可能直接 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. 常见误区
-
误区:梯度裁剪会改变优化方向。 纠正:它只按比例缩放整个梯度,方向不变(除非逐元素裁剪)。
-
误区:有了 LayerNorm 就不会梯度爆炸。 纠正:归一化只控制激活尺度;权重矩阵的谱范数、残差累积、softmax 饱和仍可能导致不稳定。
-
误区:loss spike 一定是数据问题。 纠正:可能是优化/数值问题(logits 爆炸、学习率过大、bf16 溢出),需要看梯度范数与激活统计。
-
误区:谱范数就是最大元素。 纠正:谱范数是最大奇异值,衡量算子放大能力;最大元素只是元素级指标。
7. 概念网络
- 前置:8.5 矩阵求导 —— 链式法则。
- 前置:10.2 条件数 —— 谱范数。
- 前置:11.6 归一化 —— 激活尺度。
- 关联:11.3 Attention —— logits 爆炸。
- 关联:12.2 随机矩阵 —— 初始化与谱。
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. 练习
- 概念:为什么残差连接能缓解梯度消失?用 Jacobian 解释。
- 计算:,,,求 ;若 20 层,放大多少倍?
- 计算:梯度 ,裁剪阈值 2,结果是什么?
- 场景:softmax 饱和为什么会让梯度消失?写出 Jacobian 并解释。
- 开放:如果只看 loss 曲线,能否判断是梯度爆炸还是过拟合?还需要哪些指标?
参考答案
- ,特征值含 1,梯度至少有一条「恒等通道」,不会被连乘压没。
- ;20 层 倍。
- 范数 3,缩放 ,结果 。
- 饱和时 接近 one-hot,,梯度被压为零。
- 不能;还需梯度范数、激活/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 进阶与全景
随堂测 · 4 题
源文件:11-大模型专题/12-反向传播与稳定性.md