anthology-of-the-sota

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

1 min read

Properties4
StatusActive
Tagstraining-optimization
Cited as “introduced by” bySOTA-142 — Set the peak learning rate by a power law in tokens so it transfers across batch size and training length
Cited as “source” bySOTA-142 — Set the peak learning rate by a power law in tokens so it transfers across batch size and training length

Graph View

Shen et al., IBM Research (2024) — ARXIV-2408.13359

Key takeaways

  • The optimal peak learning rate depends on batch size, training length, model size and more, and searching for it at the final scale is what makes it expensive. The paper studies these correlations under the WSD schedule (LIT-144) and finds that the optimal learning rate follows a power law in the number of training tokens and the batch size.
  • The Power scheduler applies that law directly: the learning rate decays as a power of the tokens seen, capped at a maximum, with WSD’s final decay on top — so it is agnostic to the total token count and to the batch size, and the same setting can be reused when either changes.
  • Combined with Maximal Update Parametrization (LIT-148), one set of hyperparameters carries across model sizes as well; the paper trains its 3B dense and MoE models this way.

Standing in the anthology

The source of SOTA-142 and one of the two schedule components of the Falcon-H1-Tiny recipe (LIT-119), which is the record’s independent adoption.

Backlinks