【新智元導讀】BERT是目前最強大的NLP預訓練模型，也是工業界目前最耗時的應用，計算量遠高於ImageNet。谷歌的研究人員提出新的優化器，使用1024塊TPU，將BERT的訓練時間從3天成功縮短到76分鐘，提速 65.2 倍！

去年，谷歌發布了最強預訓練模型 BERT，宣告了NLP領域的一項重大突破。

BERT 在 33 億文本的語料上訓練語言模型，再分別在不同的下游任務上微調，在11個不同的 NLP 任務均得到了目前為止最好的結果。

不過，在 33 億文本的語料上預訓練一個 BERT 模型的成本是非常大的，谷歌用了 16 個自己的 TPU 集群（一共 64 塊 TPU）來訓練大號版本的 BERT，一共花了約4天的時間。

如此巨大的訓練成本，讓普通研究者難以嘗試自己去訓練一個BERT。

有沒有辦法加快BERT的訓練呢？近日，來自Google、UC Berkeley、UCLA的幾位研究人員提出新的優化器——LAMB 優化器，將訓練的batch size推到硬體的極限，使用 TPU Pod ( 1024塊 TPUv3 晶元)，將BERT的訓練時間從3天縮短到了76分鐘！

論文地址：

「鏈接」

其中一作尤洋(Yang You)來自UC Berkeley，這項工作於他在Google Brain實習期間完成。

接下來，新智元帶來對這篇論文的譯介：

加快深度神經網路最有效的方法

大批量訓練 (large-batch training) 是加快大型分散式系統中深度神經網路訓練的關鍵。然而， large-batch 訓練是很困難的，因為它會產生一種泛化差距 (generalization gap)。直接優化通常會導致測試集的準確性下降。

BERT 是一種最先進的深度學習模型，建立在用於語言理解的深度雙向 transformers 之上。對 BERT 來說，當擴大批大小 (例如超過 8192) 時，以前的 large-batch 訓練技術效果並不好。BERT 的預訓練也需要很長時間才能完成 (使用 16 個 TPUv3 晶元大約需要 3 天)。

為了解決這個問題，我們提出了 LAMB 優化器，它幫助我們將批大小擴大到 65536，而不會丟失準確性。

LAMB 是一個通用的優化器，適用於小批量和大批量，並且除了學習率外不需要超參數調優。基線 BERT-Large 模型需要 100 萬次迭代才能完成預訓練，而 batch size 為65536/32768 的 LAMB 只需要8599 次迭代。我們將 batch size 推到 TPUv3 pod 的內存上限，可以在 76 分鐘內完成 BERT 訓練 (表 1)。

表 1：我們使用 SQuAD-v1 的 F1 score 作為精度指標。F1 的基線成績是由 BERT 的公共 github 提供的預訓練模型 (BERT- large) 實現的 (截止到 2019 年 2 月 1 日)。我們在實驗中使用 tpuv3。我們使用了與基線相同的設置：總 epochs 的前 9/10 使用序列長度128，最後 1/10 使用序列長度 512。所有的實驗運行相同數量的 epochs。

深度神經網路的訓練是十分耗時的。目前，減少訓練時間最有效的方法是使用多個晶元(如 CPU、GPU 和 TPU) 來並行化 SGD 變體的優化過程。由於前向傳播和反向傳播中不同層之間的數據依賴關係，使得跨層的並行化效率並不高。相反，研究人員在每次迭代中並行化小批量中的數據點。如果確定了訓練的 epochs 的數量，那麼線性地增大batch size 意味著會線性地減少迭代次數 (即更新權重的次數)。為了最小化訓練時間，最大化 batch size 將是理想的。

然而，大批量的訓練是困難的。例如，使用大小為 512 的 batch size 訓練在 ImageNet上訓練 AlexNet，能實現 80% 以上的 top-5 測試精度。但將 batch size 擴大到 4096之後，直接訓練可能只能獲得 50% ~ 60% 的 top 5 精度。

Keskar 等人 (10) 認為在大批量訓練中存在一個泛化差距(generalization gap)。Hoffer等人 (6) 認為，訓練時間越長，泛化差距越小。然而，訓練時間過長意味著進行大批量訓練就沒有好處了。

因此，大批量訓練的目標是在一定數量的 epochs 內達到可觀的精度。通過設計一系列的學習率計劃表，研究者已經可以將 ImageNet 訓練的 batch size 擴大到 32K，並且精度損失較小。據我們所知， Ying et al. 實現了目前最快的 ImageNet 訓練速度，並且達到了 76+% 的 top-1 精度。通過使用 LARS 優化器，將 batch size 擴展到 32K,，Ying等人使用 TPUv3 Pod，在 2.2 分鐘內完成了 ResNet-50 的 ImageNet 訓練。（最新，富士通研究院刷新了這一速度，將 ImageNet 訓練時間降到 74.7 秒）

BERT 是目前最先進的深度學習語言模型。BERT 建立在用於語言理解的深度雙向transformers 之上。對 BERT 來說，當將 batch size 擴大到非常大時 (例如超過8192)，以前的 large-batch 訓練技術效果並不好。BERT 的預訓練也需要很長時間才能完成 (使用 16 個 TPUv3 晶元大約需要 3 天)。

為了擴大 BERT 的 batch size，本文提出LAMB 優化器。LAMB 支持自適應 element-wise updating 和精確的逐層修正 (layer-wise correction)。

LAMB 是一個適用於小批量和大批量的通用優化器。用戶只需要調整學習率，不需要調其他超參數。使用 LAMB，我們可以將 BERT 預訓練的批大小擴大到 64K，而不會丟失準確性。

