基于Sherman-Morrison公式的K-FAC算法①

刘小雷,高凯新,王 勇

(天津大学 数学学院,天津 300350)

近些年来,深度学习已经在计算机视觉和自然语言处理等领域取得了重要的进展.然而随着研究的深入,模型越来越复杂,往往需要耗费大量的训练时间和计算成本.因此,采用有效的训练方法是十分有必要的.以随机梯度下降(SGD)为代表的一阶优化方法是当前深度学习中最常用的方法.近些年来,一系列SGD的改进算法被提出并也被广泛于应用深度学习中,比如,动量SGD (SGDM[1]),Adagrad[2],Adam[3].这些一阶优化方法具有更新速度快,计算成本低等优点,但是也具有收敛速度慢,需要进行复杂调参等缺点.

通过曲率矩阵修正一阶梯度,二阶优化方法可以得到更为有效的下降方向,使得收敛速度大大加快,减少了迭代次数和训练时间.对于有着上百万甚至更多参数的深度神经网络而言,其曲率矩阵的规模是十分巨大的,这样大规模的矩阵的计算,存储和求逆在实际计算中是难以实现的.因此,对曲率矩阵的近似引起了广泛的研究.其中最基本的方法是对角近似,其在实际计算中取得了较好的效果,但是在近似过程中丢失了很多曲率矩阵的信息,而且忽略了参数之间的相关性.在对角近似的基础上,一些更为精确的算法也被提出,这些算法不再局限于曲率矩阵的对角元素,同时也考虑了非对角元素的影响.这些方法对曲率矩阵的研究都取得的一定了进展[4–8].但是如何在深度学习中更加有效地利用曲率矩阵得到更有效的算法,仍然是应用二阶优化方法面临的重要挑战.

自然梯度下降可以被视为一种二阶优化方法,其中自然梯度定义为梯度与模型的Fisher 信息矩阵的乘积.该方法最初在文献[9]中被提出,其在深度学习中有着重要的应用.文献[10]中提出了一种近似全连接神经网络中自然梯度的有效方法,称之为K-FAC.K-FAC算法首先通过假设神经网络各层之间的数据是独立的,将Fisher 信息矩阵近似为块对角矩阵;将每个块矩阵近似为两个更小规模矩阵的克罗内克乘积,通过克罗内克乘积的性质可以有效计算近似后的Fisher 信息矩阵及其逆矩阵.K-FAC有效地减少了自然梯度下降的计算量并取得了很好的实验效果.这一方法也被应用到其他的神经网络中,包括卷积神经网络[11–13],循环神经网络[14],变分贝叶斯神经网络[15,16].通过设计K-FAC的并行计算框架,其在大规模问题中的有效性也得到了验证[17,18].

K-FAC 算法在众多问题中都有着很好的表现,在保持K-FAC 算法有效性的前提下,进一步降低计算成本和减少计算时间是非常值得研究的问题.在本文中,我们基于K-FAC 算法的近似思想,结合拟牛顿法的思想,提出了一种校正Fisher 信息矩阵的有效方法.该方法的主要思想是先用K-FAC 方法进行若干次迭代,保存逆矩阵的信息;在后续迭代中利用该逆矩阵以及新的迭代中产生的信息,结合Sherman-Morrison 公式进行求逆计算,大大减少了迭代时间.实验中,改进的KFAC 算法比K-FAC 算法有相同甚至更好的训练效果,同时大大减少了计算时间.

1 K-FAC 算法

神经网络的训练目标是获得合适的参数θ 来最小化目标函数h(θ),给定损失函数:

其中,x,y分别是训练输入和标签,θ是模型的参数向量,p(y|x,θ)表示预测分布的密度函数.那么Fisher 信息矩阵的定义如下:

文献[10]中给出了自然梯度的定义:F−1∇θh.在实际计算中面临的主要挑战是计算自然梯度,也就是计算逆矩阵F−1及其与∇θh的乘积.在深度学习中,由于F的规模太大,直接计算F−1是不切实际的.K-FAC 提供了一种近似F−1的有效方法,其近似过程可以分为两个步骤.首先,K-FAC 将矩阵F按网络的各层分割成块矩阵,通过假设不同层之间数据是独立的,将F近似为块对角矩阵;其次将每个块矩阵近似为两个更小规模矩阵的克罗内克乘积,根据克罗内克乘积求逆及运算的相关性质,大大减少了计算量.

