企业官网建设流程全解析

平时看自动微分框架的源码，JVP 和 VJP 这两个缩写每次都会出现，但很多刚接触的人会把它们当成两种神秘的“黑魔法”。说白了，JVP（Jacobian-vector product）是雅可比矩阵右乘一个向量，VJP（vector-Jacobian product）是一个向量左乘雅可比矩阵。名字看着对称，实际用途、计算效率和背后的数学直觉却完全不对称。这篇文章我想把这两个算子拆开，从雅可比矩阵本身讲起，再做手算推导、伪代码实现，最后聊聊工程上的坑和性能取舍。

估计不少人是先接触“梯度下降”再反向去补雅可比的，我也是这样。一开始以为所谓的自动微分就是“算个导数”，直到某次要在项目里算多输出函数的一阶灵敏度，才知道真正要面对的是 m x n 的雅可比矩阵。如果把矩阵硬建出来，再大一点就内存爆炸；但用 JVP/VJP 去算，你根本不需要建矩阵。这个思路值得好好讲透。

1. 先弄清楚雅可比矩阵到底在描述什么

1.1 从单个导数到一整张偏导表

单变量函数的导数 df/dx 我们都熟悉，它回答的是“x 变化一点点，f 跟着变化多少”。一旦输入变成 n 维向量 x = [x1, x2, ..., xn]，输出变成 m 维向量 y = [y1, y2, ..., ym]，这个问题就不再是一个数能回答的，而必须用一张表格来记录所有交叉影响。

雅可比矩阵 J 就是这张表，第 i 行第 j 列的元素是 ∂y_i / ∂x_j。J 的形状是 m 行 n 列。你可以把 J 想象成一张“变化传递责任表”：第 j 个输入动一点点，第 i 个输出会被连带放大多少倍。每一个偏导就是一条传递链路的“增益系数”。

这里要强调一下形状约定。输入维度 n 对应列数，输出维度 m 对应行数。很多人在代码里把行列搞反，后面 VJP 的形状推导就会全乱。我自己的记忆方法是：矩阵写出来必须能直接乘上输入方向的向量，所以输入维度必须排在内部维度的位置上，也就是 m x n 矩阵右乘 n 维向量，结果自然是 m 维。

1.2 一张具体的雅可比矩阵长什么样

光说概念太抽象，直接写个例子。设：

f(x1, x2) = [ x1 * x2, sin(x1) ]

输入是二维，输出也是二维，所以雅可比矩阵是 2 行 2 列。四个偏导分别是：

• ∂(x1*x2) / ∂x1 = x2
• ∂(x1*x2) / ∂x2 = x1
• ∂sin(x1) / ∂x1 = cos(x1)
• ∂sin(x1) / ∂x2 = 0

于是：

J = [[ x2, x1 ], [ cos(x1), 0 ]]

如果取具体点 (x1, x2) = (2.0, 0.0)，矩阵就变成数值形式：

J = [[ 0.0, 2.0 ], [ cos(2.0), 0.0 ]]

这个例子后面全文都会用。可以直观看到：雅可比矩阵不是存在于空泛定义里的概念，它就是一个偏导数值表。但真正要命的地方在于，这张表的规模随维度增长非常快。n 维输入、m 维输出，就要存 mn 个偏导。等 n 和 m 都是几万的时候，显式建矩阵纯属灾难。

1.3 为什么不直接建矩阵？

因为绝大多数实际场景要的根本不是整张矩阵。深度学习里常见的损失函数是一个标量，输出 m = 1，这时候雅可比矩阵退化成一行，其实我们只关心梯度向量。中间层雅可比哪怕存在，也只是在反向传播过程中被“向量穿透”了一下，并不会落地成完整矩阵。

如果你非要先把 J 完整算出来，再拿它去乘向量，那就等于白白付出了 O(mn) 的时间和 O(mn) 的内存。更尴尬的是，算完之后你可能只需要 J.v 这一个 m 维结果，矩阵里的 mn 个元素几乎全被浪费。而 JVP 和 VJP 的高明之处，就是让你绕开矩阵构造，直接得到“矩阵乘以向量”这个最终产物。

