MKOR: Momentum-Enabled Kronecker-Factor-Based Optimizer Using Rank-1 Updates

要約

本研究では、ディープニューラルネットワーク（DNN）の学習時間と収束特性を改善する、MKORと呼ばれるランク1更新を用いた運動量有効クロネッカー因子ベースのオプティマイザーを提案する。2次手法は、1次手法と比較して高い収束率を実現する一方で、モデルサイズや学習バッチサイズに対して3次的な複雑さを有しています。そのため、大規模言語モデル（LLM）のような変形モデルでは、スケーラビリティとパフォーマンスが低下します。これらのモデルでは、バッチサイズが注目メカニズムのシーケンス長によってスケールするため、モデルサイズとバッチサイズが大きくなってしまうのです。MKORの計算量はモデルサイズに対して2次関数的であり、2次手法の計算ボトルネックを緩和している。計算量が多いため、最新の2次法の実装では、2次情報の更新頻度が低く、更新による収束性の向上が期待されても、それを十分に生かすことができない。MKORは、2次更新の通信複雑度を低減し、線形通信複雑度を達成することで、2次更新の頻度を増加させます。また、MKORのハイブリッドバージョン（MKOR-H）を提案し、2次更新で収束が促進されなくなった場合、トレーニングの途中で1次オプティマイザにフォールバックする。実験によると、MKORは、64GPUのBERT-Large-Uncasedにおいて、LAMBオプティマイザなどの最新の1次手法と、KAISA/KFACなどの最高の2次手法の実装をそれぞれ最大2.57倍と1.85倍上回る性能を示しました。

要約(オリジナル)

This work proposes a Momentum-Enabled Kronecker-Factor-Based Optimizer Using Rank-1 updates, called MKOR, that improves the training time and convergence properties of deep neural networks (DNNs). Second-order techniques, while enjoying higher convergence rates vs first-order counterparts, have cubic complexity with respect to either the model size and/or the training batch size. Hence they exhibit poor scalability and performance in transformer models, e.g. large language models (LLMs), because the batch sizes in these models scale by the attention mechanism sequence length, leading to large model size and batch sizes. MKOR’s complexity is quadratic with respect to the model size, alleviating the computation bottlenecks in second-order methods. Because of their high computation complexity, state-of-the-art implementations of second-order methods can only afford to update the second order information infrequently, and thus do not fully exploit the promise of better convergence from these updates. By reducing the communication complexity of the second-order updates as well as achieving a linear communication complexity, MKOR increases the frequency of second order updates. We also propose a hybrid version of MKOR (called MKOR-H) that mid-training falls backs to a first order optimizer if the second order updates no longer accelerate convergence. Our experiments show that MKOR outperforms state -of-the-art first order methods, e.g. the LAMB optimizer, and best implementations of second-order methods, i.e. KAISA/KFAC, up to 2.57x and 1.85x respectively on BERT-Large-Uncased on 64 GPUs.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL