言語モデルの学習率スケジューリングの進展
新しい手法がモデル訓練における学習率の管理の仕方を変えてるんだ。
Yikang Shen, Matthew Stallone, Mayank Mishra, Gaoyuan Zhang, Shawn Tan, Aditya Prasad, Adriana Meza Soria, David D. Cox, Rameswar Panda
― 1 分で読む
モデルをトレーニングする時、特に大規模な言語モデルにおいては、適切な学習率を見つけるのがめっちゃ大事。学習率はモデルがデータから学ぶ速さを決めるんだ。もし学習率が高すぎると、モデルは最適な解を見逃しちゃう。逆に低すぎると、学習が遅すぎて時間を無駄にしちゃう。
モデルをトレーニングする時は、データバッチのサイズやトークンの総数、モデルのサイズ、他の設定など、いろんな要素が影響してくる。この複雑さのせいで、ベストな学習率を見つけるのはめっちゃ難しい。特に数十億のパラメータを持つ巨大なモデルの場合、この課題はさらに大きくなる。コストや時間の制約から、すべての設定をテストするのは不可能だしね。
学習率の管理
一部の研究者は、これを簡単にするために、小さなモデルやデータセットでテストして、その結果を大きなモデルのトレーニングに活かそうとしてる。この方法が効果的な場合もあるけど、小さいモデルから大きいモデルに学習率をうまく移行する方法を理解するには、もっと研究が必要なんだ。
この問題を解決するために提案されている方法の一つが、ウォームアップ・ステイブル・デケイ（WSD）スケジューラーって呼ばれる学習率スケジュール。これは、学習率を徐々に上げて、トレーニングの大半を安定させて、最後に徐々に下げるっていう三つの主要なフェーズがある。この方法の本当の利点は、事前に特定のトレーニングステップの数を設定しなくていいから、いろんなトレーニングシナリオに柔軟に対応できること。
WSDスケジューラーの理解
WSDスケジューラーでは、ウォームアップフェーズで学習率をゼロからピーク値まで上げる。安定フェーズではそのピークを大半のトレーニング時間維持する。最後にデケイフェーズで学習率を短い期間でゼロまで下げる。この柔軟性のおかげで、研究者はモデルのパフォーマンスを観察して、予め決めた条件に縛られることなくトレーニングを調整できる。
でも、これらの利点にも関わらず、研究ではWSDスケジューラーがトレーニングに使うトークンの総数に影響されることがわかってる。つまり、データの量を大きく変えると、最適な学習率が変わるかもしれなくて、その結果が最適じゃないことがあるってこと。
学習率の関係を調査
もっと深く理解するために、研究者たちは学習率、データバッチのサイズ、トレーニングトークン数の関係を調べてる。実験を通じて、これらの変数が特定の方法でつながっているというべきパワー則的な関係が見つかった。
研究者が異なるデータサイズで学習率を分析した時、パターンが見えた。例えば、トークン数が増えると最適な学習率が減少する傾向がある。一方、バッチサイズが増えると最適な学習率が増加する。つまり、ベストな学習率は固定されてなくて、トレーニングに関わる他の要素によって変わるってこと。
パワースケジューラーの導入
この発見をもとに、新しいタイプの学習率スケジューラー、パワースケジューラーが導入された。このスケジューラーはトレーニングトークンとバッチサイズに依存せずに動くことを目指してる。トレーニング中に学習率を簡単に調整できるから、事前に制限を指定する必要がない。
パワースケジューラーは学習率を柔軟に保ちながら、トレーニングの進行に応じて適応するように設計されてる。このアプローチを使うことで、研究者は事前に決めた条件にモデルを合わせることの複雑さを避けられる。
実験結果
初期テストでは、パワースケジューラーがWSDスケジューラーなどの既存の方法と同等の良いパフォーマンスを示してる。構造化された実験では、チームがパワースケジューラーを使っていろんなモデルをトレーニングした結果、一貫して強力なパフォーマンスが確認された。
制御された設定では、パワースケジューラーを使ってトレーニングされたモデルが、従来のスケジューリング方法を使ったモデルと同レベルかそれ以上の結果を出した。これから、パワースケジューラーは複雑な言語モデルに取り組む人たちにとって有用なツールになるかもしれない。
実世界での応用
もっと現実的なトレーニング環境でも、パワースケジューラーを使ったモデルは高いレベルのパフォーマンスを維持できた。このテストでは、研究者が大規模なデータセットで大きなモデルを効率的かつ効果的にトレーニングした。
強力な競合に対抗しても、パワースケジューラーを使ったモデルは有望な結果を示した。実験の結果、パワースケジューラーでトレーニングされたモデルは、過度に大きなトレーニングデータを必要とせずに素晴らしいパフォーマンスを達成できることが確認された。
結論：学習率スケジューリングの未来
系統的な研究と実験を通じて、学習率、バッチサイズ、トレーニングトークンの関係がより明確になってきてる。この理解は、トレーニング方法を改善して、大規模な言語モデルのトレーニングでより良い結果を得るために重要なんだ。
パワースケジューラーの導入は、この分野における進展を示していて、様々なトレーニング条件に適応できる柔軟な学習率管理のアプローチを提供してる。今後の研究で、これらの発見は間違いなく機械学習技術の進化に貢献し、研究者がより有能で効率的なモデルを開発できるようになるだろう。
タイトル: Power Scheduler: A Batch Size and Token Number Agnostic Learning Rate Scheduler
概要: Finding the optimal learning rate for language model pretraining is a challenging task. This is not only because there is a complicated correlation between learning rate, batch size, number of training tokens, model size, and other hyperparameters but also because it is prohibitively expensive to perform a hyperparameter search for large language models with Billions or Trillions of parameters. Recent studies propose using small proxy models and small corpus to perform hyperparameter searches and transposing the optimal parameters to large models and large corpus. While the zero-shot transferability is theoretically and empirically proven for model size related hyperparameters, like depth and width, the zero-shot transfer from small corpus to large corpus is underexplored. In this paper, we study the correlation between optimal learning rate, batch size, and number of training tokens for the recently proposed WSD scheduler. After thousands of small experiments, we found a power-law relationship between variables and demonstrated its transferability across model sizes. Based on the observation, we propose a new learning rate scheduler, Power scheduler, that is agnostic about the number of training tokens and batch size. The experiment shows that combining the Power scheduler with Maximum Update Parameterization (muP) can consistently achieve impressive performance with one set of hyperparameters regardless of the number of training tokens, batch size, model size, and even model architecture. Our 3B dense and MoE models trained with the Power scheduler achieve comparable performance as state-of-the-art small language models. We open-source these pretrained models at https://ibm.biz/BdKhLa.
著者: Yikang Shen, Matthew Stallone, Mayank Mishra, Gaoyuan Zhang, Shawn Tan, Aditya Prasad, Adriana Meza Soria, David D. Cox, Rameswar Panda
最終更新: 2024-09-11 00:00:00
言語: English
ソースURL: https://arxiv.org/abs/2408.13359
ソースPDF: https://arxiv.org/pdf/2408.13359
ライセンス: https://creativecommons.org/licenses/by/4.0/
変更点: この要約はAIの助けを借りて作成されており、不正確な場合があります。正確な情報については、ここにリンクされている元のソース文書を参照してください。
オープンアクセスの相互運用性を利用させていただいた arxiv に感謝します。