2. JVP 和 VJP 的数学身份：方向决定复杂度

2.1 两个算子的准确定义

设 J(x) 是函数 f: R^n -> R^m 在点 x 处的雅可比矩阵，形状 m x n。

JVP，jacobian-vector product，定义是 J(x) v。这里 v 是 n 维切向量，代表输入方向。结果是一个 m 维向量，它回答的问题是：如果输入沿着 v 方向移动一点点，输出会发生多大变化。相当于对每个输出分量都算一次“沿 v 方向的方向导数”。

VJP，vector-jacobian product，定义是 u^T J(x)。这里 u 是 m 维伴随向量，代表输出空间的加权方向。结果是一个 n 维向量，回答的问题是：如果输出端先按 u 加权汇总，再往回看各个输入的影响，每个输入各占多少。反向传播里的梯度，本质上就是损失函数这个标量输出对网络中每个张量的 VJP。

注意：VJP 的符号经常被人理解为“向量乘雅可比矩阵”，其实为了保证维度正确，写成 u^T J 或者 J^T u 都是同一个东西。我习惯把 VJP 想成“雅可比矩阵的转置，再右乘向量”，这样和 JVP 的统一形式比较好记。

2.2 计算复杂度完全不对称

这是标题里最值得展开的部分。JVP 和 VJP 不只是方向不同，它们的理论代价取决于 n 和 m 的对比。

• 显式构造完整 J 再乘向量：时间 O(mn)，额外内存 O(mn)。
• JVP 用前向模式计算：时间 O(n)，额外内存 O(1)（相对于原函数求值）。
• VJP 用反向模式计算：时间 O(m)，额外内存取决于是否保存中间激活。

这里的时间复杂度指的是“额外操作的规模”，不把对函数本身求值的代价算进去。实际工程中，反向传播需要存中间值，内存会高一些，但比起 O(mn) 的矩阵存储要友好太多。

所以选型的规律非常清晰：当输出维度 m 很大、输入维度 n 很小时，用 JVP 划算；当输入维度 n 很大、输出维度 m 很小时，用 VJP 划算。神经网络参数侧动辄上千万，而每个样本的标签通常只有几个数字，这正是反向传播选用 VJP 的根本原因。

2.3 从链式法则看两种“方向”

理解自动微分的时候，计算图是一条从输入到输出的有向路径。JVP 是沿着计算图从前往后逐层传播切向量，每经过一个算子，就把“输入侧的扰动”映射成“输出侧的扰动”。VJP 则是从后往前传播伴随向量，每反向经过一个算子，就把“输出侧的灵敏度”回传成“输入侧的灵敏度”。

举一个生活化类比。想象一个工厂流水线，原料从入口进，成品从出口出。JVP 相当于问：我在入口多加一批原料，最终每个成品数量怎么变？VJP 相当于问：如果我对每个成品定价，那么入口的每种原料应该分摊多少成本？两种问题都有用，但回答它们的路径正好相反。

由于链式法则本身可以按不同顺序计算乘积，JVP 和 VJP 都是在算同一堆导数，只是结合律用的顺序不同。这也是为什么很多框架同时提供两个接口：它们互相配合，才能覆盖前向模式和反向模式的全部需求。

3. 不建矩阵的实算路径：手把手实现 JVP 与 VJP

3.1 用前向模式实现 JVP

前向模式最经典的做法是引入二元数，也叫对偶数。每个变量同时携带两部分：函数值（实数部分）和导数（扰动部分）。对基本运算，比如加法、乘法、三角函数，都额外定义一套对应的导数传播规则。

回到第 1 节的例子，f(x1, x2) = [x1*x2, sin(x1)]。在点 (2.0, 0.0) 处，输入方向 v = (1.0, -1.0)。

如果我们手推，先算雅可比矩阵，再算矩阵乘：

J = [[0.0, 2.0], [cos(2.0), 0.0]]

J.v = [ 0.01.0 + 2.0(-1.0), cos(2.0)1.0 + 0.0(-1.0) ] = [ -2.0, cos(2.0) ]

如果走前向自动微分，整个过程不需要构造 J，而是把 v 当作每个输入的“伴随扰动值”同步传播：

• 设 x1 = 2.0, 切向量 1.0
• 设 x2 = 0.0, 切向量 -1.0
• 计算 y0 = x1x2，并同时算切向量：x1'x2 + x1x2' = 1.00.0 + 2.0*(-1.0) = -2.0
• 计算 y1 = sin(x1)，切向量：cos(x1)*x1' = cos(2.0)*1.0

最终 JVP 就是 [ -2.0, cos(2.0) ]，和显式矩阵计算结果完全一致。

代码层面，主流自动微分库的通用接口大概长这样：

def f(x): y0 = x[0] * x[1] y1 = sin(x[0]) return [y0, y1] x0 = (2.0, 0.0) v = (1.0, -1.0) # 通用调用方式，具体函数名不同框架略有差异 y_value, jvp_result = jvp(f, x0, v) print(y_value) # [0.0, sin(2.0)] print(jvp_result) # [-2.0, cos(2.0)]

用前向模式算 JVP 时，整个计算的代价大致等于算一次函数再加上“顺着计算图传播切向量”的额外代价，这个额外代价正比于输入维度 n。比如 n 只有几十，但输出维度 m 有几十万，JVP 就非常占便宜。

3.2 用反向模式实现 VJP

VJP 对应的自动微分模式是反向模式，也就是我们熟悉的“先算函数值，再从输出往输入回溯”。同样用例子，给定一个输出方向的伴随向量 u = (3.0, 5.0)。

如果显式算，先写出 J^T 再乘 u：

J^T = [[ x2, cos(x1) ], [ x1, 0 ]]

在 (2.0, 0.0) 处，J^T = [[0.0, cos(2.0)], [2.0, 0.0]]，所以：

u^T J = [ 3.00.0 + 5.0cos(2.0), 3.02.0 + 5.00.0 ] = [ 5*cos(2.0), 6.0 ]

反向模式手工回溯是这样的：

• 输出 y0 对本地输入 x1, x2 的偏导分别是 x2, x1，在 (2.0, 0.0) 处是 0.0 和 2.0
• 输出 y1 对本地输入 x1 的偏导是 cos(x1)，对 x2 是 0
• 伴随向量 u = (3.0, 5.0) 把两个输出通道的灵敏度汇总回输入端：
  • 对 x1 的累积：3.0 * 0.0 + 5.0 * cos(2.0) = 5*cos(2.0)
  • 对 x2 的累积：3.0 * 2.0 + 5.0 * 0.0 = 6.0

结果正是 J^T u。用通用接口写就是：

u = (3.0, 5.0) # 返回原函数值和 VJP 结果 y_value, vjp_result = vjp(f, x0, u) print(y_value) # [0.0, sin(2.0)] print(vjp_result) # [5*cos(2.0), 6.0]

这里要注意一个容易错的地方：VJP 的输入 u 必须是 m 维，因为它是从输出端过来的向量；输出是 n 维。如果你不小心传成 n 维，很多框架会直接报形状错误。我自己就在刚写接口时踩过这个，因为当时总下意识觉得“输入向量”就该是 n 维。

3.3 用两个基础算子拼出完整雅可比矩阵

有些场景确实需要完整的雅可比矩阵，比如做矩阵秩分析、绘制灵敏度热力图。这时候不建议直接用中心差分式的一列列硬算，而可以用 JVP 或者 VJP 拼出来。

如果想得到完整 J，可以对每个标准基向量 e_j 分别调一次 JVP，一共 n 次，把 n 个 m 维结果按列拼起来。想得到 J^T，就对每个标准基向量 e_i 分别调一次 VJP，一共 m 次，把结果按列拼起来。

def full_jacobian_jvp(f, x0, n, m): # 每个基向量作为切向量，得到一列 columns = [] for j in range(n): e_j = [1.0 if i == j else 0.0 for i in range(n)] _, col = jvp(f, x0, e_j) columns.append(col) # 按列拼成 m x n 矩阵 return transpose(columns)

这样拼出来的矩阵，比直接手推偏导表达式要省心得多，而且精度就是机器精度级别的。代价是 n 次前向传播，如果 n 很小，性价比很高。

3.4 组合拳：Hessian 向量积也能用 JVP/VJP 表达

JVP 和 VJP 不只是用来求一阶灵敏度，它们可以组合出二阶量。常见的需求是 Hessian 向量积（HVP），即 H v，其中 H 是标量函数 f: R^n -> R 的 Hessian 矩阵，v 是 n 维向量。

关键观察是：

H v = ∇(∇f(x) · v) = ∇(grad_f(x) · v)

也就是说，先把函数 g(x) = ∇f(x) · v 构造出来，再对它求梯度。而 ∇f(x) 本身就可以通过一次 VJP 算出来，所以 HVP 可以通过“先 VJP 得到梯度，再做一次 VJP 或 JVP”实现，完全不需要显式构造 Hessian 矩阵。

另一种等价写法是直接对 grad_f 做 JVP：JVP(grad_f, x0, v) 的结果就是 H v。这条式子在我做二阶优化时会经常用到。用 HVP 代替 Hessian 矩阵，可以把牛顿法类算法的单步成本从 O(n^2) 降到一个和反向传播同量级的水平。

3.5 永远要做的数值验证

理论推导容易，代码实现容易错。我每次写完自定义算子或自定义层，都会先用数值法验证。原理很简单，泰勒展开说明切线方向就是 JVP 的极限：

f(x0 + h*v) ≈ f(x0) + h * (J v)

所以可以用小 h，比如 1e-6，分别把 f(x0 + h*v) 和 f(x0) 算出来，两者之差除以 h，应该非常接近 JVP 结果。VJP 也可以类似验证：构造一个小随机向量 u，然后验证 u 与 J.v 的点积等于 v 与 VJP 的点积，也就是 (u^T J) v = u^T (J v)。这个恒等式是检验 VJP 实现是否正确的最快方法。

4. 工程实战里的坑位与水位

4.1 怎么判断该用 JVP 还是 VJP

很多项目卡在性能上，不是因为梯度算法不对，而是选错了方向。我给一个简单判断标准：先看输入维度和输出维度谁大。

• 输入维度 n 小，输出维度 m 大：用 JVP。典型场景是参数很少、要对大量输出通道做灵敏度分析。
• 输入维度 n 大，输出维度 m 小：用 VJP。典型场景是深度学习训练，参数上千万，损失函数只有一个标量。
• 两个都差不多：看你是想在计算图上前向传播切向量，还是反向传播伴随向量，选择更好实现的。

工程中我最常踩的坑，是在深度模型里想分析“输入图像某几个像素的扰动会如何影响网络输出”。这时候输入维度可能只有几万，但输出维度可能是几十个类别，JVP 很合适。我却一开始习惯性用了 VJP，结果每算一次都要把整条反向图跑一遍，白白慢了几十倍。

4.2 显式构建矩阵往往是性能元凶

我见过不少项目，第一步就是把中间层的雅可比矩阵整个算出来，再去做特征归因。结果矩阵大小相当于“特征数 x 特征数”，随便一个中等模型就是几千万个元素，内存直接爆炸。

用算子形式，也就是 JVP 或 VJP 形式，内存占用立刻从 O(mn) 降到 O(1) 或 O(m+n) 量级。换句话说，你要的从来不是那张矩阵，而是矩阵对一个向量的作用结果。这个“算子化”的思想，是 JVP/VJP 在工程中最值钱的地方。

下面是一组示意数据，实际数字会因框架和硬件浮动，但量级可以参考：

“适合把完整 J 显式算出来”的场景真的很少，通常只有矩阵本身不超过几万元素时才值得。

4.3 常见错误速查表

我把自己踩过以及帮同事排查过的典型错误整理了一下：

最常见也最隐蔽的错误，是自定义反向传播算子时，VJP 函数写错了雅可比转置的相乘顺序。因为反向传播库里常见的接口是 grad_outputs 输入，输出是 grad_inputs，看起来很像“梯度”，但本质上它就是一个带伴随向量的 VJP。你写的返回梯度形状如果不对，最后整个训练过程都不会报错，只是结果异常。

4.4 稀疏结构其实可以天然保留

很多函数的雅可比矩阵不是稠密的，而是带有明显结构化模式的。比如逐元素激活函数的雅可比矩阵，几乎就是对角矩阵；卷积层单个样本的雅可比，也有明显的带状结构。

如果显式建矩阵，你需要专门用稀疏矩阵去存储，否则内存和计算都会浪费在大量零元素上。而用 JVP/VJP，根本不需要关心“稀疏”，因为算子的传播过程天然绕过了无效的偏导路径。这也是为什么大型模型的反向传播不会因为参数多而退化到 O(n^2)，因为每一条雅可比路径都是按需计算，而不是全局枚举。

5. 从方阵到非方阵：那些容易想当然的场景

5.1 梯度就是 VJP 的标量特例

当输出 m = 1 时，雅可比矩阵退化成 1 x n 的行向量，VJP 的伴随向量 u 退化成一个标量。只要这个标量取 1，VJP 的结果就是梯度。所以反向传播本质上是“把损失函数对输出的标量导数，作为伴随向量逐层左乘雅可比”。

这么看，梯度下降又神秘又普通。它不过是无数次 VJP 的叠加。理解这一点后，再看自定义损失函数怎么写，思路会清晰很多：任何可微的标量损失函数，都能通过一次全局 VJP 得到所有参数梯度。

5.2 输出向量输入标量时的灵敏度分析

反过来，如果输入 n = 1，输出 m 很大，J 就是 m 行 1 列的列向量。这时候 JVP 只需要一次前向传播就能算出输出对输入水平方向的导数，而 VJP 则需要 m 次才能得到那个单元素的输出，明显不划算。

一个典型例子是，计算某组超参数连续变化时模型的多个性能指标如何变化。输入是 1 维超参，输出是几百个评价指标，用 JVP 一次就能得到整条灵敏度曲线。这类分析任务里，选对方向省下的时间不是常量级，而是几十到几百倍。

5.3 不显式建矩阵的奇异值近似

有些场景想用 Jacobian 的最大奇异值判断函数局部放大效应，这时候不需要完整 J。可以使用幂迭代法，交替调用 JVP 和 VJP：

• 先随机初始化 n 维向量 v
• 用 JVP 得到 u = J v，u 是 m 维
• 用 VJP 得到 w = J^T u，w 是 n 维
• 对 w 归一化，再重复迭代

最终向量模长的增长速率，会收敛到最大奇异值。整个过程中没有任何地方需要构造 J，完全依靠两个算子交替作用。我做嵌入空间度量的稳定性分析时用过这个思路，效果非常稳。

5.4 诊断梯度消失和爆炸的技巧

与其把网络某层雅可比矩阵拉出来看，不如直接观察梯度范数。训练过程中每个中间张量其实都对应一个梯度向量，这些梯度就是用 VJP 一层层传回来的。如果某层梯度范数异常小，说明上游信号传不透；如果异常大，说明该层灵敏度高得过份。

不需要额外写代码，只要在训练循环里把每一层的梯度范数打印出来，就能形成一张“梯度水位图”。很多时候比看损失曲线更早发现问题，尤其是训练初期就出现的爆炸或者消失。

最后再分享一个小习惯

我现在拿到任何需要求导的问题，第一反应都不会是“我要算雅可比矩阵”，而是“我要让哪个向量乘进去”。这个思维转换，就是弄懂 JVP 和 VJP 带来的最大收益。实际项目里，用前向求切线、反向传伴随、再用两者交替构造高阶算子，这三板斧基本够用。后来我做二阶优化、特征敏感性分析、正则化诊断，几乎全是围绕 JVP/VJP 组合展开的，很少再碰显式矩阵构造。把方向选对，很多性能问题根本不需要硬扛，绕一下就过去了。