パワースケジューラー：バッチサイズとトークン数に関係なく学習率をスケジューリング

Power Scheduler: A Batch Size and Token Number Agnostic Learning Rate Scheduler

August 23, 2024

著者: Yikang Shen, Matthew Stallone, Mayank Mishra, Gaoyuan Zhang, Shawn Tan, Aditya Prasad, Adriana Meza Soria, David D. Cox, Rameswar Panda

cs.AI

要旨

言語モデルの事前学習に最適な学習率を見つけることは困難な課題です。これは、学習率、バッチサイズ、トレーニングトークンの数、モデルサイズ、および他のハイパーパラメータとの複雑な相関があるだけでなく、数十億または数兆のパラメータを持つ大規模な言語モデルに対してハイパーパラメータ検索を行うのは莫大なコストがかかるためです。最近の研究では、小さなプロキシモデルと小さなコーパスを使用してハイパーパラメータ検索を行い、最適なパラメータを大規模なモデルと大規模なコーパスに移行させることが提案されています。モデルサイズに関連する深さや幅などのハイパーパラメータに関しては、ゼロショット転送性が理論的にも実証的にも証明されていますが、小規模コーパスから大規模コーパスへのゼロショット転送は未開拓の領域です。本論文では、最近提案されたWSDスケジューラの最適な学習率、バッチサイズ、およびトレーニングトークン数との相関を研究します。数千回の小規模な実験の結果、変数間にべき乗則の関係があり、その転送可能性をモデルサイズを超えて実証しました。この観察に基づき、トレーニングトークン数やバッチサイズに関係なく動作する新しい学習率スケジューラであるPowerスケジューラを提案します。実験では、Powerスケジューラを最大更新パラメータ化（muP）と組み合わせることで、トレーニングトークンの数、バッチサイズ、モデルサイズ、さらにはモデルアーキテクチャに関係なく、一連のハイパーパラメータで印象的なパフォーマンスを一貫して達成できることを示しました。Powerスケジューラでトレーニングされた3BのdenseモデルとMoEモデルは、最先端の小規模言語モデルと同等のパフォーマンスを達成します。これらの事前学習済みモデルは、https://ibm.biz/BdKhLa でオープンソース化されています。

English

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.