最近看深度学习二阶优化方面的工作，顺便做一些整理。以下内容翻译自

"Xiaoxin et al., He Large-Scale Deep Learning Optimizations: A Comprehensive Survey, 2021"(https://arxiv.org/abs/2111.00856)

 

深度学习的优化在理论和经验上主要由一阶梯度方法主导，二阶优化方法由于涉及二阶求导或数据的二阶统计远没有被流行起来，尽管其有着很强的理论性质，主要是受阻于其计算量、内存消耗和通信成本。

1. 二阶优化基础

许多机器学习问题都可以简单描述为最小化关于变量$x\in \mathbb{R}^d$的损失函数：
\begin{equation}
\min_{x\in \mathbb{R}^d} F(x)
\end{equation}
当训练一个神经网络的权重时，我们尝试尽可能降低误差曲面。在大多数情况下，我们常用SGD来更新参数向量去解决该优化问题：
\begin{equation}
x_{t+1}=x_t-\eta_t g_t
\end{equation}
在实践中，另外一种非常流行的算法是自适应优化算法（如AdaGrad，Adadelta、RMSProp、Adam等）。有一些基本算法对参数向量中每个元素进行更新。每个元素有其自己的步长，用过去的梯度进行自适应更新。
\begin{equation}
x_{t+1}^{(i)}=x_t^{(i)}-\eta_t^{(i)} g_t^{(i)}, i=1,\cdots,d
\end{equation}
这些方法也有一些动量变形有一些不同的更新的规则。 潜在的且非常强大的一类算法被称为预条件(precondition)算法，其用一些矩阵称为预条件子 (preconditioner)在取一个步长时对梯度进行变换。一般来说，二阶优化的思想是通过一个局部优化对目标函数进行建模
\begin{equation}
f(x+\delta) \approx M(x) = f(x) + \nabla f(x)^\top \delta + \frac{1}{2}\delta^\top B(x) \delta
\end{equation}
这里$B$是对称预条件子(preconditioner)，$\delta$ 是参数的改变量。在牛顿法里面，$B=H$, 或 $B=H+\lambda I$. 由于关于$\delta$的$M(x_t)$得到
\begin{equation}
\delta^*={\rm{arg}}\min_\delta M(x_t) = - B^{-1}\nabla f(x),
\end{equation}
然后应用这个更新
\begin{equation}
x_{t+1} = x_t + \delta^*.
\end{equation}
这类算法包括AdaGrad, Adam, 其中预条件子是对角的。但是，如果这些度量不是对角的而是全矩阵预条件子，这些算法将更加强大，例如，完全AdaGrad、Natural Gradient以及如牛顿方法、拟牛顿法等。 在优化中预条件方法通常可以在许多场景下得到更快的收敛速度或更好的``条件数"。但是它也有明显的问题：假设参数量为$n$, 我们需要：1）最少二次方的空间复杂度$\Omega(n^2)$来存储预条件子。2）$n^3$的时间复杂度计算预条件子的逆以应用于梯度向量。 一般来说， 这些二阶方法并不实际，因为其计算成本和内存对于深度学习的设置中是难以承受的。可替换的预条件子是对角化近似或用SGD。

最近，二阶方法取得了显著进展，通过探索全矩阵和对角矩阵之间的折中。这些方法通常以模块化的方式对梯度进行预条件，模块化可以实现与全矩阵方式的相同或近似的效果，而且可以在存储和训练时间上实现与对角情况相同。受自然梯度方法的思想启发，Martens等人提出了利用Fisher矩阵的Kronecker因子近似作为预条件矩阵，其可以应用到多层感知器（MLP）,之后又被扩展到其他架构，如卷积神经网络和循环神经网络中。之后，基于Hessian和拟牛顿方法的Kronecker因子近似也被提出。

2. 牛顿法

一个函数的梯度被定义为偏导数的向量
\begin{equation}
g_t = \nabla f(x) = [\frac{\partial f}{x_1},\ldots, \frac{\partial f}{x_n} ]
\end{equation}
这意味着我们能假设神经网络的误差曲面局部表现类似一个圆。我们将忽略曲面的所有曲率，这可能会导致我们的训练会变得非常缓慢。 为了修正该问题， 我们可以用来自函数的二阶导数的信息。牛顿法的思想是应用线性变换将椭圆变成圆。如果我们应用这个变换到梯度向量，就好像我们在圆误差曲面的下坡。形式上，牛顿法用Hessian矩阵作为预条件子，
\begin{equation}
H = [H_{ij}],\quad H_{ij} = \frac{\partial^2 f}{\partial x_i\partial x_j}
\end{equation}
该Hessian是参数的函数，我们需要取其逆并乘以梯度，然后朝该方向走一定距离
\begin{equation}
\Delta = -\eta H(x)^{-1} \nabla f(x)
\end{equation}
如果其是一个真正的二次曲面且选择正确的学习率，我们将可以通过一步到达曲面的最小点。但是，该单步会涉及一些复杂内容，如Hessian矩阵的求逆。 假设我们神经网络仅仅有1M的参数，该Hessian矩阵将有1T(1M$\times$1M)的项,其求逆是完全做不到的。

曲率矩阵. 曲率矩阵中每个元素给出的是当我们向其他方向移动时，一个方向的梯度如何变换的。曲率矩阵中的非对角线项对应于误差面中的“扭转”。一个扭转意味着当你朝着一个方向移动时，梯度在另外一个方向的变化。 如果我们有一个漂亮的圆形曲面，非对角线的所有元素将变成0。当我们朝着一个方向移动，梯度在其他方向上将不会改变。 但是当我们有一个椭圆的误差面，如果我们朝着一个放移动，其他方向的梯度将会改变。这实际上是梯度方向的问题所在。 当梯度方向更新了权重，与此同时，其也更新了其他权重，引起第一个权重的梯度变化。这意味着当我们更新权重我们可能会导致事情变得更加糟糕。由于所有其他权重的变化，梯度实际上可能有反向符号。所以我们得到的权重越多，我们需要改变他们就越谨慎，因为所有其他权重的同时变化可能会改变权重的梯度。

如何避免求解大矩阵的逆.  曲率的密集型计算限制了二阶优化方法在深度学习中的应用。 为了解决该问题，有许多方法被提出。 一种十分流行的工作是采用对角近似，如AdaGrad、RMSProp、Adadelta等。但是这些对角元素仅仅由很少比例的交互所组成，其忽略了曲率矩阵中大多数交互项。实验表明这种做法在实践中相对于调优的SGD是基本没有改善的。这些对角线方法的好处似乎主要在于易于选择学习率，但可能不会提供除此之外的任何基本好处。我们可以做的另外一种做法是用低秩矩阵来近似曲率矩阵，捕捉它的主要方面。这里代表性的方法是Limited-Memory BFGS(L-BFGS)。

最近，许多关注主要在全矩阵和对角矩阵的折中，寻求两种极端之间的平衡。一些研究者寻找方法其与全矩阵情况相同或相似效果的，且在实际应用中能够与对角矩阵在存储和运行时间相似。 一些最近的近似全矩阵的方法有K-FAC、Shampoo等。其他方法还有自动植入Hessian操作算子，如Hessian-Free方法和信赖域方法。

3. Hessian-Free方法

Hessian-Free方法是一个拟牛顿方法，其用无低秩近似。称为“Free”的原因是其没有显式的计算预条件子B取而代之的是二次模型$M(\delta)$的近似最小化。该Hessian-Free方法主要受两个观察启发。第一个是计算对于任意向量$v$的矩阵向量乘法$Hv$是相对容易的，例如利用有限差去近似极限。
\begin{equation}
Hv = \lim_{\epsilon\longrightarrow 0} \frac{\nabla f(x+\epsilon v)-\nabla f(x) }{\epsilon}
\end{equation}
第二个观测是线性共轭梯度仅用矩阵向量乘法最小化正定二次成本函数，其相对容易计算的。 共轭梯度没有像牛顿方法达通过一步到最小点。其通过梯度方向开始，在这个方向达到最小值，可能需要重新评估梯度或重新评估误差，以找到这个方向的最小值。一旦做完，共轭梯度方法找另外一个方向，按照第二个方向达到最小点。这项技术选择第二个方向时不会打乱它在第一个方向上已经做过的最小化，这被称为共轭方向。 “共轭”是指当我们走向新的方向时，我们不会改变前面方向上的梯度。共轭梯度方法对于$n$-维二次曲面可以在$n$迭代步到达全局最小点。更加重要的是，在一个典型的二次曲面上的许多小于$n$的步骤中，它将减少的误差与在最小点基本接近。如果执行全部$n$步，其计算成本将于求整个矩阵逆相差无异。阻尼Hessian-Free优化的一个简单变形如算法所示：

深度学习的二阶优化进展

但是，Hessian-free的常见变形并不能在神经网络模型中表现并不是很好。 一些强化技术被提出，如对Hessian的高斯-牛顿近似、共轭梯度的early stop、阻尼方法等。 最近研究表明，深度神经网络学习比以前认为的使用简单的方法更容易。仔细调优的动量方法足以处理深度和循环神经网络目标中的曲率问题，而不需要复杂的二阶方法。尽管SGD有或没有动量仍然是在大多数情况下最广泛使用和最好的方法，事实上Hessian-Free方法用比SGD的100-1000倍更少的迭代，原则上二阶方法可以帮助更多，只要我们可以便宜的使用这些迭代计算。

4. Kronecker因子近似曲率算法(K-FAC)

Kronecker因子近似曲率算法(Kronecker-Factored Approximate Curvature, K-FAC)是一种自然梯度近似方法，其预条件子(Preconditioner)是一种Fisher信息矩阵的高质量近似。下面首先介绍自然梯度下降，然后解释K-FAC的算法框架。

自然梯度下降(NGD)，是建立在信息几何的二阶优化方法。NGD利用Fisher信息矩阵作为损失函数的曲率，可以正确的获得损失函数全貌，并且在迭代方面比一种简单的一阶方法收敛得更快。关于网络分布$p(y|x;\theta)$的Fisher信息矩阵为
\begin{equation}
F=E[\nabla\log(p(y|x;\theta))\nabla \log(p(y|x;\theta))^\top].
\end{equation}
重要的是，$F$的一个特征是，它可以被解释为模型的对数似然的负期望Hessian，即
\begin{equation}
F=-E_{p(y|x;\theta)}[H_{\log p(x|\theta)}]
\end{equation}
从该结论中可以看到$F$作为对数似然函数曲率度量的作用。因此，$F$的直接应用是在二阶优化方法中作为$H$的直接替代。利用KL散度来衡量两种模型的不同程度，自然梯度下降NGD的更新规则为
\begin{equation}
\theta_{t+1} = \theta_t - \eta_t F^{-1} \nabla f(\theta_t).
\end{equation}

K-FAC近似。 K-FAC近似Fisher信息矩阵可以使很容易计算逆矩阵。首先，K-FAC将$F$近似为$\hat F$, 一种对角块矩阵，其中每个块对应神经网络的一层。
\begin{equation}
\hat F = {\rm{diag}} (\hat F_1, \ldots, \hat F_l,\ldots, \hat F_L).
\end{equation}
下一步，每个对角块矩阵$\hat F_l$ 被近似成Kronecker乘积,即
\begin{equation}
\hat F_l \approx A_{l-1} \otimes G_l.
\end{equation}
这被称为Kronecker的因式分解，$G_l, A_{l-1}$也称为Kronecker因子，分别表示第$l$层输出的梯度和第$l-1$层的激活输出。根据Kronecker矩阵乘积的重要性质$(A \otimes B)^{-1}=A^{-1}\otimes B^{-1}$, $\hat F_l$的逆可以通过下式子进行计算
\begin{equation}
\hat F_l^{-1} = A_{l-1}^{-1}\otimes G_l^{-1}.
\end{equation}
第$l$层的$w_l$最终迭代步如下
\begin{equation}
w_l^{(t+1)} =w_l^{(t)} - \eta^{(t)}G_l^{-1}\nabla L_l^{(t)}A_{l-1}^{-1}.
\end{equation}
在多数实现中，Tikhonov正则项被用于避免病态条件数矩阵逆，通过在$\hat F$的对角线上添加一个阻尼系数$\gamma$,
\begin{align}
V_1 &= Q_G^\top L_i(w_i^{(k)})Q_A\\
V_2 &= V_1/(v_G v_A^\top+\lambda)\\
(\hat F_l +\gamma I)^{-1} \nabla L_i(w_i^{(k)}) &= Q_G V_2 Q_A^\top.
\end{align}
在实践中，可以通过减少这些因素及其特征分解的计算频率来避免重要的计算和通信，但代价是引入隐藏的信息。

5. Shampoo算法

Shampoo是另外一种预条件算法，其是全矩阵AdaGrad的近似版本，其首先通过独立处理每一层（块对角线）来近似。它使用了小矩阵， 其 Kronecker乘积近似全AdaGrad统计。这两种近似使Shampoo比对角线预条件具有更强的表达力，而且可以避免大规模存储和计算。从数学上来讲，所寻找的预条件子Shampoo可以写成两个较小的矩阵$L$和$R$的Kronecker乘积，即
\begin{equation}
\mathrm{arg}\min_{H=L\otimes R,\\L,R>0} H^{-1}\bar G_t + Tr(H).
\end{equation}
虽然她不能解决精确的优化问题，但它有一个很好的上界，
\begin{equation}
\frac{1}{r} (\sum_{t=1}^T g_t(g_t)^\top )^{1/2} \preceq (\sum_{t=1}^T G_t G_t^\top)^{1/4}\otimes (\sum_{t=1}^T G_t^\top G_t )^{1/4}
\end{equation}
左侧是全AdaGrad预条件子，其上界是两个小矩阵$L_t, R_t$的Kronecker乘积（右边的$\otimes$的左侧项$L_t$，右侧项为$R_t$）。Shampoo的更新如下：
\begin{align}
L_t &= L_{t-1} + G_t G_t^\top,\\
R_t &= R_{t-1} + G_t^\top G_t,\\
W_{t+1} &= W_{t} - L_t^{-1}G_t R_t^{-1}.
\end{align}

Shampoo和K-FAC都采用了一种近似于Fisher信息矩阵的预条件方案，尽管他们在结构上相似，但在几个重要方面却有所不同。一个重要的不同是，Shampoo构建与层类型无关，而K-FAC十分依赖于前馈神经网络中的反向传播梯度的结构。相比之下，Shampoo实际上忽略了特定的模型结构，而只依赖于标准的梯度信息。最近，Anil等人以多种方式扩展了Shampoo算法，使其适用于更大范围的深度架构。

尽管一阶方法在近十年里都占据了主导地位，但最近二阶方法，如K-FAC和Shampoo显示出了一定的前景，它们降低了全矩阵二阶算法的空间和运行成本，并已适用于更大范围的深度架构。