1024塊TPU在燃燒！將BERT預(yù)訓(xùn)練模型的訓(xùn)練時長從3天縮減到了76分鐘

“Jeff Dean稱贊，TensorFlow官方推特支持，BERT目前工業(yè)界最耗時的應(yīng)用，計算量遠高于ImageNet。我們將BERT的訓(xùn)練時間從三天縮短到了一小時多?！盪C Berkeley大學(xué)在讀博士尤洋如是說道。

近日，來自Google、UC Berkeley、UCLA研究團隊再度合作，成功燃燒1024塊TPU，將BERT預(yù)訓(xùn)練模型的訓(xùn)練時長從3天縮減到了76分鐘。batch size技術(shù)是加速神經(jīng)網(wǎng)絡(luò)訓(xùn)練的關(guān)鍵，在“Reducing BERT Pre-Training Time from 3 Days to 76 Minutes”這篇論文中，作者提出了LAMB優(yōu)化器，它支持自適應(yīng)元素更新和分層校正。

論文傳送門：https://arxiv.org/pdf/1904.00962.pdf

論文摘要：batch size增加到很大時的模型訓(xùn)練是加速大型分布式系統(tǒng)中深度神經(jīng)網(wǎng)絡(luò)訓(xùn)練的關(guān)鍵。但是，這種模型訓(xùn)練很難，因為它會導(dǎo)致一種泛化差距。直接優(yōu)化通常會導(dǎo)致測試集上的準(zhǔn)確性下降。

BERT是一種先進的深度學(xué)習(xí)模型，它建立在語義理解的深度雙向轉(zhuǎn)換器上。當(dāng)我們增加batch size的大小（如超過8192）時，此前的模型訓(xùn)練技巧在BERT上表現(xiàn)得并不好。BERT預(yù)訓(xùn)練也需要很長時間才能完成，如在16個TPUv3上大約需要三天。

為了解決這個問題，我們提出了LAMB優(yōu)化器，可將batch size擴展到65536，且不會降低準(zhǔn)確率。LAMB是一個通用優(yōu)化器，batch size大小均使用，且除了學(xué)習(xí)率之外不需要別的參數(shù)調(diào)整。

基線BERT-Large模型需要100萬次迭代才能完成預(yù)訓(xùn)練，而batch size大小為65536/32768的LAMB僅需要8599次迭代。我們還將batch size進行內(nèi)存限制，接近TPUv3 pod，結(jié)果可在76分鐘內(nèi)完成BERT訓(xùn)練。

據(jù)悉，該論文的一作是來自UC Berkeley計算機科學(xué)部的在讀博士尤洋，同時也是Google Brain的實習(xí)生。據(jù)公開信息顯示，尤洋的導(dǎo)師是美國科學(xué)院與工程院院士，ACM/IEEE fellow，伯克利計算機系主任，以及首批中關(guān)村海外顧問James Demmel教授。他當(dāng)前的研究重點是大規(guī)模深度學(xué)習(xí)訓(xùn)練算法的分布式優(yōu)化。2017年9月，尤洋等人的新算法以24分鐘完成ImageNet訓(xùn)練，刷新世界紀錄。

在此之前，他曾在英特爾實驗室、微軟研究院、英偉達、IBM沃森研究中心等機構(gòu)實習(xí)。尤洋本科就讀于中國農(nóng)業(yè)大學(xué)計算機系，碩士保送清華大學(xué)計算機系，是一名杠杠的理工學(xué)霸！