谷歌實習生新演算法提速驚人！BERT訓練從三天三夜，縮短到一個小時
郭一璞 安妮 發自 凹非寺
量子位 報道 | 公眾號 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技術和產品新動態
