提到雅可比矩阵、JVP、VJP这三个词很多搞深度学习的人第一反应是“论文里见过公式能抄但真要写代码就分不清哪个是哪个”。尤其是想自己写自定义算子或者做高阶求导的时候经常把输入方向向量和输出方向向量传反或者说不清楚为什么反向传播只跑一次就能拿到全部梯度。这篇文章就把这件事彻底讲透雅可比矩阵到底是什么JVP和VJP各自的几何含义以及在主流自动微分框架里怎么写、怎么验证、怎么避坑。适合刚接触自动微分原理、想理解反向传播本质、或者要手写可微函数的读者。1. 先把雅可比矩阵和两种向量积的定义对齐1.1 雅可比矩阵只是“多元导数的排列”一个输入为向量、输出也为向量的函数它的导数不是单个数字而是一整张偏导表。假设输入是 (x \in \mathbb{R}^N)输出是 (y \in \mathbb{R}^M)那雅可比矩阵 (J) 就是 (M \times N) 的矩阵第 (i) 行第 (j) 列为[ J_{ij} \frac{\partial y_i}{\partial x_j} ]举一个具体的例子函数[ f(x_0, x_1) (x_0^2 x_1,\ \sin(x_0),\ x_0 x_1) ]输入维度 (N2)输出维度 (M3)那么雅可比矩阵就是[ J \begin{bmatrix} 2x_0 1 \ \cos x_0 0 \ x_1 x_0 \end{bmatrix} ]这个矩阵把输入空间里的一个小扰动“线性映射”到输出空间的变化量是一阶近似的最核心载体。理解这一点后面所有内容都顺了。1.2 JVP给输入一个方向看输出往哪走JVP 的全称是 Jacobian-vector product从名字顺序就能看出来是雅可比矩阵乘一个向量。数学表达是[ \text{JVP} J \cdot v ]这里 (v) 是输入空间里的一个方向向量形状和 (x) 一样是 (N) 维。结果是一个 (M) 维向量含义是输入沿着 (v) 方向变化时输出的一阶变化率。直觉上可以理解成“顺着推”。它在正向模式自动微分forward mode里出现一次计算就能同时得到函数值和某个方向上的导数。注意这里矩阵和向量相乘不要求把完整矩阵显式构造出来而是通过一次前向传播就能算出来。1.3 VJP给输出一个权重反推输入空间的梯度VJP 的全称是 vector-Jacobian product是向量和雅可比矩阵相乘但向量在左边[ \text{VJP} u^T \cdot J ]这里 (u) 是输出空间里的一个向量形状和 (y) 一样是 (M) 维。数学上结果是一个 (N) 维行向量实际代码中通常把它看成列向量也就是 (J^T u)。含义是把输出空间的每个分量按 (u) 加权后反向传播回输入空间得到输入空间的梯度。反向模式自动微分reverse mode的核心就是一组嵌套的 VJP神经网络的 backprop 每一层都在做这个运算。直觉上可以理解成“逆着拽”。1.4 符号方向最容易翻车很多资料里这个方向极其容易搞混我建议直接记住下面这张表名称需要的方向向量向量形状结果形状对应模式JVP(Jv)输入侧方向 (v)(N) 维(M) 维正向模式VJP(u^T J)输出侧方向 (u)(M) 维(N) 维反向模式提示你只需要记住一句话——JVP 的输入向量对着输入维度VJP 的输入向量对着输出维度。反过来基本必错。2. 从定义到效率为什么深度学习只认 VJP2.1 链式法则的两种行进方向考虑复合函数 (y f(g(x)))雅可比矩阵满足链式法则[ J J_f \cdot J_g ]正向模式从 (x) 出发先算 (J_g v)再算 (J_f (J_g v))一路从输入往输出推。反向模式从输出侧出发先拿一个方向向量 (u)算 (u^T J_f)再把这个结果继续往上游传算 ((u^T J_f) J_g)一路从输出往输入收。这两种方向本质都是链式法则只是“从哪头开始”。深度学习的损失函数是个标量输出维度 (M1)这时候反向模式只需要给输出一个 1 维的初始向量也就是常数 1然后一路 VJP 回输入一次反向传播就够了。如果换成正向模式想得到这个标量输出对每个输入的梯度就要对每个输入维度分别做一次 JVP等于要跑 (N) 次前向传播。输入维度动辄几十万、几百万这个差距是毁灭性的。2.2 一次 JVP 和一次 VJP 的成本单次 JVP 的成本大约等价于一次前向传播单次 VJP 大约等价于一次反向传播。这里的“大约”是因为除了计算函数值之外还要额外携带一个切向量或者伴随向量穿过整个计算图所以实际开销大概是普通前向的 1.5 到 2 倍。但核心在于它不依赖层的宽度做平方复杂度展开因为你始终在做“矩阵乘向量”而不是“矩阵乘矩阵”。反过来如果显式构造每一层的完整雅可比矩阵假设网络某层从 512 维映射到 512 维这个矩阵就有 26 万个元素再深几层光是存这些矩阵就能把显存吃掉大半。而 JVP 和 VJP 始终只维护 512 维的向量这就是它们能在深度学习里全面取代显式雅可比的原因。2.3 完整雅可比该用多少轮循环有时你就是想要完整雅可比矩阵比如某个分析工具需要输出输入对输出的全部偏导。这里有一个非常实用的选型规则用 JVP 拿完整矩阵需要循环 (N) 次每次喂一个第 (i) 个输入维度的单位向量 (e_i)因为 (J e_i) 正好是 (J) 的第 (i) 列。用 VJP 拿完整矩阵需要循环 (M) 次每次喂一个第 (j) 个输出维度的单位向量 (e_j)因为 (e_j^T J) 正好是 (J) 的第 (j) 行。所以当输入维度小于输出维度时用 JVP 更省当输出维度小于输入维度时用 VJP 更省。深度学习中输出是标量也就是 (M1)所以 VJP 只需要一轮这正好对应反向传播一次出梯度。3. 实操用自动微分框架计算 JVP 和 VJP3.1 先准备一个手算能验证的函数下面统一用这个函数做实验方便随时手算对照[ f(x) \begin{bmatrix} x_0^2 x_1 \ \sin x_0 \ x_0 x_1 \end{bmatrix} ]对应雅可比矩阵[ J \begin{bmatrix} 2x_0 1 \ \cos x_0 0 \ x_1 x_0 \end{bmatrix} ]取 (x (1.0, 2.0))做一次加法后四个数字都比较干净(f(x) (3.0, 0.8415, 2.0))(J [[2, 1], [0.5403, 0], [2, 1]])这套数值后面会反复出现作为验证基准。3.2 PyTorch 里直接调用 jvp 和 vjp在 PyTorch 中torch.autograd.functional包提供了现成的接口。对于单个 Tensor 输入和单个 Tensor 输出的函数用法非常直接import torch def f(x: torch.Tensor) - torch.Tensor: return torch.stack([ x[0] ** 2 x[1], torch.sin(x[0]), x[0] * x[1], ]) x torch.tensor([1.0, 2.0]) v torch.tensor([0.5, -1.0]) # 输入侧方向形状与 x 一致 # JVP: 返回 (函数输出, J v) y, jv torch.autograd.functional.jvp(f, x, v) print(y:, y) print(Jv:, jv) u torch.tensor([1.0, 0.0, 2.0]) # 输出侧方向形状与 y 一致 # VJP: 返回 (函数输出, u^T J) y, vtj torch.autograd.functional.vjp(f, x, u) print(vjp:, vtj)手算验证一下。当 (x (1,2)) 时(v (0.5, -1))[ Jv \begin{bmatrix} 2 1 \ 0.5403 0 \ 2 1 \end{bmatrix} \begin{bmatrix} 0.5 \ -1 \end{bmatrix}\begin{bmatrix} 0 \ 0.27015 \ 0 \end{bmatrix} ]而 (u (1, 0, 2)) 时[ u^T J \begin{bmatrix} 1 0 2 \end{bmatrix} \begin{bmatrix} 2 1 \ 0.5403 0 \ 2 1 \end{bmatrix}\begin{bmatrix} 6 3 \end{bmatrix} ]写代码时最容易犯的错就是把v和u的形状搞反JVP 的向量必须跟输入维度对齐VJP 的向量必须跟输出维度对齐。如果你传入的形状不对有些框架会直接报错有些框架会静默广播后果更加隐蔽。3.3 自定义反向传播的本质就是写 VJP很多人第一次接触自定义backward时一头雾水其实你写的函数就是在算 VJP。看一个输入是标量、输出是两个分量的例子class MyFn(torch.autograd.Function): staticmethod def forward(ctx, x): ctx.save_for_backward(x) return torch.stack([x * x, x * x * x]) staticmethod def backward(ctx, grad_output): x, ctx.saved_tensors # J [2x, 3x^2]^T即形状为 (2, 1) # 返回结果 J^T grad_output return grad_output[0] * 2 * x grad_output[1] * 3 * x * x这个backward返回的不是某个单一导数而是把上游传来的两个梯度加权求和。这正是 (u^T J) 的一个实例(u) 是grad_output(J) 是输出对输入的雅可比矩阵(J^T u) 就是返回的梯度。如果你接手过一个自定义算子发现 backward 里“把多个梯度加起来”那绝对不是人为设计而是 VJP 的自然结果。3.4 JAX 风格里同时看两种接口在 JAX 中JVP 用jax.jvpVJP 用jax.vjp语义和 PyTorch 完全一致但 API 风格略有不同import jax import jax.numpy as jnp def f_jax(x): return jnp.array([ x[0] ** 2 x[1], jnp.sin(x[0]), x[0] * x[1], ]) x jnp.array([1.0, 2.0]) tangent jnp.array([0.5, -1.0]) # JVP: primals_out 是 f(x)tangents_out 是 J tangent primals_out, tangents_out jax.jvp(f_jax, (x,), (tangent,)) # VJP: 先得到 vjp 函数再用输出侧向量调用 y, vjp_fn jax.vjp(f_jax, x) cotangent jnp.array([1.0, 0.0, 2.0]) x_grad, vjp_fn(cotangent) # 注意返回的是元组JAX 里vjp_fn返回一个元组即使只有一个输入也得用逗号解包。这个细节非常容易漏漏了之后拿到的就是(array,)的元组而不是数组本身后面做运算会报各种奇怪的 shape 错误。4. 完整雅可比矩阵的高效生成与数值校验4.1 用单位向量扫出完整矩阵虽然框架也提供jacobian接口但建议自己写一遍循环因为能帮你彻底理解列和行的来源。下面这段代码用 JVP 扫列、用 VJP 扫行结果应该完全相同def jacobian_by_jvp(f, x, n): cols [] for i in range(n): e torch.zeros_like(x, dtypetorch.float32) e[i] 1.0 _, jv torch.autograd.functional.jvp(f, x, e) cols.append(jv) return torch.stack(cols, dim1) # 形状: (M, N) def jacobian_by_vjp(f, x, m): rows [] for i in range(m): e torch.zeros(m, dtypetorch.float32) e[i] 1.0 _, vj torch.autograd.functional.vjp(f, x, e) rows.append(vj) return torch.stack(rows, dim0) # 形状: (M, N)为什么 JVP 是列、VJP 是行因为 (J e_i) 取出第 (i) 列(e_j^T J) 取出第 (j) 行。这个规则比任何口诀都可靠直接来自矩阵乘法定义。4.2 手算数值对照实验沿用 (x(1,2))用下面的脚本分别算两种方式得到的完整雅可比x torch.tensor([1.0, 2.0]) J1 jacobian_by_jvp(f, x, n2) J2 jacobian_by_vjp(f, x, m3) print(J1) # tensor([[2.0000, 1.0000], # [0.5403, 0.0000], # [2.0000, 1.0000]]) print(torch.max(torch.abs(J1 - J2))) # 应该是 0对比手算结果完全一致。这一步的意义在于验证你自己写的自定义算子的雅可比方向是否正确。以后写新算子时先用 JVP 和 VJP 分别扫一遍再和有限差分对照基本能堵住 90% 的梯度 bug。4.3 有限差分交叉验证有限差分虽然精度有限但谁都能看懂适合做快速 sanity check。中心差分公式def numerical_jvp(f, x, v, eps1e-6): with torch.no_grad(): f_plus f(x eps * v) f_minus f(x - eps * v) return (f_plus - f_minus) / (2 * eps)将numerical_jvp的结果和torch.autograd.functional.jvp的结果比较v torch.tensor([0.5, -1.0]) _, jv_analytic torch.autograd.functional.jvp(f, x, v) jv_numerical numerical_jvp(f, x, v) print(torch.norm(jv_analytic - jv_numerical))在我这套参数下误差大约在 (10^{-6}) 量级。如果误差到 (10^{-2}) 级别先怀疑函数写错了如果报 NaN可能是eps太小导致舍入误差如果误差只出现在某些分量上多半是某个维度的方向向量没有配对。注意中心差分对梯度验证已经很够用别用前向差分。前向差分精度差一截而且对振荡函数特别不稳定。统一用 (\epsilon 10^{-6}) 作为起步值就够了。4.4 别急着调用现成的 jacobian 接口像torch.autograd.functional.jacobian这种接口确实是方便但它内部会把你输出的完整矩阵拼起来占用的内存是 (M \times N)。如果你要分析一个 1024 维输入、1024 维输出的模块显式构造出来的雅可比矩阵就是 100 多万个元素。而如果你的算法只用到 (Jv) 或者 (J^T u)完全没必要把整张矩阵摊开。判断标准就一句话你的下游逻辑要不要逐元素访问整个矩阵不要的话就坚持用向量积形式。5. 常见事故现场与排查思路5.1 方向向量传反这是最经典的事故打算算 (Jv)结果传了一个跟输出维度一样长的向量进去。框架大概率会抛一个 shape mismatch 的异常但有些场景下因为广播机制反而不会报错只是结果莫名其妙。排查办法很简单先看一眼v.shape和x.shape是否一致再看一眼 VJP 那边的向量和y.shape是否一致。强类型习惯能避免绝大多数此类问题。5.2 高阶导计算时忘了 create_graph用autograd.functional.jvp、vjp计算高阶导数时默认是不建计算图的。很多人第一次算 Hessian-vector product 时在第一次求导后想要再对结果求导结果发现梯度全部消失。参考下面正确写法grads torch.autograd.grad(f(x), x, create_graphTrue)[0] hv torch.autograd.grad((grads * v).sum(), x)[0]注意第一次grad必须带create_graphTrue否则第二次求导拿不到计算路径。这个问题极其隐蔽因为代码不会报错只会静默返回None或者常数。5.3 自定义 backward 返回结构不匹配自定义Function时backward返回的梯度数量必须和forward的输入数量一致。如果某个输入不需要梯度要返回None而不是返回一个零张量。返回零张量虽然数值上是对的但会切断梯度流甚至导致不必要的显存占用。实践中我习惯把backward的组织方式和forward的参数列表一一对齐先在纸上写出每个输入的形状和是否参与梯度计算再写代码。别着急这个环节写错是最浪费调试时间的。5.4 JAX 元组解包遗漏前面已经提过jax.vjp返回的vjp_fn永远返回元组哪怕只有一个输入。类似地jax.jvp如果函数返回标量或数组就没问题但如果返回多个输出也要按结构解包。看到tuple类型报错时第一反应应该是在赋值那行补一个逗号。5.5 有限差分验证中把方向和数值差分混淆验证某个自定义算子的 JVP 时有人会写(f(x eps) - f(x)) / eps然后拿它和 (Jv) 比较发现完全对不上。这很正常因为 (v) 不一定等于单位向量。数值差分必须带着方向向量走也就是f(x eps * v) - f(x - eps * v)而不是沿着坐标轴逐分量猜。6. 进阶扩展HVP、牛顿法、隐函数与更多场景6.1 黑塞矩阵向量积把 JVP 和 VJP 叠加一次黑塞矩阵是梯度向量的雅可比矩阵。对于标量函数 (f(x))黑塞矩阵向量积 (Hv) 可以通过两次自动微分得到[ Hv \frac{\partial}{\partial x} \left( \nabla f(x)^T v \right) ]在 PyTorch 里手写def hvp_manual(f, x, v): grads torch.autograd.grad(f(x), x, create_graphTrue)[0] return torch.autograd.grad((grads * v).sum(), x)[0] x torch.tensor([1.0, 2.0], requires_gradTrue) v torch.tensor([0.5, -1.0]) print(hvp_manual(lambda z: (z ** 3).sum(), x, v)) # 理论值: H diag(6x)所以 Hv [3, -12]这里的技巧是把内积 ((grads \cdot v)) 当成一个新的标量函数再求一次梯度。整个过程没有显式构造黑塞矩阵只用了雅可比向量积。这对二阶优化、正则化惩罚项、扰动分析非常有用。6.2 大规模非线性方程组的牛顿法解非线性方程组 (F(x)0) 时牛顿法每一步需要解线性系统[ J(x_k) \cdot \delta -F(x_k) ]当 (J) 很大时不会有人先把它完整构造出来再求逆而是用 GMRES 这类迭代方法每次迭代只需要计算 (J \cdot \delta)。这就是 JVP 的经典应用场景。换句话说JVP 不仅仅是自动微分里的概念更是大规模数值线性代数的基石。只要算法里出现“矩阵乘以向量”并且这个矩阵是雅可比或者黑塞JVP 和 VJP 就是唯一的可行实现方式。6.3 神经 ODE、可微渲染与模型解释性带连续时间动力系统的模型里伴随方法的反向传播需要求解伴随方程本质是连续版本的 VJP。可微渲染中给相机参数一个扰动想知道图像每个像素亮度怎么变对应 JVP。算输入对 logit 的显著性图时虽然方法名五花八门底层动作仍然是 (J^T u)。我个人在实际操作中的体会是判断该用 JVP 还是 VJP永远先看你想“往哪个方向传”。从参数出发推到输出用 JVP从输出损失往回传到参数用 VJP。把维度写清楚、把形状标好再去碰框架 API能少踩至少一半的坑。这也是我做自定义算子和高阶导数相关开发时最重要的一个习惯。