深度学习基本功——自动微分的正向模式与反向模式：理解JVP与VJP

​

引言

在现代机器学习和深度学习中，自动微分（Automatic Differentiation, AD）扮演着至关重要的角色。它使得我们能够有效地计算复杂函数的导数，从而在优化过程中更新模型参数。本文将探讨自动微分的两种主要模式：正向模式（forward mode）和反向模式（inverse mode），以及它们对应的雅可比-向量乘法（JVP）和向量-雅可比乘法（VJP）。最后，我们将分析为什么当前深度学习神经网络普遍采用反向模式。

本文参考了以下的自动微分时的资料，这篇博客也算是对自己学习的总结：

• 机器学习工具（四）自动微分之VJP、JVP和JAX
• 前向微分和正向微分怎么理解？微分方式跟反向传播算法啥关系？【自动微分】系列第三篇

正向模式和反向模式

我们先考虑输出是标量的情况，假设我们有以下函数关系：
$$y=f(\mathbf{x})=\mathbf{a}(\mathbf{b}(\mathbf{c}(\mathbf{x})))$$

其中：

• $\mathbf{x}\in {\mathbb{R}}^n$ 是输入向量
• $\mathbf{a}\in {\mathbb{R}}^m$
• $\mathbf{b}\in {\mathbb{R}}^k$
• $\mathbf{c}\in {\mathbb{R}}^p$

根据链式法则，得到导数:
$$\frac{\partial y}{\partial \mathbf{c}}\in {\mathbb{R}}^{1\times p},\frac{\partial \mathbf{c}}{\partial \mathbf{b}}\in {\mathbb{R}}^{p\times k},\frac{\partial \mathbf{b}}{\partial \mathbf{a}}\in {\mathbb{R}}^{k\times m},\frac{\partial \mathbf{a}}{\partial \mathbf{x}}\in {\mathbb{R}}^{m\times n}$$

则有:
$$f^{\prime }(\mathbf{x})=\frac{\partial y}{\partial \mathbf{c}}\cdot \frac{\partial \mathbf{c}}{\partial \mathbf{b}}\cdot \frac{\partial \mathbf{b}}{\partial \mathbf{a}}\cdot \frac{\partial \mathbf{a}}{\partial \mathbf{x}}$$

正向模式自动微分是从输入向输出计算导数的方法。在这一模式中，首先计算输入
$\mathbf{x}$对各个中间变量的导数，然后利用链式法则逐步向输出传播。计算顺序是从右往左，用括号表示如下：
$$f^{\prime }(\mathbf{x})=\frac{\partial y}{\partial \mathbf{c}}\cdot \left(\frac{\partial \mathbf{c}}{\partial \mathbf{b}}\cdot \left(\frac{\partial \mathbf{b}}{\partial \mathbf{a}}\cdot \frac{\partial \mathbf{a}}{\partial \mathbf{x}}\right)\right)$$

反向模式自动微分是从输出向输入传播导数的方法。在这一模式中，首先计算输出
$y$对各个中间变量的导数，然后利用链式法则反向传播。计算顺序是从左往右，用括号表示如下：
$$f^{\prime }(\mathbf{x})=\left(\left(\left(\frac{\partial y}{\partial \mathbf{c}}\cdot \frac{\partial \mathbf{c}}{\partial \mathbf{b}}\right)\cdot \frac{\partial \mathbf{b}}{\partial \mathbf{a}}\right)\cdot \frac{\partial \mathbf{a}}{\partial \mathbf{x}}\right)$$

总之，正向模式是从输入开始，逐步计算输出的梯度；而反向模式是从输出开始，先计算输出对中间变量的梯度，再向输入传播。

这是对于$y$为标量的情况，那么对于$\mathbf{y}$为向量的情况下，正向和反向又会怎么进行呢？此时我们就需要引入“雅可比矩阵（Jacobian Matrix）”了。

雅克比矩阵以及JVP、VJP

雅可比矩阵

对于函数 $\mathbf{y}=f(\mathbf{x})$ ，其中 $f:{\mathbb{R}}^n\to {\mathbb{R}}^m$ ，那么 $\mathbf{y}$ 中关于 $\mathbf{x}$ 的梯度可以表示为雅克比矩阵 :
$${\mathbf{J}}_f=\left[\begin{array}{ccc}\frac{\partial y_1}{\partial x_1}& \cdots & \frac{\partial y_1}{\partial x_n}\\ ⋮& \ddots & ⋮\\ \frac{\partial y_m}{\partial x_1}& \cdots & \frac{\partial y_m}{\partial x_n}\end{array}\right]$$

JVP

雅可比向量积（Jacobian-Vector Product, JVP）是指雅可比矩阵与一个向量的乘积，用于描述在给定输入变化时，输出变化的灵敏度。在 JVP 中，我们从输入出发，计算其对输出的影响。形式上，对于函数 $\mathbf{y}=f(\mathbf{x})$ 及输入向量 $\mathbf{v}\in {\mathbb{R}}^n$ ，JVP 定义为：
$$\mathrm{JVP}(\mathbf{v})={\mathbf{J}}_f\cdot \mathbf{v}$$

这表示在输入 $\mathbf{x}$ 沿着向量 $\mathbf{v}$ 变化时，输出 $\mathbf{y}$ 的变化。
举个例子，输出$\mathbf{y}$关于 $x_1$ 的导数则可通过Jacobian矩阵乘以对应的one-hot向量$\mathbf{v}$来得到，
$${\partial }_{x_1}\mathbf{y}={\mathbf{J}}_f\cdot \mathbf{v}=\left[\begin{array}{c}\frac{\partial y_1}{\partial x_1}\\ ⋮\\ \frac{\partial y_m}{\partial x_1}\end{array}\right]=\left[\begin{array}{ccc}\frac{\partial y_1}{\partial x_1}& \cdots & \frac{\partial y_1}{\partial x_n}\\ ⋮& \ddots & ⋮\\ \frac{\partial y_m}{\partial x_1}& \cdots & \frac{\partial y_m}{\partial x_n}\end{array}\right]\cdot \left[\begin{array}{c}1\\ ⋮\\ 0\end{array}\right],\mathbf{v}\in {\mathbb{R}}^n$$

$\mathbf{v}$ 本来可以是任意方向的向量，但是我们这里为了方便举例，就简化为 one-hot 向量，并使得结果仅为输出向量 $\mathbf{y}$ 中第一个元素 $\mathbf{y}$ 关于 $x_1$ 的梯度。

VJP

雅可比-向量积（Jacobian-Vector Product, VJP） 是指一个向量与雅可比矩阵的乘积，用于描述在给定输出变化时，输入变化的灵敏度。在 VJP 中，我们从输出的特定元素出发，计算其对输入的梯度。形式上，对于函数 $\mathbf{y}=f(\mathbf{x})$ 及输出向量中的一个特定元素的 one-hot 向量 $\mathbf{v}\in {\mathbb{R}}^m$ ，VJP 定义为:
$$\mathrm{VJP}(\mathbf{v})={\mathbf{v}}^T\cdot {\mathbf{J}}_f$$

这表示在输出 $\mathbf{y}$ 沿着向量 $\mathbf{v}$ 变化时，输入 $\mathbf{x}$ 的变化。
举个例子，假设我们希望计算输出 $\mathbf{y}$ 中第一个元素 $y_1$ 对输入 $\mathbf{x}$ 的导数，可以通过雅可比矩阵与对应的 one-hot 向量 $\mathbf{v}$ 的乘积得到:
$${\partial }_{y_1}\mathbf{x}={\mathbf{v}}^T\cdot {\mathbf{J}}_f=[1,0,\dots ,0]\cdot \left[\begin{array}{ccc}\frac{\partial y_1}{\partial x_1}& \cdots & \frac{\partial y_1}{\partial x_n}\\ ⋮& \ddots & ⋮\\ \frac{\partial y_m}{\partial x_1}& \cdots & \frac{\partial y_m}{\partial x_n}\end{array}\right],\mathbf{v}\in {\mathbb{R}}^m$$
$\mathbf{v}$ 本来可以是任意方向的向量，但我们这里为了方便举例，就简化为 one-hot 向量，这样可以使得结果仅为输出向量 $\mathbf{y}$ 中第一个元素 $y_1$ 关于所有输入 $\mathbf{x}$ 的梯度。

总结

刚才介绍的JVP、VJP是不是听上去感觉和正向、方向模式似乎有些关系？

没错！

• 正向模式对应于利用 JVP来实现输出向量（所有输出）对单一参数的求导。特别的，正向模式适用于输入维度较小的情况，因为它可以有效地逐步计算出每个输入对输出的影响，从而得到雅可比矩阵的每一列。

• 反向模式则对应于利用 VJP 来实现输出的某个分量对参数向量（所有参数）的求导。在这一模式下，我们从输出出发，计算输出变化相对于输入变化的灵敏度，使用雅可比矩阵的转置与输出变化的向量相乘。这使得反向模式特别适合于输出维度较小的情况，因为它允许我们在一次反向传播中计算出雅可比矩阵的每一行。
  ​

为什么反向模式多用于深度学习？

在深度学习中，模型的输出通常是一个标量（如损失值）或低维向量，而输入则可能是高维（如图像、文本等）。反向模式能够高效地计算损失对所有参数的梯度，适合大规模模型训练。这种模式可以在一次前向传播后，利用一次反向传播获得所有参数的梯度，极大地提高了计算效率。因此，反向模式在深度学习中的广泛应用得益于其在处理高维输入时的优势。