📚 课程目录第3章自动微分原理解析

3. 自动微分原理解析

要真正用好 JAX 的梯度能力，不能只停留在“调用 jax.grad 就能求导”的层面。自动微分（Automatic Differentiation, AD）既不是数值差分，也不是符号推导，而是沿着程序实际执行的计算过程，逐条应用链式法则。理解这一点，才能解释为什么 JAX 的梯度函数要求纯函数、为什么 grad 默认是反向模式、以及为什么 jax.jvp 和 jax.vjp 会以不同形态出现。

一、计算图与追踪：JAX 如何“看见”你的函数

JAX 的自动微分建立在追踪（tracing）机制之上。当你把一个函数交给 jax.grad、jax.jit 或 jax.vmap 时，JAX 并不会立刻用真实数值执行它，而是传入抽象值（Tracer）作为参数。函数体照常运行，但每个运算不再产生具体数字，而是记录“这个输出是由哪些输入、经过什么原始操作得到的”，从而构建出一张计算图（jaxpr）。

举例来说，函数 f(x) = sin(x) * x + 2 被追踪时，JAX 会依次记录 sin、乘法、加法三个原始操作及其依赖关系。这张图是后续求导、编译、向量化的共同基础。正因为追踪依赖 Python 控制流在追踪时“走哪条路”，所以含有数据依赖分支（如 if x > 0）的函数在 jit 下会出问题——追踪只看到其中一条分支。这也是 JAX 强调纯函数与静态控制流（jax.lax.cond、lax.scan）的根本原因。

二、前向模式与反向模式：两种链式法则的走法

设复合函数 y = f(g(h(x)))。链式法则给出 dy/dx = f'·g'·h'。两种模式的差别在于从哪一端开始乘。

前向模式（Forward Mode）把输入端的“扰动”一路向前传播：给 x 一个单位切向量，计算 h'(x)，再乘 g'、再乘 f'，一次前向遍历即得一个输入方向对输出的导数。它的计算量与输入维度成正比，适合输入少、输出多的场景，例如对单个标量参数求整个向量输出的导数。

反向模式（Reverse Mode）则先做一次前向计算并保存中间值，再从输出端反向传播“伴随值（adjoint）”：从 dy/dy = 1 出发，依次乘 f'、g'、h'，一次反向遍历即得所有输入对一个输出的梯度。它的计算量与输出维度成正比，适合输入多、输出少的场景——而深度学习损失函数恰恰是“百万参数 → 一个标量损失”，因此反向模式成为训练的主力。

记忆要点：前向模式“一次传播一个输入方向”，反向模式“一次传播一个输出方向”。参数多、损失标量的神经网络，天然属于反向模式的甜区。

三、JAX 的 JVP 与 VJP 变换

JAX 把两种模式暴露为两个基础变换，它们都是函数到函数的变换，而非直接返回数值。

jax.jvp(f, primals, tangents) 对应前向模式。给定函数 f、输入点 primals 和切向量 tangents，返回 (f(x), J·v)，即函数值与雅可比矩阵乘切向量的结果。它一次给出一个方向的方向导数。

jax.vjp(f, *primals) 对应反向模式。它返回 (f(x), vjp_fn)，其中 vjp_fn 是一个闭包：你传入输出端的伴随向量 w，它返回 wᵀ·J，即向量转置乘雅可比。注意 VJP 是“先算函数值、再按需反向”的两段式结构，这正是反向模式需要保存中间值的体现。

import jax
import jax.numpy as jnp

def f(x):
    return jnp.sin(x) * x + 2.0

x = 1.5
# 前向模式：给定切向量 v=1.0，得到方向导数
y, y_dot = jax.jvp(f, (x,), (1.0,))
print(y, y_dot)   # 函数值 与 f'(x)

# 反向模式：先得函数值与 vjp 闭包，再传入伴随 w=1.0
y2, vjp_fn = jax.vjp(f, x)
(grad_x,) = vjp_fn(1.0)
print(y2, grad_x) # 与 y_dot 数值一致

# jax.grad 本质是 vjp 的便捷封装（标量输出时）
print(jax.grad(f)(x))

上面示例中 jvp 与 vjp 得到的导数值应当一致（示例参数 x=1.5 为教学取值，实际运行结果以环境为准）。这揭示了一个重要事实：jax.grad 默认就是反向模式，它内部借助 VJP 实现；当函数有多个输出时，grad 要求你通过 argnums 与 has_aux 明确指定标量目标，否则需改用 jax.jacrev / jax.jacfwd。

四、如何选择：JVP 还是 VJP

选择依据是雅可比矩阵的形状。设 f: Rⁿ → Rᵐ，雅可比为 m×n。前向模式每次算一列（一个输入方向），需 n 次得完整雅可比；反向模式每次算一行（一个输出方向），需 m 次。因此：

n 小 m 大（少输入、多输出）：用 jvp / jacfwd 更省。
n 大 m 小（多输入、少输出，如损失函数）：用 vjp / jacrev 更省。
需要完整雅可比且维度都不大时，jax.jacfwd 与 jax.jacrev 结果相同，仅性能路径不同。

常见误区：认为“反向模式一定更快”。它只在输出维度小于输入维度时占优；对一个 R¹ → R¹⁰⁰ 的函数，前向模式反而更划算。选型要看形状，不要背结论。

五、追踪机制带来的实践约束

由于求导依赖追踪，函数必须是可追踪的纯函数：不能有副作用（如原地修改外部数组）、不能依赖 Python 侧的可变状态、数据依赖的 Python 分支需改写为 lax.cond。此外，反向模式需要保存前向中间值，因此内存开销随计算深度增长；对超深网络可用 jax.checkpoint（即 remat）以重算换内存。理解“追踪 → 计算图 → 链式法则遍历”这条主线，就能预判大多数 JAX 求导报错的根源，并在 JVP 与 VJP 之间做出有依据的选择。