1. 从自动微分说起为什么需要JVP和VJP做机器学习或者科学计算的朋友几乎每天都要和自动微分打交道。但大多数人只熟悉整体框架比如 PyTorch 里的backward()或者 JAX 里的grad()却很少细究自动微分底层真正在算的那些向量-矩阵运算。其实无论框架怎么封装自动微分最终都跑在两类基本操作上一类叫 JVP另一类叫 VJP。它们分别对应雅可比矩阵左乘一个向量和右乘一个向量是所有前向模式、反向模式自动微分的计算基石。我最早接触这两个概念是在阅读某个框架的自定义操作扩展文档时。当时需要实现一个自定义算子要求同时提供jvp和vjp两个函数。我对着数学定义看了很久才明白框架要的到底是什么。后来在实现一个跨平台系统的自动微分模块时又亲手把这两类算子写了一遍才彻底理解它们的差异。这篇内容就围绕 JVP 和 VJP 展开不讲空泛的理论重点说清楚它们是什么为什么框架需要它们以及实际写代码时有哪些容易踩的坑。无论你是想深入理解自动微分原理还是正在写自定义算子、设计可微分程序这篇文章都能帮你少走很多弯路。先给一个最直观的定性JVP 是“给定输入扰动求输出的扰动”用于前向传播VJP 是“给定输出梯度求输入的梯度”用于反向传播。听起来简单但要真正掌握还需要理解雅可比矩阵的运算方向、维度约定以及如何处理非方阵场景。下面逐层拆解。2. 数学定义与代码落地2.1 雅可比矩阵的两种乘法假设有一个可微函数 (f: \mathbb{R}^n \rightarrow \mathbb{R}^m)输入是 (x \in \mathbb{R}^n)输出是 (y f(x) \in \mathbb{R}^m)。函数 (f) 在 (x) 处的导数是一个 (m \times n) 的矩阵称为雅可比矩阵 (J_f(x))其中第 (i) 行第 (j) 列是 (\partial y_i / \partial x_j)。有了这个矩阵接下来就可以和向量做乘法。如果有一个向量 (v \in \mathbb{R}^n)它和 (J) 的乘积是[ J \cdot v ]结果是 (\mathbb{R}^m) 的向量。这个运算就是雅可比向量积缩写为 JVP。它在计算时并不需要真的构造出完整的 (J)而是直接计算方向导数。换句话说JVP 回答的是当输入沿 (v) 方向变化一个无穷小量时输出沿哪个方向变化。反过来如果有一个向量 (u \in \mathbb{R}^m)它和 (J) 的乘积是[ u^T \cdot J ]注意这里 (u^T) 是行向量乘上 (m \times n) 的 (J)结果是 (\mathbb{R}^n) 的行向量。在编程实现里通常写成向量形式[ J^T \cdot u ]也就是矩阵的转置乘以向量得到 (\mathbb{R}^n) 的向量。这个运算叫做向量雅可比积缩写为 VJP。它回答的是给定输出侧的一个“梯度” (u)如何把梯度传回输入侧。很多初学者会混淆这两者的维度。我自己的记法是JVP 中雅可比矩阵在左边向量在右边所以它要求向量的维度等于输入维度 (n)结果的维度等于输出维度 (m)。VJP 中雅可比矩阵的转置在左边向量在右边所以向量维度等于输出维度 (m)结果维度等于输入维度 (n)。两个操作正好是“对偶”的。2.2 前向模式与反向模式JVP 是前向模式自动微分的核心。一次 JVP 计算等价于让函数做一次带方向导数的前向传播。如果你要计算 (J) 的所有元素那就需要 (n) 次 JVP分别让方向向量取单位向量 (e_1, e_2, \dots, e_n)。因为输入维度是 (n)所以前向模式计算完整雅可比矩阵的复杂度是 (O(n)) 次函数求值。VJP 是反向模式自动微分的核心。一次 VJP 计算等价于把输出侧的某个梯度向量沿着计算图反向传播。如果要计算完整的 (J)那就需要 (m) 次 VJP分别让向量取单位向量 (e_1, e_2, \dots, e_m)输出维度上的单位向量。反向模式计算完整雅可比矩阵的复杂度是 (O(m)) 次反向过程。请记住这个差别因为在实际应用中前向模式适合“输入少、输出多”的函数反向模式适合“输入多、输出少”的函数。神经网络的损失函数通常是输入是成千上万个参数、输出是一个标量即 (m1)所以反向模式 VJP 占了绝对主导。而某些科学计算场景比如求解常微分方程初值问题的敏感度可能输入维度很小但输出是一条时间序列输出维度很大这时前向模式的 JVP 就更有优势。深度学习框架的backward()本质上就是在执行一系列 VJP。每一层的反向传播函数接收从上一层传来的梯度 (u)然后计算当前的 VJP把梯度传给下一层。这就是为什么很多框架只需要你定义vjp而不需要你手动推导所有梯度公式。2.3 一个小样例手算与代码对比光说数学定义还是有点抽象我们用一个小函数来实际算一遍。假设[ f(x_1, x_2) (x_1^2 x_2, \ \sin(x_1 x_2)) ]输入是二维的输出也是二维的。雅可比矩阵是[ J \begin{bmatrix} 2x_1 1 \ x_2 \cos(x_1 x_2) x_1 \cos(x_1 x_2) \end{bmatrix} ]取 (x (2, 3))那么[ J \begin{bmatrix} 4 1 \ 3\cos 6 2\cos 6 \end{bmatrix} ]假设方向向量 (v (1, 0))那么 JVP[ J v \begin{bmatrix} 4 \cdot 1 1 \cdot 0 \ 3\cos6 \cdot 1 2\cos6 \cdot 0 \end{bmatrix} \begin{bmatrix} 4 \ 3\cos6 \end{bmatrix} ]如果取输出侧向量 (u (1, 1))那么 VJP[ J^T u \begin{bmatrix} 4 \cdot 1 3\cos6 \cdot 1 \ 1 \cdot 1 2\cos6 \cdot 1 \end{bmatrix} \begin{bmatrix} 4 3\cos6 \ 1 2\cos6 \end{bmatrix} ]看到没有JVP 和 VJP 的数值完全不同。它们虽然都利用了同一个雅可比矩阵但乘法方向不同对应的物理含义也不同。在代码层面如果用 JAX 来验证可以这样写import jax import jax.numpy as jnp def f(x): return jnp.array([x[0]**2 x[1], jnp.sin(x[0] * x[1])]) x jnp.array([2.0, 3.0]) v jnp.array([1.0, 0.0]) u jnp.array([1.0, 1.0]) # JVP 可以通过 jax.jvp 调用 y, jvp_result jax.jvp(f, (x,), (v,)) print(JVP result:, jvp_result) # VJP 可以通过 jax.vjp 调用 y_val, vjp_func jax.vjp(f, x) vjp_result vjp_func(u) print(VJP result:, vjp_result)运行结果应该与手算一致。这里jax.jvp的输入参数是(primals, tangents)jax.vjp返回一个函数把输出侧的余切向量映射到输入侧的余切向量。这种接口设计在 JAX 里非常统一其他框架比如 PyTorch 的torch.autograd.functional.jvp和torch.autograd.functional.vjp也提供了类似的功能。3. 工程实践里的关键细节3.1 自定义算子时怎么用 JVP/VJP很多深度学习框架允许用户扩展自定义算子但难点不在于前向计算而在于反向梯度的实现。以 PyTorch 为例如果你用torch.autograd.Function自定义算子核心方法是forward和backward。backward接收的就是输出侧的梯度返回的是输入侧的梯度这正好是 VJP 的实现。也就是说框架的backward就是 VJP。但如果你用的是 JAX 的自定义函数jax.custom_jvp或jax.custom_vjp那么你就需要显式地提供 JVP 或 VJP 规则。我自己在实现一个自定义算子时最初的思路是先写出数学公式然后手动推导 VJP。后来发现其实可以先实现 JVP再用“转置”的方式自动得到 VJP。因为 JVP 和 VJP 本质上互为转置如果你能把 JVP 实现为一个线性算子那这个线性算子的转置就是 VJP。在 JAX 中可以对自定义的 JVP 规则调用jax.linear_transpose来获得 VJP。这样做的好处是只需要推导一次前向导数方向上的规则不必单独手工推导反向公式而且不容易出错。举个例子如果你实现了一个自定义算子my_op定义了它的jvp规则jax.custom_jvp def my_op(x): return x ** 2 x my_op.def_jvp def my_op_jvp(primals, tangents): x, primals x_dot, tangents return my_op(x), 2 * x * x_dot x_dot在这个规则中tan gent 方向的计算就是 JVP。JAX 可以基于这个自动转置出 VJP从而让grad(my_op)正常工作。这里的关键点是JVP 规则必须对 tangent 是“线性”的否则转置结果就不正确。上面的例子2 * x * x_dot x_dot对x_dot是线性的所以没问题。如果在实现过程中发现反向梯度不对通常是因为 JVP 规则中函数输出对 tangent 并非线性或者你手动写的 VJP 维度搞错了。自定义算子的风险就在这里数学上只差一个转置但在代码里你可能把行和列弄反了。3.2 存储与计算复杂度权衡JVP 和 VJP 不只是数学概念它们还直接影响计算图和内存的构造方式。前向模式自动微分在计算过程中可以不保存中间激活值因为它从输入方向开始一路向前带出导数。当你只需要一个方向导数时内存开销很低。但反向模式需要先走一遍前向计算保存所有中间变量然后再从输出往输入反向传播所以内存开销通常大得多。不过反向模式的优势在于一次前向、一次反向就能获得所有参数的梯度而前向模式要想获得所有梯度需要对每个输入维度各做一次前向复杂度是 (O(n)) 倍。这就解释了为什么深度学习训练几乎全部采用反向模式。但在某些场景下比如要计算一个参数对多个输出的敏感度或者做逆问题求解时JVP 反而更有效率。我自己做过一个实验一个函数输入维度是 2输出维度是 10000。用前向模式求完整雅可比只需要 2 次 JVP用反向模式求完整雅可比需要 10000 次 VJP。显然前向模式更快。反过来如果是输入维度 10000、输出维度 2那反向模式更快。这个维度权衡是工程选型时的重要依据。另外还有一个细节JVP 可以通过复数和实数之间的技巧来高效计算比如使用二元数dual number或者复数运算来同时得到函数值和方向导数。在 JAX 中jax.jvp的实现会利用这些技巧自动对任何原语求切线。如果你自己实现高性能 JVP也可以用复数小技巧但一般不建议自己造轮子。3.3 与深度学习框架的衔接多数框架中VJP 是隐藏的。比如 PyTorch 自动微分引擎会根据每个张量的grad_fn构建反向图grad_fn内部的backward方法就是局部 VJP 的实现。框架已经帮你做了大部分工作你通常不需要直接碰 VJP。但在以下情况你不得不理解它自定义torch.autograd.Function时需要实现backward。在 JAX 里自定义jax.custom_vjp时需要实现jvp和vjp或者只用custom_jvp然后自动转置。在模型解释工具中比如计算积分梯度、平滑梯度会用到输入对输出的敏感度这本质上是 VJP 或 JVP。在科学计算中结合jax.experimental.sparse或自定义稀疏算子时JVP/VJP 是接口约定的核心。以 PyTorch 自定义函数为例import torch class MySquare(torch.autograd.Function): staticmethod def forward(ctx, x): ctx.save_for_backward(x) return x ** 2 staticmethod def backward(ctx, grad_output): x, ctx.saved_tensors return grad_output * 2 * x print(MySquare.apply(torch.tensor(3.0, requires_gradTrue)))这里的backward(ctx, grad_output)接收的grad_output就是输出侧梯度返回的grad_output * 2 * x其实就是 VJP 的结果。如果你把x看成输入f(x) x^2的雅可比是2x输出侧向量是grad_output那么 (J^T \cdot grad_output 2x \cdot grad_output)。注意这里雅可比是 (1 \times 1) 矩阵转置不变。如果输入输出都是多维这个逻辑就更清晰了。用一段伪代码表示多维情况假设函数 (y Wx)其中 (x \in \mathbb{R}^n)(W \in \mathbb{R}^{m \times n})。前向时输出 (y Wx)。反向时已知grad_output是 (\mathbb{R}^m) 的向量要求grad_input即 (W^T \cdot grad_output)。同时参数 (W) 的梯度是grad_output与外层输入 (x) 的外积。这整个逻辑就是 VJP 的展开。4. 常见问题与排查实录4.1 维度对不上反向结果全乱最常见的坑就是维度约定不一致。实现 VJP 时你接收的梯度向量维度应该等于输出维度返回的应该等于输入维度。很多新手会搞反导致最终结果要么尺寸不匹配要么梯度方向完全错误。我之前遇到一个案例某开发者实现了一个自定义层输入是(batch, features)输出是(batch, n_components)。他在backward里误把grad_output的features维度当成输出维度最后返回的梯度形状和输入不一致。排查了半天绕回来看数学定义才发现问题出在维度映射上。建议在写自定义算子之前先用纸笔画一下形状标注用字母表示各维度。JVP 的 tangent 形状必须和 primal 输入完全一致VJP 的 cotangent 形状必须和 primal 输出完全一致。这不仅仅是维度数量相同每个轴的长度也要对齐。4.2 实现了 JVP 但无法求 VJP这是 JAX 用户经常遇到的问题。如果你只使用jax.custom_jvp给算子定义了 JVP 规则那么jax.grad应该可以正常工作但不一定所有场景都支持。比如你在 JVP 规则里面使用了非线性操作作用于 tangent那么 JAX 的线性转置机制可能在转置时失败或者产生错误梯度。我在一个实现指数函数的自定义算子时JVP 规则写成了def exp_jvp(primals, tangents): x, primals x_dot, tangents y jnp.exp(x) return y, y * x_dot这个规则对 tangent 是线性的系数是y但它依赖于 primal不依赖于 tangent所以没问题。但如果我写上jnp.exp(x_dot)那就错了因为exp(x_dot)不是x_dot的线性函数。JAX 的转置会自动检查这种线性性但有时报错信息不够明显。你可以用jax.linear_transpose显式测试或者直接检查自定义 JVP 是否保持切线方向的线性。另外还需要注意自定义 JVP 规则中输出 tangent 必须是与 primal 输出相同形状的张量且对 tangent 输入是线性的。自定义 VJP 规则中输出 cotangent 必须是与 primal 输入相同形状的张量且对 cotangent 输入是线性的。4.3 性能坑不必要的完整雅可比构建很多人误以为 JVP 和 VJP 必须先构建雅可比矩阵但实际不是这样。两者都直接计算“矩阵-向量积”不需要显式构建矩阵。如果你在实现中用jax.jacfwd先算出完整矩阵再与向量相乘那不仅浪费内存还会大幅增加计算时间。正确做法是始终使用jax.jvp或jax.vjp这类算子它们底层会复用原始函数的计算逻辑以增量方式广播导数。我在优化一个可微分模拟器时刚开始为了调试方便先计算了完整雅可比矩阵结果发现一个中等规模的问题内存直接爆掉。改成 JVP/VJP 并行的方向传播后内存占用下降了两个数量级。对于高维函数显式雅可比矩阵是 (m \times n) 的如果 (m) 和 (n) 都是上万矩阵就有上亿个元素根本不可能存储。而 JVP/VJP 只需要存储与输入或输出同维度的向量内存代价小得多。还有一个容易被忽略的坑VJP 的反向过程会保存前向的中间结果也就是 activations。如果你在自定义规则中把中间变量从保存列表中漏掉反向时会报“尝试访问已释放的张量”之类的错误。如果反过来保存太多不需要的中间结果内存又会被浪费。所以请尽量只保存 VJP 中必要的输入或中间值其它能重新计算的就在反向时重算。我整理了一个速查表方便你快速定位问题场景推荐操作注意事项需要输出对输入的方向导数jax.jvp或torch.autograd.functional.jvptangent 维度必须等于输入维度需要输入对输出的梯度标量 lossjax.grad或 PyTorch 的backward()本质是多次 VJP 的组合求完整雅可比矩阵且输入维度小前向模式 多次 JVP复杂度正比于输入维度求完整雅可比矩阵且输出维度小反向模式 多次 VJP复杂度正比于输出维度自定义算子且只实现 JVPjax.custom_jvp保证 tangent 映射线性自定义算子且只实现 VJPjax.custom_vjp或 PyTorchFunction.backward保证 cotangent 映射线性检查自定义规则是否正确用有限差分对比数值梯度小心步长选取建议用中心差分5. 动手验证用有限差分检验 JVP/VJP 实现理论再怎么讲都不如实际跑一次来得踏实。无论你是在 JAX、PyTorch 还是 TensorFlow 中写自定义算子都建议用有限差分验证 JVP 和 VJP 的结果。这个方法很简单对输入向量分别在每个分量上加一个小扰动 (\epsilon)然后用函数输出变化除以 (\epsilon)就得到方向导数的近似值。比如验证 JVP取一个方向向量 (v)计算[ \frac{f(x \epsilon v) - f(x - \epsilon v)}{2\epsilon} ]这个结果应该和jax.jvp(f, (x,), (v,))的结果在机器精度允许的范围内一致。验证 VJP 的方法稍复杂一点因为 VJP 是反向算子你需要在输出侧随机选一个向量 (u)然后用输入侧的扰动来计算[ \frac{u \cdot (f(x \epsilon e_i) - f(x - \epsilon e_i))}{2\epsilon} \approx (J^T u)_i ]对于每个输入维度 (i) 做一次这种检查。实际操作中我不建议对每个维度都测而是随机抽几个维度测试或者使用jax.test_util.check_grads这类内置工具自动做多次随机方向的有限差分验证。JAX 的check_grads函数支持指定阶数可以验证一阶导甚至高阶导。如果你用的是 PyTorch可以结合torch.autograd.gradcheck来验证自定义算子的梯度它会自动用数值微分去对比你的反向实现。我自己在实现一个图像处理 Demo 时就用jax.test_util.check_grads检查过自定义的仿射变换算子。当时数值对不上最后发现是因为我在前向计算中使用了取整操作导致函数在局部不光滑。这个问题让有限差分出现很大误差但反向梯度本身可能是正确的。所以遇到有限差分检查不过时要判断是不是函数本身不可导而不是一味怀疑 VJP 实现。6. 实际项目中的选型心得在真实项目里你很少会单独使用 JVP 或 VJP更多是在框架的自动微分体系里间接使用它们。但理解这两者可以让你在碰到下面几类问题时游刃有余需要实现一个自定义的损失函数且其梯度无法用现有算子组合导出。需要做神经网络的敏感性分析比如计算输入对输出的梯度范数、局部 Lipschitz 常数。需要在可微物理模拟器中定义边界条件它可能是一个隐式函数必须通过隐函数定理来求 JVP 或 VJP。需要优化一个包含黑盒操作的流程黑盒只提供前向 JVP你要通过矩阵转置来得到 VJP进而接续反向传播。我曾经参与过一个跨平台系统的自动微分集成其中最棘手的部分就是一个跨语言的数值内核它本身只实现了前向模式的导数也就是 JVP。为了让这个内核能用于深度学习训练我用线性转置的方法把它的 JVP 转换成了 VJP然后包装成一个 PyTorch 可用的自定义算子。这个过程中最大的收获是JVP 和 VJP 之间的关系是线性的只要你能用矩阵表示 JVP它的转置就是 VJP。这比从零推导反向公式容易得多也不容易出错。但要注意线性转置方法要求 JVP 必须是严格线性的。如果你的 JVP 实现里包含了比如abs、clip这类分段线性函数在分段边界处的导数不唯一转置结果也可能不唯一。这种情况下建议在 JVP 规则里显式定义次梯度或者用光滑近似替代否则梯度可能在某些点不稳定。另一个经验是在实现自定义 JVP 时尽量不要依赖中间结果的全局状态最好把所有依赖都作为 primal 参数传入否则转置时状态管理会很头疼。我最后的建议是在写任何自动微分相关代码之前先花十分钟画一张“维度流向图”。左边写输入维度右边写输出维度中间写函数名。JVP 是从输入维度方向来的向量流向输出维度VJP 是从输出维度来的向量流向输入维度。把这张图贴在屏幕旁边调试自定义算子时就不会迷失方向。如果你现在正准备实现自定义算子请务必先确认框架要求的接口是 JVP 还是 VJP再用我上面提到的有限差分法验证这样能省下大量的排查时间。