伴随方法：线性方程的伴随方程（Adjoint Equation）

伴随方法是 Neural-ODE 中十分重要的一个方法，它让一个计算量复杂到基本无法求解的问题变得有可能。在神经网络中嵌套线性方程或者非线性方程也会遇到同样的问题，这篇文章从最简单的例子线性方程中的网络参数求解中，表达一下伴随方法的思想以及一些公式的推导。

假设现在有一个线性系统 $\mathbf{A}\boldsymbol{x}=\boldsymbol{b}$ ，其中矩阵 $\mathbf{A}$ 和 $\boldsymbol{b}$ 都是参数 $\theta$ 的函数，那么线性系统可以表示为 $\mathbf{A}(\theta)\boldsymbol{x}=\boldsymbol{b}(\theta)$ 。在机器学习领域， $\mathbf{A}(\theta)$ 和 $\boldsymbol{b}(\theta)$ 可以看做是神经网络， $\theta$ 是神经网络的参数，那么自然而然地，我们的目标就是想要求得损失函数关于网络参数 $\theta$ 的导数，然后利用梯度下降以及优化算法来训练网络。

对于一个线性方程，有许多的方法来求解得到 $\boldsymbol{x}$ ，假设 $\boldsymbol{x}$ 会作为模型最后的预测结果，那么最终它会输入到一个损失函数 $J(\boldsymbol{x})$ 中，可能会有真实标签与其对应。因此，我们最终要求的就是损失函数关于参数的导数 ${\text{d}J}/{\text{d}\theta}$ 。

因为 $\mathbf{A}(\theta)$ 和 $\boldsymbol{b}(\theta)$ 都是由 $\theta$ 决定的，因此 $\boldsymbol{x}$ 实际上也是 $\theta$ 的隐式函数，所以可以写成 $\boldsymbol{x}(\theta)$ 。我们假设参数 $\theta$ 的维度为 $P$ ，即 $\theta\in\mathbb{R}^{P}$ ，其他的矩阵以及向量的维度分别为 $\mathbf{A}(\theta)\in\mathbb{R}^{N\times N}$ ， $\boldsymbol{x}(\theta)\in\mathbb{R}^N$ ， $\boldsymbol(\theta)\in\mathbb{R}^N$ 。有得时候损失函数也会是 $\theta$ 的函数，因此具体地写出来损失函数就是 $J(\boldsymbol{x}(\theta);\theta)$ .

注意：为了方便各种符号的简化，下面继续表示这些变量的时候，会省略后面的 $\theta$ ，但是读者应该记住这些变量依旧是 $\theta$ 的函数，在求导的时候要一直考虑这一项。

我们想要得到的是 $\text{d}J/\text{d}\theta$ ，要注意的是这里表达的是全微分，因此有：

\underbrace{\frac{\text{d}J}{\text{d}\theta}}_{\mathbb{R}^{1\times P}} = \underbrace{\frac{\partial J}{\partial \theta}}_{\mathbb{R}^{1\times P}} + \underbrace{\frac{\partial J}{\partial \boldsymbol{x}}}_{\mathbb{R}^{1\times N}} \times \underbrace{\frac{\text{d}\boldsymbol{x}}{\text{d}\theta}}_{\mathbb{R}^{N\times P}}\tag{1},

在每一个变量的下面都标上了各自的维度。因为 $\boldsymbol{x}$ 和 $\theta$ 都是一个向量，因此 $\text{d}\boldsymbol{x}/\text{d}\theta$ 是一个雅可比矩阵，在这式子当中， $\text{d}\boldsymbol{x}/\text{d}\theta$ 是最难求的。

我们对于线性系统 $\mathbf{A}\boldsymbol{x}=\boldsymbol{b}$ 的两端，都对 $\theta$ 进行求导，可以得到：

\frac{\text{d}}{\text{d}\theta}(\mathbf{A}\boldsymbol{x}) = \frac{\text{d}}{\text{d}\theta}(\boldsymbol{b})

\frac{\text{d} \mathbf{A}}{\text{d}\theta}\boldsymbol{x}+\mathbf{A} \underbrace{\frac{\text{d}\boldsymbol{x}}{\text{d}\theta}}_{\text{target}} = \frac{\text{d}\boldsymbol{b}}{\text{d}\theta}

我们的目标是求出 ${\text{d}\boldsymbol{x}}/{\text{d}\theta}$ 这一项，对其进行简单的变换：

\mathbf{A}\frac{\text{d}\boldsymbol{x}}{\text{d}\theta} = \frac{\text{d}\boldsymbol{b}}{\text{d}\theta}-\frac{\text{d}\mathbf{A}}{\text{d}\theta}\boldsymbol{x},\quad\text{（移项）}

方程两边同时左乘 $\mathbf{A}$ 的逆，得到：

\underbrace{\frac{\text{d}\boldsymbol{x}}{\text{d}\theta}}_{\mathbb{R}^{N\times P}} = \underbrace{\mathbf{A}^{-1}}_{\mathbb{R}^{N\times N}} \left( \underbrace{\frac{\text{d}\boldsymbol{b}}{\text{d}\theta}}_{\mathbb{R}^{N\times P}} - \underbrace{\frac{\text{d}\mathbf{A}}{\text{d}\theta}}_{\mathbb{R}^{N\times N\times P}} \underbrace{\boldsymbol{x}}_{\mathbb{R}^{N}} \right)\tag{2},

