谷歌實習生新算法提速驚人！BERT訓練從三天三夜，縮短到一個小時

郭一璞 安妮 2019-04-03 15:36:33 來源：量子位

郭一璞 安妮 發自 凹非寺

量子位 報道 | 公眾號 QbitAI

從頭開始訓練一遍當今最強的語言AI需要多久？

現在，這個時間從三天三夜，一下縮短到一個多小時！

帶來這個進步的，是Google Brain實習生，尤洋。這位小哥來自中國河南，曾經的清華計算機碩士第一，目前在加州大學伯克利分校讀博。

他最近完成的一項研究，把預訓練BERT的速度足足提高了64倍。訓練時間從4860分鐘，變成了76分鐘11秒。

而且訓練完成后，在機器問答數據集SQuAD-v1上測試一下，F1得分比原來的三天三夜版還要高一點點。

別人家的實習生，究竟又使出了什么神技？

費時到省時

要縮短神經網絡的訓練時間，本來有成熟的方法可用：

雙管齊下，一是堆上大量CPU、GPU或者TPU增加算力，二是增加批大小減少迭代次數。

在計算機視覺研究中，這種方法就很常用。前兩天，富士通的研究人員還用這種方法，74.7秒在ImageNet上訓練完ResNet-50。

但是，視覺領域的這些訓練方法，搬到BERT上是行不通的。BERT是目前工業界訓練起來最耗時的應用，計算量遠高于ImageNet。

此外，大批量訓練還有個“通病”，就是會產生泛化誤差（Generalization Gap），導致網絡泛化能力下降，如此直接優化往往會導致測試集上準確度下降。

怎么辦？

為了能用大批量訓練BERT，尤洋和他的同事們提出了LAMB優化器。這是一個通用的神經網絡優化器，無論是大批量還是小批量的網絡都可以使用，也無需在學習率之外調試超參數。

靠超大批量也適用的LAMB，他們將批大小由512擴展到了65536。

65536是什么概念呢？這已經達到了TPU內存的極限，也是第一次有研究用2000以上的超大批量來訓練BERT。

于是，迭代次數大大降低。此前，BERT-Large模型需要1000000次迭代才能完成預訓練過程，耗時81.4小時。有了LAMB加持用上大批量，只需要進行8599次迭代，預訓練時間直接縮短到76分鐘。

這樣看來，加速64倍！

因吹斯聽，這LAMB到底是何方神器？

LAMB優化器

它的全稱是Layer-wise Adaptive Moments optimizer for Batch training，和大家熟悉的SGD、Adam屬于同類，都是機器學習模型的優化器（optimizer）。

原本，三天三夜的BERT訓練，用的是權重衰減的Adam優化器。

而這一次的新型優化器LAMB，是在論文一作尤洋2017年的一項研究啟發下產生的。當時，他提出了一種用于大批量卷積神經網絡的優化器LARS。

LARS使用系數eeta控制信任率（trust ratio），但是這種做法可能會導致一些問題，造成一些差異。

因此，在LAMB里，研究團隊刪除了eeta，對于0|w|或者0|g|的層，直接把信任率設置成1.0，消除了BERT訓練的差異。

另外，LARS里用權重衰減來計算信任率：

考慮到權重衰減，LAMB里的信任率公式改成了這樣：

另外，LARS雖然在ImageNet上運行的很好，但是在ImageNet上用的模型參數比BERT少得多，因此，這一次，研究團隊把LARS中的

改成了：

通過一系列改動，LAMB優化器的最大批量被提升到了32K。

真·優化

是騾子是馬，也該拉出來溜溜了。

研究人員用常規訓練與混合批訓練兩種方法，測試LAMB優化器的優化成果，實測效果不錯。

在測試中，他們大大增加了算力，選擇1024核的TPUv3 Pod進行訓練，其中的1024個TPU核心，每秒可以提供超過10億次浮點（100 petaflops）的混合精度運算。

最后研究人員敲定，和原版BERT模型一樣，用Wikipedia和ooksCorpus數據集預訓練，然后利用斯坦福的SQuAD-v1數據集進行測試，測試得出的F1 Score的值則用來衡量準確度。

結果顯示，隨著批量的增大，迭代次數逐減少，F1 Score的波動并不顯著，F1值維持在90以上，但訓練時間明顯縮短。

△?測試結果

當批大小大于65536、序列長度達到128時，訓練時間就沒有明顯的減少了。

當用了1024塊TPU，批大小為32768或65536、迭代8599次時，訓練時間縮減到最小，預訓練過程只需要76.19分鐘。

并且，最終達到了101.8%的弱縮放效率（weak scaling efficiency）。

學霸一作：本碩都是第一名

這項研究的作者是UC伯克利計算機科學部的在讀博士尤洋，同時他也是Google Brain的實習生。

小哥哥是一位大學霸，他本科就讀于中國農業大學計算機專業，是班里的第一名，碩士保送了清華計算機系，在134名入讀清華的佼佼者中，他依然是第一名。

作為第一名的學霸，在申請博士的時候，尤洋一下子喜提了UC伯克利、CMU、芝加哥大學、UIUC、佐治亞理工、西北大學六所名校的全獎offer，簡直是名校任挑。

于是，他從六所名校里選擇了UC伯克利，UC伯克利剛好位于灣區，尤洋也因此有機會游走于Google Brain、英特爾實驗室、微軟研究院、英偉達、IBM沃森研究中心等知名企業、研究院實習，趁實習的機會為TensorFlow、英偉達GPU上部署caffe、英特爾CPU部署caffe等大型知名開源項目做出了貢獻。

甚至，還有機會趁實習去皮衣哥黃仁勛家里開爬梯~真是讓人羨慕啊。

△?今天沒穿皮衣

另外，尤洋也是一位論文高產選手，光是一作的頂會論文就有十幾篇，其中還包含去年ICPP的最佳論文，還會有IPDPS 2015的最佳論文。而且，他還拿到了2014年的西貝爾學者獎（Siebel Scholar）。

傳送門

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

Yang You, Jing Li, Jonathan Hseu, Xiaodan Song, James Demmel, Cho-Jui Hsieh

https://arxiv.org/abs/1904.00962

—?完?—

誠摯招聘

量子位正在招募編輯/記者，工作地點在北京中關村。期待有才氣、有熱情的同學加入我們！相關細節，請在量子位公眾號(QbitAI)對話界面，回復“招聘”兩個字。

量子位?QbitAI · 頭條號簽約作者

?’?’ ? 追蹤AI技術和產品新動態