MARS: Unleashing the Power of Variance Reduction for Training Large Models

要約

ディープ ニューラル ネットワーク (さらに最近では大規模モデル) のトレーニングには、効率的でスケーラブルなオプティマイザーが必要です。
Adam、AdamW、およびそのバリアントなどの適応勾配アルゴリズムがこのタスクの中心となっています。
過去 10 年間に、凸型と非凸型の両方の設定における確率的最適化の加速を目的とした多数の分散削減アルゴリズムが開発されたにもかかわらず、ディープ ニューラル ネットワークや大規模な言語モデルのトレーニングにおいて分散削減は広く成功を収めていません。
その結果、現代の AI では依然としてあまり好まれないアプローチとなっています。
この論文では、大規模モデルの効率的なトレーニングのために分散削減の力を解き放つために、事前条件付き勾配法とスケーリングされた確率的再帰運動量手法による分散削減を調和させる統合最適化フレームワーク MARS (Make vAriance Reduction Shine) を提案します。
私たちのフレームワーク内では、それぞれ AdamW、Lion、Shampoo に基づく事前条件付き勾配更新を利用する MARS の 3 つのインスタンスを導入します。
また、アルゴリズムと既存のオプティマイザーとの間の接続も描画します。
GPT-2 モデルのトレーニングに関する実験結果は、MARS が一貫して AdamW を大幅に上回っていることを示しています。

要約(オリジナル)

Training deep neural networks–and more recently, large models–demands efficient and scalable optimizers. Adaptive gradient algorithms like Adam, AdamW, and their variants have been central to this task. Despite the development of numerous variance reduction algorithms in the past decade aimed at accelerating stochastic optimization in both convex and nonconvex settings, variance reduction has not found widespread success in training deep neural networks or large language models. Consequently, it has remained a less favored approach in modern AI. In this paper, to unleash the power of variance reduction for efficient training of large models, we propose a unified optimization framework, MARS (Make vAriance Reduction Shine), which reconciles preconditioned gradient methods with variance reduction via a scaled stochastic recursive momentum technique. Within our framework, we introduce three instances of MARS that leverage preconditioned gradient updates based on AdamW, Lion, and Shampoo, respectively. We also draw a connection between our algorithms and existing optimizers. Experimental results on training GPT-2 models indicate that MARS consistently outperforms AdamW by a large margin.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google