Gradcheck 机制

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

它将涵盖实值和复值函数的前向和后向 AD（自动微分），以及高阶导数。本文档还涵盖了 gradcheck 的默认行为以及传递 fast_mode=True 参数（下文称为 fast gradcheck）的情况。

符号和背景信息

在本文档中，我们将使用以下约定

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)$。

对于简单的实数到实数情况，我们将 $J_f$ 表示与 $f$ 相关的 Jacobian 矩阵，其大小为 $M \times N$。该矩阵包含所有偏导数，其中位置 $(i, j)$ 的元素为 $\frac{\partial y_i}{\partial x_j}$。后向模式 AD 计算的是给定大小为 $M$ 的向量 $v$ 的量 $v^T J_f$。另一方面，前向模式 AD 计算的是给定大小为 $N$ 的向量 $u$ 的量 $J_f u$。

对于包含复数值的函数，情况要复杂得多。此处仅提供要点，完整描述请参阅 复数的 Autograd。

满足复数可微性（Cauchy-Riemann 方程）的约束对于所有实值损失函数来说都过于严格，因此我们转而使用 Wirtinger 微积分。在 Wirtinger 微积分的基本设置中，链式法则需要访问 Wirtinger 导数（下文称为 $W$）和共轭 Wirtinger 导数（下文称为 $CW$）。 $W$ 和 $CW$ 都需要传播，因为通常情况下，尽管它们的名字如此，但一个并不是另一个的复共轭。

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

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

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

Default backward mode gradcheck behavior

Real-to-real functions

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

Default real input numerical evaluation

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

$$\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$ 矩阵。

Default real input analytical evaluation

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

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

Complex-to-real functions

为了测试函数 $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`.

默认复数输入解析评估

由于反向模式AD已经精确计算出 $CW$ 导数的两倍，我们只需在此处使用与实数到实数情况相同的技巧，并在存在多个实数输出时逐行重建矩阵。

具有复数输出的函数

在这种情况下，用户提供的函数不符合autograd的假设，即我们进行反向AD计算的函数是实数值的。这意味着直接在这种函数上使用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}$。然后，我们根据 $\mathcal{P}$，使用上面描述的实数到实数情况或复数到实数情况，对 $hr$ 和 $hi$ 进行基本gradcheck。

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

快速反向模式gradcheck

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

这里的高层策略是找到一个标量量，它可以由数值方法和解析方法高效计算，并且能够很好地代表慢速gradcheck计算出的完整矩阵，以确保它能捕获雅可比矩阵中的任何差异。

实数到实数函数的快速gradcheck

我们想要在此处计算的标量量是 $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$ 进行点积，以得到期望值。

复数到实数函数的快速gradcheck

与实到实情况类似，我们想对全矩阵进行约化。但 $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}$$

这将需要进行四次实数到实数的有限差分评估（是上述方法的两倍）。由于这种方法没有更多的自由度（实值变量的数量相同），并且我们在这里尝试获得最快的评估速度，因此我们使用了上面提到的另一种公式。

对具有复杂输出的函数的快速 gradcheck

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

Gradgradcheck 实现

PyTorch 还提供了用于验证二阶梯度的实用工具。这里的目标是确保反向实现也具有适当的可微性并计算正确的结果。

此功能通过考虑函数 $F: x, v \to v^T J_f$ 并使用上述定义的 gradcheck 对此函数进行检查。请注意，本例中的 $v$ 只是一个与 $f(x)$ 类型相同的随机向量。

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