企业官网建设流程全解析

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)) 次反向过程。

请记住这个差别，因为在实际应用中，前向模式适合“输入少、输出多”的函数，反向模式适合“输入多、输出少”的函数。神经网络的损失函数通常是输入是成千上万个参数、输出是一个标量，即 (m=1)，所以反向模式 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_grad=True)))

这里的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 中必要的输入或中间值，其它能重新计算的就在反向时重算。

我整理了一个速查表，方便你快速定位问题：

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，再用我上面提到的有限差分法验证，这样能省下大量的排查时间。