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

TL;DR

The paper tackles the challenge of learning-rate selection for large-scale pretraining, where traditional cosine schedules require predefined total steps. It demonstrates that under the Warmup-Stable-Decay regime, the optimal LR follows a power-law with respect to batch size and training tokens, and introduces the Power scheduler, which operates independently of token count and batch size when paired with Maximum Update Parametrization (μP). By revealing that η_opt ≈ β a T^b and that coefficients are largely model-size invariant, the authors enable cross-size transfer and develop a token-agnostic scheduling approach that performs robustly in both dense and mixture-of-experts models. Empirically, Power scheduler with μP achieves strong results on 3B models and is open-sourced, offering practical benefits for continual pretraining, early stopping, and scalable transfer across architectures.

Abstract

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.

Figures (6)

• Figure 1: Illustration of learning rate curves for Cosine, WSD, and our Power schedulers.
• Figure 2: Left: Learning Rate v.s. Test Perplexity for different numbers of training tokens. The optimal learning rate decreases with respect to the number of training tokens. Right: Learning Rate v.s. Test Perplexity for different batch sizes. The optimal learning rate increases with respect to the batch size.
• Figure 3: Batch Size v.s. Learning Rate for different numbers of training tokens. The optimal learning rate increases with respect to the batch size.
• Figure 4: a, b, c: The $\gamma$ of best three batch sizes v.s. The number of training tokens. (Dots could overlap.) d: The average $\gamma$ of the three best batch sizes v.s. The number of training tokens. A power-law function can model the correlation between $\gamma$ and $T$.
• Figure 5: Hyperparameter search results for 32B training tokens. Left: Test perplexity vs $a$. The test perplexity is not sensitive to different choices of $a$ within the range of $[3,5]$. Right: Test perplexity vs $b$. The test perplexity is more sensitive to the choices of $b$. The optimal $b$ is between $-0.52$ and $-0.51$.