Gradcheck 機制

本筆記概述了 gradcheck() 和 gradgradcheck() 函數的工作原理。

這份文件將涵蓋實數和複數值函數的前向和反向模式自動微分 (AD)，以及更高階的導數。本說明也涵蓋了 gradcheck 的預設行為，以及傳遞 fast_mode=True 參數的情況（以下稱為快速 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$ 相關聯的 Jacobian 矩陣寫成 $f$，其大小為 $M \times N$。 此矩陣包含所有偏導數，使得位置 $(i, j)$ 上的條目包含 $\frac{\partial y_i}{\partial x_j}$。 然後，反向模式 AD 計算的是，對於給定的向量 $v$ (大小為 $M$)，數量 $v^T J_f$。 另一方面，前向模式 AD 計算的是，對於給定的向量 $u$ (大小為 $N$)，數量 $J_f u$。

對於包含複數值的函式，情況就複雜得多了。 我們僅在此處提供要點，完整的描述可以在 複數的自動微分 中找到。

對於所有實數值的損失函數來說，要滿足複可微性（柯西-黎曼方程式）的限制過於嚴苛，因此我們選擇使用 Wirtinger 微積分。在 Wirtinger 微積分的基本設定中，鏈鎖律需要存取 Wirtinger 導數（以下稱為 $W$）和共軛 Wirtinger 導數（以下稱為 $CW$）。$W$ 和 $CW$ 都需要傳播，因為一般來說，儘管它們有這樣的名稱，但彼此並非複共軛。

為了避免需要傳播這兩個值，對於反向模式 AD，我們始終假設正在計算導數的函數是實數值函數，或是較大的實數值函數的一部分。這個假設意味著，在反向傳遞期間計算的所有中介梯度也與實數值函數相關。在實務上，這個假設在進行最佳化時沒有太多限制，因為這類問題需要實數值的目標（因為複數沒有自然的排序）。

在這個假設下，使用 $W$ 和 $CW$ 的定義，我們可以證明 $W = CW^*$ （這裡我們用 $*$ 表示複共軛），因此實際上只需要將其中一個值「反向傳播通過圖」，因為另一個值可以輕易地恢復。為了簡化內部計算，PyTorch 使用 $2 * CW$ 作為它反向傳播的值，並在使用者要求梯度時返回。與實數情況類似，當輸出實際上在 $\mathcal{R}^M$ 中時，反向模式 AD 不會計算 $2 * CW$，而僅計算給定向量 $v \in \mathcal{R}^M$ 的 $v^T (2 * CW)$。

對於前向模式 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$，我們以解析和數值兩種方式重建完整 Jacobian 矩陣 $J_f$，其大小為 $M \times N$。解析版本使用我們的反向模式 AD，而數值版本使用有限差分。然後，逐一比較兩個重建的 Jacobian 矩陣是否相等。

預設實數輸入數值評估

如果我們考慮一維函數的基本情況（$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)$)即可近似完整的 Jacobian 矩陣。

處理多個輸入的情況（$N \gt 1$）在計算上更為昂貴。在這種情況下，我們依序遍歷所有輸入，並對 $x$ 的每個元素依序應用 $eps$ 擾動。這使我們能夠逐列重建 $J_f$ 矩陣。

預設實數輸入解析評估

對於解析評估，我們使用如上所述的事實：反向模式 AD 計算 $v^T J_f$。對於具有單個輸出的函數，我們只需使用 $v = 1$ 即可通過單次反向傳遞恢復完整的 Jacobian 矩陣。

對於有多個輸出的函數，我們採用 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 章）得知：

請注意，在上述方程式中，$\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}$$

where $q \in \mathcal{P}$。然後，我們對 $hr$ 和 $hi$ 執行基本的梯度檢查（gradcheck），根據 $\mathcal{P}$，使用上面描述的實數到實數或複數到實數的情況。

請注意，在撰寫本文時，程式碼並未明確建立這些函數，而是透過將 $\text{grad\_out}$ 參數傳遞給不同的函數，來手動執行與 $real$ 或 $imag$ 函數的鏈式法則。當 $\text{grad\_out} = 1$ 時，我們正在考慮 $hr$。當 $\text{grad\_out} = 1j$ 時，我們正在考慮 $hi$。

快速反向模式梯度檢查

雖然上述的 gradcheck 公式很棒，但為了確保正確性與可除錯性，它的速度非常慢，因為它會重建完整的 Jacobian 矩陣。本節介紹一種更快執行 gradcheck 的方法，且不會影響其正確性。可除錯性可以透過在偵測到錯誤時加入特殊的邏輯來恢復。在這種情況下，我們可以執行預設版本，重建完整的矩陣，為使用者提供完整的詳細資訊。

這裡的高層次策略是找到一個純量，可以使用數值方法和解析方法有效地計算，並且能夠充分代表慢速 gradcheck 計算出的完整矩陣，以確保它能夠捕捉到 Jacobian 矩陣中的任何差異。

實數到實數函數的快速 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$ 向量進行點積。

快速複數輸入解析評估

對於解析的情況，事情更簡單，我們將公式改寫為

因此，我們可以利用反向模式自動微分 (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 的快速版本來實作。