从零发明数值最优化 5：共轭梯度法

在之前的梯度下降法中，我们提到，即使在精确线搜索下，梯度下降法也大概率走出一条全是直角拐弯的优化路径。牛顿法虽然解决了这个问题，然而它的存储开销和计算开销实在太大了。我们最希望的当然是既只用一阶梯度，又能快速下降。

在通常的函数中这可能比较难以做到，那让我们退而求其次，来考察一下严格正定的二次型函数：

$$f(x) = \frac{1}{2} x^\top A x - b^\top x, \quad A \in \mathbb{R}^{n \times n}, \quad A \succ 0, \quad A^\top = A$$

求这个函数的最小值当然就是求 $\nabla f(x) = Ax - b = 0$。

其实，共轭梯度法最初在 1952 年就是为了求出 $Ax = b$ 而设计的。不过这里我们为了逻辑上的通畅，选择先有函数再算梯度的逻辑。

直观感受

在一个圆上进行下降，不管从什么地方出发都是一步的事。此时，$A = \mu I$，$x^\star = \frac{1}{\mu} b$。

在一个椭圆上就不然。椭圆上的下降曲线要不断拐弯，逼近得很慢。而且即使每步都用精确线搜索优化到了最佳，下一步又会破坏已知的最优。

那我们的想法当然是，能不能第一步最优，第二步在把不破坏第一步最优的基础上也最优，第三步在不破坏前两步最优的基础上也最优……

如何下降

那么，应该怎么走呢？

比较直观的想法是，矩阵是一个 $n \times n$ 的东西，而梯度是一个 $n \times 1$ 的东西，所以大概 $n$ 次就能从任意点收敛到最优值。而且，因为要走 $n$ 步，每步都走“尽” $\mathbb{R}^n$ 的一个维度，所以每走一步都要走到当前步最优的位置，而且肯定不能损害前一步的最优性。

我们经常算便宜的梯度 $\nabla f(x_k) = Ax_k - b$。另外，我们还知道 $\nabla f(x^\star) = A x^\star - b = 0$。我们希望最终 $x_k = x^\star$，即误差 $e_k = x_k - x^\star = 0$，所以会想把上面这两个式子做个差 $\nabla f(x_k) - \nabla f(x^\star) = \nabla f(x_k) = A(x_k - x^\star) = A e_k$，得到 $\nabla f(x_k) = A e_k$。这个式子看上去很有意思，我们先记下来。

把我们沿着走的方向记为 $p_1, p_2, \dots, p_n$。我们希望最后满足 $-e_1 = \sum_{i=1}^n \alpha_i p_i$。

那么当走了 $k$ 步时，我们应该怎么知道第 $k+1$ 步往哪儿走，走多远呢？

走多远

假设上帝已经给了我们一组能够铺满全空间的搜索方向 $p_1, p_2, \dots, p_n$。

在走完前 $k-1$ 步后，当前的累积位置是 $x_k = x_1 + \sum_{i=1}^{k-1} \alpha_i p_i$，当前的剩余误差是 $e_k = x_k - x^\star$。从全局来看，我们还剩下的总差距是：

$$-e_k = x^\star - x_k = \sum_{i=k}^n \alpha_i p_i$$

我们希望在第 $k$ 步就把这个式子中的 $\alpha_k$ 解出来，并且不用知道未来还要走多远，怎么走。总的来说，我们希望 $p_i$ 和 $p_j$ 之间满足某种联系，使得它们能进行顺滑的消消乐。

最常规的联系就是两两垂直，即 $p_i^\top p_j = 0 \quad (i \neq j)$。然而，即使满足这个条件，把上式两边都左乘一个 $p_k$，我们也只能得到

$$-p_k^\top e_k = \alpha_k \Vert p_k \Vert^2 + \sum_{i=k+1}^n \alpha_i (p_k^\top p_i) = \alpha_k \Vert p_k \Vert^2$$

问题是很明显的：我们不知道终末状态，所以 $e_k$ 也无从得知。

然后你突然想起来上面那个神秘的式子：$\nabla f(x_k) = A e_k$。如果我们左乘一个 $p_k^\top A$，又会得到什么呢？

$$- p_k^\top A e_k = \alpha_k (p_k^\top A p_k) + \sum_{i = k+1}^n \alpha_i (p_k^\top A p_i)$$

化简，

$$- p_k^\top \nabla f(x_k) = \alpha_k (p_k^\top A p_k) + \sum_{i = k+1}^n \alpha_i (p_k^\top A p_i)$$

刚刚上帝给了你 $p_1$ 到 $p_n$，所以 $p_k$ 是已知的。$A$ 和 $\nabla f(x_k)$ 当然也是已知的。这么看来，你要求 $\alpha_k$，最方便的办法当然就是让后面那一大串求和等于零，也就是

$$p_i^\top A p_j = 0, \quad i \neq j$$

这就是大名鼎鼎的共轭条件。而在这个条件下，每步的步长就是

$$\alpha_k = \frac{- p_k^\top \nabla f(x_k)}{p_k^\top A p_k}$$

往哪儿走

