入门神经网络优化算法（六）：二阶优化算法K-FAC

优化算法系列文章索引：

• 入门神经网络优化算法（一）：Gradient Descent，Momentum，Nesterov accelerated gradient
• 入门神经网络优化算法（二）：Adaptive Optimization Methods：Adagrad，RMSprop，Adam
• 入门神经网络优化算法（三）：待定
• 入门神经网络优化算法（四）：AMSGrad，Radam等一些Adam变种
• 入门神经网络优化算法（五）：二阶优化算法Natural Gradient Descent（Fisher Information）
• 入门神经网络优化算法（六）：二阶优化算法K-FAC
• 入门神经网络优化算法（七）：二阶优化算法Shampoo

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

A block-wise Kronecker-factored Fisher approximation

自然梯度的算法关键就是计算Fisher矩阵的逆，$F^{-1}$。首先表示对数似然loss对参数的梯度为（score function）：

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

$W_i$表示第$i$层的参数矩阵，很容易知道，$D{W}_i=g_i{\bar{a}}_{i-1}^T$，$g_i$表示第$i$层的梯度，${\bar{a}}_{i-1}$表示input。

介绍Kronecker product：

以及一个小性质：$vec(u{v}^T)=v\otimes u$。因此上面的Fisher matrix中的每一个小项可以写成：

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

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

Approximating $\tilde{F}$ as block-diagonal

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

Kronecker product还有一个很好用的性质：$(A\otimes B{)}^{-1}=A^{-1}\otimes B^{-1}$，因此我们最后可以得到（近似）Fisher Matrix的逆为：

所以，最后我们只要计算${\bar{A}}_{i,j}$和$G_{i,j}$的逆就好了。但是，这样还是要去算Kronecker product，这个感觉还是有点复杂。那么还需要一个很好用的性质：$(A\otimes B)\text{vec}(X)=\text{vec}(BX{A}^T)$

考察任意其中一层$i$，我们实际要算的是$u_i={\tilde{F}}_{ii}^{-1}{g}_i$，我们记$g_i=\text{vec}(V_i)$，$V_i$是梯度矩阵形式，size类比于$W_i$。我们需要算的最终形式是：$({\bar{A}}_{i-1,i-1}^{-1}\otimes G_{i,i}^{-1})\text{vec}(V_i)=\text{vec}(G_{i,i}^{-1}{V}_i{\bar{A}}_{i-1,i-1}^{-1})$。不考虑等式右边的vec，我们可以得到第$i$层的自然梯度：

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

参考资料

[1] Optimizing Neural Networks with Kronecker-factored Approximate Curvature, 2016
[2] 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