梯度检查机制

本笔记概述了 gradcheck() 和 gradgradcheck() 函数的工作原理。

它将涵盖前向和后向模式自动微分，包括实值和复值函数以及高阶导数。本笔记还涵盖了 gradcheck 的默认行为以及传递 fast_mode=True 参数的情况（以下称为快速梯度检查）。

符号和背景信息

在本笔记中，我们将使用以下约定

1. $x$, $y$, $a$, $b$, $v$, $u$, $ur$ 和 $ui$ 是实值向量，而 $z$ 是复值向量，可以用两个实值向量表示为 $z = a + i b$.

2. $N$ 和 $M$ 是两个整数，我们将分别用于输入和输出空间的维度。

3. $f: \mathcal{R}^N \to \mathcal{R}^M$ 是我们基本的实数到实数函数，使得 $y = f(x)$。

4. $g: \mathcal{C}^N \to \mathcal{R}^M$ 是我们基本的复数到实数函数，使得 $y = g(z)$。

对于简单的实数到实数情况，我们将与 $f$ 相关的雅可比矩阵写为 $J_f$，大小为 $M \times N$。此矩阵包含所有偏导数，使得位置 $(i, j)$ 处的元素包含 $\frac{\partial y_i}{\partial x_j}$。然后反向模式自动微分计算的是，对于给定的一个大小为 $M$ 的向量 $v$，数量 $v^T J_f$。另一方面，前向模式自动微分计算的是，对于给定的一个大小为 $N$ 的向量 $u$，数量 $J_f u$。

对于包含复数值的函数，情况要复杂得多。我们这里只提供要点，完整的描述可以在 复数的自动微分 中找到。

满足复可微性（柯西-黎曼方程）的约束条件对于所有实值损失函数来说过于严格，因此我们选择使用Wirtinger微积分。在Wirtinger微积分的基本设置中，链式法则需要访问Wirtinger导数（在下文中称为$W$）和共轭Wirtinger导数（在下文中称为$CW$）。$W$和$CW$都需要进行传播，因为一般来说，尽管它们的名字，一个不是另一个的复共轭。

为了避免必须传播这两个值，对于反向模式自动微分，我们始终假设正在计算导数的函数要么是实值函数，要么是更大实值函数的一部分。这个假设意味着我们在反向传播过程中计算的所有中间梯度也与实值函数相关联。在实践中，当进行优化时，这个假设并不严格，因为此类问题需要实值目标（因为复数没有自然的排序）。

在这个假设下，使用$W$和$CW$的定义，我们可以证明$W = CW^*$（我们在这里使用$*$表示复共轭），因此这两个值中只有一个实际上需要“反向传播到图中”，因为另一个可以很容易地恢复。为了简化内部计算，PyTorch 使用$2 * CW$作为它反向传播的值，并在用户请求梯度时返回。类似于实数情况，当输出实际上在$\mathcal{R}^M$中时，反向模式自动微分不会计算$2 * CW$，而只计算$v^T (2 * CW)$，对于给定的向量$v \in \mathcal{R}^M$.

对于前向模式AD，我们使用类似的逻辑，在这种情况下，假设该函数是更大函数的一部分，其输入位于$\mathcal{R}$中。在此假设下，我们可以做出类似的断言，即每个中间结果对应于一个输入位于$\mathcal{R}$中的函数，在这种情况下，使用$W$和$CW$的定义，我们可以证明对于中间函数，$W = CW$。为了确保前向和后向模式在一维函数的基本情况下计算相同的量，前向模式还计算$2 * CW$。与实数情况类似，当输入实际上位于$\mathcal{R}^N$中时，前向模式AD不会计算$2 * CW$，而只计算给定向量$u \in \mathcal{R}^N$的$(2 * CW) u$。

默认后向模式gradcheck行为

实数到实数函数

为了测试函数$f: \mathcal{R}^N \to \mathcal{R}^M, x \to y$，我们以两种方式重建大小为$M \times N$的完整雅可比矩阵$J_f$：解析法和数值法。解析版本使用我们的后向模式AD，而数值版本使用有限差分。然后逐元素比较这两个重建的雅可比矩阵以检查是否相等。

默认实数输入数值评估

如果我们考虑一维函数的基本情况（$N = M = 1$），那么我们可以使用来自维基百科文章的基本有限差分公式。我们使用“中心差分”以获得更好的数值特性

$$\frac{\partial y}{\partial x} \approx \frac{f(x + eps) - f(x - eps)}{2 * eps}$$

这个公式可以很容易地推广到多个输出的情况（$M \gt 1$），方法是将$\frac{\partial y}{\partial x}$ 变为大小为 $M \times 1$ 的列向量，类似于 $f(x + eps)$。在这种情况下，上述公式可以原样重用，并且仅使用用户函数的两次评估（即 $f(x + eps)$ 和 $f(x - eps)$）来近似完整的雅可比矩阵。

