2017 LARS:LARGE BATCH TRAINING OF CONVOLUTIONAL NETWORKS (訓練大Batch的卷積神經網絡)

0. 寫在前面

前一段時間調研AI芯片在儘量保證準確率的情況下，如何快速訓練大規模數據集並分析其可行性。UC Berkeley的大佬尤洋[1]的這篇<LARS：LARGE BATCH TRAINING OF CONVOLUTIONAL NETWORKS> 提供了一個很好的保證在大Batch Size的分佈式訓練情況下，精度損失很小的策略。下面， 讓我們開始~

1. 摘要 (Abstract)

加速訓練大規模CNN的場景策略是分佈式訓練（即加多計算節點）. 由於data-parallel的SGD的使用和計算節點的增大（成百上千個計算節點），訓練的模型精度會劇烈降低。我們認爲，目前對大batch的訓練策略不夠泛化而且可能導致訓練無法收斂。

爲了克服這個問題，作者提出了一種基於Layer-Wise Aaptive Rate Scaling (LARS) 的學習算法，通過使用LARS， 作者在Alexnet的batch size = 8K，Resnet-50 的batch size = 32K的時候可以讓模型收斂而無精度損失。

2. 介紹 (Introduction)

本部分是刪節版介紹，一些related work基本不會提及. 如感興趣， 請查閱英文原版.

我們都知道，訓練一個大型的CNN網絡是非常耗時的，最直接的加速方式就是增加算力(比如更多的GPU節點)，通過data-parallel 的SGD來加速訓練。其中， 每個GPU節點（或者AI芯片）接收整體batch （比如8K，32K）中的一部分，爲了充分利用計算資源，每個計算節點接收的batch應儘可能大。但是，天下沒有免費的午餐，這種通過加計算節點分佈式訓練的策略，對模型的精度有顯著的負面影響。

容易理解，當Epoch保持固定(數據集也不變的時候), 大的batch size意味着權重調整的次數下降/減少. 一個直觀的解決方式是(加大)學習率來compensate. 然而，使用更大的學習率會使得網絡在初始階段就發散（這也是爲啥近幾年warm up策略火的原因）。很多研究人員針對這個問題提出自己的思路，這裏不詳細介紹，總之：線性調整學習率 + warm up是大batch 訓練的state-of-the-art （17年前最好的方法）。

但是，作者使用這種策略來在ImageNet上訓練Alexnet，發現batch size在2K後，隨着batch size的增大，模型開始出現精度下降甚至發散的情況。爲了解決這種情況，作者將Alexnet裏面的歸一化策略Local Response Normalization （LRN）替換爲Batch Normalization （BN），通過將LRN替換爲BN，作者發現精度的損失變小了。

爲了分析在大學習率下訓練的穩定性，作者分析了模型中每層的權值範數和更新的梯度範數。一個有意思的現象被觀察到：當這個比例很高時，訓練就變得不穩定；當比例很低時，權值的更新變得很慢。因爲這個比例在不同層之間差異很大，所以有必要對神經網絡中的每一層都使用其獨特的學習率。這就是LARS被提出的基礎。

與ADAM和RMSProp等學習算法不同的是，LARS：

• 網絡中的每層都使用不同的學習率，使得訓練更穩定。
• 權值更新的量級與權值的範數有關，這樣更好的控制訓練的速度。

通過使用LARS，作者聲稱他們訓練的Alexnet-BN和Resnet-50在batch size過萬的情況下，精度不出現損失的情況。

3. 背景 (Background)

深度學習中，權值更新（或者說模型學習）的策略如下：

在每一步$t$時, mini-batch個$x_i$被從訓練集中挑選出來，通過計算該mini-batch的損失函數的梯度$\nabla L(x_i, w)$， 網絡的權值更新通過隨機梯度法如下：

假設我們有$N$個計算節點，總共的batch size爲$B$, 那麼每個計算節點分到的mini-batch爲 $\frac{B}{N}$個樣本。但隨着batch加大，我們發現訓練就變得愈發困難（模型精度下降乃至發散）。之前的思路爲了避免這種問題，會不斷的調整一系列超參數（學習率，動量等）。

Krizhevsky[2]指出一種Linear Scaling策略：當我們對$B$增大$k$倍時，應該同樣對學習率增大$k$倍（並保持其它超參數不變）。

