爱笑的小姐姐 · 2021年02月02日

MindSpore高阶优化器系列(1)

文章转载于:知乎
作者:于璠

深度概率学习系列已经介绍完啦,接下来给大家分享下高阶优化器系列,想要看前一个系列的小伙伴们戳下面的地址哦。

于璠:MindSpore贝叶斯应用工具箱​zhuanlan.zhihu.com图标

本文先跟大家分享下常用的一二阶优化器。深度学习训练过程可以看成是loss下降过程,而每一次loss如何下降,就是优化器做的事情。合适的优化器可以让深度学习训练时间大大减少。优化器可以分为一阶优化器和二阶优化器,深度学习中常用的优化器是一阶优化器,如SGD(Stochastic Gradient Descent)和其算法变种,二阶优化器近年来也被用于深度学习训练,并取得了不错的结果。接下来我们就具体展开来看一下。

背景

首先我们来看一下深度学习训练过程是在做什么事。先假设训练样本数据集: 参数 表述的深度神经网络模型为: 定义在模型输出 和真实标签 之间的损失函数为: 网络参数学习的过程是最小化损失函数的过程: 给定数据集、模型、损失函数后,深度学习训练问题归结为优化问题,深度神经网络训练优化问题参数规模巨大,需要显著大量的计算,难以计算出解析解。因此该过程也常常被比喻成下山,如下图所示,一个人站在山顶的时候如何在有限视距内寻找最快路径下山呢?

这个时候就是优化器大展身手了,好的优化器可以以更少的时间下山。业界的优化算法可分为一阶优化算法和二阶优化算法。下面我们就一起来看下常用的一阶优化器和二阶优化器。

一阶优化器

简介

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

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

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

常用一阶优化器

Momentum

Momentum是一阶优化算法,它对SGD进行了改进,在梯度下降的过程中加入了惯性,使得梯度在方向不变的维度上速度变快,梯度方向有所改变的维度上的更新速度变慢,与SGD相比可以加快收敛并减少震荡。其更新规则如下:

是动量超参, 为在第 个step时的学习率, 是需要更新的参数, 是损失函数对于参数的梯度。

Adam

Adam也是一种自适应学习率的方法,该算法使用了动量变量 和一阶梯度的指数加权移动平均变量 并在初始时将他们置为0,其更新规则如下:

其中 和 为超参,通常取为0.9和0.999, 为在第 个step时的学习率, 是需要更新的参数, 是损失函数对于参数的梯度, 是一个很小的数,一般为1e-8,为了避免分母为0。

二阶优化器

简介

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

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

从数学公式上来看,与一阶优化算法相比,二阶优化算法则是先将 与一个矩阵 相乘,参数更新规则如下: 其中 即为二阶信息矩阵,不同的二阶优化算法中的 定义是不同的,常见的二阶优化算法有牛顿法,自然梯度法等,分别对应的二阶信息矩阵 为 矩阵, 矩阵。

牛顿法有着很好的局部收敛性质,当函数 在最优值点 满足 是正定矩阵且 矩阵在极值点附近是李普希兹连续时,牛顿法二次收敛到最优值点。 矩阵是一个由多变量实值函数的所有二阶偏导数组成的方块矩阵。 矩阵可以表示为: 其中 即为损失函数, 是需要更新的参数。

在SGD中,参数空间和函数空间的度量用的都是欧式距离,但欧式距离在一些情况下不能作为函数空间准确的距离度量。例如神经网络中,参数引起的目标函数变化是概率的变化,这并不适合在欧几里得空间度量,它不是概率属性变化的合理表征。KL散度是分布之间距离的合理度量。当使用KL散度作为概率分布之间距离的度量时。此时参数更新时,用到的梯度就是自然梯度。自然梯度法中的 矩阵可以表示为: 其中 是网络模型的预测分布, 是其概率密度, 是需要网络模型的参数。

二阶优化算法虽然收敛速度快,但是计算二阶矩阵的逆的时间复杂度为 其中 为二阶信息矩阵维度,当模型参数量为 时,对应的二阶信息矩阵的大小为 。在深度学习模型中, 常常在数百万的量级,此时二阶信息矩阵的逆无法计算。因此如何降低二阶信息矩阵求逆的计算复杂度成为关键问题。

深度学习中的二阶优化器

KFAC

KFAC(Kronecker-factored Approximate Curvature)是由Martens等学者于2015年提出的基于自然梯度法的二阶优化算法,该算法高效地近似了自然梯度法中的Fisher矩阵,并且给出了充分的理论证明,给后续许多二阶优化算法带来了启发。这里重点跟大家分享下他是如何做近似的,具体的理论推导建议大家看论文哦。

第一步,KFAC假设不同层之间的参数是独立的,从而将二阶信息矩阵按网络层解耦,近似成块对角矩阵,然后对二阶信息矩阵求逆就变为了对这些块对角矩阵求逆。

第二步将按层解耦后的每个块矩阵又近似为两个小得多的矩阵的Kronecker乘积。我们一起来回顾下Kronecker乘积公式:

由于两个矩阵的Kronecker乘积的逆等于它们逆的Kronecker乘积,且这两个较小矩阵会比整个块矩阵更容易计算和求逆,通过这样的近似和分解,KFAC大大简化了Fisher矩阵的计算。

上面这段话还不够直观,不急,我们通过一个栗子来看一下。比如对于一个3-4-2的神经网络,