考虑一个有L层的神经网络,al−1,sl分别表示第l层的输入和输出,那么sl=Wlal−1,其中Wl为第l层的权重矩阵,l∈{1,2,···,L}.方便起见,定义如下的符号:

那么Fisher 信息矩阵可以被表示为:

即为:

因此F可以被看作一个L×L的块矩阵.定义Fij=E[vec(Dθi)vec(Dθi)T],i,j∈{1,2,···,L}.通过假设神经网络各层之间数据的独立性,也就是Fij=0(i≠j),那么F可以被近似为:

然而由于每个块矩阵Fll的规模仍然很大,因此需要进一步近似.每个块矩阵Fll可以被写作:

因此Fisher 信息矩阵的逆矩阵可以被近似为:

为了保持训练的稳定性,需要对克罗内克因子Al−1,l−1和Gll添加阻尼如下:

其中,λ是阻尼参数,πl理论上可以是一个任意的正数,但在实验中发现根据以下公式计算的πl是一个更好的选择.

其中,dl−1和dl分别是矩阵Al−1,l−1和Gll的维度.因此由上述内容可以得到训练中第l层参数的更新规则如下:

其中,η是学习率,m表示迭代次数.

图1是K-FAC 算法近似过程示意.

图1 K-FAC 算法近似过程示意

2 算法改进

在实际计算中,Fisher 信息矩阵F的主对角线上的每个块矩阵的规模仍然很大,直接对矩阵{F11,F22,···,FLL}进行求逆的计算成本很高.K-FAC 算法采用的方法是将Fisher 信息矩阵近似为两个小规模矩阵的克罗内克乘积,从而将求解大规模矩阵的逆转化为求解两个小规模矩阵的逆,大大降低了求逆的成本.在本文中我们结合K-FAC 方法采取的这种近似方法,给出了Fisher信息矩阵F的另一种近似,其求逆的费用更低.我们基于下面的Sherman-Morrison 公式来改进K-FAC 算法.

Sherman-Morriso 公式:假设X∈Rm×m是可逆矩阵,p,q∈Rm为任意列向量,则可逆当且仅当1+qTX−1p≠0,而且如果X+pqT可逆,逆矩阵可以由以下公式得到:

由上述公式可以看出,如果矩阵X的逆已知(或者很容易求),那么利用Sherman-Morrison 公式可以将矩阵的求逆运算转化为矩阵向量乘积,从而可以减少大量的计算时间.因此我们根据Sherman-Morrison 公式,结合拟牛顿的算法思想,提出了一种K-FAC 算法的改进算法.在实际计算中,每次迭代均更新逆矩阵需要很高的计算成本,因此实验中一般设置若干次迭代更新一次逆矩阵.在下文中,我们用k表示逆矩阵更新的次数.我们提出的算法主要是对K-FAC 算法的求逆运算进行了进一步的改进.其主要思想是用K-FAC 算法先进行k次求逆运算,保存第k次求逆得到的逆矩阵信息;后续迭代中利用该逆矩阵的信息以及在新的迭代中产生的信息,结合Sherman-Morrison 公式进行求逆运算.下面我们以矩阵A00,G11为例说明改进的方法,对于矩阵All(l∈{1,2,···,L−1})和Gll(l∈{2,3,···,L})均采用相同的近似方法.为方便表示,我们在下文中省略下标.

(1)首先,按照K-FAC 算法进行k次求逆运算,在求出逆矩阵(A(k))−1和(G(k))−1后保留逆矩阵的信息.

(2)其次,矩阵A(k+1)和G(k+1)表示的是在第k+1 次更新逆矩阵时计算得到的矩阵,利用矩阵A(k+1)和G(k+1)构造向量u(k+1)和v(k+1).从现有文献中可以看出,矩阵A(k+1)和G(k+1)的主对角线上的元素占主导性的信息,因此我们选取其主对角线上的元素来构造向量u(k+1)和v(k+1).我们选取:

那么u(k+1)u(k+1)T和v(k+1)v(k+1)T就是保留矩阵A(k+1)和G(k+1)主对角线元素的对角矩阵.

(3)最后,利用Sherman-Morrison 公式可以得到:

其中,α,β是两个合适的正数.

