从零实现 autograd：graph 怎么建？
backward 按什么顺序？为什么用 reverse mode？
先弄清楚：每个概念是什么
- Computational graph
- forward 时每个结果都记下自己的输入和 op，最后得到一个 DAG。leaf 是输入和 parameter，root 是 loss。例：\(L = xy + x\) 有 4 个节点：\(x\)、\(y\)、\(u = xy\)、\(L = u + x\)。
- Local gradient 与 chain rule
- 记 \(\bar v = \partial L/\partial v\)。每个 op 只需要自己输出对输入的导数，收到 \(\overline{\text{out}}\) 后乘上 local derivative 传给输入。
- Topological sort
- 让每个节点都排在它的输入之后。backward 按逆序走，一个节点收齐所有下游的 gradient 之后才往上游传。
- Fan-out 与 +=
- 一个变量被用了多次，它的 gradient 是各条路径之和，所以代码里写 +=。例：\(x = 2,\ y = 3\) 时 \(\partial L/\partial x = y + 1 = 4\)，写成 = 会被覆盖成 3。
- Reverse mode vs forward mode
- forward mode 从输入往后推，一遍只得到对一个输入的导数，\(n\) 个 parameter 要 \(n\) 遍。reverse mode 从 scalar loss 往回推，一遍得到全部 gradient，代价是 forward 的中间结果要存到 backward 用完（见 activation recomputation）。
看图：L = xy + x 的 graph 和 backward 顺序
30 秒回答
Each value stores its data, grad, and the op and inputs that produced it, so the forward pass builds a DAG. backward() topologically sorts the graph from the loss, sets its grad to 1, and visits nodes in reverse order; each node multiplies the incoming grad by its local derivative and adds it into its inputs' grads. It must be += because a value used several times gets the sum over all paths; the reverse order guarantees a node's grad is complete before it propagates. Reverse mode fits a scalar loss: one backward pass, a few times the forward cost, gives every parameter's gradient.
推导 / 关键公式
这个求和就是代码里的 +=。对 \(f:\mathbb{R}^n\to\mathbb{R}^k\)，求整个 Jacobian 的代价：
\[ \text{forward mode: } n\,c\,\mathrm{ops}(f) \] \[ \text{reverse mode: } k\,c\,\mathrm{ops}(f),\qquad c < 6 \]\(c\) 通常 2–3。loss 是 scalar，\(k = 1\)，所以 reverse mode 一遍就够。
| op | forward | backward（\(\bar y\) 是收到的 gradient） |
|---|---|---|
| add | \(y = a + b\) | \(\bar a\) += \(\bar y\)，\(\bar b\) += \(\bar y\)；被 broadcast 的维度要 sum 回原 shape |
| mul | \(y = ab\) | \(\bar a\) += \(b\,\bar y\)，\(\bar b\) += \(a\,\bar y\) |
| matmul | \(C = AB\) | \(\bar A\) += \(\bar C B^\top\)，\(\bar B\) += \(A^\top \bar C\) |
| ReLU | \(y = \max(x, 0)\) | \(\bar x\) += \(\bar y\,[x>0]\) |
| softmax + CE | \(L = -\log p_{y}\) | \(\bar z = p - \text{onehot}(y)\) |
手撕代码（scalar 版）
import math
class Value:
def __init__(self, data, parents=()):
self.data, self.grad = data, 0.0
self.parents = parents
self._backward = lambda: None # 把 out.grad 传给 parents
def __add__(self, o):
out = Value(self.data + o.data, (self, o))
def _backward():
self.grad += out.grad # +=：fan-out 时各路径相加
o.grad += out.grad
out._backward = _backward
return out
def __mul__(self, o):
out = Value(self.data * o.data, (self, o))
def _backward():
self.grad += o.data * out.grad
o.grad += self.data * out.grad
out._backward = _backward
return out
def tanh(self):
t = math.tanh(self.data)
out = Value(t, (self,))
def _backward():
self.grad += (1 - t * t) * out.grad
out._backward = _backward
return out
def backward(self):
topo, seen = [], set()
def build(v): # DFS 后序 = topological order
if v not in seen:
seen.add(v)
for p in v.parents:
build(p)
topo.append(v)
build(self)
self.grad = 1.0 # dL/dL = 1
for v in reversed(topo): # 从 loss 往输入走
v._backward()
x, y = Value(2.0), Value(3.0)
L = x * y + x
L.backward() # x.grad = 4.0, y.grad = 2.0必答点
- forward 建 DAG：每个节点存 data、grad、parents 和 local backward
- backward：topological sort，从 loss（grad = 1）逆序调用 local backward
- gradient 用 +=：同一个变量被用多次时各路径相加
- reverse mode：scalar loss 一遍 backward 拿到全部 gradient；代价是存 activation
- 用 finite difference 做 gradient check
追问
matmul、bias、softmax + CE 的 backward 怎么推？
matmul \(Y = XW\)，X 是 [n, in]，W 是 [in, out]。按元素写 \(Y_{ij} = \sum_k X_{ik} W_{kj}\)，所以 \(\partial Y_{ij}/\partial W_{kl} = X_{ik}\,[l = j]\)，\(\partial Y_{ij}/\partial X_{ab} = W_{bj}\,[a = i]\)。用 chain rule 对所有 \(Y_{ij}\) 求和：
\[ \frac{\partial L}{\partial W_{kl}} = \sum_{i,j} \frac{\partial L}{\partial Y_{ij}} \frac{\partial Y_{ij}}{\partial W_{kl}} = \sum_i X_{ik}\, dY_{il} = (X^\top dY)_{kl} \] \[ \frac{\partial L}{\partial X_{ab}} = \sum_j dY_{aj}\, W_{bj} = (dY\,W^\top)_{ab} \]bias \(Y = XW + b\)：\(db_j = \sum_i dY_{ij}\)，即 dY 按 token 求和。elementwise \(H = f(A)\)：\(dA = dH \odot f'(A)\)。
softmax + CE（合成一个 op，数值更稳）：\(L = -\log p_t = -z_t + \log\sum_k e^{z_k}\)，对 \(z_j\) 求导得 \(-[j = t] + e^{z_j}/\sum_k e^{z_k}\)，即 \(dz = p - \text{onehot}(t)\)。例：z = [2, 1, 0]、t = 0 → p ≈ [0.665, 0.245, 0.090]，dz ≈ [−0.335, 0.245, 0.090]。
怎么验证 gradient 写对了？
central difference \(\big(f(x+\epsilon) - f(x-\epsilon)\big)/2\epsilon\)，用 float64、\(\epsilon \approx 10^{-6}\)，逐个输入比 relative error（应远小于 \(10^{-6}\)）。也可以和 PyTorch 的 .grad 对比。
为什么不直接从 loss 递归调 parents 的 backward？
节点没收齐下游 gradient 就往上传，被用 \(k\) 次就传 \(k\) 次，路径一多就指数爆炸，按 topological order 每个 local backward 只跑一次。图很深时递归 DFS 会超 Python recursion limit，改用显式 stack。
PyTorch 里实际是怎么做的？
tensor 的 grad_fn 指向 graph 里的 Function 节点，graph 每个 iteration 重新建。backward 完默认释放 graph，再 backward 一次要传 retain_graph=True。只有 requires_grad=True 的 leaf tensor 把 gradient 累加进 .grad（见 gradient accumulation）。
什么时候反而用 forward mode？
输入少、输出多（\(n \ll k\)），或者只要 Jacobian-vector product 时。reverse mode 一遍给出的是 vector-Jacobian product \(J^\top r\)（从 \(\bar y = r\) 往回推）。
来源Baydin et al., Automatic differentiation in machine learning: a survey §3.1（p.10）、§3.2 + Table 3（pp.12–13）· PyTorch, Autograd mechanics（How autograd encodes the history、Saved tensors、Setting requires_grad）· torch.Tensor.backward · Karpathy, micrograd（README：scalar Value 上的 reverse-mode autodiff）