loss.backward()函数工作原理详解

背景

训练神经网络的核心是优化一个损失函数（loss），该函数用一个标量数值衡量模型当前预测的误差。通过梯度下降法，我们反复微调每一个参数，使得损失下降。为了执行一次梯度下降，对于每个参数 ( p )，我们需要知道：如果对 ( p ) 施加微小扰动，损失会如何变化（方向和大小）。这个量就是导数。计算所有参数的导数，沿反方向迈出一小步，就完成了一次学习更新：p->data -= learning_rate * p->grad。

因此，整个问题归结为：如何一次性获得所有参数的导数？一个真实网络可能有成百上千乃至数万亿个参数，但只有一个损失值。这种“多输入，单输出”的形状，正是后续一切机制运作的根本原因。

核心内容

局部导数与链式法则

我们不需要知道整个庞大网络的导数，只需要知道每个独立操作的局部导数。例如：

• 加法：( z = x + y )，对 ( x ) 的局部导数为 1（因为 ( x ) 增加 1，结果增加 1）。
• 乘法：( z = x \cdot y )，对 ( x ) 的局部导数为 ( y )（因为 ( x ) 变化 ( \Delta x )，结果变化 ( y \cdot \Delta x )）。

其他操作（减法、双曲正切等）类似。要组合这些局部导数，需要链式法则：如果 ( z ) 依赖于 ( y )，( y ) 依赖于 ( x )，那么 ( \frac{dz}{dx} = \frac{dz}{dy} \cdot \frac{dy}{dx} )。当 ( x ) 通过多条路径影响 ( z ) 时，各路径贡献相加。

运行示例

考虑以下代码片段（来自 microcrad，一个标量自动微分引擎，灵感来自 Andrej Karpathy 的 micrograd）：

Value *a = value_create_leaf(2.0);
Value *b = value_create_leaf(3.0);
Value *e = value_mul(a, b);   // e = a * b = 6
Value *L = value_mul(e, a);   // L = e * a = 12

数学上：先计算 ( e = a \cdot b )，再计算 ( L = e \cdot a )。代入得 ( L = a \cdot b \cdot a = a^2 \cdot b )。当 ( a = 2.0, b = 3.0 ) 时，( L = 12 )。

注意 ( a ) 被使用了两次：一次用于生成 ( e )，一次直接参与 ( L )。这正是多条路径的情况。

手工计算导数：

• ( L = e \cdot a )，局部导数：( \frac{\partial L}{\partial e} = a = 2 )，( \frac{\partial L}{\partial a} = e = 6 )。
• ( e = a \cdot b )，局部导数：( \frac{\partial e}{\partial a} = b = 3 )，( \frac{\partial e}{\partial b} = a = 2 )。

应用链式法则：

• ( \frac{dL}{db} = \frac{\partial L}{\partial e} \cdot \frac{\partial e}{\partial b} = 2 \times 2 = 4 )。
• ( a ) 有两条路径：通过 ( e ) 的路径：( \frac{\partial L}{\partial e} \cdot \frac{\partial e}{\partial a} = 2 \times 3 = 6 )；直接路径：( \frac{\partial L}{\partial a} = 6 )。总和：( 6 + 6 = 12 )。

为什么要反向传播

有两种方向应用链式法则：

• 前向模式：选定一个输入（如 ( a )），从该输入向前传播影响，计算 ( \frac{\partial e}{\partial a} )、( \frac{\partial L}{\partial a} )。一次扫描只能得到关于这一个输入的导数。要得到 ( b ) 的导数，必须再做一次完整扫描。每个输入一次扫描，成本随参数数量线性增长。

• 反向模式：从输出 ( L ) 开始，种子设为 1，向后传播到所有输入。一次扫描就能同时得到 ( \frac{dL}{da} )、( \frac{dL}{db} ) 等所有输入的导数。

问题的形状是“多参数，单损失”，前向模式需要每个参数一次扫描——当参数百万级时不可接受。反向模式一次扫描即可给出所有参数梯度。这种不对称性正是神经网络可训练的根本原因。这就是反向模式自动微分，在机器学习中称为“反向传播”。microcrad 在构建计算图后所做的所有工作就是一次反向扫描。

注意：在节点的梯度馈送到其前驱节点之前，该梯度必须已经完成累加。例如 ( a ) 的梯度，必须等经过 ( e ) 的路径和直接路径的贡献都合并后，才算完成。如果节点按错误顺序访问，可能会传播部分累加的梯度。因此，我们需要对图进行拓扑排序，使每个节点出现在其依赖节点之后，然后反向遍历该列表——这是反向传播算法的第一步。

计算图的自动构建

反向传播需要一个图——记录每个值由哪些操作数和什么操作产生。但开发者无需显式构建图，只需执行前向计算，图作为副作用自动生成。

秘密在于：Value 不仅仅是一个数字，它还是一个记住来源的数字：

typedef struct Value {
    double data;               /* 节点持有的标量值 */
    double grad;               /* dLoss/dThisValue，反向传播时填充 */
    struct Value **prev;       /* 操作数（图中的前驱节点） */
    int32_t op_code;           /* 产生该值的操作 */
    /* 内存管理字段省略 */
} Value;

由操作产生的 Value 通过 prev 指针指回其操作数，并通过 op_code 标记自身由何种操作产生。沿着 prev 指针从任意节点出发，即可反向遍历计算图。

以乘法操作为例（简化版，去掉了错误处理）：

Value *value_mul(Value *v1, Value *v2) {
    Value **prev = malloc(2 * sizeof(Value *));
    prev[0] = v1; value_retain(v1);
    prev[1] = v2; value_retain(v2);
    Value *result = value_create(v1->data * v2->data, 2, prev);
    result->op_code = MUL_OP_CODE;
    return result;
}

每个操作都做三件事：

1. 计算结果并存储在新节点中。
2. 记录操作数到 prev。
3. 用操作码标记该节点，以便反向传播时知道应使用哪个局部导数规则。

因此，运行示例中的四行代码不仅计算了 12，每个操作都悄悄留下了指向其输入的节点。当得到 L 时，实际上也持有了一个图的根节点，该图记录了 L 的完整历史：a 是一个节点，有两个箭头分别指向 e 和 L，这正是我们手工求导时遇到的“两条路径”，也是为什么 a 的梯度会有两个贡献相加。

反向传播的代码实现

value_backward 函数执行两个步骤：

1. 构建图的拓扑排序，确保每个节点出现在其依赖节点之后。自然，损失节点成为图的输出节点。
2. 将输出节点的梯度初始化为 1，然后反向遍历拓扑排序列表，对每个节点应用其局部导数规则，将梯度累加到其前驱节点上。

关键要点

• 反向传播的本质：利用链式法则从输出向输入反向传播梯度，一次扫描即可获得所有参数相对于损失函数的导数，