[論文レビュー] AdaScale SGD: A User-Friendly Algorithm for Distributed Training

Tyler B. Johnson, Pulkit Agrawal+2 more|arXiv (Cornell University)|Jul 9, 2020

Advanced Neural Network Applications被引用数 20

ひとこと要約

AdaScale SGD は、勾配の分散に基づいて学習率を動的にスケーリングすることで、大バッチ学習に自動的に適応する使いやすいアルゴリズムであり、ハイパーパrameterのチューニングが不要な近似的な線形スケーリングを実現する。モデルの品質を多様なタスクやバッチサイズで維持し、線形または平方根スケーリングといった固定スケーリングルールを上回り、明示的なウォームアップ段階を必要としない自然なウォームアップ的挙動を示す。

ABSTRACT

When using large-batch training to speed up stochastic gradient descent, learning rates must adapt to new batch sizes in order to maximize speed-ups and preserve model quality. Re-tuning learning rates is resource intensive, while fixed scaling rules often degrade model quality. We propose AdaScale SGD, an algorithm that reliably adapts learning rates to large-batch training. By continually adapting to the gradient's variance, AdaScale automatically achieves speed-ups for a wide range of batch sizes. We formally describe this quality with AdaScale's convergence bound, which maintains final objective values, even as batch sizes grow large and the number of iterations decreases. In empirical comparisons, AdaScale trains well beyond the batch size limits of popular "linear learning rate scaling" rules. This includes large-batch training with no model degradation for machine translation, image classification, object detection, and speech recognition tasks. AdaScale's qualitative behavior is similar to that of "warm-up" heuristics, but unlike warm-up, this behavior emerges naturally from a principled mechanism. The algorithm introduces negligible computational overhead and no new hyperparameters, making AdaScale an attractive choice for large-scale training in practice.

研究の動機と目的

• 大バッチ分散学習において、固定の学習率スケーリングルールが性能を低下させるという課題に対処すること。
• 大規模学習における手動のハイパーパrameterチューニングやヒューリスティックなウォームアップスケジュールの必要性を排除すること。
• 確率的勾配の固有の分散に基づいて学習率を原理的かつ適応的にスケーリングする方法を開発すること。
• 画像分類、物体検出、機械翻訳、音声認識を含む多様な機械学習タスクにおいて、信頼性があり使いやすい大バッチ学習を可能にすること。
• 線形または平方根スケーリングの理論的裏付けのある代替手段を提供し、スケールに応じた収束性とモデルの正確性を維持すること。

提案手法

• AdaScale は、各イテレーションで確率的勾配の推定分散に逆比例して学習率を動的に調整することで、バッチサイズにかかわらず期待される更新量の大きさが安定するようにする。
• このアルゴリズムは、$\eta_t \mathbb{E}[\|\bar{\mathbf{g}}_t\|^2]$ の不変性を維持する。これは収束境界の中心的役割を果たし、一貫した最適化の進行を保証する。
• 新たなハイパーパrameterを導入せず、計算オーバーヘッドも最小限に抑えられ、大規模な学習システムにおいて実用的である。
• 適応的スケーリングメカニズムから自然にウォームアップ的挙動が生じ、明示的なウォームアップ段階を必要としない。
• AdaScale は標準的な学習率スケジュール（例：指数関数的減衰）と互換性があり、それらを異なるバッチサイズやスケールに自動的に適応する。
• アルゴリズムは、バッチサイズが増加しイテレーション回数が減少する場合でも、最終的な目的関数値を維持することを保証する理論的収束境界に基づいて導出される。

実験結果

リサーチクエスチョン

• RQ1 手動の再チューニングが不要な学習率適応戦略を設計でき、広範なバッチサイズにおいてモデルの品質を維持できるか？
• RQ2 勾配の分散に基づく適応的学習率スケーリングは、線形または平方根スケーリングといった固定ルールと比較して、モデルの正確性と学習速度の面でどのように異なるか？
• RQ3 明示的なウォームアップハイパーパramーターやスケジュール変更なしに、自然にウォームアップ的挙動を示す適応的メカニズムが存在するか？
• RQ4 AdaScale は、最終的なモデル性能を維持しながら、大バッチ分散学習で線形スケーリングをどの程度達成できるか？
• RQ5 AdaScale は、画像分類、物体検出、機械翻訳、音声認識を含む多様な機械学習タスクに効果的に適用可能か？

主な発見

• AdaScale は、ImageNet（バッチサイズ32,000）やTransformer（1バッチあたり最大262,000トークン）を含む複数のベンチマークで、大バッチ学習においてほぼ完全な線形スケーリングを達成した。
• CIFAR-10では、スケール16における指数関数的減衰スケジュールの13×13グリッドにおいても、AdaScale はスケーリングされたSGDの直接的なハイパーパラメータ探索を上回るモデル品質を維持した。
• 線形スケーリングにウォームアップを組み合わせた手法とは異なり、バッチサイズが増加するにつれてモデルの正確性が低下するのに対し、AdaScale は極端なスケールでもバリデーション正確性を維持した。
• AdaScale の収束境界により、バッチサイズが増大しイテレーション回数が減少する場合でも、最終的な目的関数値が維持されることを保証し、その頑健性に理論的根拠を与える。
• アルゴリズムは、明示的なウォームアップハイパーパラメーターやスケジュール変更なしに、自然にウォームアップ的学習率挙動を生成した。
• 実験的結果から、AdaScale は多様なタスクとスケールにおいて一貫してモデル品質を維持することが示され、実世界の分散学習における信頼性と実用性を示した。