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分鐘！

論文地址：

https://arxiv.org/pdf/1904.00962.pdf

其中一作尤洋(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優化器的理論分析。