(五点五) 自动微分 2.0
迟到的承诺
欢迎回到零基础深度学习之旅。上一篇文章的结尾我承诺过一篇”自动微分 2.0”，讲如何通过张量级操作提升性能，其实代码早就写好了，只是一直没有写文章，现在终于兑现了。
先复习上一篇的内容：（五）里的标量版自动微分实现叫 System A，本文的张量版叫 System B（与《The Little Learner》附录对应）。 System A中每个标量被包装成一个 Dual，每次基本运算都会创建一个新的 Dual，并用一个闭包（link 函数）记住”梯度该如何向上游传播”。整个函数调用链就是计算图。
本文实现了两个版本的System B，一个坚持零依赖，一个交给 NumPy。两个版本共享同一个核心思想，把它们放在一起看，能更清楚地看到性能提升到底来自哪里。
System A 慢在哪里
问题不在 Python 本身，而在计算图的粒度。
System A 的图节点是标量。假设参数是一个 30×30 的权重矩阵，做一次矩阵-向量乘法：每个输出元素需要 30 次乘法和 29 次加法，30 个输出元素就是约 1800 个基本运算，每个运算都创建一个 Dual 对象和一个闭包。一次前向传播创建的 Python 对象数量 = 数值个数 × 运算次数，而训练循环的每一步都要重建这张图。
也就是说，就算每个闭包本身再快，数量也会压垮我们。
关键一步：把节点从标量提升到张量
两个 System B 的全部要点是同一句话：Dual 的实部不再是一个数，而是一整个张量。
一次矩阵-向量乘法，System A 要建 1800 个图节点，System B 只建 1 个。图的大小从”数值个数 × 运算次数”变成了”运算次数”，与张量多大无关。
但这一步立刻带来一个新问题。System A 里，ext1/ext2 只需要扩展前向计算，而梯度是由标量层面的 link 闭包自动携带的。现在节点是张量，前向在张量层面走，反向也必须在张量层面走：每个张量级运算都需要知道”整个张量的梯度怎么算”。前向和反向，需要被分别扩展。
两个 System B 对这个问题给出了两种回答。
System B 之一：纯 Python 版
第一个版本坚持零依赖：张量还是嵌套列表，元素级计算还是 Python 循环，只是减少图节点。
它的核心技巧是让每个基本运算同时携带两个函数，一个前向（rho）和一个反向（nabla），并且可以被”问出来”：
def prim(rho_fn: Callable, nabla_fn: Callable) -> Callable:
def primitive(*args):
if len(args) == 1:
arg = args[0]
if arg is rho_function: # 问它：你的前向是什么？
return rho_fn
elif arg is nabla_function: # 问它：你的反向是什么？
return nabla_fn
else:
return prim1_dual(rho_fn, nabla_fn, arg) # 正常计算
elif len(args) == 2:
return prim2_dual(rho_fn, nabla_fn, args[0], args[1])
return primitive
add_00 = prim(lambda ra, rb: ra + rb, lambda ra, rb, z: (z, z))
mul_00 = prim(lambda ra, rb: ra * rb, lambda ra, rb, z: (rb * z, ra * z))
rho_function 和 nabla_function 是两个哨兵值：把它们当参数传给一个运算，运算就交出自己的前向或反向部分。这样 ext2 就可以把一个标量运算拆开，前向和反向各自扩展到张量层面，再组装回一个新的运算：
def ext2(f: Callable, n: int, m: int) -> Callable:
rho_fn = rho_function(f) # 拆出前向
nabla_fn = nabla_function(f) # 拆出反向
extended_rho = ext2_rho(rho_fn, n, m) # 前向：和 System A 的 ext2 一样递归下降
extended_nabla = ext2_nabla(nabla_fn, n, m) # 反向：结构相同的递归下降
return prim(extended_rho, extended_nabla) # 组装回一个"可再扩展"的运算
最微妙的部分在反向的递归下降里。前向广播时，一个低阶张量会和高阶张量的每个元素分别运算，也就是说它参与了多条路径。（五）里讲过多路径法则：多条路径的梯度要相加。纯 Python 版把这条规则写得明明白白：
def desc_nabla_u(g, n, t, m, u, z):
zs = z if not is_scalar(z) else [z] * len(u)
results = [g(t, ui, zi) for ui, zi in zip(u, zs)]
gt = reduce(add_rho, [gi for gi, _ in results]) # ← t 被广播了，梯度求和！
gu = [ui for _, ui in results]
return gt, gu
那句 reduce(add_rho, ...) 就是整个 System B 最重要的一行代码：前向传播中的广播（复制），对应反向传播中的求和。一个被广播出去的偏置 b，它的梯度是它参与的每条路径的梯度之和。
用这个引擎拟合平面（真实参数 w=[3,2]、b=1），训练后得到 [3.0000164, 2.0000164], 1.0000164。
System B 之二：NumPy 版，把循环交给 C
第二个版本在同一个思想上再进一步：既然节点已经是张量了，元素级的计算为什么还要用 Python 循环？交给 NumPy。
def make_primitive(forward_fn, backward_fn):
def primitive(*args):
input_arrays = [get_r(arg).data for arg in args]
result = FlatTensor(forward_fn(*input_arrays)) # 前向：一次 NumPy 调用
if not any(is_dual(arg) for arg in args):
return result # 不求导就不建图
def link(d, upstream_grad, sigma):
input_grads = backward_fn(*input_arrays, upstream_grad.data)
for i, arg in enumerate(args):
if is_dual(arg):
grad = input_grads[i]
if grad.shape != arg.r.shape: # 前向广播过？
grad = sum_to_shape(grad, arg.r.shape) # 反向就求和
get_k(arg)(arg, FlatTensor(grad), sigma)
return Dual(result, link)
return primitive
tadd = make_primitive(lambda a, b: a + b, lambda a, b, grad: (grad, grad))
tmul = make_primitive(lambda a, b: a * b, lambda a, b, grad: (grad * b, grad * a))
结构和纯 Python 版一模一样：先算值，再定义 link。区别只是嵌套列表换成了 np.ndarray，递归下降换成了 NumPy 的向量化运算。
而纯 Python 版里那句 reduce(add_rho, ...)，在这里变成了 sum_to_shape，同一条”广播的反向是求和”规则，从结构递归换成了按形状求和：
def sum_to_shape(grad, target_shape):
"""把梯度求和回参数原本的形状"""
ndims_added = grad.ndim - len(target_shape)
if ndims_added > 0: # 多出来的前导维度
grad = np.sum(grad, axis=tuple(range(ndims_added)))
for i, (g_dim, t_dim) in enumerate(zip(grad.shape, target_shape)):
if t_dim == 1 and g_dim > 1: # 被广播的维度（1 → n）
grad = np.sum(grad, axis=i, keepdims=True)
return grad.reshape(target_shape)
两个实现、同一条规则的两种写法：如果只看其中一个，很容易把它当成实现细节记住就算了；两个放在一起，才看出这是自动微分里的一条定律。顺带一提，求和的反向恰好是广播（tsum 的梯度把上游梯度复制给每个输入元素）：广播与求和，在前向和反向里互为镜像。
NumPy 版还多了一个真正的矩阵乘法，它的梯度规则值得单独记住：
matmul = make_primitive(
lambda x, y: np.matmul(x, y),
lambda x, y, grad: (
np.matmul(grad, np.swapaxes(y, -1, -2)), # ∂L/∂x = grad @ yᵀ
np.matmul(np.swapaxes(x, -1, -2), grad), # ∂L/∂y = xᵀ @ grad
),
)
不想背公式可以用形状倒推：x 是 (m,k)，y 是 (k,n)，grad 是 (m,n)；x 的梯度必须是 (m,k)，用 grad 和 y 拼出 (m,k) 的唯一方式就是 grad @ yᵀ。形状对了，公式基本就对了，这是（五）里”梯度形状永远等于参数形状”规则的又一次胜利。
性能对比：两次提速，各自来自哪里
三个引擎跑完全相同的任务（可以自己跑tiny-learner 仓库 examples/04_engines_benchmark.py）：
| 引擎 | 小任务：平面拟合（2 个参数），Adam 1000 步 | 大任务：拟合 30×30 线性映射、20 个样本，Adam 100 步 |
|---|---|---|
| System A（标量节点） | 0.24 s | 13.9 s |
| System B（张量节点，纯 Python） | 0.37 s | 8.7 s |
| System B（NumPy） | 0.35 s | 0.03 s |
这张表比我动手之前预想的诚实得多，有三个值得停下来看的地方：
- 小任务上 System A 反而最快。 张量级节点有自己的管理开销；参数只有两三个标量时，”少建图”省下的还不够付出的。优化是针对规模的，优化是为规模准备的，规模不够，省下的还不够付开销的。
- 大任务上，A → B（纯 Python）提速约 1.6×。 这就是”图粒度”这个思想单独的价值：不靠 NumPy，只靠少建节点。有了明显的提升，但是提升幅度不够大，因为元素级计算仍然是 Python 循环，而它才是开销的大头。
- B → B（NumPy）提速约 300×。 数量级的差距来自向量化：同样的图结构，把逐元素循环交给 C。
所以”深度学习框架为什么快”的完整答案是两层：张量级的计算图（结构上的），加上向量化的底层运算（实现上的）。两层缺一不可，但贡献极不成比例，这正是把两个 System B 分开实现才能测出来的东西。
优化器层：一行都不用改
最让我满意的部分：（四）里写的整个优化器体系——gradient_descent 构建器、动量、RMSProp、Adam——在三个引擎上原封不动地工作。
def velocity_u(vp, g, h):
v = tsub(tmul(h.momentum, vp.velocity), tmul(h.lr, g))
p = tadd(vp.parameter, v)
return VelocityP(p, v)
这段代码和（四）里的一模一样。它当初操作的是标量级 Dual 的嵌套列表，现在底下可能是张量级 Dual、也可能是 NumPy 数组，但它只通过 tadd/tmul/tsub 这些接口说话，底层表示换掉了它也毫不知情。当初坚持统一的张量运算接口、把 inflate/deflate/update 分离出来，回报在这一刻兑现。
在配套的 tiny-learner 仓库里，三个引擎被做成了可互换的组件（tiny_learner/engines/），同一个优化器工厂喂进任何一个引擎都能训练。想亲手验证本文的每一个说法，clone 下来跑 examples/04_engines_benchmark.py 即可。
总结
System B 用一个动作解决了 System A 的性能问题：把计算图的节点从标量提升到张量。由此得到的知识点：
- 图的粒度决定开销——节点数从”数值个数 × 运算数”降到”运算数”，这一步不需要 NumPy 就能兑现；
- 广播的反向是求和，求和的反向是广播——纯 Python 版的
reduce(add_rho)和 NumPy 版的sum_to_shape是同一条定律的两种写法；- 前向与反向需要分别扩展——rho/nabla 拆分是张量级自动微分的结构性要求；
- matmul 的梯度规则——grad @ yᵀ 与 xᵀ @ grad，用形状就能推出来。
而损失函数和优化器完全没有改动，这是对我们接口设计的最好检验。
下一篇，我们用这套引擎兑现整个系列最初的承诺：训练一个真正的神经网络。