[論文レビュー] Reducing BERT Pre-Training Time from 3 Days to 76 Minutes

Yang You, Jing Li+4 more|arXiv (Cornell University)|Apr 1, 2019

Advanced Neural Network Applications被引用数 73

ひとこと要約

この論文では、TPUv3ポッド上で最大32,868の非常に大きなミニバッチサイズを可能にすることで、BERTの事前学習を3日間からわずか76分に短縮する、新しい大規模ミニバッチ最適化手法LAMB（Layer-wise Adaptive Moments）を紹介する。LAMBは、層ごとの適応的学習率と形式的な収束保証を用いることで、BERTおよびResNet-50の両方で先行手法を上回る性能を達成する。

ABSTRACT

Training large deep neural networks on massive datasets is very challenging. One promising approach to tackle this issue is through the use of large batch stochastic optimization. However, our understanding of this approach in the context of deep learning is still very limited. Furthermore, the current approaches in this direction are heavily hand-tuned. To this end, we first study a general adaptation strategy to accelerate training of deep neural networks using large minibatches. Using this strategy, we develop a new layer-wise adaptive large batch optimization technique called LAMB. We also provide a formal convergence analysis of LAMB as well as the previous published layerwise optimizer LARS, showing convergence to a stationary point in general nonconvex settings. Our empirical results demonstrate the superior performance of LAMB for BERT and ResNet-50 training. In particular, for BERT training, our optimization technique enables use of very large batches sizes of 32868; thereby, requiring just 8599 iterations to train (as opposed to 1 million iterations in the original paper). By increasing the batch size to the memory limit of a TPUv3 pod, BERT training time can be reduced from 3 days to 76 minutes. Finally, we also demonstrate that LAMB outperforms previous large-batch training algorithms for ResNet-50 on ImageNet; obtaining state-of-the-art performance in just a few minutes.

研究の動機と目的

• BERTのような大規模なディープニューラルネットワークの学習時間が長いため、膨大な計算リソースと長時間の学習を要するという課題に対処すること。
• ディープラーニングにおける大規模ミニバッチ確率的最適化の効率を向上させること。現状、この分野は理論的裏付けが乏しく、実務では多くの手動チューニングに依存している。
• 非常に大きなミニバッチで安定かつ高速に学習できる、一般化可能で適応的な最適化戦略を開発すること。
• LAMBおよび先行のLARS最適化手法について、一般非凸設定下での形式的収束解析を提供し、理論的堅牢性を保証すること。
• ImageNet や GLUE の標準ベンチマークにおいて、顕著に短縮された学習時間と少ない反復回数で、BERT および ResNet-50 で最先端の性能を達成すること。

提案手法

• 各層ごとに勾配ノルムとパラメータノルムの比に基づいて学習率をスケーリングする、層ごとの適応的学習率を導入した大規模ミニバッチ最適化手法LAMBを提案する。
• Adam最適化手法を拡張し、層ごとの適応的学習率を導入することで、大規模ミニバッチ学習下での安定性と収束性を向上させる。
• 勾配とパラメータの大きさを各層でバランスさせる正規化機構を導入し、学習の安定性を高める。
• LAMBおよびLARSの両者について、一般非凸設定下で局停留点への収束を示す形式的収束解析を実施する。
• TPUv3ポッドのメモリ容量を活用して、ミニバッチサイズを最大32,868までスケールアップし、学習反復回数を著しく削減する。
• 極めて大きなミニバッチサイズでもモデルの安定性を維持できる学習率スケジュールを採用し、標準的大規模ミニバッチ手法で一般的に見られる発散を回避する。

実験結果

リサーチクエスチョン

• RQ1 大規模ミニバッチ学習のための一般化可能な適応戦略は、モデル品質を損なわせることなく、BERTの事前学習時間を著しく短縮できるか？
• RQ2 層ごとの適応的学習率は、深層ネットワークの大規模ミニバッチ設定下での最適化の安定性と収束性をどのように向上させるか？
• RQ3 LAMBは、LARSなどの先行手法と同様に、非凸最適化設定下で形式的な収束保証を達成できるか？
• RQ4 LAMBは、TPUv3ポッドのメモリ制限に近いミニバッチサイズでBERTの学習を可能にするか？
• RQ5 ImageNet や GLUE の標準ベンチマークにおいて、LAMBは既存の大規模ミニバッチ最適化手法と比較して、精度と学習速度の両面で優れているか？

主な発見

• LAMBは、TPUv3ポッド上でミニバッチサイズ32,868を可能にすることで、BERTの事前学習を3日間からわずか76分に短縮する。
• LAMBを用いたBERTの学習では、1回の反復が8,599回にまで減少し、元のBERT論文の100万回に比べて著しく削減される。
• LAMBは数分間でResNet-50/ImageNetで最先端の性能を達成し、以前の大規模ミニバッチ手法を上回る。
• 標準的大規模ミニバッチ最適化で一般的に見られる発散問題を回避しながら、極めて大きなミニバッチサイズでも安定した学習が可能である。
• 形式的収束解析により、LAMBが一般非凸設定下で局停留点への収束を示すことが確認され、理論的根拠が得られた。
• LAMBは、BERTおよびResNet-50の両方で、先行の大規模ミニバッチ学習アルゴリズムを上回り、速度と精度の両面で一貫した向上を示した。