同样的，我们在变量下面标上对应的维度。要注意的是，这里 $\text{d}\mathbf{A}/\text{d}\theta$ 和 $\boldsymbol{x}$ 的维度是不匹配的，但是我们不拘泥于这里，我们关注的点在于如果要通过最直接的方式去求解 ${\text{d}\boldsymbol{x}}/{\text{d}\theta}$ 所需要的时间是有多大。这里只需要记住，无论如何，括号里面最终得到的矩阵维度为 $N\times P$ 的大小。同时也不用去过度的关注矩阵 $\mathbf{A}$ 要如何求逆（因为这里是一个神经网络的输出，所以求逆会使得问题变得更为复杂），因为在后面会发现其实没有必要对 $\mathbf{A}$ 求逆。

将式子 (2) 与线性方程 $\mathbf{A}\boldsymbol{x}=\boldsymbol{b}$ 进行对比可以发现，其实这就是由 $P$ 个线性方程组成的更大的线性方程。求解一个线性方程可以用 LU 分解或者 QR 分解，它们的时间复杂度为 $\mathcal{O}(N^3)$ ，时间花费太过于大，对于神经网络来说，参数一多基本无法求解。因此，我们要使用另外一种更为高效的方法 —— 伴随方法，来求解这个问题。

伴随方法（Adjoint Method）

我们观察 (1) 式子以及 (2) 式，会发现实际上 (1) 式的最后一项就是我们想要求的「目标」，那么我们可以将 (2) 代入到 (1) 式中，得到 (3) 式：

\underbrace{\frac{\text{d}J}{\text{d}\theta}}_{\mathbb{R}^{1\times P}} = \frac{\partial J}{\partial \theta} + \underbrace{\frac{\partial J}{\partial \boldsymbol{x}}}_{\mathbb{R}^{1\times N}} \underbrace{\mathbf{A}^{-1}\left( \frac{\text{d}\boldsymbol{b}}{\text{d}\theta} - \frac{\text{d}\mathbf{A}}{\text{d}\theta}\boldsymbol{x}\right)}_{\mathbb{R}^{N\times P}}\tag{3},

我们发现最后括号里面的那一整块维度是 $N\times P$ 的，而我们最终需要的只是一个 $1\times P$ 的向量，这说明，实际上我们不需要额外求解 $P$ 个线性方程，而只需要额外求解 1 个线性方程就能行了。

我们重新把 (3) 式分块来看：

\frac{\text{d}J}{\text{d}\theta} = \frac{\partial J}{\partial \theta} + \underbrace{\left( \frac{\partial J}{\partial \boldsymbol{x}} \mathbf{A}^{-1}\right)}_{\lambda^\top} \left( \frac{\text{d}\boldsymbol{b}}{\text{d}\theta} - \frac{\text{d}\mathbf{A}}{\text{d}\theta}\boldsymbol{x} \right)\tag{4},

我们令 $\lambda^\top = \frac{\partial J}{\partial \boldsymbol{x}} \mathbf{A}^{-1}$ ，称 $\lambda\in\mathbb{R}^N$ 为伴随变量（adjoint variable），然后对这个方程进行如下变换：

\lambda^\top \mathbf{A} = \frac{\partial J}{\partial \boldsymbol{x}},\quad\text{(两边右乘 $\mathbf{A}$)}

\left( \lambda^\top \mathbf{A} \right)^\top = \left( \frac{\partial J}{\partial \boldsymbol{x}} \right)^\top,\quad\text{(两边进行转置）}

最后我们得到 (5) 式：

\underbrace{\mathbf{A}^\top}_{\mathbb{R}^{N\times N}} \underbrace{\lambda}_{\mathbb{R}^{N}} = \underbrace{\left( \frac{\partial J}{\partial \boldsymbol{x}} \right)^\top}_{\mathbb{R}^{N}}\tag{5}

观察 (5) 式不难发现，这其实与 $\mathbf{A}\boldsymbol{x}=\boldsymbol{b}$ 的形式是完全一样的，而且我们不用计算矩阵 $\mathbf{A}$ 的逆，而是直接用它的转置，关于 $\frac{\partial J}{\partial \boldsymbol{x}}$ 这一项，利用自动微分可以很简单地计算出来。

这种求解方法就很好地规避了求逆，并且使得问题的维度大大地减小了。对于伴随方法，可以通过以下三步来计算：

第一步：前向求解 $\mathbf{A}\boldsymbol{x}=\boldsymbol{b}$ ，得到 $\boldsymbol{x}$ 的解；

第二步：后向求解伴随方程 $\mathbf{A}^\top \lambda = \left( \frac{\partial J}{\partial \boldsymbol{x}} \right)^\top$ ，得到伴随变量 $\lambda$ ；

第三步：代回原式：

\frac{\text{d}J}{\text{d}\theta} = \underbrace{\frac{\partial J}{\partial \theta}}_{\text{may be zero in many problems}} + \lambda^\top \left( \frac{\text{d}\boldsymbol{b}}{\text{d}\theta} - \frac{\text{d}\mathbf{A}}{\text{d}\theta} \boldsymbol{x} \right)

利用这样的伴随方法，只需要求解两个线性系统就可以得到 $\frac{\text{d}J}{\text{d}\theta}$ 。而对于 $\frac{\partial J}{\partial \boldsymbol{x}}, \frac{\partial J}{\partial\theta}, \frac{\text{d}\boldsymbol{b}}{\text{d}\theta}, \frac{\text{d}\mathbf{A}}{\text{d}\theta}$ ，这几个矩阵利用自动微分可以更为简单地求得。

参考：

[1] Machine Learning & Simulation. Adjoint Equation of a Linear System of Equations - by implicit derivative. YouTube