76分鐘訓練BERT！谷歌大腦新型優化器LAMB加速大批量訓練

選自arXiv

作者：Yang You、Jing Li等

機器之心編輯部

去年，谷歌發布了基於雙向 Transformer 的大規模預訓練語言模型 BERT並開源。該模型參數量非常大——3 億，訓練時間也很長。近日，來自谷歌大腦的研究者提出了一種新型優化器 LAMB，有效降低 BERT 預訓練時間，76 分鐘即可完成 BERT 預訓練！

儘管BERT效果驚人，但它所需的計算量非常大，原作者在論文中也表示每次只能預測 15% 的詞，因此模型收斂得非常慢。如果我們想保留這種 Mask 機制，那麼就需要尋找另一種加速方法了。

當時，BERT 的作者在 Reddit 上也表示預訓練的計算量非常大，Jacob 說：「OpenAI 的 Transformer 有 12 層、768 個隱藏單元，他們使用 8 塊 P100 在 8 億詞量的數據集上訓練 40 個 Epoch 需要一個月，而 BERT-Large 模型有 24 層、2014 個隱藏單元，它們在有 33 億詞量的數據集上需要訓練 40 個 Epoch，因此在 8 塊 P100 上可能需要 1 年？16 Cloud TPU 已經是非常大的計算力了。」

而在谷歌大腦的這篇新論文中，研究者提出新型優化器LAMB，通過使用 65536/32768 的批量大小，他們只需要 8599 次迭代、76 分鐘就能完成 BERT 預訓練。總體而言，相比基線 BERT-Large 用 16 塊 TPU 晶元，LAMB 訓練 BERT-Large 用了一個 TPU v3 Pod（1024 塊 TPU 晶元），因此時間也由 3 天降低為 76 分鐘。

不過在 BERT 原論文中，訓練 BERT-Large 使用了 64 塊 TPU 晶元（16 Cloud TPU），它需要花四天時間完成訓練。

論文：Reducing BERT Pre-Training Time from 3 Days to 76 Minutes

論文地址：https://arxiv.org/abs/1904.00962

摘要：大批量訓練是加速大型分散式系統中深度神經網路訓練的關鍵。但是，大批量訓練難度很大，因為它會產生泛化差距（generalization gap），直接優化通常會導致測試集準確率受損。BERT [4] 是當前最優的深度學習模型，它基於用於語言理解的深度雙向 transformer 而構建。當我們擴展批量大小時（比如批量大小超過 8192），之前的大批量訓練技術在 BERT 上的執行性能並不好。BERT 預訓練需要大量時間（使用 16 個 TPUv3 訓練 3 天）。

為了解決這個問題，來自谷歌大腦的研究者提出了一種新型優化器 LAMB，可在不損害準確率的情況下將批量大小擴展至 65536。LAMB 是一款通用優化器，它適用於小批量和大批量，且除了學習率以外其他超參數均無需調整。基線 BERT-Large 模型的預訓練需要 100 萬次迭代，而 LAMB 使用 65536/32768 的批量大小，僅需 8599 次迭代。研究者將批量大小擴展到 TPUv3 pod 的內存極限，在 76 分鐘內完成了 BERT 的訓練。

具體來講，LAMB 優化器支持自適應元素級更新（adaptive element-wise updating）和準確的逐層修正（layer-wise correction）。LAMB 可將 BERT 預訓練的批量大小擴展到 64K，且不會造成準確率損失。BERT 預訓練包括兩個階段：1）前 9/10 的訓練 epoch 使用 128 的序列長度，2）最後 1/10 的訓練 epoch 使用 512 的序列長度。

基線 BERT-Large 模型的預訓練需要 100 萬次迭代，研究者僅使用 8599 次迭代就完成了預訓練，將訓練時間從 3 天縮短到 76 分鐘。該研究使用的訓練批量大小接近 TPUv3 pod 的內存極限。LAMB 優化器可以將批量大小擴展到 128k 甚至更大，由於硬體限制，研究者在序列長度為 512 時使用的批量大小為 32768，在序列長度為 128 時使用的批量大小為 65536。該論文中的 BERT 模型均指 BERT-Large。為公平起見，研究中所有實驗均運行同樣數量的 epoch（即固定數量的浮點運算）。實驗結果見下表。

表 1：該研究使用 SQuAD-v1 的 F1 分數作為準確率度量。

