当梯度变成玄学

当梯度变成玄学

你大概率经历过这样的时刻:模型loss突然炸成NaN,或者训练了三天发现权重纹丝不动。你检查了数据、调小了学习率、换了初始化方法,最后在一篇论坛帖子里看到有人说“试试把某个算子换成数值稳定的版本”,问题莫名其妙地解决了。

你松了一口气,但心里清楚:自己并没有真正理解为什么。

这种无力感的根源,是我们把 .backward() 当成了一个黑盒。深度学习框架把微积分封装得太过优雅,以至于我们忘记了——计算机其实不懂极限,不懂链式法则,更不懂什么偏导数。它只会做加减乘除和内存读写。当我们写下 loss.backward() 时,机器执行的并不是数学意义上的“求导”,而是一套精心设计的工程流程。

这套流程不是凭空出现的。在它成为工业标准之前,人们尝试过人肉推导公式、用差分逼近斜率、用符号系统代数求导。每一种实践的局限,都精确地指向了下一个突破的方向。

这篇文章不会给你更多调试技巧。它会带你重走这条从“笨办法”到“自动微分”的演进之路。当你理解了机器求导的真实面貌,那些曾经神秘的梯度异常,将不再是玄学,而是可追溯、可解释的工程现象。

起点与困境:从人肉实践到机器算法

.backward() 成为默认选项之前,实践者面对三种选择。第一种是人肉推导,也就是手动微分;后两种是真正的计算机求导算法:数值微分和符号微分。人肉推导是自动微分出现之前的默认实践,后两者是算法尝试。它们各自的局限共同指向了自动微分的必要性。

手动微分:默认实践的不可扩展天花板

# 两层MLP的前向与手写梯度
h = torch.relu(x @ W1 + b1)
y = h @ W2 + b2
loss = ((y - target) ** 2).mean()

# 手动推导的梯度(仅展示W1部分)
dy = 2 * (y - target) / y.size(0)
dh = dy @ W2.T
dW1 = x.T @ (dh * (h > 0))  # ReLU导数硬编码,漏掉就静默错误

严格来说,这不是“让计算机求导”,而是人完成全部数学推导后,机器仅负责执行写好的梯度代码。它是没有算法时的默认状态——当参数只有几个时完全可行,但当模型膨胀到数十亿参数、网络结构频繁迭代时,其不可扩展的天花板便暴露无遗。任何结构调整都意味着重推整条链式法则,工程上已不可持续。

这条路径至今没有消失:教学里用它讲清链式法则,实现自定义算子时用它推导局部导数,极端部署场景下甚至直接手写梯度以省去 autograd 依赖。它的问题从来不是“错误”,而是“不可扩展”。

既然人推不动,能不能让机器自己“算”出梯度?于是有了真正意义上的计算机求导算法:数值微分与符号微分。先看最朴素的一种。

数值微分:2n次前向传播的代价

def numerical_grad(param, loss_fn, eps=1e-5):
    grad = torch.zeros_like(param)
    for i in range(param.numel()):      # 遍历每个参数
        param.data[i] += eps
        loss_plus = loss_fn()
        param.data[i] -= 2 * eps
        loss_minus = loss_fn()
        param.data[i] += eps             # 恢复原值
        grad.view(-1)[i] = (loss_plus - loss_minus) / (2 * eps)
    return grad

# 对含百万参数的模型调用一次 ≈ 200万次前向传播
grad_W1 = numerical_grad(W1, lambda: loss_fn(x, target))

数值微分对任意黑盒函数都适用,无需解析推导。这意味着百万参数模型训练一步就要跑两百万次前向——这在GPU时代仍是不可接受的开销。同时,eps 的选择永远在精度与浮点噪声之间走钢丝。

有没有办法既精确又高效?

符号微分:被误解的膨胀,与真实的瓶颈

import sympy as sp
x, W1, b1, W2, b2 = sp.symbols('x W1 b1 W2 b2')
h = sp.Max(0, x * W1 + b1)
y = h * W2 + b2
dy_dW1 = sp.diff(y, W1)

# SymPy默认输出未优化的表达式树,字符串长度随层数增长
# 但这并非符号微分的固有缺陷,而是朴素实现的产物
print(len(str(dy_dW1)))  # 简单两层MLP已超500字符

一个常见的误解是:符号微分必然导致表达式指数级膨胀,因此不可用。学术研究早已澄清,这其实是“神话”。如果符号系统在求导时复用中间结果(比如把 x*W1+b1 当作一个整体而非每次重新展开),其输出规模其实和前向模式自动微分一样,都是线性的。SymPy 之所以膨胀,是因为它没做这件事。

