入門神經網絡優化算法（六）：二階優化算法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$層的參數矩陣，很容易知道，$DW_i = g_i \bar{a}_{i-1}^T$，$g_i$表示第$i$層的梯度，$\bar{a}_{i-1}$表示input。

介紹Kronecker product：

以及一個小性質：$vec(uv^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}(BXA^T)$

考察任意其中一層$i$，我們實際要算的是$u_i = \tilde{F}_{ii}^{-1}g_i$，我們記$g_i = \text{vec}(V_i)$，$V_i$是梯度矩陣形式，size類比於$W_i$。我們需要算的最終形式是：$(\bar{A}^{-1}_{i-1,i-1}\otimes G^{-1}_{i,i})\text{vec}(V_i) = \text{vec}(G^{-1}_{i,i} V_i \bar{A}^{-1}_{i-1,i-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