在 PyTorch 中，尽管我们可以（链接未公开），我们一般难以计算一个多维函数的 Jacobi 矩阵。然而，Jacobi 矩阵与一个向量的乘积，也就是 Jacobi 向量积（Jacobian-Vector Product, JVP）却可以通过前向传播进行快速计算。

方向导数与偏导

给定一个 \mathbb{R}^n \to \mathbb{R}^m 的连续可导函数 f 和一个单位向量 v \in \mathbb{R}^n，其在 x 位置沿 v 的方向导数定义为：

\frac{\partial f}{\partial v} = \lim_{h \to 0} \frac{f(x + hv) - f(x)}{h}.

可以证明，方向导数与 Jacobi 矩阵 J 有如下的关系：

\frac{\partial f}{\partial v} = Jv = \operatorname{JVP}(f, x, v).

而偏导数是方向导数的特例，也就是说 f 对第 i 个分量的偏导数 f_i' 就有

\frac{\partial f}{\partial x_i} = J e_i = \operatorname{JVP}(f, x, e_i),

其中 e_i 是 \mathbb{R}^n 的第 i 个标准基 (\underbrace{0, \dots, 0}_{i - 1 \text{个}}, 1, \underbrace{0, \dots, 0}_{n - i \text{个}})^\mathrm{T}。

Jacobi 向量积