BERT 預訓練包括兩個階段：

(1) 前 9/10 的 epochs 使用 128 的序列長度；

(2) 後 1/10 epochs 使用 512 的序列長度。

baseline 需要 100 萬次迭代來完成 BERT 預訓練，但我們只需要 8599 次迭代，這使我們能夠將 BERT 訓練時間從 3 天減少到 76 分鐘。

我們將批大小推到了 TPU Pod 的硬體極限。批大小大於 32768(序列長度為 512) 的話將耗盡內存。批大小大於 65536(序列長度為 128) 則不會帶來任何加速。我們的優化器可以將批大小擴大到 128k，甚至更大。由於硬體限制，序列長度為 512 的設置下，我們在批大小達到 32768 時停下，在序列長度為 128 的設置下，批大小達到 65536 時停止。

本文中所有的 BERT 模型都指 BERT-Large 模型。為了進行公平的比較，本文所有的實驗都運行相同數量的 epochs(即固定數量的浮點運算)。我們的結果如表 1 所示。

LAMB優化器

LAMB的全稱是Layer-wise Adaptive Moments optimizer for Batch training。

BERT 訓練的基線使用權重衰減的 Adam 作為優化器，這是 Adam 優化器的一個變體。另一個成功應用於大批量卷積神經網路訓練的自適應優化器是 LARS。

這些優化器啟發我們提出了新的優化器，用於大批量 BERT 訓練。我們提出的 LAMB 優化器的概述如演算法 1 所示。

實驗和結果

常規訓練

TPU 是浮點運算的強大計算硬體。我們在所有的實驗中都使用了 TPUv3。TPUv3 Pod有 1024 個晶元，可以為混合精度計算提供超過 100 petaflops 的性能。我們的結果如表 1 所示。基線模型在預訓練時使用 Wikipedia 和 BooksCorpus 數據集。

我們使用了與原始 BERT 模型相同的數據集，即 Wikipedia 和 BooksCorpus，分別有2.5B 和 8 億單詞。原始 BERT 模型的作者首先以 128 的序列長度進行了 900k 次迭代訓練，然後以 512 的序列長度進行了 100k 迭代訓練。

16 個 TPUv3 晶元的總訓練時間約為 3 天。我們使用 SQuAD-v1 的 F1 分數作為精度指標。F1 得分越高，準確度越高。斯坦福問答數據集 (SQuAD) 是一個閱讀理解數據集，包含眾包工作者從維基百科的文章中提出的問題，每一個問題的答案都是對應閱讀文章的一段文字，或者該問題無法回答。我們從 BERT 的公開 GitHub 庫上下載了預訓練好的模型。

使用作者提供的腳本，baseline 的 F1 得分為 90.395。在我們的代碼中，我們使用了BERT 的作者提供的數據集和基線模型，只修改了優化器。通過使用 LAMB 優化器，我們能夠在批大小為 32768 的 15625 次迭代中獲得 91.460 的 F1 分數 (序列長度為 128的 14063 次迭代和序列長度為 512 的 1562 次迭代)。

我們把訓練時間從 3 天減少到100 分鐘左右。我們將批大小推到了 TPU Pod 的硬體極限。批大小大於 32768 時 (序列長度為 512) 將導致 TPU Pod 耗盡內存。

我們實現了76.7% 的弱擴展效率 (49.1 倍的加速，64 倍的計算資源)。由於我們在 TPU Pod 上使用同步數據並行來進行分散式訓練，因此在互連上傳輸梯度會帶來通信開銷。梯度的大小與訓練後的模型相同。

Mixed-Batch 訓練

如前所述，BERT 預訓練包括兩個階段：

(1) 前 9/10 的 epoch 使用 128 的序列長度，

(2) 最後 1/10 的 epoch 使用 512 的序列長度。

對於第二階段，由於內存限制，TPUv3 Pod 上的最大批大小為 32768，因此我們將第二階段在批大小達到 32768 時停止。

對於第一階段，由於內存限制，TPUv3 Pod 上的最大批大小是 131072。但是，當我們將批大小從 65536 增加到 131072 時，並沒有看到加速，因此我們在第一階段批大小達到 65536 時停止。

此前，Smith 等人也研究了混合批訓練。但是，他們在訓練中增大了批大小，而我們減小了批大小。

我們能夠從頭到尾充分利用硬體資源。Smith 等人的研究只在最後階段充分利用了硬體資源。增加批大小可以 warm-up 和穩定優化過程，但是減小批大小會給優化過程帶來混亂，導致訓練不收斂。

在實驗中，我們發現了一種有助於穩定第二階段優化的方法。由於我們切換到一個不同的優化問題，有必要重新 warm-up 優化過程。在第二階段，我們沒有降低學習率，而是將學習率從零開始增加 (re-warm-up)。

通過這些改變，我們只需要 8599 次迭代，可以在 76 分鐘左右完成 BERT 訓練，實現了 101.8% 的弱縮放效率 (weak scaling efficiency)，提速 65.2 倍，利用了 64 倍的計算資源。

結論

Large batch 技術是加快神經網路深度訓練的關鍵。在本文中，我們提出了支持adaptive element-wise updating 和 layer-wise correction 的 LAMB 優化器。LAMB是一個通用的優化器，適用於小批量和大批量。通過使用 LAMB，我們可以將 BERT 預訓練的 batch size 擴展到 64K，而不會丟失準確性。我們將 BERT 的訓練時間從 3 天減少到 76 分鐘左右，並將批大小推到了 TPU Pod 的硬體極限。我們正在研究 LAMB優化器的理論分析。

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