二階優化！訓練ImageNet僅需35個Epoch

選自arXiv

作者：Kazuki Osawa 等

機器之心編譯

在「x 分鐘訓練 ImageNet」問題上，人們通常採用的方法是增加批大小並加大算力。隨著 ResNet-50 在 ImageNet 上的訓練時間已用秒計，人們開始轉向其他研究方向。來自東京工業大學的研究者近日採用二階方法，實現了和優化 SGD 類似的準確率和效率。

隨著神經網路的尺寸和訓練數據的持續增長，人們對分散式計算的需求也逐漸增大。在深度學習中實現分散式並行的常用方式是使用數據並行方法，其中數據被分配進不同進程中，而模型在這些進程中重複。當每個模型的 mini-batch 大小保持不變，以增加計算/通信比時，整個系統上的 mini-batch 大小會隨著進程數量成比例增長。

在 Mini-batch 大小超過某一點之後，驗證準確率就會開始下降。這一大尺寸 mini-batch 泛化限制廣泛見於學習不同的模型和數據集的情況中，Hoffer 等人曾將這種限制歸因於更新的限制，同時建議進行更長的訓練。這引出了增大學習率，並成比例增加 mini-batch 的方向，同時在訓練的前幾個 epoch 里逐漸增加學習率。這種方法可以讓訓練 mini-batch 達到 8k，在使用 ResNet-50 訓練 ImageNet 時可以達到訓練 90 epoch，達到 76.3% 的 top-1 驗證準確率。將這種學習速率方法與其他的一些技術，如 RMSprop warm-up、無動平均的批歸一化，以及緩速啟動學習率策略等技術結合，Akiba 等人曾經實現在相同的數據集和模型上，batch size 32k 的情況下，在 15 分鐘時間裡實現 74.9% 的準確率。

當然，人們在這之後還找到了更加複雜的方法，如 LARS，其中人們通過對不同層的權重和梯度的歸一化實現不同的學習速率。這使得我們在 mini-batch 大小為 32k 的情況下，無需特別的修正就可以在 14 分鐘的訓練後（64 epoch）達到 74.9% 的準確率。另外，結合 LARS 和批歸一化反直覺修正結合，我們可以在 mini-batch 大小 65k 的情況下實現 75.8% 的準確率。

使用小 batch size 幫助前幾個 epoch 的快速收斂，隨後逐步增加批大小的方法是另外一個被證明成功的思路。使用這種自適應批大小的方法，Mikami 等人曾實現在 224 秒鐘內達到 75.03% 的準確率。人們還提出了 mini-batch 的分層同步，但這些方法目前還沒有經過人們的廣泛驗證。

在東京工業大學最近的研究中，人們通過使用二階優化，採用數學上更加嚴謹的方法來處理大 mini-batch 的問題。研究人員認為在大 mini-batch 訓練中，每一個 mini-batch 都會更具統計穩定性，通過二階優化方法可能會展現優勢。新方法的另一個獨特性是與其他的二階方法相比，前者可以接近 Hessian 的準確性。與 TONGA 等粗近似 Hessian 的方法、非 Hessian 方法不同，新方法採用 Kronecker 係數近似曲率（K-FAC）方法，K-FAC 的兩個主要特徵是它要比一階隨機梯度下降方法（SGD）要快，而且它可以容忍相對較大的 mini-batch 大小而無需任何特定修正。K-FAC 已經被成功用在了卷積神經網路、ImageNet 分散式內存訓練、循環神經網路、貝葉斯深度學習和強化學習上。

在該論文中，研究人員的貢獻包括：

使用同步 all-worker 的方式實現了分散式 K-FAC 優化器。研究人員使用半精度浮點數進行計算，並利用 Kronecker 係數的對稱性減少算力消耗。

通過在 ImageNet 上訓練 ResNet-50 作為基準，研究人員首次展示了二階優化方法與高度優化的 SGD 相比可以實現類似的泛化能力。通過僅僅 35 個 epoch 的訓練，研究人員即實現了 75% 的 top-1 準確率，其中 mini-batch 大小不到 16,384——而即使 mini-batch 達到了 131,072，準確度也為 75%。

表 1：對帶有 K-FAC 的 ImageNet，ResNet-50 的訓練 epoch（迭代）和 top-1 單季驗證準確率

論文展示了在經過數百次迭代後，我們能夠減少更新 K-FAC Fisher 矩陣的頻率。這麼做我們能夠減少 K-FAC 的成本。研究人員展示了在 10 分鐘內，使用 1024 塊 Tesla V100 GPU，訓練 ResNet-50 的 Top-1 準確率準確率達到 74.9% 的結果。

論文顯示，Batch 歸一化層的 Fisher 矩陣可被近似為對角矩陣，進一步減少計算和內存消耗。

表 2：在 ImageNet 上訓練 ResNet-50 的不同方法性能對比。

圖 1:新的分散式 K-FAC 設計概覽。在三層模型上有兩個訓練流程（GPU)。

論文：Second-order Optimization Method for Large Mini-batch: Training ResNet-50 on ImageNet in 35 Epochs

論文鏈接：https://arxiv.org/abs/1811.12019

摘要：深度神經網路的大規模分散式訓練會遭受由有效 mini-batch 大小增加所帶來的泛化差距。先前，曾有一些方法嘗試使用改變 epochs 和層上的學習率來解決這一問題，或者對一些 batch 歸一化做特別的修改。我們提出一種使用二階優化方法的替代性方法，達到了類似於一階方法的泛化能力，但其收斂更快且能處理更大的 mini-batch。為了在基準上測試我們的方法，我們在 ImageNet 上訓練了 ResNet-50。在 35 個 epoch 內，我們的方法把低於 16,384 的 mini-batch 收斂到了 75% 的 Top-1 驗證準確率，而即使是 mini-batch 大小為 131,072 時，我們花費 100 個 epoch 也只能取得 75% 的準確率。

本文為機器之心編譯，轉載請聯繫本公眾號獲得授權。

------------------------------------------------

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！