在上表中，基線 F1 分數來自於 BERT 開源 GitHub 中 BERT-Large 預訓練模型的對應分數。實驗所用硬體為 TPUv3，實驗設置與基線 BERT-Large 相同：前 9/10 的訓練 epoch 使用 128 的序列長度，最後 1/10 的訓練 epoch 使用 512 的序列長度。所有實驗均運行同樣數量的 epoch。

LAMB (Layer-wise Adaptive Moments optimizer for Batch training)

BERT 基線模型的訓練使用 Adam with weight decay（Adam優化器的變體）作為優化器 [15]。另一個成功用於大批量卷積神經網路訓練的自適應優化器是 LARS [21]。這些優化器啟發了研究者提出用於大批量 BERT 訓練的新型優化器 LAMB。LAMB 優化器詳見演算法 1。

實驗

常規訓練

TPU 是強大的浮點運算計算硬體。研究者在所有實驗中均使用 TPUv3。每個 TPUv3 pod 擁有 1024 個晶元，可提供超過 100 petaflops 的混合精度計算。實驗結果見表 1。基線模型預訓練過程中使用了 Wikipedia 和 BooksCorpus 數據集。研究者使用與開源 BERT 模型相同的數據集進行預訓練，即包含 2.5B 單詞的 Wikipedia 數據集和包含 800M 單詞的 BooksCorpus 數據集。

BERT 作者首先以 128 的序列長度進行 900k 次迭代，然後以 512 的序列長度進行 100k 次迭代。在 16 塊 TPUv3 上訓練的總時間大約為 3 天。該研究使用 SQuAD-v1 的 F1 分數作為準確率度量。較高的 F1 分數意味著較好的準確率。研究者下載了 BERT 開源項目中提供的預訓練模型。使用 BERT 作者提供的腳本，基線模型得到了 90.395 的 F1 分數。

該研究中，研究者使用 BERT 作者提供的數據集和基線模型，僅改變了優化器。使用新型優化器 LAMB 後，研究者以 32768 的批量大小進行了 15625 次迭代，得到了 91.460 的 F1 分數（用於序列長度 128 的迭代有 14063 次，用於序列長度 512 的迭代有 1562 次）。研究者將 BERT 訓練時間從 3 天縮短到大約 100 分鐘。

該研究取得了 76.7% 的弱可擴展性效率（weak scaling efficiency）。研究者在 TPU Pod 上使用了分散式訓練的同步數據並行化，因此梯度遷移會產生通信開銷。這些梯度與訓練後的模型大小一樣。在 ImageNet 數據集上訓練 ResNet-50 時的弱可擴展性效率可達到 90 %，因為 ResNet-50 比 BERT 的參數少得多（25 million vs 300 million）。LAMB 優化器無需用戶調整超參數，用戶只需輸入學習率即可。

混合批次訓練（Mixed-Batch Training）

如前所述，BERT 預訓練主要分為兩部分，1）前面 9/10 的 Epoch 使用 128 的序列長度；2）最後 1/10 的 Epoch 使用 512 的序列長度進行訓練。對於第二階段而言，因為內存限制，TPUv3 Pod 上最大的批量大小為 32768，因此第二階段使用的批大小為 32768。對於第一階段，受限於內存，TPUv3 Pod 上最大的批量大小為 131072。然而，研究者將批大小從 65536 增加到 131072 時，並沒有發現明顯的加速，因此研究者將階段 1 的批大小定為 65536。

此前，也有其它研究者討論了混合批次訓練，然而他們是在訓練中增大批量大小；而本研究是降低批大小，這樣他們能從開始到結束充分利用硬體資源。增大批大小能起到 warm-up 初始訓練和穩定最優化過程的效果，但是降低批大小可能會帶來最優化過程的紊亂，並導致訓練的不收斂。

在研究者的實驗中，他們發現有一些技術能穩定第二階段的最優化過程。因為這兩階段會切換到不同的最優化問題，因此有必要重新 warm-up 最優化過程。研究者沒有在第二階段進行學習率衰減，而是將學習率從零開始增加（re-warm-up）。和第一階段一樣，研究者在 re-warm-up 之後執行學習率衰減。因此，研究者能夠以 8599 次迭代完成 BERT 訓練，且僅用時 76 分鐘，達到了 101.8% 的弱可擴展性效率。

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

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

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