原始二阶信息矩阵维度为20*20,经过第一步近似变为了一个12*12和8*8的矩阵,再经过第二步近似,分别变为了3*3和4*4;4*4和2*2维度的矩阵。如下图所示:

该算法在CIFAR-10上做了实验,使用cuda-convnet网络,该网络包含3层卷积层和一层全连接层,下图中的蓝线为SGD with Momentum,红线为KFAC(图中的KFC-pre即为KFAC),实线是测试误差,虚线是训练误差。KFAC达到19%的测试误差仅需3分钟,SGD需要9分钟;在训练误差上KFAC的优势更加明显,KFAC达到6%仅需4分钟,SGD需要30分钟。

KFAC使自然梯度法应用到了深度学习训练中,并取得了很好的结果,但是KFAC未考虑大型神经网络计算复杂度,例如对于大型神经网络ResNet50,全连接层维度为2048*1000,按照KFAC算法分解后,仍然要计算一个2048维和一个1000维的逆矩阵,求逆时间虽然大大降低,但是依然很慢,在ResNet50模型,ImageNet全集上,单次迭代所需时间是秒级别,单迭代所需时间是一阶的33倍。

SP-NGD

SP-NGD(Scalable and Practical Natural Gradient Descent)是由Osawa等学者于2020年提出的基于自然梯度法的二阶优化器,并在大数据集ImageNet上训练ResNet50,使用1024块V100,仅需5.5min即可收敛到75.4%。

SP-NGD受KFAC启发,在KFAC基础上主要做了以下几点改进。

第一点,SP-NGD使用了过时的二阶信息矩阵,核心思想是,在训练过程中不是每个迭代都更新二阶信息矩阵的逆,而是使用上一次更新保存下来的二阶信息矩阵的逆,当二阶信息矩阵变化量超过一定阈值时减少二阶信息矩阵逆的更新间隔,当变化量持续不变,增大二阶信息矩阵逆的更新间隔从而减少计算量。

大致过程是判断当前更新的二阶信息矩阵 与上一次的 是否相似,其中文中定义两个矩阵 相似为: 其中 是 范数, 是设定的阈值,若相似再判断与上上次的二阶信息矩阵 是否相似,若不相似,更新间隔保持不变,若相似则增大下次更新间隔。具体计算规则见下图:

第二点,为了更高效的进行分布式训练,提出了数据和模型混合并行,减少了通讯时间,具体做法是:

Stage1. 给每块GPU分配不同的数据集,做数据并行,每块GPU根据分配到的数据集计算前向过程和前向过程中二阶信息矩阵需要的信息A。

Stage2. 对A进行归并,归并时使用ReduceScatterV,而不是常规使用的AllReduce,这个操作是指每个GPU只归并指定层中的A,将数据并行转为模型并行,并且在A通讯的时候同时计算反向过程,进一步减少计算时间。

Stage3. 对反向过程中计算得到的相关信息也使用ReduceScatterV做归并。这时每块GPU有指定层的完整信息。

Stage4. 每块GPU分别计算指定层的二阶信息矩阵的逆,并更新网络参数。

Stage5. 最后对网络参数进行AllGatherV,同步所有GPU中的参数信息。

文中还有其他的改进点,比如对BN层的二阶梯度的特殊处理,如何做数据增强,如何更好地调参等, 这里我就不具体展开啦,有兴趣的小伙伴可以去看论文哦。接下来,我们来看下这个算法的实验结果,在ImageNet上训练ResNet50,使用1024块V100,仅需5.5min即可收敛到75.4%,该算法使二阶优化器在大数据集和大型网络中与一阶优化器比较也取得了有竞争力的结果。

到这里,这篇文章想要分享的内容就结束啦,如果有什么不对之处欢迎大家批评指正。

最后再做个预告,下一篇跟大家分享MindSpore自研二阶优化器THOR,该优化器在8块昇腾芯片下仅需66.7分钟就可以让ImageNet+ResNet50收敛到75.9%,当使用256块时,仅需2.7分钟。

参考文献:

[1] Qian N. On the momentum term in gradient descent learning algorithms[J]. Neural networks, 1999, 12(1): 145-151.

[2] Kingma D P, Ba J. Adam: A method for stochastic optimization[J]. arXiv preprint arXiv:1412.6980, 2014.

[3] Martens J, Grosse R. Optimizing neural networks with kronecker-factored approximate curvature[C]//International conference on machine learning. 2015: 2408-2417.

[4] Grosse R, Martens J. A kronecker-factored approximate fisher matrix for convolution layers[C]//International Conference on Machine Learning. 2016: 573-582.

[5] Osawa K, Tsuji Y, Ueno Y, et al. Scalable and Practical Natural Gradient for Large-Scale Deep Learning[J]. arXiv preprint arXiv:2002.06015, 2020.

- The End -

推荐阅读

更多嵌入式AI技术相关内容请关注嵌入式AI专栏。
推荐阅读
关注数
18790
内容数
1342
嵌入式端AI,包括AI算法在推理框架Tengine,MNN,NCNN,PaddlePaddle及相关芯片上的实现。欢迎加入微信交流群,微信号:aijishu20(备注:嵌入式)
目录
极术微信服务号
关注极术微信号
实时接收点赞提醒和评论通知
安谋科技学堂公众号
关注安谋科技学堂
实时获取安谋科技及 Arm 教学资源
安谋科技招聘公众号
关注安谋科技招聘
实时获取安谋科技中国职位信息