入门神经网络优化算法:二阶优化算法K-FAC

Posted 大饼博士X

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了入门神经网络优化算法:二阶优化算法K-FAC相关的知识,希望对你有一定的参考价值。

文章目录


优化算法系列文章索引:

上一篇介绍了二阶优化算法Natural Gradient Descent(自然梯度算法),虽然可以避免计算Hessian,但是依然在计算代价上极高,对于大型的神经网络参数规模依然不可能直接计算。本篇继续介绍自然梯度算法后续的一个近似计算方法K-FAC[1],让自然梯度算法可以(近似)实现。如果还不清楚自然梯度算法,可以回看:入门神经网络优化算法(五):二阶优化算法Natural Gradient Descent(Fisher Information)

A block-wise Kronecker-factored Fisher approximation

自然梯度的算法关键就是计算Fisher矩阵的逆, F − 1 F^-1 F1。首先表示对数似然loss对参数的梯度为(score function):


vec()表示把一个矩阵向量化,D表示梯度计算。所以可以把Fisher矩阵写成:


W i W_i Wi表示第 i i i层的参数矩阵,很容易知道, D W i = g i a ˉ i − 1 T DW_i = g_i \\bara_i-1^T DWi=giaˉi1T g i g_i gi表示第 i i i层的梯度, a ˉ i − 1 \\bara_i-1 aˉi1表示input。

介绍Kronecker product:

以及一个小性质: v e c ( u v T ) = v ⊗ u vec(uv^T) = v\\otimes u vec(uvT)=vu。因此上面的Fisher matrix中的每一个小项可以写成:


这里引入第一次近似:Kronecker product期望近似成期望的Kronecker product。目的是要把 F F F每一个小项表示成两个矩阵的Kronecker product。而 A ˉ i j \\barA_ij Aˉij G i j G_ij Gij可以用batch的平均来计算。但是这样的近似还是很难算出整个Fisher矩阵。观察Fisher矩阵非对角块,我们可以发现Fisher矩阵建立了各种两两层之间的参数梯度关系。这个和一般的一阶梯度方法就是很大的不同了(类似Hessian矩阵,计算复杂度很大)。


推到这里,还是不能很容易算,接下来引入第二次近似

Approximating F ~ \\tildeF F~ as block-diagonal

只考虑对角块元素,也就是只在每一个layer内考虑计算FIsher矩阵,这样就变成一个块对角矩阵。我们知道块对角矩阵的逆就很容易求了,只要求每一个块的逆就行了。

Kronecker product还有一个很好用的性质: ( A ⊗ B ) − 1 = A − 1 ⊗ B − 1 (A \\otimes B)^-1 = A^-1 \\otimes B^-1 (AB)1=A1B1,因此我们最后可以得到(近似)Fisher Matrix的逆为:


所以,最后我们只要计算 A ˉ i , j \\barA_i,j Aˉi,j G i , j G_i,j Gi,j的逆就好了。但是,这样还是要去算Kronecker product,这个感觉还是有点复杂。那么还需要一个很好用的性质: ( A ⊗ B ) vec ( X ) = vec ( B X A T ) (A \\otimes B)\\textvec(X) = \\textvec(BXA^T) (AB)vec(X)=vec(BXAT)

考察任意其中一层 i i i,我们实际要算的是 u i = F ~ i i − 1 g w i u_i = \\tildeF_ii^-1g_wi ui=F~ii1gwi,我们记 g w i = vec ( V i ) g_wi = \\textvec(V_i) gwi=vec(Vi) V i V_i Vi是梯度矩阵形式,size类比于 W i W_i Wi,是i层的outsize*insize。(注意,这里的 g w i g_wi gwi表示 W i W_i Wi权重的梯度,前面 g i g_i gi表示的是top-diff)。

我们需要算的最终形式是: ( A ˉ i − 1 , i − 1 − 1 ⊗ G i , i − 1 ) vec ( V i ) = vec ( G i , i − 1 V i A ˉ i − 1 , i − 1 − 1 ) (\\barA^-1_i-1,i-1\\otimes G^-1_i,i)\\textvec(V_i) = \\textvec(G^-1_i,i V_i \\barA^-1_i-1,i-1) (Aˉi1,i11Gi,i1)vec(Vi)=vec(Gi,i1ViAˉi1,i11)。不考虑等式右边的vec,我们可以得到第 i i i层的自然梯度,size类比于 W i W_i Wi


那么自然梯度就可以算了,但是其中有两个比较大的矩阵逆怎么办呢?复杂度也很高啊!这里就没有进一步近似了,参考我的博客:三十分钟理解:矩阵Cholesky分解,及其在求解线性方程组、矩阵逆的应用,用Cholesky分解来求解矩阵逆。所以到这里,我们可以发现,二阶算法即使在一系列近似以后,计算复杂度依然很大,但是K-FAC已经是相对比较容易计算的二阶算法了,有研究工作[2]就利用了K-FAC,并在分布式计算环境下实现算法。只需要35epoch,16K Batchsize下,可以训练ResNet50在ImageNet下达到75%的Top1准确率,效果非常不错。


另外最新的一个工作是[7],也是用到了K-FAC来训练卷积神经网络。主要参考了[9]中关于preconditioned
gradient计算方案,用特征分解来代替了矩阵求逆。在这里我们再复习一下上面的流程,用[7]中的截图





上面公式(13)有一点小问题,应该是 ∇ L i ( w i ( k ) ) \\nabla L_i(w_i^(k)) Li(wi(k))。用特征分解的方法来替换原来对A和G的求逆。

参考资料

[1] Optimizing Neural Networks with Kronecker-factored Approximate Curvature, 2016
[2] 2019 Large-Scale Distributed Second-Order Optimization Using Kronecker-Factored Approximate Curvature for Deep Convolutional Neural Networks
[3] DISTRIBUTED SECOND-ORDER OPTIMIZATION USING KRONECKER-FACTORED APPROXIMATIONS
[4] 2013 Revisiting natural gradient for deep networks
[5] 2014 New insights and perspectives on the natural gradient method
[6] Phd Thesis, James Martens, SECOND-ORDER OPTIMIZATION FOR NEURAL NETWORKS
[7] 2020 Convolutional Neural Network Training with Distributed K-FAC
[8] 2018 an_evaluation_of_fisher_approximations_beyond_kronecker_factorization
[9] 2016 A Kronecker-factored approximate Fisher matrix for convolution layers

以上是关于入门神经网络优化算法:二阶优化算法K-FAC的主要内容,如果未能解决你的问题,请参考以下文章

入门神经网络优化算法:一文看懂二阶优化算法Natural Gradient Descent(Fisher Information)

神经网络优化算法综述

入门神经网络优化算法:Gradient Descent,Momentum,Nesterov accelerated gradient

PSO算法解决带约束条件的优化问题

优化算法 - 特点

数值优化基础