大规模神经网络优化：二阶优化器_天富娱乐-天富注册装修站

在本节中，我们将学习一系列二阶优化器。首先我们会比较全面的了解与二阶优化器相关的基础概念。在充分理解这些概念的基础下，我们将很轻松的理解各类二阶优化器。(系列目录）

考虑函数的二阶泰勒展开

$f(x+\\Delta x)\\approx f(x)+\ abla f(x)^T\\Delta x+\\frac{1}{2}\\Delta x^T\ abla^2f(x)\\Delta x$

当上式取得最小值时，有 $\\Delta x=-[\ abla^2f(x)]^{-1}\ abla f(x)$ 。牛顿法即使用更新： $\ heta_{n+1}=\ heta_n-\\epsilon H^{-1}\ abla f(\ heta_n)$ ，其中 $\ heta$ 为网络参数，黑塞阵 $H=\ abla^2f(\ heta_n)$ ，这里我们简单考虑 $H$ 可逆。在凸函数上对比梯度下降法收敛率 $O(1/\\sqrt{\\epsilon})$ ，当初值接近收敛值时，牛顿法有更快的收敛率 $O(\\log\\log(1/\\epsilon))$ 。

牛顿法计算黑塞阵（ $O(n^2)$ ）和矩阵求逆（朴素 $O(n^3)$ ，当 $H$ 不可逆时需求解 $Hs=-\ abla f(x)$ ）的计算复杂度高。

在前面的章节中，我们已经讨论过高斯-牛顿阵是黑塞阵很好的近似。记神经网络为 $F=\\sigma\\circ f,f(x,\\cdot):\\mathbb{R}^n\\rightarrow\\mathbb{R}^m,\\sigma:\\mathbb{R}^m\\rightarrow\\mathbb{R}$ ， $n$ 为参数量。则 $H\\approx G=J^TH_\\sigma J$ 。其中 $H_\\sigma$ 是函数 $\\sigma$ 对 $f(x)$ 的黑塞阵（形状为 $m\ imes m$ ）， $J$ 是 $f(x)$ 的雅可比矩阵（形状为 $m\ imes n$ ）。通过该方法可以计算雅克比矩阵避免对黑塞阵的计算，同时保证 $G$ 是半正定的。

另一条思路是将 $H(\ heta_t)$ 估计为 $B(\ heta_t)$ ， $B(\ heta_t)$ 利用递推的方式从 $B(\ heta_{t-1})$ 更新。我们可以让 $B$ 为对称正定阵等良好性质便于求得更新值 $s$ ，使得 $B(\ heta_t)s(\ heta_t)=-\ abla f(\ heta_t)$ 。

联想割线法（secant method）中对梯度的估计，我们有 $B(\ heta_k)\\approx\\frac{\ abla f(\ heta_{k}) - \ abla f(\ heta_{k-1})}{\ heta_{k}-\ heta_{k-1}}$ ，注意到 $s(\ heta_t)=\ heta_{k}-\ heta_{k-1}$ ，整理可得 $B(\ heta_k)s(\ heta_{k-1})=\\Delta \ abla\ heta$ 。在该公式下，根据不同的假设可以得到不同的拟牛顿法。这包括 SR1 方法（symmetric rank-one，假设 $\\Delta B$ 秩为一），Sherman-Morrison 方法（SR1 基础上直接对逆进行递推），BFGS 方法（假设 $\\Delta B$ 秩为二），DFP 方法（BFGS 基础上直接对逆进行递推）。其中，BFGS 方法可以保证 $B$ 正定，总时间复杂度为 $O(n^2)$ ，记 $B(\ heta_{k+1})$ 为 $B_+$ ， $B(\ heta_{k})$ 为 $B$ ，更新公式为

$B_{+}=B-\\frac{Bss^TB}{s^TBs}+\\frac{(\\Delta\ abla \ heta)(\\Delta\ abla \ heta)^T}{(\\Delta \ abla\ heta)^Ts}$

然而 $O(n^2)$ 的还是难以接受。L-BFGS 通过储存最近的 k 个 $\\{s_i, \\Delta(\ abla\ heta)_i\\}_{i=t-k}^k$ ，通过每次从初值开始，进行 k 次迭代用 IFQN 算法直接计算 $H^{-1}g$ ，从而降低复杂度到 $O(kn)$ 。但对于神经网络而言还是太大。

黑塞向量积 Hessian-vector-product 是关于 $H$ 的一个基本操作。虽然优化中我们更关注 $H^{-1}v$ ，而不是 $Hv$ （ $v$ 是[ $n\ imes 1$ ]的向量）。但该方法在分析神经网络曲面，以及凸函数的共轭梯度法中有应用（虽然在深度学习中梯度存在噪声，共轭梯度法不能保证方向的共轭性，难以应用）。

虽然计算 $H$ 需要 $O(n^2)$ 的复杂度，计算 $Hv$ 只需要 $O(n)$ 。在第一次反向传播获得梯度 $g$ 后，将梯度和向量 $v$ 做乘积得到的值 $g^Tv$ 再次进行反向传播，就可以得到 $Hv$ 。其原因在于 $\\frac{d(g^Tv)}{d\ heta}=\\frac{dg}{d\ heta}^Tv+g^T0=Hv$ 。这里 $Hv$ 的数学含义是方向梯度 $\ abla_v f$ 。PyTorch 中的实现为F.hvp。

目前的深度学习框架往往都提供了高效的 vector-Jaccobian-product 和 Jaccobian-vector-product，例如 PyTorch 中的 F.vjp 和 F.jvp。为了计算高斯-牛顿阵与向量的乘积 Gauss-Newton-vector-product（ $Gv=(J^TBJ)v$ ），第一步可以先计算 $g(x)=\\frac{\\partial \\mathcal{L}}{\\partial f(x)}$ 。通过计算 $g(x)$ 雅克比矩阵与向量 $v$ 的乘积（jvp），可得 $J_{g(x)}v=\\frac{\\partial L}{(\\partial f(x))^T\\partial f(x)}J_{f(x)}v=BJv$ 。继而计算向量 $(BJv)^T$ 与 $f(x)$ 雅克比矩阵的乘积（vjp），即可得到答案 $(BJv)^TJ=(J^TBJv)^T=(Gv)^T$ 。

利用 $Hv$ ，我们可以对黑塞阵的对角线元素进行估计。Hutchinson 方法是对实对称矩阵对角线元素的估计方法。令 $\\hat h=u\\cdot(Hu)$ ，其中 $u$ 是一个服从于高斯分布 $N(0,I_d)$ 的随机向量，则有 $\\mathbb{E}_{u}[u\\cdot(Hu)]=\\mathrm{diag}(H)$ 。

数学推导可参考该中文资料

在第一篇中我们已经定义过费雪信息矩阵（Fisher Information Matrix, FIM）。这里我们从神经网络的视角重新定义一次。一般而言，神经网络旨在建模一个分布 $p(x|\ heta)$ （对于判别式神经网络，建模条件分布 $p(y|x;\ heta)$ ），其中 $\ heta$ 是神经网络的参数。使用交叉熵损失与最大似然等价，目标为最大化 $\\log p(x|\ heta)$ 。定义评分函数 $s(\ heta)=\ abla_\ heta\\log p(x|\ heta)$ 为对数似然函数的梯度（神经网络梯度的负数）。

评分函数有性质 $\\mathbb{E}_{p(x|\ heta)}[s(\ heta)]=0$ 。这里要格外注意一点，期望是在分布 $p(x|\ heta)$ 下计算的，而不是在数据集上计算的。记数据分布为 $x\\sim \\mathcal{D}$ ， $\\mathcal{D}$ 与 $p(x|\ heta)$ 之间存在差距，随着训练的进行不断减小（训练就是 $p$ 对 $\\mathcal{D}$ 的拟合）。故 $\\mathbb{E}_{x\\sim D}[s(\ heta)]$ 不一定为零。类似的期望取值之后也要注意。

此时费雪信息阵可以有一个比较直观的定义：FIM 即评分函数（网络导数）的二阶矩。即 $F=\\mathbb{E}_{p(x|\ heta)}[\ abla_\ heta\\log p(x|\ heta)\ abla_\ heta\\log p(x|\ heta)^T]$ 。由于 $\\mathbb{E}_{p(x|\ heta)}[\ abla_\ heta\\log p(x|\ heta)]=0$ ，此时二阶矩就是评分函数的协方差。

在之前我们已经知道 FIM 与 Hessian 的关系，即 $F=-\\mathbb{E}_{x\\sim p(x|\ heta)}[H]$ 。但此黑塞阵不是在数据集上计算的。虽然当训练充分，网络拟合的分布接近实际分布时两者应该接近，但一般而言， $F$ 与实际的 $\\mathbb{E}_{x\\sim\\mathcal{D}}[H]$ 的差距有多少呢？

这里我们借助[LLH'23+]中的推导分析判别式网络，证明其 FIM 就是该网络的高斯-牛顿阵，即 $F=G$ 。首先，记网络输出为 $p(y|x,\ heta)=\\vec p=\ extrm{softmax}(\\vec z),\\vec z=f(\ heta,x)$ 。对 $\\vec z$ 求导可得 $\\frac{\\partial L}{\\partial \\vec z}=\\vec p - \\vec y$ 。故在拟牛顿法开头提到的 $H_\\sigma=H_{\ extrm{softmax}}=\\frac{\\partial L}{\\partial^2 \\vec z}=\\frac{\\partial \\vec p - \\vec y}{\\partial \\vec z}=\ extrm{diag}(\\vec p)-\\vec p\\vec p^T$ 与真实类别无关。其中 $\ extrm{diag}(\\vec p)$ 是一个对角阵，对角线上的元素为 $\\vec p$ 。

由于 $H_\\sigma$ 与真实类别无关，我们无需按照数据集进行采样，而是可以直接在网络输出的分布 $p(y|x,\ heta)$ 上采样。如此 $H_\\sigma=F_\\sigma=\\mathbb{E}_{y\\sim p(\ heta,x)}[\\frac{\\partial L}{\\partial z}\\frac{\\partial L}{\\partial z}^T]$ 。由于 $J$ 是 $z$ 对 $\ heta$ 的雅克比矩阵，与真实标签无关，故 $G=J^TH_\\sigma J=\\mathbb{E}_{y\\sim p(\ heta,x)}[J\\frac{\\partial L}{\\partial z}(\\frac{\\partial L}{\\partial z}J)^T]=\\mathbb{E}_{y\\sim p(\ heta,x)}[\ abla L\ abla L^T]=F$ 。

由于在神经网络中， $G$ 是 $H$ 很好的近似，故 $F$ 也是 $H$ 很好的近似。因此，可以用 $F$ 代替 $H$ 进行优化，即 $\ heta_{n+1}=\ heta_n-\\epsilon F^{-1}\ abla f(\ heta_n)$ 。但 $F$ 仍然需要 $O(n^2)$ 的储存开销，同时还有按网络分布采样计算梯度，以及计算 $F^{-1}$ 的计算开销。

FIM 阵除了可以看作是对 Hessian 的估计以外， $\ heta_{n+1}=\ heta_n-\\epsilon F^{-1}\ abla f(\ heta_n)$ 本身还具有一个直观的数据解释。单步优化中，我们希望更新方向使损失最小：

$d^*=\\arg\\min_{d,D(\ heta,\ heta+d)=\\epsilon}\\mathcal{L}(\ heta+d)$

其中 $D$ 是一个空间度量，衡量两个参数的距离。在（最速）梯度下降方法中使用了欧氏距离，可得 $d^*=-\\frac{\\epsilon}{{||\ abla_\ heta\\mathcal{L}(\ heta)||}}{\ abla_\ heta\\mathcal{L}(\ heta)}$ ，即负梯度方向。当更新量为 $\\epsilon$ 时，学习率 $\\alpha=\\frac{\\epsilon}{{||\ abla_\ heta\\mathcal{L}(\ heta)||}}$ 。

然而我们的目标是最大化似然函数，参数的改变导致似然函数 $p$ 在函数空间的改变。故我们希望使用一个度量 $D$ 衡量函数空间的距离，此时我们可以选择 KL 散度（虽然其不是一个严格的距离），即 $D(\ heta,\ heta+d)=\\mathrm{KL}(p(\ heta)||p(\ heta+d))=\\mathbb{E}_{x\\sim p(x|\ heta)}[\\log p(x|\ heta)-\\log p(x|\ heta+d)]$ 。用泰勒展开近似到二阶项可得 $\\mathrm{KL}(p(\ heta)||p(\ heta+d))\\approx\\frac{1}{2}d^TFd$ ，进一步可得 $d^*$ 的方向为 $\ ilde{\ abla}\\mathcal{L}(\ heta)=-F^{-1}\ abla_\ heta\\mathcal{L}(\ heta)$ 。该项被称作自然梯度，沿该方向的梯度下降被称作自然梯度下降。

我们计算当 KL 散度的更新量为定值 $\\epsilon$ 时，学习率 $\\alpha$ 应该为多少。令 $\\Delta \\mathrm{KL}=\\frac{1}{2}(\\alpha\ ilde{\ abla}_\ heta\\mathcal{L})^TF(\\alpha\ ilde{\ abla}_\ heta\\mathcal{L})=\\epsilon$ ，解得 $\\alpha=\\sqrt{\\frac{2\\epsilon}{\ ilde{\ abla}_\ heta\\mathcal{L}^TF\ ilde{\ abla}_\ heta\\mathcal{L}}}$ 。我们有：

$({\ ilde{\ abla}_\ heta\\mathcal{L}^TF\ ilde{\ abla}_\ heta\\mathcal{L}})^{1/2}=((F^{-1}g)^TF(F^{-1}g))^{1/2}=(g^TF^{-1}g)^{1/2}\\dot=\\lambda(x)$

记 $v=-F^{-1}g$ 为自然梯度更新方向，又有：

$({\ ilde{\ abla}_\ heta\\mathcal{L}^TF\ ilde{\ abla}_\ heta\\mathcal{L}})^{1/2}=((F^{-1}g)^TF(F^{-1}g))^{1/2}=(v^TFv)^{1/2}=||v||_{F}$

其中 $\\lambda(x)$ 被称作牛顿减小项（Newton decrement， $F$ 是 $H$ 的近似），通过计算可以得到牛顿下降法单步更新损失的最大更新值为 $\\frac{1}{2}\\lambda(x)^2$ 。而 $||v||_{H}$ 为牛顿更新步长在黑塞阵下的范数（一些论文中也叫做方向锐度 directional sharpness）。故条件 KL 散度变化为定值，也可以改写为。更新值 $-\\alpha\ ilde{\ abla}_\ heta\\mathcal{L}=-\\frac{\\sqrt{2\\epsilon}}{||v||_F}\ ilde{\ abla}_\ heta\\mathcal{L}$ 。

回忆梯度下降的中更新项为 $-\\frac{\\epsilon}{{||\ abla_\ heta\\mathcal{L}(\ heta)||}}{\ abla_\ heta\\mathcal{L}(\ heta)}$ ，如果在分母上保留 $||\ abla_\ heta\\mathcal{L}(\ heta)||$ ，那么该梯度下降为归一化梯度下降法（normalized gradient descent）。而带学习率项的自然梯度下降与归一化梯度下降有相同的形式，或许可以称作“归一化自然梯度下降”或“归一化牛顿下降”。

上述各类算法中，我们都再梯度前引入了一个逆矩阵得到形如 $P^{-1}\ abla\\mathcal{L}$ 的更新方向。我们可以称 $P$ 为预条件子（preconditioner）。目前对优化器的设计，主要是对预条件子的设计。当预条件子为对角阵时， $P^{-1}$ 可以很方便的获得。此时也可以看作在每个维度上进行自适应更新（adaptive） $\\frac{g_i}{p_i}$ ，比如 Adam 类的优化器以及 Sophia 优化器。预条件子可能能起到一下作用：1.降低方差 2.对高条件数问题鲁棒 3.提供高阶（曲率）信息。

给定列向量 $u,v$ ，其外积可以定义为 $uv^T=u\\otimes v^T$ 。其中前一项是两个向量的乘积，后一项是两个向量的 Kronecker 积。Kroneckor 积的定义如下：

在神经网络中，参数具有结构性质（第 $i$ 层的参数记为 $W_i$ ）。在计算 FIM 阵时，我们把所有的参数平等对待。定义 $\ extrm{vec}(\\cdot)$ 算符将 Tensor $W_i$ 转为为向量（flatten）。Kronecker 积可以方便的表示外积向量化的结果，即 $\ extrm{vec}(uv^T)=v \\otimes u$ 。

Kronecker 积有一系列良好的性质：

$(A\\otimes B)(C\\otimes D)=(AC)\\otimes(BD)$
对于正定矩阵 $A,B$ ， $(A\\otimes B)^{-1}=A^{-1}\\otimes B^{-1}$
$(A\\otimes B)\\otimes \ extrm{vec}(X)=\ extrm{vec}(BXA^T)$

K-FAC 算法是自然梯度下降的一种近似实现。定义 MLP 网络中第 $i$ 层为 $a_i=\\phi_i(z_i),z_i=W_ia_{i-1}$ （为简化不考虑 bias），其中 $\\phi$ 为激活函数。梯度 $\ abla_\ heta\\log p(x|\ heta)=[\ extrm{vec}(\\frac{\\partial\\mathcal{L}}{\\partial W_i})^T]^T$ 。K-FAC 算法使用三个近似以高效的实现自然梯度下降。

近似一：作者利用服从数据集分布的数据来计算二阶矩，进而近似计算 FIM 阵，即 $F=\\mathbb{E}_{x\\sim\\mathcal{D}}[\ abla_\ heta\\log p(x|\ heta)\ abla_\ heta\\log p(x|\ heta)^T]\\approx\\frac{1}{N}\\sum_{i=1}^N[\ abla_\ heta\\log p(x_i|\ heta)\ abla_\ heta\\log p(x_i|\ heta)^T]$ 。这里 $x_i$ 是从数据集中采样的数据。但是只有当网络分布接近真实分布时两者才能近似。实际上一个好的估计是经验 FIM 阵： $F=\\frac{1}{N}\\sum_{i=1}^N[p(x_i|\ heta)\ abla_\ heta\\log p(x_i|\ heta)\ abla_\ heta\\log p(x_i|\ heta)^T]$ 。然而由于梯度计算是批量进行的，我们只能得到一个批量样本中作为一个整体的 $\\frac{1}{B}\\sum_{i=1}^B\ abla_\ heta\\log p(x_i|\ heta)$ 。作者用该近似避免了额外的采样计算开销。

近似二：将每层的参数作为整体，可以把 FIM 阵分块。 $F_{ij}=\\mathbb{E}[\ extrm{vec}(\\frac{\\partial\\mathcal{L}}{\\partial W_i})\ extrm{vec}(\\frac{\\partial\\mathcal{L}}{\\partial W_i})^T]$ 。又根据链式法则中有 $\ extrm{vec}(\\frac{\\partial\\mathcal{L}}{\\partial W_i})=\ extrm{vec}(\\frac{\\partial L}{\\partial a_i}a_{i-1}^T)=a_{i-1}\\otimes \\frac{\\partial L}{\\partial a_i}$ 。结合 Kronecker 积的性质 1 我们有 $F_{ij}=\\mathbb{E}[(a_{i-1}\\otimes \\frac{\\partial L}{\\partial \\vec a_i})(a_{i-1}\\otimes \\frac{\\partial L}{\\partial \\vec a_i})^T]=\\mathbb{E}[a_{i-1}a_{j-1}^T\\otimes \\frac{\\partial L}{\\partial a_i}\\frac{\\partial L}{\\partial a_i}^T]$ 。

如果我们维护 $F_{ij}$ 的移动平均作为估计，则需要 $O(n^2)$ 的储存空间。K-FAC 利用近似 $\\mathbb{E}[a_{i-1}a_{j-1}^T\\otimes \\frac{\\partial L}{\\partial a_i}\\frac{\\partial L}{\\partial a_i}^T]=\\mathbb{E}[a_{i-1}a_{j-1}^T]\\otimes \\mathbb{E}[\\frac{\\partial L}{\\partial a_i}\\frac{\\partial L}{\\partial a_i}^T]$ 从而只需维护两个向量的移动平均，即 $O(n)$ 的储存空间。[MG20]主要根据实验说明其效果不错。下图中左图是 $F$ ，中间是 $F$ 的近似，右图是两者的差。

近似三：在前向和反向传播中， $a_i$ 和 $\\frac{\\partial L}{\\partial a_i}$ 都是已有的值，故 $F$ 阵已经可以计算了。为了方便计算 $F$ 的逆，作者进一步对 $F$ 阵的结构做了两种假设，分别是块对角（block-diagnoal）矩阵和块三对角矩阵（block-tridiagonal）。这里我们只讨论前者。在该假设下，我们只关心 $F_{i,i}$ ，即层内的相关性（假设不同层间弱相关）。根据 Kronecker 积的性质二，我们有 $F_{i,i}^{-1}=(a_{i-1}a_{i-1}^T)^{-1}\\otimes (\\frac{\\partial L}{\\partial a_i}\\frac{\\partial L}{\\partial a_i}^T)^{-1}\\dot=A_{i-1,i-1}^{-1}\\otimes G_{i,i}^{-1}$ 。根据性质三，可得自然梯度 $F_{i,j}^{-1}\ extrm{vec}(\\frac{\\partial\\mathcal{L}}{\\partial W_i})=\ extrm{vec}(G_{i,i}^{-1}\\frac{\\partial\\mathcal{L}}{\\partial W_i}A_{i-1,i-1}^{-1})$ 。

学习率调整：通过上式计算出自然梯度后，我们还要对自然梯度的学习率进行修正。与之前推导自然梯度下降学习率类似，当 KL 散度的近似值 $\ u=\\frac{1}{2}(\\alpha\ ilde{\ abla}_\ heta\\mathcal{L})^TF(\\alpha\ ilde{\ abla}_\ heta\\mathcal{L})$ 大于阈值 $C$ 时（[GM16]中选择 0.3）将其裁剪。即学习率 $\\alpha\\leftarrow\\alpha\\min(1, \\sqrt{\\frac{C}{\ u}})$ 。如此可以限制每次更新时 KL 变化的幅度。其中， $\ u$ 的计算可以通过”黑塞向量积“一节的技术快速计算。作者实际每个几个批量计算一次，进一步降低均摊开销。

其它实现细节：在该方法的一个实现中，作者基于 K-FAC 添加动量时，动量应用在自然梯度上。权重衰减也应该发生在函数空间，故应使用 $(F+\\lambda I)^{-1}$ 作为预条件子，实际中仍然使用参数空间的衰减 $\\frac{1}{2}\\lambda\ heta^T\ heta$ 。计算正定矩阵 $A$ 的逆时，可以通过特征值分解 $A=Q\\Lambda Q^T$ ，其中 $Q$ 是正交矩阵， $\\Lambda$ 是对角阵，其对角线上的元素为特征值。故 $A^{-1}=Q\\Lambda^{-1}Q^T$ 。为了获取前向传播中的 $a$ 和反向传播中的 $\\frac{\\partial L}{\\partial a}$ ，可以通过 pytorch 的 hook 功能快捷实现。

缺点：虽然通过三个近似，可以一定程度高效的实现自然梯度下降，然而 K-FAC 算法具有不少缺点。

通过带噪声的梯度信息近似曲率信息会使高阶信息估计不准确。即使是在大批量（低噪声）的情况下，[MMY+19]通过实验表面 K-FAC 效果不如 SGD。
三个近似一定程度也损失了精度。
对于不同的网络结构（如卷积等）需要重新推导，实现较为复杂。

受到 K-FAC 的启发，[GRB21]将拟牛顿法 (L-)BFGS 推广到神经网络上。与 K-FAC 类似，作者在 Kronecker 积以及块对角阵假设的基础下设计了 K-BFGS 算法。由于使用不广，这里不再赘述。

Adaptive 类算法（Adagrad[DHS11], Adam[KB14]等）中的自适应学习率，可以看作是使用了对角阵的预条件子。在该篇内容中，我们还会从别的视角探究这类算法。本篇则从曲率信息的角度进行探究，为了方便这里不讨论权重衰减、动量修正等技术。记梯度为 $g_t$ ，预条件子为 $P$ ，更新量为 $U$ （梯度下降即 $U_t=-\\alpha g_t$ ）， $A^{1/2}=X,XX=A$ 。

Adagrad 算法在一位参数上的自适应学习率为 $-\\frac{1}{\\sqrt{\\sum_{i=1}^tg_i^2}}$ 。而在原始论文中，实际上 Adagrad 的原始形式是 $U=-\\alpha({\\sum_{i=1}^tg_tg_t^T})^{-1/2}g_t=-\\frac{\\alpha}{\\sqrt{t}}(\\frac{1}{t}\\sum_{i=1}^tg_tg_t^T)^{-1/2}g_t$ 。这里我们采取与原论文不同的分析方法。注意到 K-FAC 中的近似一，即用数据分布近似计算 FIM 阵。我们同样可以将 Adagrad 看作是近似 FIM 阵的预条件子加上随时间的梯度衰减，即 $P^{-1}=-\\frac{1}{\\sqrt{t}}(\\frac{1}{t}\\sum_{i=1}^tg_tg_t^T)^{-1/2}=-\\frac{1}{\\sqrt{t}}F^{-1/2}$ 。这里预条件子使用了 $F^{-1/2}$ 而非 $F^{-1}$ ，这与曲率之间能否建立联系呢？

在归一化自然梯度下降中，我们通过设定学习率 $\\alpha=\\sqrt{\\frac{2\\epsilon}{g^TF^{-1}g}}$ 让每步的更新幅度相等。在学习率与动量一章中，我们讨论过高条件数对优化的不利影响。考虑 $F\\approx H$ 的特征值 $\\lambda_i$ 对应的特征向量 $e_i$ ，如果我们希望对每个特征向量进行伸缩变换，使 $||s_ie_i||_H$ 相等，从而让不同特征值对应的更新幅度相等。令 $||s_ie_i||_H=\\frac{1}{2}(s_ie_i)^TF(s_ie_i)=\\epsilon$ ，解得 $s_i=\\sqrt{2\\epsilon\\lambda_i}$ 。记梯度 $g$ 在 $e_i$ 上的分量为 $d_i$ ，此时 $-U=F^{-1}\\sum_i s_i\\cdot d_ie_i=\\sum_i \\sqrt{2\\epsilon\\lambda_i}\\cdot \\frac{d_i}{\\lambda_i}e_i=\\sqrt{2\\epsilon}\\sum_{i}\\frac{1}{\\sqrt{\\lambda_i}}d_ie_i$ 。对比 $F^{-1}g=\\sum_i\\frac{1}{{\\lambda_i}}d_ie_i$ ，我们发现特征值进行了开根，故我们可以选择 $F^{-1/2}$ 作为预条件子以满足 $s_i$ 的选择。由此可见，Adagrad 算法（以及不包括 $\\frac{1}{\\sqrt{t}}$ 学习率衰减的 Adam 算法）可以看作是”特征值归一化的自然梯度下降“。

Adagrad 中继续将 $P$ 限制为对角阵，这样可以减小 $P$ 的储存空间、求逆以及开方的复杂度。由于 $g_tg_t^T$ 的对角线构成的矩阵为 $\\mathrm{diag}(\\sum_{i=1}^tg_{i,j}^2)_{j=1}^n$ ，其中 $g_{t,i}$ 是第 $i$ 个参数的导数。故 $P^{-1}=-\\frac{1}{\\sqrt{t}}F^{-1/2}=-\\frac{1}{\\sqrt{t}}\\mathrm{diag}(\\frac{1}{\\sqrt{g_{t,i}^2}})$ 。

从曲率的角度看，Adaptive 类算法实际上是利用数据分布近似的对角 FIM 阵（牛顿阵）下的特征值归一化的自然梯度下降（或牛顿梯度下降）。

Shampoo 算法[GKS18]是对 Adagrad 更新量的近似计算。记 Adagrad 的更新量为 $V^{-1/2}g,V=\\sum_{t=1}^Tg_tg_t^T$ 是对 $\\sqrt{t}F$ 的近似计算。 $V$ 矩阵是一个[nxn]的矩阵，储存、求逆、求方根开销都很大。如果使用对角近似，我们可以看成在每个位子上使用了[1x1]的矩阵，大量减小了开销，然而丢失了大量信息。那能否在这两者之间取得一个平衡呢？

一个简单的思路是类似 K-FAC 一样使用块对角矩阵，即只考虑每层参数内的相关性。对于第 $l$ 层参数 $W_l$ ， $g_l$ 为该层导数， $G_l$ 是 $g_l$ 重整（reshape）成与 $W_l$ 形状相同的矩阵。则更新公式为 $W_{t+1}=W_t-\\alpha (\\sum_{t=1}^Tg_lg_l^T)^{-1/2}g_l\\dot=W_t-\\alpha V_l^{-1/2}g_l$ 。但这样做需要维护一个 $O(m^2n^2)$ 矩阵以及对其求逆和求方根，仍然是不小的开销。

Shampoo 算法将 $V_l^{-1/2}g_l$ 进行近似，这里我们只做直观理解，理论证明可以参考原论文。对于形如[ $m\ imes n$ ]的二维矩阵 $W_l$ ，令 $L_t=\\sum_{t=1}^TG_lG_l^T$ ([ $m\ imes m$ ])， $R_t=\\sum_{t=1}^TG_l^TG_l$ [ $n\ imes n$ ]，更新为 $W_{t+1}=W_t-\\alpha L_t^{-1/4}G_tR_t^{-1/4}$ 。其中 $-\\frac{1}{4}+-\\frac{1}{4}=-\\frac{1}{2}$ ，与 Adagrad 中的次数相同。如此一来，我们只需要维护 $O(m^2)$ 和 $O(n^2)$ 的矩阵，以及对其求逆和求方根。这样大大减小了计算开销。

对于一般的张量 $W_t$ 有形状[ $n_1\ imes n_2\ imes\\cdots\ imes n_k$ ]，总大小为 $n$ 。维护 $k$ 个矩阵，每个矩阵是 $H_l^{(i)}=\\sum_{t=1}^TG_l^{(i)}G_l^{(i)T}$ 。这里 $G_l^{(i)}$ 是 $g_l$ 重整成[ $n_i\ imes \\frac{n}{n_i}$ ]形状的矩阵。如此一来， $H_l^{(i)}$ 有形状[ $n_i\ imes n_i$ ]。为了计算最后的更新量，将 $G_t$ 轮流重整成[ $\\frac{n}{n_i}\ imes n_i$ ]形状的矩阵，计算 $G_t\\leftarrow H_l^{(i)-1/2k}G_t$ ，再将其 reshape 成原状。此时总次数为 $k\ imes-\\frac{1}{2k}=-\\frac{1}{2}$ ，与 Adagrad 中的次数相同。维护的矩阵大小为 $O(n_i^2)$ ，与直积计算整层的 $A_t$ 相比减小了计算开销。

[YGS+21]放弃了 K-FAC 和 Adagrad 中的近似一，并估计为对角矩阵。其通过上文提到的 Hutchinson 方法估计 $H$ 的对角阵 $D$ ，维护梯度动量 $m_t$ 和对角阵动量 $D$ 。使用与 Adam 一致的更新方式，即 $D^{-1/2}m_t$ ，在每个维度上更新方向为 $\\frac{m_t^i}{\\sqrt{d^i_t}}$ 。

在此基础上，AdaHessian 还用平均池化将 $H$ 的值进行平滑，从而使训练更稳定。

Sophia 算法[LLH'23]在 AdaHessian 算法的基础上进行了改进。首先，其通过真实的 $F$ 函数对 $H$ 进行了估计。根据“费雪信息阵与牛顿法”一节的内容，计算 $F$ 时损失函数应使用按照网络输出概率采样的真实标签，即

$F=\\mathrm{diag}(\\mathbb{E}_{y\\sim p(y|x;\ heta)}[\ abla_\ heta\\mathcal{L}_\ heta(x,y)\ abla_\ heta\\mathcal{L}_\ heta(x,y)^T])=\\mathbb{E}_{y\\sim p(y|x;\ heta)}[{\ abla_\ heta\\mathcal{L}_\ heta(x,y)\\cdot\ abla_\ heta\\mathcal{L}_\ heta(x,y)}]$

注意到实际训练中我们无法得到单个样本的 $\\mathbb{E}_y[\ abla_\ heta\\mathcal{L}_\ heta(x,y)]$ ，只能得到一个批量的 $\\hat G_B=\\frac{1}{B}\\sum_{b=1}^B\\mathbb{E}_y[\ abla_\ heta\\mathcal{L}_\ heta(x_b,y)]$ 。不过幸好由于 $\\mathbb{E}_y[\ abla_\ heta\\mathcal{L}_\ heta(x_b,y)]=0$ ，我们有

$F=\\mathbb{E}_y[\\frac{1}{B}\\sum_{b=1}^B\ abla_\ heta\\mathcal{L}_\ heta(x_b,y)\\odot\ abla_\ heta\\mathcal{L}_\ heta(x_b,y)]=\\mathbb{E}_y[\\frac{1}{B}\\sum_{b=1}^B (B\\cdot \\hat G_B)\\odot(B\\cdot\\hat G_B)]=B\\cdot\\mathbb{E}_y[\\hat G_B\\odot \\hat G_B]$

由于该操作需要额外的一次反向传播，作者通过每隔 $k$ 步采样一次的方式进行近似，并做动量更新，从而降低了计算开销。

此外，Sophia 没有使用 Adam 中 -1/2 次的“特征值归一化”，而是直接使用自然梯度下降，即 $F^{-1}g$ 。为了应对自然梯度下降训练的不稳定性，作者对更新值进行裁剪，使其不超过给定的阈值。

[DHS11]Adaptive Subgradient Methods for Online Learning and Stochastic Optimization
[KB14]Adam: A Method for Stochastic Optimization
[MG15]Optimizing Neural Networks with Kronecker-factored Approximate Curvature
[GM16]A Kronecker-factored approximate Fisher matrix
for convolution layers
[GKS18]Shampoo: Preconditioned Stochastic Tensor Optimization
[MMY+19]Inefficiency of K-FAC for Large Batch Size Training
[GRB21]Practical Quasi-Newton Methods for Training Deep Neural Networks
[YGS+21]AdaHessian: An Adaptive Second Order Optimizer for Machine Learning
[LLH'23+]Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training

北京装修公司知名品牌电话：13988889999

大规模神经网络优化：二阶优化器

内容搜索 Related Stories

推荐内容 Recommended

天富注册 SERVICE

天富登录 CASE

现在致电 13988889999 OR 查看更多联系方式 →

现在致电 13988889999 OR 查看更多联系方式 →