前言

反向传播(Backpropagation)是训练神经网络的核心算法。本文将深入探讨其数学本质——多元函数链式法则的高效实现。

计算图视角下的链式法则

考虑一个简单的三层网络:

$$ \begin{aligned} z_1 &= W_1 x + b_1 \ a_1 &= \sigma(z_1) \ z_2 &= W_2 a_1 + b_2 \ \hat{y} &= \text{softmax}(z_2) \ \mathcal{L} &= -\sum_{k} y_k \log \hat{y}_k \end{aligned} $$

对应的计算图如下:

前向计算图

标量反向传播

对于标量函数 $f(g(h(x)))$,链式法则给出:

$$ \frac{df}{dx} = \frac{df}{dg} \cdot \frac{dg}{dh} \cdot \frac{dh}{dx} $$

Sigmoid 激活函数为例,其导数的简洁形式是一个经典技巧点:

1def sigmoid(x: float) -> float:
2    s = 1.0 / (1.0 + math.exp(-x))
3    return s
4
5def sigmoid_backward(d_out: float, s: float) -> float:
6    # sigmoid 导数: sigma(x) * (1 - sigma(x))
7    return d_out * s * (1.0 - s)

矩阵/张量反向传播

当推广到矩阵时,维度一致性是关键。对于矩阵乘法 $Y = XW$:

$$ \frac{\partial \mathcal{L}}{\partial X} = \frac{\partial \mathcal{L}}{\partial Y} \cdot W^T $$

$$ \frac{\partial \mathcal{L}}{\partial W} = X^T \cdot \frac{\partial \mathcal{L}}{\partial Y} $$

这里的 · 表示矩阵乘法。在实际框架中,这些梯度由 autograd 引擎 自动计算。

PyTorch 实战:自定义 Autograd Function

 1import torch
 2from torch.autograd import Function
 3
 4class MyLinear(Function):
 5    """自定义线性层的前向与反向传播"""
 6
 7    @staticmethod
 8    def forward(ctx, input, weight, bias=None):
 9        # 保存反向所需变量
10        ctx.save_for_backward(input, weight, bias)
11        output = input.mm(weight.t())
12        if bias is not None:
13            output += bias.unsqueeze(0).expand_as(output)
14        return output
15
16    @staticmethod
17    def backward(ctx, grad_output):
18        input, weight, bias = ctx.saved_tensors
19        grad_input = grad_weight = grad_bias = None
20
21        # 根据链式法则计算各梯度
22        if ctx.needs_input_grad[0]:
23            grad_input = grad_output.mm(weight)        # dL/dX = dL/dY · W
24        if ctx.needs_input_grad[1]:
25            grad_weight = grad_output.t().mm(input)    # dL/dW = Y_grad^T · X
26        if bias is not None and ctx.needs_input_grad[2]:
27            grad_bias = grad_output.sum(0)             # dL/db = sum(dL/dY, dim=0)
28
29        return grad_input, grad_weight, grad_bias

调试梯度时,数值梯度检查(torch.autograd.gradcheck)是必不可少的工具:

1# 使用 gradcheck 验证自定义反向传播
2$ python -c "
3import torch
4from my_linear import MyLinear
5input = torch.randn(20, 10, dtype=torch.double, requires_grad=True)
6weight = torch.randn(5, 10, dtype=torch.double, requires_grad=True)
7torch.autograd.gradcheck(MyLinear.apply, (input, weight), eps=1e-6, atol=1e-4)
8print('Gradient check passed!')
9"

激活函数对比

激活函数公式导数值域优点缺点
Sigmoid$\sigma(x) = \frac{1}{1+e^{-x}}$$(0, 0.25]$平滑、有界梯度消失
Tanh$\tanh(x) = \frac{e^x-e^{-x}}{e^x+e^{-x}}$$(0, 1]$零中心化梯度消失
ReLU$f(x)=\max(0,x)$${0, 1}$计算快、缓解梯度消失Dead Neuron
GELU$x \cdot \Phi(x)$$\approx (0, 1]$平滑、概率解释计算略贵
SwiGLU$x \cdot \sigma(\beta x)$LLM 标配参数量翻倍

Why GELU? GELU(Gaussian Error Linear Unit)在 Transformer 架构中表现优异,因为它是 ReLU 的平滑近似——对每个输入应用一个由标准正态 CDF $\Phi(x)$ 决定的"门控概率"。

常见问题排查清单

在实践中,反向传播的调试常有以下痛点:

  1. 梯度爆炸(Gradient Explosion):损失突然变成 NaN
    • 使用 torch.nn.utils.clip_grad_norm_ 进行梯度裁剪
    • 检查学习率是否过大
  2. 梯度消失(Vanishing Gradient):loss 几乎不下降
    • 换用 ReLU / GELU 替代 Sigmoid
    • 添加 BatchNorm / LayerNorm
  3. 形状不匹配(Shape Mismatch):RuntimeError: shape '[64, 10]' is invalid for input of size 1280
    • 仔细追踪 forward 中每一步的 tensor shape
    • 善用 .view() / .reshape() / .flatten()

现代框架的计算图范式对比

框架图构建方式动态图分布式策略
PyTorchDefine-by-RunDDP / FSDP
TensorFlow 2.xEager + tf.functionMirroredStrategy
JAX函数式 + grad()/vmap()pmap() / pjit
Flux.jl (Julia)原生 Julia 自动微分

总结

我们用一张表回顾核心概念:

概念符号直觉理解
前向传播$x \mapsto f(x)$从输入计算输出
反向传播$\frac{\partial \mathcal{L}}{\partial x}$回答"改变 $x$ 会如何影响最终损失?"
梯度累积$\sum \nabla$用 micro-batch 模拟大 batch
链式法则$\frac{\partial z}{\partial x} = \frac{\partial z}{\partial y} \frac{\partial y}{\partial x}$将复杂梯度分解为简单梯度的乘积

反向传播的本质,就是把一个复杂函数的梯度计算分解为有向无环图上每个节点的局部雅可比乘积。掌握了这一点,也就掌握了所有现代深度学习框架的核心抽象。


推荐阅读:《The Autodiff Cookbook》by JAX team