处理具有多个输入的情况($N \gt 1$)在计算上更昂贵。在这种情况下，我们依次遍历所有输入，并对$x$的每个元素依次应用$eps$扰动。这使我们能够逐列重建$J_f$矩阵。

默认实数输入解析评估

对于解析评估，我们使用前面描述的事实，即反向模式自动微分计算$v^T J_f$。对于具有单个输出的函数，我们只需使用$v = 1$即可通过一次反向传递恢复完整的雅可比矩阵。

对于具有多个输出的函数，我们采用一个for循环，该循环遍历输出，其中每个$v$都是一个与每个输出依次对应的one-hot向量。这使得能够逐行重建$J_f$矩阵。

复数到实数函数

为了测试一个函数$g: \mathcal{C}^N \to \mathcal{R}^M, z \to y$，其中$z = a + i b$，我们重建包含$2 * CW$的（复数）矩阵。

默认复数输入数值评估

首先考虑一个简单的例子，其中$N = M = 1$。根据（第 3 章）这篇研究论文，我们知道

$$CW := \frac{\partial y}{\partial z^*} = \frac{1}{2} * (\frac{\partial y}{\partial a} + i \frac{\partial y}{\partial b})$$

注意，上式中的$\frac{\partial y}{\partial a}$ 和 $\frac{\partial y}{\partial b}$ 是 $\mathcal{R} \to \mathcal{R}$ 导数。为了数值计算这些导数，我们使用上面描述的实数到实数情况的方法。这使得我们能够计算 $CW$ 矩阵，然后将其乘以 $2$。

请注意，在撰写本文档时，代码以一种稍微复杂的方式计算此值。

# Code from https://github.com/pytorch/pytorch/blob/58eb23378f2a376565a66ac32c93a316c45b6131/torch/autograd/gradcheck.py#L99-L105
# Notation changes in this code block:
# s here is y above
# x, y here are a, b above

ds_dx = compute_gradient(eps)
ds_dy = compute_gradient(eps * 1j)
# conjugate wirtinger derivative
conj_w_d = 0.5 * (ds_dx + ds_dy * 1j)
# wirtinger derivative
w_d = 0.5 * (ds_dx - ds_dy * 1j)
d[d_idx] = grad_out.conjugate() * conj_w_d + grad_out * w_d.conj()

# Since grad_out is always 1, and W and CW are complex conjugate of each other, the last line ends up computing exactly `conj_w_d + w_d.conj() = conj_w_d + conj_w_d = 2 * conj_w_d`.

默认复数输入解析评估

由于反向模式自动微分已经精确地计算了两倍的 $CW$ 导数，因此我们在这里使用与实数到实数情况相同的技巧，并在有多个实数输出时逐行重建矩阵。

具有复数输出的函数

在这种情况下，用户提供的函数不符合自动微分（autograd）的假设，即我们对其进行反向自动微分计算的函数是实值函数。这意味着直接对该函数使用 autograd 没有明确定义。为了解决这个问题，我们将用两个函数替换对函数 $h: \mathcal{P}^N \to \mathcal{C}^M$ 的测试（其中 $\mathcal{P}$ 可以是 $\mathcal{R}$ 或 $\mathcal{C}$）：$hr$ 和 $hi$，使得

$$\begin{aligned} hr(q) &:= real(f(q)) \\ hi(q) &:= imag(f(q)) \end{aligned}$$

其中 $q \in \mathcal{P}$。然后，我们对 $hr$ 和 $hi$ 都进行基本的梯度检验（gradcheck），根据 $\mathcal{P}$ 使用上面描述的实数到实数或复数到实数的情况。

请注意，截至编写本文时，代码没有显式地创建这些函数，而是通过将 $\text{grad\_out}$ 参数传递给不同的函数来手动执行链式法则，并使用 $real$ 或 $imag$ 函数。当 $\text{grad\_out} = 1$ 时，我们正在考虑 $hr$。当 $\text{grad\_out} = 1j$ 时，我们正在考虑 $hi$。

快速反向模式梯度检验

虽然上述梯度检验的公式很好，但为了确保正确性和可调试性，它非常慢，因为它会重建完整的雅可比矩阵。本节介绍了一种以更快的方式执行梯度检验的方法，而不会影响其正确性。可以通过在检测到错误时添加特殊的逻辑来恢复可调试性。在这种情况下，我们可以运行重建完整矩阵的默认版本，以便向用户提供完整的详细信息。

这里的核心策略是找到一个标量量，该标量量可以通过数值方法和分析方法有效地计算，并且能够足够好地表示慢速梯度检验计算的完整矩阵，以确保它能够捕获雅可比矩阵中的任何差异。

实数到实数函数的快速梯度检验

