谷歌實(shí)習(xí)生新算法提速驚人！BERT訓(xùn)練從三天三夜，縮短到一個(gè)小時(shí)
安妮極致利用1024個(gè)TPU
郭一璞 安妮 發(fā)自 凹非寺
量子位 報(bào)道 | 公眾號 QbitAI
從頭開始訓(xùn)練一遍當(dāng)今最強(qiáng)的語言AI需要多久？
現(xiàn)在，這個(gè)時(shí)間從三天三夜，一下縮短到一個(gè)多小時(shí)！
帶來這個(gè)進(jìn)步的，是Google Brain實(shí)習(xí)生，尤洋。這位小哥來自中國河南，曾經(jīng)的清華計(jì)算機(jī)碩士第一，目前在加州大學(xué)伯克利分校讀博。
他最近完成的一項(xiàng)研究，把預(yù)訓(xùn)練BERT的速度足足提高了64倍。訓(xùn)練時(shí)間從4860分鐘，變成了76分鐘11秒。
而且訓(xùn)練完成后，在機(jī)器問答數(shù)據(jù)集SQuAD-v1上測試一下，F(xiàn)1得分比原來的三天三夜版還要高一點(diǎn)點(diǎn)。
別人家的實(shí)習(xí)生，究竟又使出了什么神技？
費(fèi)時(shí)到省時(shí)
要縮短神經(jīng)網(wǎng)絡(luò)的訓(xùn)練時(shí)間，本來有成熟的方法可用：
雙管齊下，一是堆上大量CPU、GPU或者TPU增加算力，二是增加批大小減少迭代次數(shù)。
在計(jì)算機(jī)視覺研究中，這種方法就很常用。前兩天，富士通的研究人員還用這種方法，74.7秒在ImageNet上訓(xùn)練完ResNet-50。
但是，視覺領(lǐng)域的這些訓(xùn)練方法，搬到BERT上是行不通的。BERT是目前工業(yè)界訓(xùn)練起來最耗時(shí)的應(yīng)用，計(jì)算量遠(yuǎn)高于ImageNet。
此外，大批量訓(xùn)練還有個(gè)“通病”，就是會產(chǎn)生泛化誤差（Generalization Gap），導(dǎo)致網(wǎng)絡(luò)泛化能力下降，如此直接優(yōu)化往往會導(dǎo)致測試集上準(zhǔn)確度下降。
怎么辦？
為了能用大批量訓(xùn)練BERT，尤洋和他的同事們提出了LAMB優(yōu)化器。這是一個(gè)通用的神經(jīng)網(wǎng)絡(luò)優(yōu)化器，無論是大批量還是小批量的網(wǎng)絡(luò)都可以使用，也無需在學(xué)習(xí)率之外調(diào)試超參數(shù)。
靠超大批量也適用的LAMB，他們將批大小由512擴(kuò)展到了65536。
65536是什么概念呢？這已經(jīng)達(dá)到了TPU內(nèi)存的極限，也是第一次有研究用2000以上的超大批量來訓(xùn)練BERT。
于是，迭代次數(shù)大大降低。此前，BERT-Large模型需要1000000次迭代才能完成預(yù)訓(xùn)練過程，耗時(shí)81.4小時(shí)。有了LAMB加持用上大批量，只需要進(jìn)行8599次迭代，預(yù)訓(xùn)練時(shí)間直接縮短到76分鐘。
這樣看來，加速64倍！
因吹斯聽，這LAMB到底是何方神器？
LAMB優(yōu)化器
它的全稱是Layer-wise Adaptive Moments optimizer for Batch training，和大家熟悉的SGD、Adam屬于同類，都是機(jī)器學(xué)習(xí)模型的優(yōu)化器（optimizer）。
原本，三天三夜的BERT訓(xùn)練，用的是權(quán)重衰減的Adam優(yōu)化器。
而這一次的新型優(yōu)化器LAMB，是在論文一作尤洋2017年的一項(xiàng)研究啟發(fā)下產(chǎn)生的。當(dāng)時(shí)，他提出了一種用于大批量卷積神經(jīng)網(wǎng)絡(luò)的優(yōu)化器LARS。
LARS使用系數(shù)eeta控制信任率（trust ratio），但是這種做法可能會導(dǎo)致一些問題，造成一些差異。
因此，在LAMB里，研究團(tuán)隊(duì)刪除了eeta，對于0|w|或者0|g|的層，直接把信任率設(shè)置成1.0，消除了BERT訓(xùn)練的差異。
另外，LARS里用權(quán)重衰減來計(jì)算信任率：
考慮到權(quán)重衰減，LAMB里的信任率公式改成了這樣：
另外，LARS雖然在ImageNet上運(yùn)行的很好，但是在ImageNet上用的模型參數(shù)比BERT少得多，因此，這一次，研究團(tuán)隊(duì)把LARS中的
改成了：
通過一系列改動，LAMB優(yōu)化器的最大批量被提升到了32K。
真·優(yōu)化
是騾子是馬，也該拉出來溜溜了。
研究人員用常規(guī)訓(xùn)練與混合批訓(xùn)練兩種方法，測試LAMB優(yōu)化器的優(yōu)化成果，實(shí)測效果不錯(cuò)。
在測試中，他們大大增加了算力，選擇1024核的TPUv3 Pod進(jìn)行訓(xùn)練，其中的1024個(gè)TPU核心，每秒可以提供超過10億次浮點(diǎn)（100 petaflops）的混合精度運(yùn)算。
最后研究人員敲定，和原版BERT模型一樣，用Wikipedia和ooksCorpus數(shù)據(jù)集預(yù)訓(xùn)練，然后利用斯坦福的SQuAD-v1數(shù)據(jù)集進(jìn)行測試，測試得出的F1 Score的值則用來衡量準(zhǔn)確度。
結(jié)果顯示，隨著批量的增大，迭代次數(shù)逐減少，F(xiàn)1 Score的波動并不顯著，F(xiàn)1值維持在90以上，但訓(xùn)練時(shí)間明顯縮短。
△?測試結(jié)果
當(dāng)批大小大于65536、序列長度達(dá)到128時(shí)，訓(xùn)練時(shí)間就沒有明顯的減少了。
當(dāng)用了1024塊TPU，批大小為32768或65536、迭代8599次時(shí)，訓(xùn)練時(shí)間縮減到最小，預(yù)訓(xùn)練過程只需要76.19分鐘。
并且，最終達(dá)到了101.8%的弱縮放效率（weak scaling efficiency）。
學(xué)霸一作：本碩都是第一名
這項(xiàng)研究的作者是UC伯克利計(jì)算機(jī)科學(xué)部的在讀博士尤洋，同時(shí)他也是Google Brain的實(shí)習(xí)生。
小哥哥是一位大學(xué)霸，他本科就讀于中國農(nóng)業(yè)大學(xué)計(jì)算機(jī)專業(yè)，是班里的第一名，碩士保送了清華計(jì)算機(jī)系，在134名入讀清華的佼佼者中，他依然是第一名。
作為第一名的學(xué)霸，在申請博士的時(shí)候，尤洋一下子喜提了UC伯克利、CMU、芝加哥大學(xué)、UIUC、佐治亞理工、西北大學(xué)六所名校的全獎(jiǎng)offer，簡直是名校任挑。
于是，他從六所名校里選擇了UC伯克利，UC伯克利剛好位于灣區(qū)，尤洋也因此有機(jī)會游走于Google Brain、英特爾實(shí)驗(yàn)室、微軟研究院、英偉達(dá)、IBM沃森研究中心等知名企業(yè)、研究院實(shí)習(xí)，趁實(shí)習(xí)的機(jī)會為TensorFlow、英偉達(dá)GPU上部署caffe、英特爾CPU部署caffe等大型知名開源項(xiàng)目做出了貢獻(xiàn)。
甚至，還有機(jī)會趁實(shí)習(xí)去皮衣哥黃仁勛家里開爬梯~真是讓人羨慕啊。
△?今天沒穿皮衣
另外，尤洋也是一位論文高產(chǎn)選手，光是一作的頂會論文就有十幾篇，其中還包含去年ICPP的最佳論文，還會有IPDPS 2015的最佳論文。而且，他還拿到了2014年的西貝爾學(xué)者獎(jiǎng)（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
—?完?—
誠摯招聘
量子位正在招募編輯/記者，工作地點(diǎn)在北京中關(guān)村。期待有才氣、有熱情的同學(xué)加入我們！相關(guān)細(xì)節(jié)，請?jiān)诹孔游还娞?QbitAI)對話界面，回復(fù)“招聘”兩個(gè)字。
量子位?QbitAI · 頭條號簽約作者
?’?’ ? 追蹤AI技術(shù)和產(chǎn)品新動態(tài)
相關(guān)閱讀