真正的瓶颈不在大小,而在形态:符号微分给出一张需要翻译的图纸,自动微分直接给出可执行的指令。前者多了一层“图纸→指令”的转换,在训练循环里这层开销足以抵消数学上的优雅。

这也解释了为何 JAX 等现代框架选择融合路线:XLA 编译器在自动微分生成的计算图上执行符号化简,既保留了AD的即时执行优势,又获得了符号方法的代数优化能力。符号微分没有失败,只是换了个位置继续发挥作用。

人肉推导的不可扩展、数值微分的效率瓶颈、符号微分的执行形态局限,共同指向了一个核心洞察:如果把“求导”不再当作独立于“计算”的事后动作,而是看作一系列已执行的基本运算,并在执行过程中同步累积梯度呢?

自动微分由此诞生。

符号微分在训练主战场之外依然活跃。DiffTaichi 和 Dr.Jit 把它和数值执行混在一起做物理仿真,研究者拿 SymPy 验证自定义算子的梯度公式,JAX/XLA 则用它优化计算图。它不再是端到端的求导方案,而是嵌在工具链里的一个零件。

自动微分的本质:链式法则的工程化

自动微分不是数值近似,也不生成符号表达式。它的核心洞察极其朴素:把复杂函数拆解为基本运算序列,并在执行这些运算的同时,按链式法则即时累积梯度。

下面这段代码用一个含分支汇聚的小计算片段,完整暴露了这一机制:

class AddMul(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, y, w):
        ctx.save_for_backward(x, y, w)
        return x * y + y * w          # y被使用两次,形成分支汇聚

    @staticmethod
    def backward(ctx, grad_z):
        x, y, w = ctx.saved_tensors
        # 两条路径对y的局部导数必须相加
        grad_x = grad_z * y
        grad_y = grad_z * x + grad_z * w   # ← 梯度在此处累积
        grad_w = grad_z * y
        print(f"grad_z={grad_z.item():.1f} | "
              f"grad_x={grad_x.item():.1f}, "
              f"grad_y={grad_y.item():.1f}(={grad_z.item():.1f}*{x.item():.1f}+{grad_z.item():.1f}*{w.item():.1f}), "
              f"grad_w={grad_w.item():.1f}")
        return grad_x, grad_y, grad_w

x = torch.tensor(2.0, requires_grad=True)
y = torch.tensor(3.0, requires_grad=True)
w = torch.tensor(4.0, requires_grad=True)
z = AddMul.apply(x, y, w)
z.backward()
# 输出: grad_z=1.0 | grad_x=3.0, grad_y=6.0(=1.0*2.0+1.0*4.0), grad_w=3.0
# x.grad==3.0, y.grad==6.0, w.grad==3.0

这段代码揭示了自动微分与前述实践的根本区别:

  • 相比人肉推导的默认实践:开发者只需定义每个算子的局部导数(backward中的三行乘法),框架负责按计算顺序自动组合所有局部导数。当y被复用十次时,累加逻辑仍由同一行代码处理,无需人工展开链式法则。
  • 对比数值微分:梯度是精确的解析结果(x、y、w本身),无任何ε扰动;且无论参数多少,反向传播仅需一次遍历。
  • 对比符号微分:不生成任何中间表达式。ctx仅保存数值张量,梯度在反向执行时动态计算、即时释放,省去了“图纸→指令”的转换开销,内存与时间成本均与原始前向同阶。

更重要的是,print语句暴露了后向模式的核心特征:梯度是从输出向输入反向流动的,且在分支汇聚点必须累加。它并非简单地把前向倒过来算,梯度流动有严格的依赖顺序。

但注意:上述代码展示的是后向模式。链式法则同样支持从输入向输出的前向模式,数学上完全等价。既然两者都正确,为什么PyTorch、TensorFlow、JAX无一例外地只暴露了 .backward()

答案藏在下一节的计算量对比里。

从前向模式到后向模式:为什么是 .backward()

链式法则有两种等价的执行顺序。数学上它们同样正确,工程上却天差地别。

前向模式:符合直觉,但扛不住多输入

# 前向模式:用(value, tangent)元组手动模拟对偶数
# 每次只能计算一个输入方向的梯度,n个输入需n次独立前向
def forward_jvp(x, y, w, seed_x=0, seed_y=0, seed_w=0):
    # 每个变量携带(值, 切线),切线初始化为seed方向
    x_dual = (x, seed_x)
    y_dual = (y, seed_y)
    w_dual = (w, seed_w)

    # z = x*y + y*w 的对偶运算(乘法/加法按对偶数规则展开)
    xy_val = x_dual[0] * y_dual[0]
    xy_tan = x_dual[0] * y_dual[1] + y_dual[0] * x_dual[1]

    yw_val = y_dual[0] * w_dual[0]
    yw_tan = y_dual[0] * w_dual[1] + w_dual[0] * y_dual[1]

    z_val = xy_val + yw_val
    z_tan = xy_tan + yw_tan
    return z_val, z_tan

x, y, w = 2.0, 3.0, 4.0
# 求dz/dx: seed=(1,0,0) → 一次前向只得一个梯度
_, dz_dx = forward_jvp(x, y, w, seed_x=1, seed_y=0, seed_w=0)
# 求dz/dy: seed=(0,1,0) → 第二次前向
_, dz_dy = forward_jvp(x, y, w, seed_x=0, seed_y=1, seed_w=0)
# 求dz/dw: seed=(0,0,1) → 第三次前向
_, dz_dw = forward_jvp(x, y, w, seed_x=0, seed_y=0, seed_w=1)
# 3输入1输出 → 前向模式需3次完整前向传播
# n输入m输出 → 通用代价 O(n · ops),ops为基本运算次数

前向模式完全符合人类“从输入到输出”的直觉,每个变量的切线随值一同向前流动。但当输入维度n远大于输出维度m时(这正是神经网络的常态),它需要n次独立前向才能拿到完整梯度——与数值微分陷入同样的O(n)困境。

后向模式:反直觉,却适合训练场景

# 后向模式:第三部分的AddMul.backward已展示
# 一次反向传播得到所有输入的梯度
# z.backward() 单次调用即得 x.grad, y.grad, w.grad
# 3输入1输出 → 后向模式仅需1次反向遍历
# n输入m输出 → 通用代价 O(m · ops),与输入维度无关
# 神经网络: m=1(标量loss), n=百万参数 → 后向比前向快百万倍

后向模式要求先完成整个前向、再逆序累积梯度,心智模型更复杂。但它将计算代价锚定在输出维度m上。对于训练场景(m=1),无论网络有多少参数,梯度获取成本恒定。

工程选择:PyTorch为何只暴露 .backward()

这个计算量差异直接决定了框架设计。PyTorch、TensorFlow、JAX的训练API全部围绕后向模式构建,因为深度学习本质上是“高维输入→标量输出”的优化问题。自定义算子必须实现 backward 而非前向梯度方法,正是因为框架默认且仅高效支持后向模式。

但前向模式并未消亡。当场景反转为“低维输入→高维输出”(如生成模型的雅可比计算)、或需要Hessian向量积、梯度检查时,JAX的 jvp、PyTorch的 torch.autograd.functional.jvp 仍是首选。理解两种模式的适用边界,才能在自定义算子、调试梯度异常、选择框架特性时做出正确判断。

这也解释了引言中的痛点:当你看到NaN梯度时,问题几乎总出在后向模式的某个局部导数实现上——因为那是唯一被高频执行的路径。

总结:超越 .backward() 的确定性

回顾这条演进之路,数值微分和符号微分作为算法各有归宿:前者是梯度检查的可靠基准,后者转型为编译器内的优化 pass。而人肉推导严格说不是一种算法,它是自动微分出现之前的默认实践,至今仍是教学、自定义算子验证和极端部署场景的起点。自动微分并非“淘汰”了它们,而是将链式法则工程化,解决了大规模训练的核心瓶颈;而后向模式又因多输入单输出的场景特性,成为了深度学习的默认选择。

理解这些底层逻辑,最终是为了在日常开发中获得确定性。当你在自定义 autograd.Function 时,本质上是在手写局部导数——漏掉分支汇聚处的梯度累加、写错正负号、或对不可导点做了错误假设,就是引言中那些NaN和静默错误的根源。当你看到官方文档警告“不要在forward里调用 .item().numpy()”时,现在能立刻反应过来:这会切断计算图,让后向模式的流水线从中间断裂,梯度无法回传。这些不再是死记硬背的规则,而是可追溯的工程因果。

求导的工程化仍在继续。JAX支持对梯度函数再次求导,可微编程框架将物理仿真纳入优化循环,XLA等编译器正将符号化简与自动微分深度融合。但无论上层如何演进,底层的心智模型始终稳定:把复杂函数拆解为基本运算,在执行中按链式法则累积梯度。掌握这一点,你就掌握了与所有现代深度学习框架对话的共同语言。