4. 自动微分原理与 grad 实战

自动微分（Automatic Differentiation，AD）既不是数值差分，也不是符号推导，而是把程序拆成一连串基本运算，再对每个基本运算套用已知的求导规则，最后按链式法则把这些局部导数组合起来。JAX 的 grad 正是这一思想的直接实现：它追踪函数在追踪过程中发生的每一次加、乘、指数、三角函数等基本操作，构造出一张计算图，然后沿图反向传播梯度。理解这一点，才能解释为什么 grad 得到的是"精确到浮点误差"的导数，而不是差分近似。

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

链式法则的组合顺序有两种，对应两种自动微分模式。设复合函数 y = f(g(h(x)))，前向模式从输入出发，一边计算函数值一边携带"当前量对输入的导数"，即计算 dy/dx = (dy/dg)·(dg/dh)·(dh/dx) 时从左往右乘；反向模式先完成一次前向计算得到输出，再从输出往回逐层乘上局部导数。

关键差异在于代价。前向模式一次传播只能得到"一个输入方向"的导数，若输入有 n 维、输出有 m 维，需要 n 次前向传播才能得到完整雅可比矩阵；反向模式一次反向传播就能得到"一个输出方向"对所有输入的导数，需要 m 次。机器学习中损失函数是标量（m = 1），参数却是成千上万维（n 很大），因此反向模式一次就能拿到全部参数的梯度，这正是深度学习框架默认采用反向模式的原因。JAX 中 grad 对应反向模式，jvp 对应前向模式，vjp 则是反向模式的底层原语。

用 grad 计算标量函数梯度

最基本的用法是给 grad 传入一个返回标量的函数。下面以 f(x) = x³ + 2x² − 5x + 1 为例，其解析导数为 f′(x) = 3x² + 4x − 5。在 x = 2.0 处，解析值应为 12 + 8 − 5 = 15。

import jax
import jax.numpy as jnp

def f(x):
    return x**3 + 2*x**2 - 5*x + 1

df = jax.grad(f)
print(df(2.0))          # 15.0
print(jax.grad(f)(0.0)) # -5.0

注意 grad 返回的是一个新函数，而不是在某个点上求出的数值。这个新函数可以像普通函数一样被调用、被 jit 编译、被再次 grad 求二阶导。例如 jax.grad(jax.grad(f))(2.0) 会得到二阶导数 6x + 4 = 16。这种"函数变换"的视角是 JAX 与命令式自动微分库的重要区别：grad、jit、vmap 都是可组合的高阶函数变换。

一个常见误区是把 grad 用在返回非标量的函数上。JAX 要求目标函数输出为标量，否则会报错，因为反向模式需要一个"种子"标量来启动回传。若确实要对向量输出求导，应使用 jax.jacobian，或显式构造标量化（如与一个余切向量做内积）后用 vjp。

多输入函数与 argnums 参数

实际模型往往有多个输入，例如 loss(params, x, y)。默认情况下 grad 只对第 0 个参数求导，这正是"对参数求梯度"的常见需求。若想指定其他位置，用 argnums：

def loss(w, b, x):
    pred = w * x + b
    return jnp.sum((pred - 3.0)**2)

# 只对 w（第 0 个参数）求导
g_w = jax.grad(loss, argnums=0)

# 同时对 w 和 b 求导，返回元组
g_wb = jax.grad(loss, argnums=(0, 1))

w, b, x = 1.5, 0.5, jnp.array([1.0, 2.0, 3.0])
print(g_w(w, b, x))
gw, gb = g_wb(w, b, x)
print(gw, gb)

argnums 接受整数、元组或列表。传元组时返回值是与 argnums 顺序一致的元组，这一点在写训练循环时非常实用：grads = jax.grad(loss, argnums=0)(params, x, y) 得到的 grads 与 params 具有相同的树结构（PyTree），可以直接交给优化器更新。若 argnums 指定的参数不是浮点类型（如整型索引），JAX 会报错，因为整数不可微；此时应确认传入的是 jnp.float32 数组。

验证梯度正确性

自动微分虽然精确，但代码写错（例如误用了 jnp 之外的 NumPy 运算、在可微路径上做了 int() 截断、或用 stop_gradient 意外切断了梯度）都会导致梯度错误。工程上应养成验证习惯，最直接的方法是与中心差分对比：

def numeric_grad(fn, x, eps=1e-4):
    return (fn(x + eps) - fn(x - eps)) / (2 * eps)

x0 = 2.0
ad = jax.grad(f)(x0)
fd = numeric_grad(f, x0)
print(ad, fd, abs(ad - fd))   # 15.0  15.0000...  误差约 1e-6 量级

中心差分的截断误差为 O(eps²)，在 eps = 1e-4 时通常能给出 6 位左右有效数字，足以发现量级错误或符号错误。判断标准不是"完全相等"，而是相对误差在合理范围内（例如 1e-5 以内）。若误差明显偏大，优先检查：函数中是否混入了非 JAX 的 NumPy 调用、是否有 float() 或 int() 强制转换、是否在 jit 内部使用了 Python 控制流导致追踪失败。

对于多输入函数，验证时应对每个参数分别做差分，并注意 argnums 与差分变量的对应关系。JAX 还提供了 jax.test_util.check_grads，可对函数自动做前向、反向与数值差分的一致性检查，适合在单元测试中使用。把"解析式手算 + 中心差分 + check_grads"三者结合，基本可以覆盖绝大多数梯度实现错误。

小结本节的关键判断链：先确认函数输出是标量，再确认要对哪个输入求导（argnums），然后确认梯度结构与参数结构一致，最后用中心差分或 check_grads 做一次数值对照。这四步走完，grad 的结果才可信，也才能安全地交给后续的优化器使用。