大力再出奇蹟，1024 張TPU，65536 batch size，僅76分鐘訓練完BERT！

大資料文摘出品

作者：Andy

BERT 作為目前工業界中訓練最耗時的應用，計算量甚至遠大於機器視覺中的 ImageNet 訓練。在 BERT 原論文中，Jacob Devlin 也是用了 16 臺雲TPU（64 個 TPU 晶片）花了整整四天，才訓練完了 BERT-large 模型。

正因如此難以訓練，導致大家也都幾乎是直接拿 Google 開源出的模型，用 GPU 在上面 finetune 一下就來用。Github上的 BERT 開源庫也因為其在 GPU 機器上預訓練的困難和耗時，往往沒給出預訓練相關指令碼。

論文下載地址：

而在這篇愚人節當天於arxiv發出的 《Reducing BERT Pre-Training Time from 3 Days to 76 Minutes》論文中，正如標題所說，作者將BERT的訓練時間縮減到了僅僅 76 分鐘。

而所用的計算機器數也是驚人的，一個 TPUv3 pod （1024張TPU晶片）。不要以為這只是個靠堆積機器就能堆積上去的成果，這還涉及到具體的訓練最佳化問題解決，比如說如何在增大batch size來提高計算通訊比的同時，又能保證其收斂。這篇論文中就是提出了一套訓練方案，在新最佳化器 LAMB 的幫組下來解決這個問題。

這篇論文中就提出了一套訓練方案，還有一個新的最佳化器LAMB來解決這個問題。

這是一作加大伯克利分校博士尤洋作為 Google Brain 實習生時完成的專案，其現在的主攻方向便是“高速並行分佈深度學習演算法”，這也是與這次論文密切相關的主題。

看他之前的論文也能看到類似的大規模訓練專案，比如之前在英偉達實習時的技術報告《LARGE BATCH TRAINING OF CONVOLUTIONAL NETWORKS》，怎麼用大 batch size 訓練CNN網路。

他在裡面提出的 LARS，也正是這篇論文中提出最佳化器 LAMB 的最初原型。論文中也有拿 LARS 做了實驗。而說到在Google Brain實習也是有跡可循，之前他便已在這實習過一段時間研究關於TPU上使用Tensorflow的課題，也為此次專案打下鋪墊。

論文下載地址：

看他之前的論文也能看到類似的大規模訓練專案，比如他之前在英偉達實習時的技術報告《LARGE BATCH TRAINING OF CONVOLUTIONAL NETWORKS》，怎麼用大 batch size 訓練CNN網路。

他在裡面提出的 LARS，也正是這篇論文中提出最佳化器 LAMB 的最初原型。論文中也拿LARS做了實驗。而說到在Google Brain實習也是有跡可循，之前他便已在這實習過一段時間研究關於TPU上使用Tensorflow的課題。

下面就來介紹論文吧。

首先在導言部分，作者稍稍介紹了大批次訓練的困難，還有此次面對的訓練物件BERT。為處理大批次 BERT 的訓練，作者提出了 LAMB 最佳化器，透過這個最佳化器將BERT訓練的batch size推到了64k（具體65536）的量級，而同時不損失精度，此外LAMB最佳化器還有一個優點就是隻用調學習率。

最後的預訓練包括兩個階段：前九個epoch用128的序列長度和64k的batch size；最後一個epoch用512的序列長度和32k的batch size，只用了驚人的8599個迭代便訓練完了BERT。而它和基線模型還有其他batch size訓練對比，如下圖。

之後，在背景部分，作者給我們分享了關於大批次訓練的知識。目前大批次訓練常常會遇到的幾個問題，還有如何對其中一些進行處理。

• 大批次會導致測試精度丟失，所以需要調節超參，比如學習率。所以需要隨著batch size的增大，線性或平方根級地增大學習率；

• 但大學習率會導致訓練初始不穩定，因此需要使用學習率預熱技巧（learning rate warmup），先用一個小學習率然後慢慢增大，到一定點切換到正常的學習率策略；

• 還有大批次訓練裡的泛化間距（generalization gap）問題，大批次訓練模型會傾向於收斂到比較尖銳的區域性最小點，這會導致訓練很容易達到一個比較高的訓練精度，卻很難獲得比較好的測試精度，也就是存在比較大的間距（gap）。目前還沒找到合適方法解決。

然後再來詳細瞭解一下本文最重要的創新，LAMB （Layer-wise Adaptive Moments optimizer for Batch training）最佳化器吧。它是基於作者之前提出的 LARS （Layer-wise Adaptive Rate Scaling）最佳化器，同時又參考了BERT原有最佳化器進行了改進。

首先為了理解 LAMB 我們需要對 LARS 有個粗略的瞭解。LARS 提出的背景是，作者發現對於網路各層，其權重和梯度的L2-norm的比值|w|/|g|變化會非常大，比如5.76和1345。

這導致的一個問題就是，一個學習率並不適應於所有層訓練，它對於有些層或許適合，但對於有些層卻可能太大。因此作者提出應該要按照層，來獲得每層的學習率（Local LR），而這個每層學習率的計算則需要之後 LAMB 中多次提到的一個trust ratio，置信比，有多大的可能我們相信當前層會在這次更新中改變它的權重。

於是 LAMB 對 LARS 最大的改進有三點，都是與trust ratio相關。

1.透過在Tensorflow具體的LARS最佳化器實現中，移除其中一個當某層的|w|和|g|都非零時用於計算 trust ratio 的係數 eeta，從而避免了BERT大批次訓練中的發散；

2.LARS 自身的trust ratio在一些自適應最佳化器，比如BERT裡用到的ADAM with weight decay 或 ADAM中，會導致不準確的學習率矯正，因為它們用了元素級的更新策略。於是作者們將trust_ratio做了些修正，保證了它和自適應最佳化器的結合。

3.最後一點，在計算梯度的L2-norm的時候，還加入了梯度的一階和二階慣量的資訊。

LAMB 演算法具體如下：

最後具體實驗部分便不累述，值得一提的一個細節是，在混合batch訓練的時候，因為中間有一個階段轉換過程，將batch size從64k降到了32k，這裡會給最佳化過程帶來噪音，進而導致訓練發散。為解決這個問題，在第二個階段的時候作者又重新進行了一次學習率預熱（re-warm-up）。