title: Power Scheduler: A Batch Size and Token Number Agnostic Learning Rate
    Scheduler
summary: 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.
paper_authors._id: 66cd41353bd2c33f87769b14
paper_authors.name: Yikang Shen
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2024-08-27T07:22:48.045Z
paper_authors._id: 6463995939359568c63d740c
paper_authors.avatarUrl: /avatars/11caacefc68c1259dbc8764cc4340481.svg
paper_authors.fullname: Yikang Shen
paper_authors.type: user
paper_authors.user: YikangS
paper_authors._id: 66cd41353bd2c33f87769b15
paper_authors.name: Matthew Stallone
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2024-08-27T07:38:27.212Z
paper_authors._id: 63b5dfd50d5913eee487b1b7
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/63b5dfd50d5913eee487b1b7/Av-QwMbfUScp0JpzoGGpK.jpeg
paper_authors.fullname: Matt
paper_authors.type: user
paper_authors.user: stallone
paper_authors._id: 66cd41353bd2c33f87769b16
paper_authors.name: Mayank Mishra
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2024-08-27T20:42:30.301Z
paper_authors._id: 62cd5057674cdb524450093d
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/62cd5057674cdb524450093d/f67rlrdsKPRTLdXCXoa_X.jpeg
paper_authors.fullname: Mayank Mishra
paper_authors.type: user
paper_authors.user: mayank-mishra
paper_authors._id: 66cd41353bd2c33f87769b17
paper_authors.name: Gaoyuan Zhang
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2024-08-28T08:01:00.248Z
paper_authors._id: 652431dc0e7792fa3aa2adc7
paper_authors.avatarUrl: /avatars/b865d99ff5267184b37ed11cd2bd5626.svg
paper_authors.fullname: Gaoyuan Zhang
paper_authors.type: user
paper_authors.user: gaoyuanz
paper_authors._id: 66cd41353bd2c33f87769b18
paper_authors.name: Shawn Tan
paper_authors._id: 66cd41353bd2c33f87769b19
paper_authors.name: Aditya Prasad
paper_authors._id: 66cd41353bd2c33f87769b1a
paper_authors.name: Adriana Meza Soria
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2024-09-04T07:25:29.310Z
paper_authors._id: 639bcaa2445b133a4e942436
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/639bcaa2445b133a4e942436/h4Oa1ljMQvKXpGCHhTJY5.jpeg
paper_authors.fullname: Adriana Meza Soria
paper_authors.type: user
paper_authors.user: amezasor
paper_authors._id: 66cd41353bd2c33f87769b1b
paper_authors.name: David D. Cox
paper_authors._id: 66cd41353bd2c33f87769b1c
paper_authors.name: Rameswar Panda
paper_authors.status: extracted_pending
paper_authors.statusLastChangedAt: 2024-08-27T03:00:06.283Z
paper_authors._id: 648a7542ac1fbb8a1a654645
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/648a7542ac1fbb8a1a654645/jLPnY11EOuOa6ZHE7dVEX.jpeg
paper_authors.fullname: Rameswar Panda
paper_authors.type: user
paper_authors.user: rpand002
paper_title: Power Scheduler: A Batch Size and Token Number Agnostic Learning Rate
    Scheduler
paper_summary: 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.
paper_ai_keywords: learning rate batch size training tokens model size hyperparameters proxy models corpus zero-shot transferability WSD scheduler power-law relationship learning rate scheduler Power scheduler Maximum Update Parameterization (muP) dense models MoE models state-of-the-art small language models