昇思MindSpore进阶教程--二阶优化器THOR

大家好，我是刘明，明志科技创始人，华为昇思MindSpore布道师。
技术上主攻前端开发、鸿蒙开发和AI算法研究。
努力为大家带来持续的技术分享，如果你也喜欢我的文章，就点个关注吧

正文介绍

深度学习训练过程可以看成损失函数损失值下降过程，合适的优化器可以让深度学习训练时间大大减少。优化器可以分为一阶优化器和二阶优化器，目前业界主流使用的仍然是一阶优化器，二阶优化器因为单步训练时间过久而没有被广泛应用，而近年来，将二阶优化应用到深度学习训练中有了理论突破，并取得了不错的结果。

本文会介绍下优化器的背景，以及MindSpore团队自研二阶优化器THOR。

优化器背景介绍

假设训练样本数据集：$D=(x_1,y_1),...,(x_i,y_i),...,(x_N,y_N),x_i\in X,y_i\in Y$，参数θ表述的深度神经网络模型为： $\hat{y}=f(x;\theta ),x\in X$，定义在模型输出和真实标签y之间的损失函数为：$L(y,\hat{y}),y\in Y$，网络参数学习的过程是最小化损失函数的过程：$\underset{\theta }{\min }L(y,\hat{y})$ 。给定数据集、模型、损失函数后，深度学习训练问题归结为优化问题，深度神经网络训练优化问题参数规模巨大，需要大量的计算，难以计算出解析解。因此该过程也常常被比喻成下山，如图1 所示，一个人站在山顶的时候如何在有限视距内寻找最快路径下山呢？

而优化器就是在做这件事情，业界的优化算法可分为一阶优化算法和二阶优化算法。下面简单介绍下业界的优化器情况。

一阶优化器

梯度下降算法（Gradient Descent, GD）是机器学习中最经典的一阶优化算法，也是众多机器学习算法中最常用的优化算法。常用的一阶优化算法（比如SGD算法）中对参数的更新采用如下规则：
，其中
是需要更新的参数，
是学习率，
是损失函数对于参数的梯度。

但是主流随机梯度下降方法有以下问题：太小的学习率会导致网络收敛过于缓慢；学习率太高可能会影响收敛，并导致损失函数在最小值上波动，甚至出现发散，对参数比较敏感；容易收敛到局部最优，难以跳出鞍点。

因此业界提出了很多随机梯度下降方法的改良算法，例如Momentum、Nesterov、AdaGrad、RMSprop、Adadelta和Adam等。这些改进后的优化算法可以利用随机梯度的历史信息来自适应地更新步长，使得它们更容易调参，而且方便使用。

二阶优化器

二阶优化算法利用目标函数的二阶导数进行曲率校正来加速一阶梯度下降。与一阶优化器相比，其收敛速度更快，能高度逼近最优值，几何上下降路径也更符合真实的最优下降路径。

例如，二阶优化算法中的牛顿法就是用一个二次曲面去拟合你当前所处位置的局部曲面，而梯度下降法是用一个平面去拟合当前的局部曲面，通常情况下，二次曲面的拟合会比平面更好，所以牛顿法选择的下降路径会更符合真实的最优下降路径。如图2 所示，左边下降路径表示牛顿法的下降曲线，右边表示一阶梯度的下降曲线，二阶算法与一阶算法相比，可以更快的走到目的地，从而加速收敛。

THOR的介绍

当前业界已有的二阶优化算法计算量较大，与一阶相比优势不明显或者使用场景较为单一。MindSpore提出了自研算法THOR(Trace-based Hardware-driven layer-ORiented Natural Gradient Descent Computation)， 算法已被AAAI录用，THOR在多个场景中均有明显收益，如在BERT和ResNet50中，收敛速度均有明显优势。THOR主要做了以下两点创新：

降低二阶信息矩阵更新频率

通过实验观察费雪矩阵的F范数（Frobenius norm）在前期变化剧烈，后期逐渐变稳定，从而假设$\{F^k{\}}_{k=1}^n$是一个马尔可夫过程，可以收敛到一个稳态分布π，其中$F^k$代表第k个迭代时的费雪矩阵。因此，在训练过程中逐步增大费雪矩阵的更新间隔，可以在不影响收敛速度的情况下，减少训练时间。例如在ResNet50中，更新间隔步数随着训练的进行越来越大，到后期每个epoch只需更新一次二阶信息矩阵。

