來源：arXiv
新智元編輯部
【新智元導讀】愛也BERT，恨也BERT！BERT是史上最強的NLP模型之一，但卻也是工業界目前最耗時的應用，計算量遠高于ImageNet！谷歌的研究人員提出新的優化器，將BERT的訓練時間從3天成功縮短到76分鐘！該論文已被ICLR 2020接收。戳右邊鏈接上 新智元We站公開課 了解更多！
谷歌發布的史上最強NLP模型之一BERT，是NLP領域的一項重大突破。
隨著大規模數據集的出現，在海量數據集上訓練大型深度神經網絡，甚至使用隨機梯度下降（Stochastic Gradient Descent，SGD）等計算效率高的優化方法，都已變得尤為具有挑戰性。例如，BERT和ResNet-50等最先進的深度學習模型在16個TPUv3芯片上訓練需要3天，在8臺Tesla P100 GPU上訓練需要29小時。
人們對于使用大批量隨機優化方法來解決此問題的興趣激增。該研究領域中最杰出的算法是LARS，它通過采用逐層自適應學習率，在幾分鐘內在ImageNet上訓練RESNET。然而，LARS對于像BERT這樣的注意力模型表現不佳，這表明其性能增益在各個任務之間并不一致。
這該如何是好？且來看看近日 Google、UC Berkeley、UCLA研究團隊的被ICLR 2020接收的力作—— 《Large Batch Optimization for Deep Learning: Training BERT in 76 minutes》。他們 首先研究了一種原則性的逐層自適應策略，來加速large mini-batches的深度神經網絡的訓練。利用該策略， 研究團隊開發了一種新的逐層自適應大批量優化技術LAMB，并給出了LAMB和LARS的收斂分析，表明在一般非凸情形下收斂到駐點（Stationary Point）。實驗結果表明，LAMB在BERT和RESNET-50訓練等任務中具有很好的性能，且超參數調整很少。特別是對于BERT訓練，提出的優化器支持使用非常大（32868）的批處理，而不會降低性能。 通過將批處理大小增加到TPUv3 Pod的內存限制， BERT訓練時間可以從3天減少到76分鐘！
https://arxiv.org/pdf/1904.00962.pdf
主要貢獻
受LARS的啟發，研究了一種特別適合大批量學習的通用適應策略，并為該策略提供了直覺。
基于自適應策略，提出了一種新的優化算法（LAMB）來實現SGD中學習速率的自適應。此外，還提供了LARS和LAMB的收斂分析，以獲得非凸情況下的一個駐點。我們將重點介紹在大型批處理設置中使用這些方法的好處。
展示了LAMB在多個挑戰性任務中的強大實證性能。使用LAMB，將訓練BERT的批量大小擴展到32k以上，而不會降低性能；時間從3天減少到76分鐘。 這是目前將BERT訓練經過時間減少到幾個小時以內的第一項工作。
還展示了LAMB在訓練RESNET這樣最先進的圖像分類模型方面的效率。 本項研究提出的自適應解算器是第一個能夠為RESNET-50實現最先進準確性的自適應解算器。
LAMB算法
LAMB的全稱是Layer-wise Adaptive Moments optimizer for Batching training。
BERT 訓練的基線使用權重衰減的 Adam 作為優化器，這是 Adam 優化器的一個變體。另一個自適應優化器LARS，早前被提出用于ImageNet上RESNET的大批量學習。
為了得到進一步的改進，研究團隊對LAMB使用混合批處理訓練。BERT訓練包括兩個階段：總epochs的前9/10使用128的序列長度，而總epochs的后1/10使用512的序列長度。由于內存限制，第二階段的訓練需要更長的序列長度，因此在TPUv3 Pod上最多只能使用32768個批處理大小。
LAMB 算法的概述如下所示：
實驗和結果
本研究在兩個重要的大批量訓練任務：BERT和 RESNET-50訓練上，比較LAMB與現有優化器的實證結果。還比較了在小批量（<1k）和小數據集（如CiFAR、MNIST）上，LAMB與現有優化器的結果。
BERT訓練
首先是加速BERT訓練的實證結果。本實驗使用與Devlin et al.相同的數據集，由Wikipedia和BooksCorpus拼接而成，分別為2.5B和8億詞。在本文中，特別關注的是SQuAD task。在實驗中，以SQuAD-v1上的F1分數作為精度指標，所有的比較都是基于Devlin等人的基線BERT模型。
為了訓練BERT, Devlin等人首先使用序列長度為128的900k迭代訓練模型，然后在最后的100k迭代中轉換為512的序列長度。這導致了在16個TPUv3芯片上大約需要3天的訓練時間。
文本首先的實驗，除了將訓練優化器更改為LAMB之外，保持與基線相同的訓練過程，使用與基線相同數量的epochs運行，但批量大小從512擴展到32K（選擇32K大小（序列長度512）主要是由于TPU Pod的內存限制）。通過使用LAMB優化器，能夠在批量大小為32768的15625次迭代（序列長度為128的14063次迭代和序列長度為512的1562次迭代）中獲得91.460的F1分數。對于批量大小為32K，本文將BERT訓練時間從3天縮短到約100分鐘。
本文在表中報告F1分數為91.345，這是未調諧版本的分數。
表 1：使用SQuAD-v1的F1 score作為精度指標。F1的基線成績是由BERT的公共知識庫提供的預訓練模型(BERT-large)實現的 (截止到2019年2月1日)。本文在實驗中使用TPUv3s。本文使用了與基線相同的設置：總epochs的前9/10使用序列長度128，最后 1/10使用序列長度512。所有的實驗運行相同數量的epochs。Dev set表示測試數據。值得注意的是，本文可以通過手動調整超參數來獲得更好的結果。
研究結論
大批量處理技術是加快深度神經網絡訓練的關鍵。本文提出了LAMB優化器，它支持自適應的元素更新（Adaptive Elementwise Updating）和逐層學習率（Layerwise Learning Rates）。此外，LAMB是一個通用優化器，適用于小批量和大批量。
本文還為LAMB優化器提供了理論分析，重點介紹了其性能優于標準SGD的情況。對于廣泛的應用程序，LAMB的性能要優于現有的優化器。通過使用LAMB，我們可以將BERT預訓練的批處理規模擴展到64K而不失準確性，從而將BERT訓練時間從3天縮短到76分鐘左右。LAMB也是第一個能夠在RESNET-50的ImageNet訓練中獲得最先進精度的大批量自適應解算器。
