Keyboard shortcuts

Press or to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

Autograd 自动微分

引言

深度学习训练的目标,是找到一组使损失尽可能小的模型参数。梯度给出了损失相对于各个参数的变化率,优化器据此更新参数。为了得到这些梯度,训练框架必须解决两个问题:如何记录前向计算,以及如何沿着计算过程反向传播梯度。

Zero 使用反向模式自动微分。它在前向计算时动态构建计算图;调用 Backward() 后,从损失的梯度开始,按照与前向相反的顺序依次执行各个算子的反向规则。每个算子接收输出端传来的梯度,计算各个输入的梯度并继续向前传递;如果同一个张量参与了多条计算路径,来自不同路径的梯度还需要累加。这个过程一直进行到模型参数,得到参数更新所需的梯度。

链式法则

考虑复合函数:

$$ y=f(g(x)) $$

它对 \(x\) 的导数可以拆成两个局部导数的乘积:

$$ \begin{aligned} \frac{\partial y}{\partial x} &= \frac{\partial y}{\partial g} \frac{\partial g}{\partial x} \end{aligned} $$

训练引擎不需要一次求出整个模型的导函数。每个算子只需知道自己的局部导数,然后将上游梯度乘上局部导数即可。

例如 \(z=x\cdot y\) 的上游梯度为 \(g\),乘法算子返回:

$$ \frac{\partial L}{\partial x}=g\cdot y, \qquad \frac{\partial L}{\partial y}=g\cdot x $$

如果一个张量被多个后续算子使用,那么它的总梯度是每条路径贡献之和:

$$ \begin{aligned} \frac{\partial L}{\partial x} &= \sum_i \frac{\partial L}{\partial y_i} \frac{\partial y_i}{\partial x} \end{aligned} $$

这也是反向传播中梯度必须累加,而不能直接覆盖的原因。

计算图

表达式 \(z=x\cdot y+x\) 可以写成一张有向无环图:

x ──┐
    ├── Mul ──┐
y ──┘          ├── Add ── z
x ────────────┘

图中的张量保存数值,算子保存局部的前向和反向规则。前向计算得到 z 时,还需要记住 Add 的输入;否则反向阶段无法继续找到 Mulxy

设 \(x=3\)、\(y=2\),前向计算得到 \(x\cdot y=6\)、\(z=9\)。反向计算从 \(\partial z/\partial z=1\) 开始:

                     ∂z/∂z = 1
                          │
                    Add::Backward
                    ┌─────┴─────┐
              传给 Mul:1    传给 x:1
                    │
                    │
               Mul::Backward
               ┌────┴────┐
       传给 x:1 × y   传给 y:1 × x
               │
               └──→ x 的两条梯度相加

Add 将上游梯度分别传给两个输入,Mul 再根据局部导数计算 \(x\) 和 \(y\) 的梯度。由于 \(x\) 同时经过 MulAdd 到达 \(z\),它需要累加两条路径的贡献。

最终得到:

$$ \frac{\partial z}{\partial x}=y+1=3, \qquad \frac{\partial z}{\partial y}=x=3 $$

这个例子同时说明了两件事:算子只计算局部梯度,训练引擎负责按照计算图组合并累加它们。

Zero 在 Tensor 中使用 Back 保存这条连接:

struct Back {
    std::unique_ptr<TensorOp> op;
    std::vector<Tensor> inputs;
};

op 是产生当前结果的算子,inputs 是该算子的输入。Back 被当前逻辑张量的所有副本共享,算子的生命周期则由 unique_ptr 明确管理。

实现

动态建图

每次运算都会创建一个新的算子实例,再通过 TensorOp::Apply 执行:

Tensor Tensor::operator*(const Tensor& other) const {
    return RunOp<OpMul>({*this, other});
}

Apply 先计算结果,然后检查输入是否需要梯度:

Tensor result = op->OpValue_(inputs);

bool need_grad = std::any_of(
    inputs.begin(), inputs.end(),
    [](const Tensor& t) { return t.RequiresGrad(); });

if (need_grad) {
    result.RequiresGrad(true);
    result.back_->op = std::move(op);
    result.back_->inputs = std::move(inputs);
}

只有至少一个输入需要梯度时,结果才会记录计算图。推理计算和普通数据处理因此不会携带不必要的反向状态。优化器更新参数时使用 RunOpNoGrad,也是为了防止更新过程自身进入计算图。

反向传播

一个节点的梯度可能来自多条路径。在执行该节点的反向函数之前,必须先收集完所有下游贡献。Zero 首先从损失张量出发进行深度优先遍历,得到拓扑序,再逆序执行每个算子的 Backward

对上面的例子,前向的依赖顺序是:

x, y → Mul → Add → z

反向传播则按照相反顺序处理:

z → Add → Mul → x, y

拓扑排序的意义不只是“倒着遍历”,而是保证处理某个节点时,它的所有下游梯度已经准备完成。

根节点的起始梯度为全 1 张量,对应:

$$ \frac{\partial L}{\partial L}=1 $$

反向遍历的核心过程可以概括为:

for (auto it = topo.rbegin(); it != topo.rend(); ++it) {
    Tensor* current = *it;
    if (!current->back_ || !current->back_->op) continue;

    Tensor grad = RowMajorGradOf(*current);
    auto input_grads = current->back_->op->Backward(
        grad, current->back_->inputs);

    for (size_t i = 0; i < current->back_->inputs.size(); ++i) {
        auto& input = current->back_->inputs[i];
        if (input.RequiresGrad()) {
            AccumulateGrad(input, input_grads[i]);
        }
    }
}

第一次写入某个梯度槽时,AccumulateGrad 直接接管算子产生的存储;如果该梯度已经存在,则就地累加。这既满足多路径的链式法则,又避免了每次累加都分配新张量。

张量身份

Back::inputs 保存的是 Tensor 值副本。如果单纯使用对象地址识别计算图节点,同一个逻辑张量的不同副本就会被当成多个节点,导致上游梯度重复计算。

Zero 为每个逻辑张量设置一个共享的梯度槽:

struct GradHolder {
    std::shared_ptr<TensorNode> grad_node;
};

普通 Tensor 副本共享 GradHolder,因此反向传播通过图中副本写入的梯度,用户持有的原张量也能立即看到。拓扑排序也以 GradHolder 的身份去重,而不是使用 Tensor* 或底层数据地址。

视图算子是一个特例。TransposeView 可以共享数据存储,但形状和步长已经改变,因此必须使用独立的 GradHolder。张量的存储共享与梯度共享并不是同一件事,这一点会在下一篇详细讨论。

示例

有了动态计算图和反向传播,最小训练循环已经可以表达:

auto w = Parameter(Tensor::From<float>({0.1f}, {1}));
auto b = Parameter(Tensor::From<float>({0.0f}, {1}));

for (int step = 0; step < 200; ++step) {
    w.ZeroGrad();
    b.ZeroGrad();

    auto pred = w * x + b;
    auto diff = pred - y;
    auto loss = diff * diff;
    loss.Backward();

    w.Update(w.Grad() * learning_rate);
    b.Update(b.Grad() * learning_rate);
}

前向表达式在计算数值的同时建立计算图,Backward() 将损失梯度传回 wb,最后再由优化器或参数更新函数修改数值。这条链路是训练引擎的核心,后续各篇将在它之上逐步扩展数据规模、模型结构和执行设备。