在改进的K-FAC 算法中,主要是结合Sherman-Morriso 公式对Fisher 信息矩阵进行了近似,因此改进的算法在计算矩阵A,G及其逆矩阵部分与K-FAC 算法有所区别,其余部分与K-FAC 算法相同.文献[1]中对K-FAC 算法整体计算复杂度进行了详细分析,因此在表1中,我们主要给出了两种方法在计算矩阵A,G及其逆矩阵的计算复杂度对比.

表1 K-FAC和改进的K-FAC 算法的求逆计算复杂度对比

对于改进的K-FAC 算法,前k次与K-FAC 算法一致.在实际计算中,k的取值远远小于t.比如在实验中我们选择k=10,t=39 100.因此前t次的计算成本占比很低.在后续的更新过程中,改进的K-FAC 算法一方面矩阵求逆部分都转化成了矩阵乘法,计算复杂度由O(n3)降为O (n2);另一方面,在后续迭代中,我们仅需要计算矩阵A和G的主对角线元素,可以直接将向量元素对应相乘,不再需要进行矩阵乘法,计算复杂度由O(n2)降为O(n).因此,通过上述两个方面改进的KFAC 算法可以减少大量的计算时间.算法1 总结了改进的K-FAC 算法的流程.

算法1.改进的K-FAC 算法η λ TFIM,Tinv k输入:训练集T,学习率,阻尼参数,Fisher 信息矩阵及其逆矩阵的更新频率,逆矩阵更新次数θ输出:模型参数All(l∈{0,1,···,L−1})Gll(l∈{1,2,···,L})m=0,t=0初始化参数和,;While 未达到终止条件do m≡0(mod TFIM)if then t<=k if then{All}L−1 l=0,{Gll}Ll=1根据式(2)计算因子m≡0(mod Tinv)if then t<=k if then{A−1 ll}L−1l=0,{G−1ll}Ll=1根据式(4)计算逆矩阵else u,v{A−1 ll}L−1l=0,{G−1ll}Ll=1计算向量,根据式(7)和式(8)计算逆矩阵end if t=t+1 end if{θl}Ll=1根据式(5)更新参数m=m+1 end While

3 实验

为了说明改进的K-FAC 算法的有效性,我们在常用的图像分类数据集上进行了实验.实验中,数据集选取的是CIFAR-10和CIFAR-100 数据集[19].这两个数据集都是由60000 张分辨率为3 2×32的彩色图像组成,训练集和测试集分别有50000和10000 张彩色图像.CIFAR-10 数据集中图像有10个不同的类,每类有6000 张图像.CIFAR-100 数据集中图像有100个不同的类,每类有600 张图像.实验中我们对两个数据集的图像都采用数据增强技术,包括随机裁剪和水平翻转.我们选择动量随机梯度下降(SGDM)和K-FAC 算法作为对比标准,在ResNet20[20]上比较了这两个方法和改进的K-FAC 算法的表现.

实验中我们采用的深度学习框架是TensorFlow,训练的硬件环境为单卡 NVIDIA RTX 2080Ti GPU.实验中批量大小(batch-size)设置为128,动量为0.9,最大迭代次数为39100,初始学习率SGDM 设置为0.03,K-FAC 算法和改进的K-FAC 算法设置为0.001,学习率每16000 次迭代衰减为原来的0.1.对于K-FAC 算法和改进的K-FAC 算法,Fisher 信息矩阵及其逆矩阵的更新频率分别为TFIM=10,Tinv=100,阻尼为0.001.在改进的K-FAC 算法中,我们令 α=β=0.1.对于所有的方法,我们均没有采用权重衰减.

在表2中,我们给出了在CIFAR-10 数据集上SGDM,K-AFC 算法和改进的K-FAC 算法的训练精度及时间比较,其中,K-FAC 算法给出了每次迭代均更新逆矩阵(1:1)和100 次迭代更新逆矩阵(100:100)的实验结果.表中第一行给出了各种算法的训练精度比较,其余各行别给出了各种方法每次迭代的平均训练时间以及测试精度首次达到89%,90%,91%,92%,93%的训练时间,表中最后一列给出了改进的K-FAC 算法(100:100)比K-FAC 算法(100:100)减少的训练时间.因为CIFAR-100 数据集和CIFAR-10 数据集图像数量和分辨率相同,这两个数据集上每次迭代的训练时间几乎相同,所以我们仅给出了在CIFAR-10 数据集上的结果,在CIFAR-100 数据集也有类似的结果.

