[論文レビュー] Scaling SGD Batch Size to 32K for ImageNet Training.

Yang You, Igor Gitman+1 more|arXiv (Cornell University)|Aug 13, 2017

Advanced Neural Network Applications参考文献 7被引用数 225

ひとこと要約

この論文では、重みと勾配のノルムに応じて各レイヤーごとに学習率を適応的に調整する、Layer-wise Adaptive Rate Scaling (LARS) と呼ばれる、学習率スケジューリング手法を提案している。これにより、安定した大バッチ学習が可能となり、ImageNet-1k において ResNet-50 では 32,768、AlexNet では 8,192 のバッチサイズにまで SGD のバッチサイズを拡大することに成功。精度の低下を伴わず、例えばバッチサイズ 4,096 ではバッチサイズ 512 と比較して 3 倍の高速化を達成した。

ABSTRACT

The most natural way to speed-up the training of large networks is to use data-parallelism on multiple GPUs. To scale Stochastic Gradient (SG) based methods to more processors, one need to increase the batch size to make full use of the computational power of each GPU. However, keeping the accuracy of network with increase of batch size is not trivial. Currently, the state-of-the art method is to increase Learning Rate (LR) proportional to the batch size, and use special learning rate with "warm-up" policy to overcome initial optimization difficulty. <p>By controlling the LR during the training process, one can efficiently use large-batch in ImageNet training. For example, Batch-1024 for AlexNet and Batch-8192 for ResNet-50 are successful applications. However, for ImageNet-1k training, state-of-the-art AlexNet only scales the batch size to 1024 and ResNet50 only scales it to 8192. The reason is that we can not scale the learning rate to a large value. To enable large-batch training to general networks or datasets, we propose Layer-wise Adaptive Rate Scaling (LARS). LARS LR uses different LRs for different layers based on the norm of the weights and the norm of the gradients. By using LARS algoirithm, we can scale the batch size to 32768 for ResNet50 and 8192 for AlexNet. Large batch can make full use of the system's computational power. For example, batch-4096 can achieve 3x speedup over batch-512 for ImageNet training by AlexNet model on a DGX-1 station (8 P100 GPUs).

研究の動機と目的

• ImageNet 学習において、バッチサイズを 8,192 を超えて拡大してもモデルの精度を維持する課題に対処すること。
• ある閾値を超えてバッチサイズを増加すると通常観察される不安定性や精度の低下を克服すること。
• 収束性とパフォーマンスを保持したまま、バッチサイズを 32,768 まで拡大することで、マルチ GPU システムの効率的利用を可能にすること。
• 重みと勾配のノルムに応じてレイヤー固有の学習率を適応的にスケーリングする学習率スケジューリング戦略を開発し、安定した大バッチ最適化を実現すること。

提案手法

• 各レイヤーの重みの L2 ノルムと勾配の L2 ノルムの比に基づくスケーリング要因を用いて、グローバルなベース学習率との積としてレイヤー固有の学習率を計算する Layer-wise Adaptive Rate Scaling (LARS) を提案。
• バッチサイズに比例してスケーリングされるグローバルなベース学習率を採用し、一般的な線形スケーリング則に従う。
• 特に大バッチの場合に初期エポックでの学習の不安定性を軽減するために、学習率のウォームアップ段階を導入。
• 大バッチ学習中の一般化性能の向上を図るために、重み減衰成分を導入。
• ハイパーパramータを学習率スケジューリングを除き、固定したまま、AlexNet や ResNet-50 といったモデル間で一貫したトレーニングプロトコルを採用。
• スケーラビリティとスループットの評価のため、DGX-1 システム上での 8 枚の P100 GPU を用いたデータ並列トレーニングを実施。

実験結果

リサーチクエスチョン

• RQ1 新規の学習率スケジューリング戦略を用いることで、ImageNet-1k における大バッチ SGD 学習（最大 32,768）を安定化・高精度化できるか？
• RQ2 大バッチサイズにおける収束性と精度を考慮した場合、標準的な線形スケーリングとウォームアップに比べ、レイヤーごとの適応的学習率スケーリング（LARS）はどのように優れているか？
• RQ3 バッチサイズを 512 から 4,096 まで増加させることで、モデルの精度を維持したままどの程度のトレーニングスピードアップが達成できるか？
• RQ4 LARS は、極めて大きなバッチサイズにおいて、より深い（ResNet-50）および浅い（AlexNet）ネットワークの両方で安定した学習を可能にするか？

主な発見

• LARS を用いることで、ImageNet-1k における ResNet-50 に対してバッチサイズ 32,768 での安定した学習が可能となり、スケールにおける最先端のパフォーマンスを達成した。
• AlexNet においても、LARS によりバッチサイズを 8,192 まで拡大でき、従来の最先端の限界（1,024）を著しく超えた。
• バッチサイズ 4,096 において、8 GPU の DGX-1 システム上での AlexNet モデルを用いて、LARS を用いた学習はバッチサイズ 512 と比較して 3 倍の高速化を達成した。
• ベース学習率の線形スケーリングと LARS のレイヤー固有の適応的スケーリングの組み合わせにより、極めて大きなバッチサイズでもモデルの精度が維持された。
• LARS におけるウォームアップ段階は、大バッチ学習における初期段階の最適化の不安定性を効果的に軽減した。
• LARS は、ResNet-50 や AlexNet といった異なるネットワークアーキテクチャにおいても、大バッチ設定で一貫したパフォーマンスを発揮した。