我们在这里想要计算的标量量是 $v^T J_f u$，对于给定的随机向量 $v \in \mathcal{R}^M$ 和一个随机单位范数向量 $u \in \mathcal{R}^N$.

为了进行数值评估，我们可以有效地计算

$$J_f u \approx \frac{f(x + u * eps) - f(x - u * eps)}{2 * eps}.$$

然后我们执行该向量与 $v$ 之间的点积，以获得我们感兴趣的标量值。

对于分析版本，我们可以使用反向模式自动微分 (AD) 直接计算 $v^T J_f$。然后我们执行与 $u$ 的点积以获得期望值。

复杂到实数函数的快速梯度检查

类似于实数到实数的情况，我们希望对整个矩阵进行降维。但是 $2 * CW$ 矩阵是复数值的，因此在这种情况下，我们将与复数标量进行比较。

由于我们在数值情况下可以有效计算的内容存在一些限制，并且为了将数值评估的数量降到最低，我们计算以下（虽然令人惊讶）标量值

$$s := 2 * v^T (real(CW) ur + i * imag(CW) ui)$$

其中 $v \in \mathcal{R}^M$, $ur \in \mathcal{R}^N$ 以及 $ui \in \mathcal{R}^N$.

快速复数输入数值评估

我们首先考虑如何使用数值方法计算 $s$。为此，请记住我们正在考虑 $g: \mathcal{C}^N \to \mathcal{R}^M, z \to y$，其中 $z = a + i b$，并且 $CW = \frac{1}{2} * (\frac{\partial y}{\partial a} + i \frac{\partial y}{\partial b})$，我们将其改写如下。

$$\begin{aligned} s &= 2 * v^T (real(CW) ur + i * imag(CW) ui) \\ &= 2 * v^T (\frac{1}{2} * \frac{\partial y}{\partial a} ur + i * \frac{1}{2} * \frac{\partial y}{\partial b} ui) \\ &= v^T (\frac{\partial y}{\partial a} ur + i * \frac{\partial y}{\partial b} ui) \\ &= v^T ((\frac{\partial y}{\partial a} ur) + i * (\frac{\partial y}{\partial b} ui)) \end{aligned}$$

在这个公式中，我们可以看到 $\frac{\partial y}{\partial a} ur$ 和 $\frac{\partial y}{\partial b} ui$ 可以像实数到实数情况下的快速版本一样进行评估。一旦计算出这些实数值量，我们就可以重建右侧的复向量，并与实数值 $v$ 向量进行点积。

快速复数输入解析评估

对于解析情况，事情变得更简单，我们将公式改写为

$$\begin{aligned} s &= 2 * v^T (real(CW) ur + i * imag(CW) ui) \\ &= v^T real(2 * CW) ur + i * v^T imag(2 * CW) ui) \\ &= real(v^T (2 * CW)) ur + i * imag(v^T (2 * CW)) ui \end{aligned}$$

因此，我们可以利用反向模式自动微分 (AD) 提供了一种计算 $v^T (2 * CW)$ 的高效方法，然后对实部与 $ur$ 和虚部与 $ui$ 进行点积，然后重构最终的复数标量 $s$。

为什么不使用复数 $u$

在这一点上，您可能想知道为什么我们没有选择一个复杂的$u$ 并且只是执行了简化$2 * v^T CW u'$. 为了深入探讨这个问题，在本段中，我们将使用$u$ 的复数版本，记为$u' = ur' + i ui'$. 使用这种复数形式的$u'$，问题在于，在进行数值计算时，我们需要计算

$$\begin{aligned} 2*CW u' &= (\frac{\partial y}{\partial a} + i \frac{\partial y}{\partial b})(ur' + i ui') \\ &= \frac{\partial y}{\partial a} ur' + i \frac{\partial y}{\partial a} ui' + i \frac{\partial y}{\partial b} ur' - \frac{\partial y}{\partial b} ui' \end{aligned}$$

这将需要对实数到实数的有限差分进行四次评估（与上面提出的方法相比多一倍）。由于这种方法没有更多的自由度（相同数量的实值变量），并且我们试图在这里获得尽可能快的评估，因此我们使用上面提到的另一种公式。

具有复数输出的函数的快速梯度检查

就像在慢速情况下一样，我们考虑两个实值函数，并对每个函数使用上面相应的规则。

Gradgradcheck 实现

PyTorch 还提供了一个实用程序来验证二阶梯度。这里的目标是确保反向传播实现也是可微的，并计算正确的结果。

此功能的实现考虑了函数 $F: x, v \to v^T J_f$，并使用上面定义的 gradcheck 对该函数进行检查。请注意，在这种情况下，$v$ 只是一个与 $f(x)$ 类型相同的随机向量。

gradgradcheck 的快速版本通过对相同函数 $F$ 使用 gradcheck 的快速版本来实现。