技术博客

数值梯度检验与梯度消失的诊断

第 14 / 147 篇 · 2025-03-22 · ai · 梯度检验有限差分调试

前面几篇推了那么多梯度公式,凭什么相信推对了、代码写对了?解析推导有符号游戏出错的可能,框架 autograd 也不能验证你自定义算子里的手写生灭。答案是一个笨办法:有限差分。它不依赖任何推导,直接从定义出发逼近导数,是反向传播代码的唯一独立裁判。这篇讲怎么把它用对——步长选多大、误差怎么比,以及梯度消失时同一套工具怎么做诊断。

这套工具的历史比深度学习长得多:数值分析里用它校验手推导数是几十年的老传统,自动微分库自己的测试套件(PyTorch 的 gradcheck、JAX 的 jax.test_util.check_grads)内核同样是它。学会手搓一遍,再看框架的实现就没有黑盒了。

中心差分:$O(h)$ 和 $O(h^2)$ 的一步之遥

导数定义给的是单侧差分:

$$ f'(x)\approx\frac{f(x+h)-f(x)}{h}, $$

但它的截断误差是 $O(h)$——泰勒展开 $f(x+h)=f(x)+hf'(x)+\frac{h^2}{2}f''(x)+\cdots$ 代入就知道,误差主导项是 $\frac{h}{2}f''(x)$。改成中心差分:

$$ f'(x)\approx\frac{f(x+h)-f(x-h)}{2h}, $$

$f(x+h)$ 与 $f(x-h)$ 的展开相减,$h^2$ 项奇偶相消,主导误差变成 $\frac{h^2}{6}f'''(x)$,即 $O(h^2)$。代价只是多算一次前向,精度却整整高一阶。做梯度检验没有理由用单侧。

多变量情形逐坐标套用同一公式:对每个分量 $x_i$ 各做两次前向。这正是数值检验 $O(n)$ 次前向的成本来源,也是它只能做"裁判"不能做"选手"的原因——autodiff 一次反向拿全部梯度,差分要为每个分量单独付费。

那 $h$ 是不是越小越好?不是,这里有一场截断与舍入的拉锯:

  • 截断误差 $\sim\frac{h^2}{6}|f'''|$,随 $h$ 减小按平方下降;
  • 舍入误差:$f(x+h)$ 与 $f(x-h)$ 在 float64 下各有约 $\epsilon_{\text{mach}}\approx2.2\times10^{-16}$ 量级的表示误差,相减再除以 $2h$,贡献 $\sim\epsilon_{\text{mach}}|f|/h$,随 $h$ 减小反比上升。

总误差呈 V 形,最优点在两股误差相等处:$h^2\sim\epsilon_{\text{mach}}/h$,即 $h\sim\epsilon_{\text{mach}}^{1/3}\approx6\times10^{-6}$。经验值 float64 取 $h=10^{-5}\sim10^{-6}$。float32 的 $\epsilon_{\text{mach}}\approx1.2\times10^{-7}$,最优 $h$ 挪到 $\sim10^{-3}$,而此时最小可达相对误差只有 $10^{-3}$ 量级——梯度检验必须在 float64 下做,float32 里检验等于没检,这是我把无数"检验不通过"归因到 dtype 之后的结论。

中心差分总误差随 h 的 V 形曲线(log-log):截断误差下行,舍入误差上行,交点附近是最优步长 步长 h(对数轴,向右减小)→ 误差(log) 截断误差 ~ h²/6·|f‴| 舍入误差 ~ ε|f|/h 最优 h ≈ ε^(1/3) float64: ~1e-5 h 太大:公式不准 h 太小:减法灾难

怎么比:相对误差指标

逐分量比较 $|a_i-b_i|$ 没有意义,梯度分量量级差异可能跨数量级。标准做法(也是 torch.autograd.gradcheck 用的)是全量相对误差:

$$ \mathrm{rel}=\frac{\|a-b\|_2}{\|a\|_2+\|b\|_2}. $$

分母用和而不是 $\max$ 或单边,是为了 $a\approx b\approx0$ 时指标仍然良性。判据经验值:float64 + 中心差分下,$\mathrm{rel}<10^{-7}$ 放心,$10^{-7}\sim10^{-5}$ 可疑(调 $h$ 再试),$>10^{-5}$ 基本就是反向写错了。注意 gradcheck 要求你的函数对输入是确定性的——dropout 不关、随机 mask 不重置种子,检验必挂,而且挂得像梯度错了。