Linear Scaling背後的邏輯很直接：在保持epoch數量不變的情況下，如果你將batch size加大k倍, 那麼更新權重的頻次因此降低k倍。所以，加大學習率是一種很自然的方式。

但這種方式對batch size大於2K的情況, 訓練變得很難收斂(training diverges), 將Alexnet中的LRN替換爲BN後, 效果會好不少~，但是精度還是有損失的。

我們發現，主要的障礙在於高學習率帶來的不穩定性，基於此，Goyal[3]等人提出LR+warm up的策略：

training starts with small LR, and then LR is gradually increased to
the target. 

這種策略，在$B$=8K的時候，Resnet-50在無精度損失的情況下收斂，是目前的state-of-the-art。

另一個由大batch帶來的問題是：“generalization gap” （泛化差距）. Keskar[4]等人指出：large-batch的方法傾向於收斂於訓練函數的精確最小化. 但是對泛化性能並沒有什麼提升…

4. 分析Alexnet在大batch下的訓練效果

下表是作者在不同batch size, 不同的基準學習率下，對Alexnet原版和Alexnet-BN版本進行試驗的結果統計。

由右表(Alexnet-BN)可以發現，$B$=512和$B$=8192時，精度掉了2.2個百分點。爲了檢查這個是不是"generalization gap" . 根據下圖的loss圖，我們並沒有發現$B$=512和$B$=8192有顯著的差別。

我們的結論是，在這種情況下，精度損失與泛化差距無關，這是由低訓練造成的。

5. Layer-wise Adaptive Rate Scaling LARS

首先, 在第一次迭代後，分析每層權值的L2範數和梯度更新的L2範數以及其相應的比值:
$||w|| / ||\nabla L(w_t)||$.

接着，對每層$l$使用其獨特的學習率$\lambda^{l}$ (Local LR)，則其權值更新的值從原來的$\Delta w_{t}^{l} = \lambda * \nabla L(w_t)$變爲$\Delta w_{t}^{l} =\gamma *\lambda^{l} *\nabla L(w_t^{l})$. 其中，$\gamma$是全局學習率 (global LR)，Local LR的定義爲:
$\lambda^{l} = \eta * \frac {||w^l||} {||\nabla L(w^{l})||}$

$\eta$這個參數決定了一種置信度：即我們有多相信當前層(第$l$層)會在一次更新中改變它的權重。

需要注意的是，現在，模型中每層權值更新的幅度與梯度的幅度無關，所以這有助於部分消除梯度消失和梯度爆炸等現象。我們的LARS可以擴展到SGD (平衡local LR與weight decay，$\beta$是weight decay係數)：

$\lambda^{l} = \eta \times \frac {||w^l||} {||\nabla L(w^{l})|| + \beta * ||w^l||}$

下表就是LARS + SGD的算法流程圖：

6. 試驗結果

訓練在Nvidia DGX1上進行，可以看出，試驗結果非常不錯，比起之前的結果好上不少。

由下圖可以看出，帶LARS的訓練策略，$B$=8K的表現遠遠比不帶LARS的測試集要好。

雖然帶LARS的SGD比普通SGD的效果更好，但是，也沒有達到baseline(即達到 $B$=512的Alexnet的水平)。 但是，這可以通過訓練更長的時間來達到。

作者認爲: 當batch非常大的時候, 隨機梯度將與真實梯度非常接近. 所以, bs大的時候相比bs小的時候並沒有更多的額外梯度信息. 下表的試驗結果驗證了這一觀點。

7. 結論

大Batch Size是擴大訓練卷積網絡的核心.

現有的方法通常是採用更大的學習率, 但這不可避免的帶來了網絡的不收斂. (即便使用warm up策略)

爲了解決這個優化問題, 作者提出了LARS, 通過使用LARS, 作者成功的將Resnet50和Alexnet-BN擴展到BS=32k的情況(訓練ImageNet).

參考資料

[1]: You Yang, Igor Gitman, Boris Ginsburg: <LARGE BATCH TRAINING OF CONVOLUTIONAL NETWORKS>
[2]: Alex Krizhevsky: <One weird trick for parallelizing convolutional neural networks>
[3]: Priya Goyal, Piotr Dollár and et al. <Accurate, large minibatch sgd: Training imagenet
in 1 hour >
[4]: Nitish Shirish Keskar, Dheevatsa Mudigere, and et al: <On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima>