大規模バッチ最適化による深層学習の訓練：BERTを76分で訓練する

大規模データセット上で大規模な深層ニューラルネットワークを訓練することは、計算上非常に困難です。この問題に対処するために、最近では大バッチ確率的最適化手法の使用に対する関心が高まっています。この研究分野で最も注目されているアルゴリズムはLARSであり、レイヤーごとの適応学習率を用いてImageNet上でResNetを数分で訓練することができます。しかし、BERTのようなアテンションモデルに対してLARSは性能が低く、その性能向上がタスク間で一貫していないことを示しています。本論文では、まず大バッチを使用して深層ニューラルネットワークの訓練を加速するための理論的なレイヤーごとの適応戦略を研究します。この戦略に基づいて、新しいレイヤーごとの適応的大バッチ最適化技術であるLAMBを開発しました。さらに、LAMBおよびLARSの収束解析を行い、一般的な非凸設定における定常点への収束を示しています。我々の実験結果は、BERTやResNet-50などの様々なタスクにおいて、わずかなハイパーパラメータ調整でLAMBが優れた性能を発揮することを示しています。特にBERTの訓練において、我々のオプティマイザはパフォーマンス低下なしに最大32868の大バッチサイズを使用可能とします。TPUv3 Podのメモリ限界までバッチサイズを増加させることにより、BERTの訓練時間は3日から76分に短縮されます（表1）。LAMBの実装は以下のURLから利用可能です: https://github.com/tensorflow/addons/blob/master/tensorflow_addons/optimizers/lamb.py