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.
66cd41353bd2c33f87769b14
Yikang Shen
claimed_verified
2024-08-27T07:22:48.045Z
6463995939359568c63d740c
/avatars/11caacefc68c1259dbc8764cc4340481.svg
Yikang Shen
user
YikangS
66cd41353bd2c33f87769b15
Matthew Stallone
claimed_verified
2024-08-27T07:38:27.212Z
63b5dfd50d5913eee487b1b7
https://cdn-avatars.huggingface.co/v1/production/uploads/63b5dfd50d5913eee487b1b7/Av-QwMbfUScp0JpzoGGpK.jpeg
Matt
user
stallone
66cd41353bd2c33f87769b16
Mayank Mishra
claimed_verified
2024-08-27T20:42:30.301Z
62cd5057674cdb524450093d
https://cdn-avatars.huggingface.co/v1/production/uploads/62cd5057674cdb524450093d/f67rlrdsKPRTLdXCXoa_X.jpeg
Mayank Mishra
user
mayank-mishra
66cd41353bd2c33f87769b17
Gaoyuan Zhang
claimed_verified
2024-08-28T08:01:00.248Z
652431dc0e7792fa3aa2adc7
/avatars/b865d99ff5267184b37ed11cd2bd5626.svg
Gaoyuan Zhang
user
gaoyuanz
66cd41353bd2c33f87769b18
Shawn Tan
66cd41353bd2c33f87769b19
Aditya Prasad
66cd41353bd2c33f87769b1a
Adriana Meza Soria
claimed_verified
2024-09-04T07:25:29.310Z
639bcaa2445b133a4e942436
https://cdn-avatars.huggingface.co/v1/production/uploads/639bcaa2445b133a4e942436/h4Oa1ljMQvKXpGCHhTJY5.jpeg
Adriana Meza Soria
user
amezasor
66cd41353bd2c33f87769b1b
David D. Cox
66cd41353bd2c33f87769b1c
Rameswar Panda
extracted_pending
2024-08-27T03:00:06.283Z
648a7542ac1fbb8a1a654645
https://cdn-avatars.huggingface.co/v1/production/uploads/648a7542ac1fbb8a1a654645/jLPnY11EOuOa6ZHE7dVEX.jpeg
Rameswar Panda
user
rpand002
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.
Research reveals a power-law relationship for optimal learning rate, batch size, and training tokens, leading to the development of a Power scheduler that achieves consistent performance across various model sizes and architectures.
learning rate batch size training tokens hyperparameter search proxy models small corpus WSD scheduler power-law relationship Power scheduler Maximum Update Parameterization dense model MoE model