表2 SGDM,K-FAC和改进的K-FAC 算法的训练精度及时间比较

从表2可以看出,K-FAC在不同的逆矩阵更新频率下((1:1)和(100:100))的测试精度相差不大,但每次迭代均更新逆矩阵耗费了大量的计算时间(每次迭代平均增加了2.07 s).结合之前的相关工作,在本文中我们更多关注若干次迭代更新逆矩阵的实验结果.因此,在后文中,我们主要基于K-FAC 算法(100:100)的实验结果进行讨论.

从测试精度看,改进的K-FAC 算法与K-FAC 算法相差不大.在CIFAR-10 数据集上,改进的K-FAC 算法的测试精度略低于K-FAC 算法,但在CIFAR-100 数据集上,改进的K-FAC 算法的测试精度高于K-FAC算法.从训练时间看,SGDM 从89%到90%,K-FAC 算法从91%到92%以及改进的K-FAC 算法从从91%到92%的训练时间差距较大,这是因为在学习率衰减之前,测试精度在较多的迭代中变化不大,衰减后才达到了相应的测试精度.改进的K-FAC 算法每个迭代的平均训练时间与SGDM 相比,仅增加了0.006 s,比KFAC 减少了0.023 s.从到达各个测试精度的时间看,改进的K-FAC 算法均比K-FAC 算法减少了大量的训练时间.比如在测试精度达到91%时,K-FAC 算法比SGDM 多花费了8 s,而我们改进的K-FAC 算法比SGDM 减少了356 s.从表格最后一行看,SGDM 最终的测试精度达不到93%,K-FAC 算法和改进的KFAC 算法都可以达到93%,而且改进的K-FAC 算法减少了237 s.从这些结果可以看出,我们改进的K-FAC算法可以达到与K-FAC 算法相近的训练精度,同时减少了大量的训练时间,而且与一阶优化方法相比在速度与精度上都具有一定的优势.

图2给出了在CIFAR-10 数据集上的实验结果,分别给出了SGDM,K-FAC 算法(100:100)和改进的KFAC 算法(100:100)的训练损失,训练精度和测试精度随迭代的变化曲线.在图中可以看出二阶优化方法(KFAC 算法和改进的K-FAC 算法)收敛速度明显快于SGDM,改进的K-FAC 算法与K-FAC 收敛速度相近.从训练损失看,所有的方法都可以达到较低的训练损失,SGDM的训练损失略高;从精度看,所有的方法都可以达到很高的测试精度,我们改进的K-FAC 算法在前期表现好于K-FAC.

图3分别给出了SGDM,K-FAC 算法和改进的KFAC 算法在CIFAR-100 数据集上的训练损失,训练精度和测试精度随迭代的变化曲线.从图中可以看出,CIFAR-100 数据集和CIFAR-10 数据集有着相似的实验结果.但从测试精度看,改进的K-FAC 算法好于KFAC.从这些结果我们可以看出,我们改进的K-FAC算法与K-FAC 算法相比,有着相似甚至更好的实验效果,说明我们提出的Fisher 信息矩阵的逆矩阵进一步近似的方法是有效的.

图2 CIFAR-10 数据集上的实验结果

4 结论

在深度学习中应用二阶优化问题面临的一个重要挑战是计算曲率矩阵的逆矩阵,由于深度神经网络拥有海量的参数导致其曲率矩阵的规模巨大而难以求逆.在本文中,我们基于K-FAC 算法对Fisher 信息矩阵的近似方法,结合拟牛顿方法的思想,在前期少量迭代中利用原方法训练,后续迭代利用新计算的矩阵信息构造秩–1 矩阵进行近似.利用Sherman-Morrison 公式大大降低了计算复杂度.实验结果表明,我们改进的KFAC 算法与K-FAC 算法有着相似甚至更好的实验效果.从训练时间看,我们的方法比原方法减少了大量的计算时间,与一阶优化方法相比我们改进的方法仍具有一定的优势.但如何在深度学习中更加有效地利用曲率矩阵的信息,得到更有效更实用的算法,仍然是在深度学习中应用二阶优化方法面临的重要挑战.

图3 CIFAR-100 数据集上的实验结果