THOR受KFAC启发，将费雪矩阵按层解耦来降低矩阵复杂度，分别针对每一层的费雪矩阵做实验，发现有些层的费雪矩阵趋于稳态的速度更快，因此在统一的更新间隔上，更加细粒度的去调整每一层的更新频率。THOR使用矩阵的迹作为判断条件，当迹的变化情况大于某一阈值时更新该层的二阶信息矩阵，否则沿用上一个迭代的二阶信息矩阵，并且引入了停止更新机制，当迹的变化量小于某个阈值时停止更新该层二阶信息矩阵，具体更新公式如下：$$\begin{array}{r}\{\begin{array}{l}update{F}_i^k,if{\Delta }^k\in ({\omega }_1,+\infty )\\ donotupdate{F}_i^{k}andset{F}_i^k=F_i^{k-1},if{\Delta }^k\in [{\omega }_2,{\omega }_1]\\ stopupdate{F}_i^{k}andset{F}_t^{k+t}\equiv F_i^{k-1}forallt=1,2,...if{\Delta }^k\in [0,{\omega }_2)\end{array}\end{array}$$

其中：

$${\Delta }^k=\frac{||tr(F_i^k+\lambda I)|-|tr(F_i^{k-1}+\lambda I)||}{|tr(F_i^{k-1}+\lambda I)|}$$

硬件感知矩阵切分

THOR在将费雪矩阵按层解耦的基础上，进一步假设每个网络层中的输入和输出块之间也是独立的，例如将每层网络的输入输出切分为n个块，这n个块之间即是独立的，根据该假设对二阶信息矩阵做进一步的切分，从而提高了计算效率。THOR结合矩阵信息损失数据和矩阵性能数据确定了矩阵分块维度，从而大大提升费雪矩阵求逆时间。

那么如何确定矩阵分块维度的呢。具体方法为：

（1）根据费雪矩阵中维度最大的那一层，确定矩阵切分维度，拿ResNet50举例，网络层中的最大维度为2048，确定矩阵切分维度为[1,16,32,64,128,256,512,1024,2048]。

（2）根据确定的矩阵维度，根据谱范数计算每个维度下的矩阵损失，具体公式为：
$$L=1-\sqrt{\frac{{\lambda }_{max}(\hat{A}{\hat{A}}^T)}{{\lambda }_{max}(A{A}^T)}}$$
3）根据确定的矩阵维度，计算每个维度下的矩阵求逆时间，再通过公式

得到每个维度下标准化后性能数据，其中
表示维度最小的矩阵的性能数据，
表示第n个维度下的性能数据。

（4）根据标注化后的矩阵损失信息和标准化后的性能数据绘图，如以ResNet50为例，如图3所示，图中下降曲线为性能曲线，上升曲线表示矩阵损失曲线，图中交叉点为106，与128最接近，最后确定矩阵切分维度为128。

实验结果

图4展示了THOR在ResNet50+ImageNet，batchsize为256时一二阶上的训练曲线图，其中train loss表示训练误差，test accuracy表示测试精度，epoch表示迭代数，wall-clock time表示时间，其中下降较快的曲线和上升较快的曲线是本算法曲线，另外差距较明显的曲线是momentum的训练曲线。

THOR还测试了在不同batchsize下ResNet50+ImageNet的收敛结果，结果见下图5，其中Hardware表示硬件平台，Software是指使用的深度学习框架，Batch size是每次训练的图片数量，Optimizer表示使用的优化器，Time指总体训练时间，Accuracy是指最后收敛精度。当batchsize为8192，使用256块Atlas训练系列产品时，只需2.7分钟精度即可收敛到75.9%。

在BERT+WIkipedia中，THOR也有不错的表现效果，以MLPerf为标准，精度达到71.2%，与一阶相比端到端提升30%，实验结果见图6，图中横坐标表示训练时间，纵坐标表示测试精度，上升较快的曲线是THOR的训练曲线，另一条为Lamb的训练曲线。