检验要覆盖的不是"一个随机点",而是边界点:输入含零(ReLU 的折点)、含相同值(max 的 ties)、接近饱和区(tanh 大输入)。梯度在光滑点大家都对,错的是折点处次梯度的约定。gradcheck 挂在这些点时要判断是你的错还是不可微点的固有歧义。

另一个容易被忽略的前提是损失必须是标量。如果检验的是一个多输出的中间层,先给它接一个固定的随机线性泛函 $w^\top y$(等价于伪造上游梯度),再对标量结果做差分——这正是第一篇矩阵求导里 VJP 的角色:检验的从来不是"层的导数"这种不存在的东西,而是"层输出与固定上游梯度内积之后的标量函数"的导数。

复步长技巧:绕开减法灾难的另一条路

有限差分的舍入误差来自"$f(x+h)$ 与 $f(x-h)$ 几乎相等却相减"的灾难性抵消。有一个漂亮的绕法——复步长微分(complex-step):对实解析函数, $$ f(x+ih)=f(x)+ihf'(x)-\frac{h^2}{2}f''(x)+\cdots\ \Rightarrow\ f'(x)=\frac{\mathrm{Im}\,f(x+ih)}{h}+O(h^2). $$

虚部里只有 $f'$ 一项,没有减法,因此不存在抵消,$h$ 可以放心取到 $10^{-50}$ 这样的量级,精度直逼机器精度。代价是被测函数必须支持复数输入且内部只用解析运算(abs、比较分支、max 这类不行),PyTorch 张量运算大部分天然支持 complex128。对自写的标量函数做终极核验时,这是我压箱底的手段;对含 max/裁剪分支的网络则老老实实用中心差分。

与自动微分的关系也值得摆正:数值差分是裁判,autodiff 是选手。差分 $O(h^2)$ 的截断决定了它永远只是近似,所以不能反过来要求 autodiff 向差分"对齐"到零;判据的容差正是为这段不可消除的距离留的。

20 行手写 gradcheck

import torch

def rel_check(f, x, h=1e-5, tol=1e-7):
    x = x.detach().clone().to(torch.float64).requires_grad_(True)
    analytic = torch.autograd.grad(f(x), x)[0]
    numeric = torch.zeros_like(x)
    it = __import__("numpy").nditer(x.detach(), flags=["multi_index"])
    for _ in it:
        i = it.multi_index
        xp, xm = x.detach().clone(), x.detach().clone()
        xp[i] += h; xm[i] -= h
        numeric[i] = (f(xp) - f(xm)) / (2 * h)
    rel = (analytic - numeric).norm() / (analytic.norm() + numeric.norm())
    return rel.item() < tol, rel.item()

# 用法:f 必须是标量输出的确定性函数
ok, rel = rel_check(lambda t: (torch.tanh(t) ** 2).sum(), torch.randn(6))
print(ok, rel)   # True, ~1e-11 量级

对张量求梯度时要逐分量扰动,所以这套东西只适合小输入——它是单元测试工具,不是大模型验证工具。大图上的正确姿势是把可疑子图拆出来单独喂小输入。

再验一个容易写错的场景——带 clamp 折点的函数,顺带演示折点附近判据的松动:

ok, rel = rel_check(
    lambda t: (torch.sigmoid(t).clamp(min=0.2) ** 2).sum(),
    torch.linspace(-3, 3, 6), tol=1e-6,
)
print(ok, rel)

clamp 的折点上,中心差分跨过不可微点得到两侧斜率的平均,autograd 给的是次梯度的单侧约定——这两个答案本来就不该相等,容差放宽到 $10^{-6}$ 正是吸收这类固有分歧。把 tol 卡死在 $10^{-9}$ 然后对折点较劲,是 gradcheck 使用中最常见的自我折磨。

梯度消失的诊断流程

同一套有限差分工具换个用途:诊断深层网络梯度消失。机制上,深层网络的梯度是 $L$ 个雅可比的连乘 $\prod_l J_l$,范数大致按奇异值连乘衰减或爆炸——平均奇异值 $\bar\sigma<1$ 时梯度范数 $\sim\bar\sigma^L$ 指数归零。tanh 网络饱和后 $J_l$ 的对角元 $\tanh'\le1$ 且饱和区趋零,是最经典的来源。

