[論文レビュー] Large Batch Optimization for Deep Learning: Training BERT in 76 minutes

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

Advanced Neural Network Applications参考文献 36被引用数 164

ひとこと要約

Lambという層ごとに適応する大規模バッチ最適化アルゴリズムを導入し、BERTの学習を最大32Kのバッチサイズで可能にし、TPUv3ポッド上でBERTの事前学習時間を3日から76分に短縮し、LambとLarsの収束保証を提供。

ABSTRACT

Training large deep neural networks on massive datasets is computationally very challenging. There has been recent surge in interest in using large batch stochastic optimization methods to tackle this issue. The most prominent algorithm in this line of research is LARS, which by employing layerwise adaptive learning rates trains ResNet on ImageNet in a few minutes. However, LARS performs poorly for attention models like BERT, indicating that its performance gains are not consistent across tasks. In this paper, we first study a principled layerwise adaptation strategy to accelerate training of deep neural networks using large mini-batches. Using this strategy, we develop a new layerwise adaptive large batch optimization technique called LAMB; we then provide convergence analysis of LAMB as well as LARS, showing convergence to a stationary point in general nonconvex settings. Our empirical results demonstrate the superior performance of LAMB across various tasks such as BERT and ResNet-50 training with very little hyperparameter tuning. In particular, for BERT training, our optimizer enables use of very large batch sizes of 32868 without any degradation of performance. By increasing the batch size to the memory limit of a TPUv3 Pod, BERT training time can be reduced from 3 days to just 76 minutes (Table 1). The LAMB implementation is available at https://github.com/tensorflow/addons/blob/master/tensorflow_addons/optimizers/lamb.py

研究の動機と目的

• 大規模ミニバッチでの高性能な深層ネットワークのトレーニングの非効率性を動機づけ、対処する。
• 大規模バッチSGDに適した層ごとに適応する最適化手法（Lamb）を開発する。
• 非凸設定におけるLarsとLambの理論的収束解析を提供する。
• BertとResNet-50の大規模バッチ領域でのLambの優れた性能を経験的に示す。

提案手法

• 更新を層ごとに単位L2ノルムへ正規化し、層ノルムの関数で学習率をスケールして大規模バッチ学習を安定化させる一般戦略を提案する。
• この戦略の下で2つのアルゴリズムを具体化する：Lars（モメンタムベース）とLamb（層ごとの正規化を備えたAdamベース）。
• LarsとLambの両方の非凸目的に対して定常点収束を示す収束解析を提供する。
• Bert（SQuAD, 2.5B+0.8B語コーパス）およびImageNet/ResNet-50を、非常に大きなバッチサイズと最小限のハイパーパラメータ調整で評価する。
• 平方根学習率スケーリングと線形ウォームアップを使用し、混合シーケンス長を持つマルチステージBert訓練のリワームアップを適用する。

実験結果

リサーチクエスチョン

• RQ1 層ごとの更新正規化と適応学習率スケーリングは、非常に大きなミニバッチで安定かつ効率的な訓練を可能にするか？
• RQ2 LambとLarsはSGDと同様の非凸設定で理論的収束保証を提供し得るか、BertとResNet-50などのタスクで実証的にどのように比較されるか？
• RQ3 BertとResNet-50はどのバッチサイズまで精度を失うことなく処理できるのか、TPUv3ポッドでの実測のウォールクロック速度アップはどの程度か？
• RQ4 Lambは多様なアーキテクチャとデータ規模（言語モデル対視覚モデル）を横断してLarsよりも頑健か？

主な発見

ソルバー	バッチサイズ	ステップ	F1スコア（dev）	TPU	時間
Baseline	512	1000k	90.395	16	81.4h
Lamb	512	1000k	91.752	16	82.8h
Lamb	1k	500k	91.761	32	43.2h
Lamb	2k	250k	91.946	64	21.4h
Lamb	4k	125k	91.137	128	693.6m
Lamb	8k	62500	91.263	256	390.5m
Lamb	16k	31250	91.345	512	200.0m
Lamb	32k	15625	91.475	1024	101.2m
Lamb	64k/32k	8599	90.584	1024	76.19m

• Lambは Bert のバッチサイズを32Kまで安定した訓練を可能にし、SQuADでのF1をベースラインと比較して維持または向上させる。
• 混合バッチ訓練とリワームアップを用いたLambにより、Bert訓練時間は約3日からTPUv3ポッド上で76分へ低下（64xのリソーススケーリングを伴う）。
• Lambは ImageNet の ResNet-50 で競争力のあるまたは優れた top-1 精度を提供し、大規模バッチサイズでAdam/AdamWを上回り、モーメンタム/Larsベースラインに匹敵または超える。
• Lars は tested バッチサイズで Bert に対して一般に Lamb より劣るが、Lamb はバッチサイズの増加に伴い高い F1 を維持する（表1と表2）。
• 理論的結果は、非凸設定において Lars と Lamb の両方が定常点へ収束することを示し、Lamb は平均層 Lipschitz 定数（L_avg）に対する依存が L_infty より有利であることを示す。