现在我们知道走多远了，那么问题就转化为往哪儿走。很遗憾，现实生活中可没有上帝，我们还是得解决已知 $p_1, \dots, p_{k-1}$，如何求 $p_k$ 的问题。

幸好，在上面的推理中，我们还顺便得到了 $p_i^\top A p_j = 0, \quad i \neq j$ 的共轭关系。

最简单的想法当然是用施密特正交化。

我们拿手头上现成的向量 $\nabla f(x_k)$ 来做每步的施密特正交化。定义残差 $r_k = A x^\star - A x_k = b - A x_k = - \nabla f(x_k)$。那么理论上，下一个方向就是

$$p_k = r_k - \sum_{i=1}^{k-1} \frac{p_i^\top A r_k}{p_i^\top A p_i} p_i$$

不过，施密特正交化是 $O(n^3)$ 的算法，和矩阵求逆一样麻烦。我们看看能不能找到一点好的性质避免这么高的时间复杂度。

我们在第 $k$ 步所处的位置 $x_k$，实际上是在前 $k-1$ 个共轭方向张成的子空间内达到了全局最优。这意味着当前的梯度与已有的所有搜索方向都正交，也就是说：

$$r_k^\top p_i = 0, \quad \forall i < k$$

第一个搜索方向没有前面的信息，我们干脆定义 $p_1 = r_1$，而每次 $p_k$ 都是用之前的 $p_1, \dots, p_{k-1}$ 和 $r_k$ 来更新的。所以我们显然有 $\mathrm{span}\{p_1, \dots, p_k\} = \mathrm{span}\{r_1, \dots, r_k\}$。

在第 $k$ 个点处，我们有 $r_k^\top p_i = 0, \quad \forall i < k$ 和 $\mathrm{span}\{p_1, \dots, p_{k-1}\} = \mathrm{span}\{r_1, \dots, r_{k-1}\}$。所以，我们得到：

$$r_k^\top r_i = 0, \quad \forall i < k$$

回顾残差的更新方式：

$$r_{i+1} = b - A x_{i+1} = b - A (x_i + \alpha_i p_i) = r_i - \alpha_i A p_i$$

$$A p_i = \frac{1}{\alpha_i} (r_i - r_{i+1})$$

代入到施密特正交化的分子上：

$$p_i^\top A r_k = r_k^\top A p_i = \frac{1}{\alpha_i} r_k^\top (r_i - r_{i+1})$$

又根据前面得到的 $r_k^\top r_i = 0, \quad \forall i < k$，我们知道

$$p_i^\top A r_k = \frac{1}{\alpha_i} r_k^\top (r_i - r_{i+1}) = \begin{cases} 0 & i < k-1 \\ - \frac{1}{\alpha_{k-1}} \Vert r_k \Vert^2 & i = k-1 \end{cases}$$

因此，

$$p_k = r_k - \frac{- \frac{1}{\alpha_{k-1}} \Vert r_k \Vert ^2}{p_{k-1}^\top A p_{k-1}} p_{k-1}$$

其实这个式子的分母还可以用 $A p_{k-1} = \frac{1}{\alpha_{k-1}} (r_{k-1} - r_k)$、$p_{k-1}^\top r_{k-1} = r_{k-1}^\top r_{k-1}$ 和 $p_{k-1}^\top r_k = 0$ 继续化简，得到

$$p_k = r_k + \frac{\Vert r_k \Vert ^2}{\Vert r_{k-1} \Vert ^2} p_{k-1}$$

这样一来，站在当前点的下降方向只需要用当前点的梯度、上一个点的梯度和上一个点的下降方向就可以得出。这大大节省了存储空间。

总览

初步的共轭梯度法先设定 $r_1 = b - A x_1$，$p_1 = r_1$，接下来进行迭代：

$$\begin{aligned} \alpha_k &= \frac{p_k^\top r_k}{p_k^\top A p_k} \\ x_{k+1} &= x_k + \alpha_k p_k \\ r_{k+1} &= b - A x_{k+1} \\ \beta_{k+1} &= - \frac{p_k^\top A r_{k+1}}{p_k^\top A p_k} \\ p_{k+1} &= r_{k+1} + \beta_{k+1} p_k \end{aligned}$$

而更进一步化简后的共轭梯度法同样先设定 $r_1 = b - A x_1$，$p_1 = r_1$，不过接下来则是：

$$\begin{aligned} \alpha_k &= \frac{\Vert{}r_k\Vert{}^2}{p_k^\top A p_k} \\ x_{k+1} &= x_k + \alpha_k p_k \\ r_{k+1} &= r_k - \alpha_k A p_k \\ \beta_{k+1} &= \frac{\Vert{}r_{k+1}\Vert{}^2}{\Vert{}r_k\Vert{}^2} \\ p_{k+1} &= r_{k+1} + \beta_{k+1} p_k \end{aligned}$$

化简后，矩阵-向量乘法的次数从每轮 2-3 次减少为严格 1 次，向量点积每轮迭代仅计算两次，且保证了 $O(n)$ 的时间复杂度，是实际场景下的高效实现。