实操排查顺序(我按这个顺序省过大量时间):

  1. 逐层记录 grad 范数:for n,p in model.named_parameters(): log(p.grad.norm()),看曲线从输出层到输入层是衰减、爆炸还是断裂;
  2. 断裂点定位:范数在某一层断崖式归零,先查该层激活是否饱和(记录激活直方图,tanh/sigmoid 看 $|$输出$|$ 是否贴 1),再查权重谱范数;
  3. 区分消失与死亡:ReLU 网络里梯度归零可能是整片神经元死亡(激活恒为 0),不是连乘衰减——激活均值恒零、对应权重梯度也恒零,用有限差分在单样本上验证一下就能区分:差分也零,是真死亡;差分非零而反向零,是路径断了;
  4. 数值而非符号:float16/bf16 下梯度下溢($<10^{-8}$ 量级)也会伪装成梯度消失,先转 float32 再下结论。

这套流程的每一步都可以压缩成日志面板上的一个数字或一张直方图,总成本不到一次训练重启的零头。真正的成本在"没装监控时靠猜"——梯度问题的根因分散在初始化、归一化、dtype、数据四个方向,猜错方向的试错按天计,装好监控的定位按小时计。

梯度爆炸的对称诊断

消失的每一步排查都有爆炸的镜像版,机制同样是雅可比连乘,只是平均奇异值 $\bar\sigma>1$:范数随深度指数增长,表现为浅层 grad norm 天文数字、loss 一步冲成 NaN。诊断顺序照搬——逐层范数曲线从输出端往回看是"增长"而不是"衰减",断裂点定位到首次爆炸的层。处置比消失更直接:全局范数裁剪先把步子压住(可信域那篇会展开),再回头查根因——常见的引爆点是某层初始化谱范数过大、残差分支无缩放地堆叠、或 bf16 里中间激活溢出后梯度跟着污染。两类问题的共同教训:先看曲线形态再动手调参,形态已经把故障类型告诉你了。

爆炸还有一个消失没有的急迫性:它会当场冲毁训练状态。消失可以"带病训练"几周才被发现,爆炸一步到位把参数打进 NaN 区域,只能靠 checkpoint 回滚。所以监控面板上 grad norm 的告警阈值该按爆炸场景设置——宁误报,不漏报。

检验的组织方式:放在哪、跑多勤

梯度检验不该是一次性动作,值得收进项目的常驻测试。我的做法是给每个手写算子配一个小尺寸 gradcheck 用例,输入维度压到个位数、dtype 固定 float64、种子固定;CI 里跑,改动算子或其下游时必过。大模型训练侧则把"逐层 grad norm"接进训练日志面板,作为常备监控而不是出事才查的救火工具——曲线形态(随深度衰减的斜率、突发尖峰)本身就是训练健康的指纹。

还有一个反直觉的经验:检验失败的梯度不一定不能用。差分失败但失败模式是"所有分量同比例偏差"时,多半是前向里藏了一个与输入无关的缩放(比如 softmax 的 log-sum-exp 常数化简写错),方向对了训练照样能跑——但这类"碰巧能用"的状态极其脆弱,换个 loss 组合就露馅,不值得保留。

组织层面再补两条细节。一是检验点的采样:随机点之外,刻意构造结构化输入——全零、全同值、按位递增——很多广播/归约类 bug 只在有重复值时触发。二是把检验写进算子开发的流程而不是事后补票:先写前向、随手写反向、立刻 gradcheck,错误在刚写完的半小时内修掉成本最低,隔一周再查就要重读自己的代码。

小结

中心差分把截断误差压到 $O(h^2)$,配合 float64 与 $h\approx10^{-5}$、相对误差 $<10^{-7}$ 的判据,构成手写反向的独立裁判;检验要特意覆盖折点与饱和区。梯度消失的诊断则是同一工具的逆向使用:逐层范数定断裂点,直方图分饱和与死亡,dtype 下溢先排除。下一篇换个尺度看问题:既然一阶信息这么好用,为什么二阶的海森矩阵在大模型里几乎没人用。