[論文レビュー] Accelerated Large Batch Optimization of BERT Pretraining in 54 minutes

Shuai Zheng, Haibin Lin+2 more|arXiv (Cornell University)|Jun 24, 2020

Topic Modeling参考文献 33被引用数 10

ひとこと要約

本稿では、微分ごとの勾配正規化と、ウォームアップ後に一定学習率を維持する新しい学習率スケジューラーを組み合わせることで、安定した大ミニバッチ最適化を実現する LANS を提案する。96K および 33K のミニバッチサイズで安定した学習を可能にし、192 台の AWS EC2 P3dn.24xlarge インスタンスを用いて BERT プレトレーニングを 54 分で完了させ、精度を損なわずにクラウド環境で新たな速度記録を樹立した。

ABSTRACT

BERT has recently attracted a lot of attention in natural language understanding (NLU) and achieved state-of-the-art results in various NLU tasks. However, its success requires large deep neural networks and huge amount of data, which result in long training time and impede development progress. Using stochastic gradient methods with large mini-batch has been advocated as an efficient tool to reduce the training time. Along this line of research, LAMB is a prominent example that reduces the training time of BERT from 3 days to 76 minutes on a TPUv3 Pod. In this paper, we propose an accelerated gradient method called LANS to improve the efficiency of using large mini-batches for training. As the learning rate is theoretically upper bounded by the inverse of the Lipschitz constant of the function, one cannot always reduce the number of optimization iterations by selecting a larger learning rate. In order to use larger mini-batch size without accuracy loss, we develop a new learning rate scheduler that overcomes the difficulty of using large learning rate. Using the proposed LANS method and the learning rate scheme, we scaled up the mini-batch sizes to 96K and 33K in phases 1 and 2 of BERT pretraining, respectively. It takes 54 minutes on 192 AWS EC2 P3dn.24xlarge instances to achieve a target F1 score of 90.5 or higher on SQuAD v1.1, achieving the fastest BERT training time in the cloud.

研究の動機と目的

• 従来数日を要する BERT プレトレーニングのトレーニング時間を短縮するため、効率的な大ミニバッチ最適化を可能にする。
• LAMB などの既存の大ミニバッチ最適化手法で 64K を超えるミニバッチサイズにスケーリングすると生じる不安定性および精度の低下を克服する。
• 発散を防ぎつつ、長期間にわたり高い学習率を維持できる学習率スケジューラーを設計し、収束を高速化する。
• モデル性能を損なわずに、クラウド環境における BERT のトレーニング速度を最先端の水準にまで向上させる。

提案手法

• LANS は、LAMB のブロックワイズな適応的学習率を拡張するため、ブロック単位の勾配正規化を導入し、異なる層間での安定した更新を保証する。
• 古典的モーメンタムとネステロフのモーメンタムの関係を活用して、モーメンタム項を改変し、収束性を向上させる。
• ウォームアップ後に一定学習率を維持する新しい学習率スケジューリングを採用し、初期ウォームアップフェーズを越えて高い学習率を長期間使用可能にする。
• 学習率は平方根則に従ってスケーリングされるが、大規模バッチサイズでの発散を回避するように調整されており、実験的チューニングにより安定性を維持している。
• データシャーディングを用いた分散トレーニングをサポートし、勾配の分散を最小限に抑え、重複のないサンプリングを保証する。
• クラウドインfraに最適化された実装であり、192 台の AWS EC2 P3dn.24xlarge インスタンスを EFA を使用して高効率な通信を実現している。

実験結果

リサーチクエスチョン

• RQ1 大ミニバッチサイズを 96K および 33K にまで拡大した場合、修正された大ミニバッチ最適化手法がモデル精度を維持できるか？
• RQ2 ウォームアップ後に一定学習率フェーズを導入することで、線形ウォームアップスケジュールと比較して、トレーニングの安定性と収束速度が向上するか？
• RQ3 LANS は、精度の低下を伴わず、クラウドベースの GPU クラスタ上で LAMB を上回るトレーニング速度を達成できるか？
• RQ4 高速収束を実現するための最適なハイパーパramータ設定（学習率、ウォームアップ、一定フェーズ比）は何か？

主な発見

• 1536 個の V100 GPU を用いて、SQuAD v1.1 デベロップメントセットで F1 スコア 90.60 を 53.6 分で達成し、クラウド上での BERT プレトレーニングの最速記録を樹立した。
• 本手法は、BERT プレトレーニングのフェーズ 1 で 96K、フェーズ 2 で 33K のミニバッチサイズにスケーリングしても、性能の低下を認めなかった。
• LAMB は 96K/33K のバッチサイズで収束しなかったが、LANS ははるかに少ないイテレーション数（4301 vs. 8599）で同じ F1 スコアを達成した。
• 提案された学習率スケジューリングにより、学習率 0.007 と 0.01 の間の性能差が AUC ポints で 5.28 から 1.91 に縮小され、より高いロバストネスを示した。
• ウォームアップ後に一定学習率フェーズを導入することで、高い学習率を長期間維持でき、収束が著しく高速化された。
• 192 台の AWS EC2 P3dn.24xlarge インスタンスを用いてトレーニング時間を 54 分に短縮し、TPU を使用した従来手法の 76.2 分を上回った。