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

arXiv: https://arxiv.org/abs/2408.13359 · HTML: https://mlsanity.com/paper/2408.13359/

## Key facts

- Cited by 28 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 23 references that are in the map.
- First posted on arXiv 23 Aug 2024.

## Builds on (most cited first, 18 of 18)

- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [PIQA: Reasoning about Physical Commonsense in Natural Language](https://mlsanity.com/paper/1911.11641/) (2019)
- [BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questions](https://mlsanity.com/paper/1905.10044/) (2019)
- [Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering](https://mlsanity.com/paper/1809.02789/) (2018)
- [Gemma 2: Improving Open Language Models at a Practical Size](https://mlsanity.com/paper/2408.00118/) (2024)
- [Solving Quantitative Reasoning Problems with Language Models](https://mlsanity.com/paper/2206.14858/) (2022)
- [Scaling Language Models: Methods, Analysis & Insights from Training Gopher](https://mlsanity.com/paper/2112.11446/) (2021)
- [MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies](https://mlsanity.com/paper/2404.06395/) (2024)
- [Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer](https://mlsanity.com/paper/2203.03466/) (2022)
- [Feature Learning in Infinite-Width Neural Networks](https://mlsanity.com/paper/2011.14522/) (2020)
- [Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations](https://mlsanity.com/paper/2405.18392/) (2024)
- [Cerebras-GPT: Open Compute-Optimal Language Models Trained on the Cerebras Wafer-Scale Cluster](https://mlsanity.com/paper/2304.03208/) (2023)
- [Tensor Programs VI: Feature Learning in Infinite-Depth Neural Networks](https://mlsanity.com/paper/2310.02244/) (2023)
- [JetMoE: Reaching Llama2 Performance with 0.1M Dollars](https://mlsanity.com/paper/2404.07413/) (2024)
- [Scattered Mixture-of-Experts Implementation](https://mlsanity.com/paper/2403.08245/) (2024)

## Built on by (most cited first, 28 of 28)

- [Power Lines: Scaling Laws for Weight Decay and Batch Size in LLM Pre-training](https://mlsanity.com/paper/2505.13738/) (2025)
- [Straight to Zero: Why Linearly Decaying the Learning Rate to Zero Works Best for LLMs](https://mlsanity.com/paper/2502.15938/) (2025)
- [Predictable Scale: Part I, Step Law -- Optimal Hyperparameter Scaling Law in Large Language Model Pretraining](https://mlsanity.com/paper/2503.04715/) (2025)
- [EuroBERT: Scaling Multilingual Encoders for European Languages](https://mlsanity.com/paper/2503.05500/) (2025)
- [A Multi-Power Law for Loss Curve Prediction Across Learning Rate Schedules](https://mlsanity.com/paper/2503.12811/) (2025)
- [Mid-Training of Large Language Models: A Survey](https://mlsanity.com/paper/2510.06826/) (2025)
- [Hyperparameter Transfer Enables Consistent Gains of Matrix-Preconditioned Optimizers Across Scales](https://mlsanity.com/paper/2512.05620/) (2025)
- [Towards Precise Scaling Laws for Video Diffusion Transformers](https://mlsanity.com/paper/2411.17470/) (2024)
- [Scaling Stick-Breaking Attention: An Efficient Implementation and In-depth Study](https://mlsanity.com/paper/2410.17980/) (2024)
- [Training Dynamics of the Cooldown Stage in Warmup-Stable-Decay Learning Rate Scheduler](https://mlsanity.com/paper/2508.01483/) (2025)
- [LoopMoE: Unifying Iterative Computation with Mixture-of-Experts for Language Modeling](https://mlsanity.com/paper/2606.04438/) (2026)
- [Time Transfer: On Optimal Learning Rate and Batch Size In The Infinite Data Limit](https://mlsanity.com/paper/2410.05838/) (2024)
- [Scaling with Collapse: Efficient and Predictable Training of LLM Families](https://mlsanity.com/paper/2509.25087/) (2025)
- [Not All Models Suit Expert Offloading: On Local Routing Consistency of Mixture-of-Expert Models](https://mlsanity.com/paper/2505.16056/) (2025)
- [Optimal Scaling Needs Optimal Norm](https://mlsanity.com/paper/2510.03871/) (2025)
- [Beyond Cosine Decay: On the effectiveness of Infinite Learning Rate Schedule for Continual Pre-training](https://mlsanity.com/paper/2503.02844/) (2025)
- [A Partition Cover Approach to Tokenization](https://mlsanity.com/paper/2501.06246/) (2025)
- [Occult: Optimizing Collaborative Communication across Experts for Accelerated Parallel MoE Training and Inference](https://mlsanity.com/paper/2505.13345/) (2025)
- [AdaLRS: Loss-Guided Adaptive Learning Rate Search for Efficient Foundation Model Pretraining](https://mlsanity.com/paper/2506.13274/) (2025)
- [Scale or Reason? A Compute-Equivalent Analysis of Reasoning Distillation](https://mlsanity.com/paper/2509.22193/) (2025)
- [SiLQ: Simple Large Language Model Quantization-Aware Training](https://mlsanity.com/paper/2507.16933/) (2025)
- [xGen-small Technical Report](https://mlsanity.com/paper/2505.06496/) (2025)
- [Puro-2B: Poor Lab's Qwen2-1.5B Trained on RTX 5090 within 5090](https://mlsanity.com/paper/2608.27370/) (2026)
- [Deriving Scaling Laws for OpenEuroLLM Models: Learning Rate, Batch Size and Loss](https://mlsanity.com/paper/2608.28308/) (2026)
- [Predictable Scaling Laws of Optimal Hyperparameters for LLM Continued Pre-training](https://mlsanity.com/paper/2606.05610/) (2026)
- [WSqD: A Horizon-Free Learning Rate Schedule for Large Model Training](https://mlsanity.com/paper/2607.10959/) (2026)
- [Blog: Survey of Optimizers](https://mlsanity.com/paper/2608.28557/) (2026)
- [To Explore The Strange New World Beyond Data Distribution: System Behavior, Causality Tax, and Non-causal Base Model](https://mlsanity.com/paper/2610.02839/) (2026)

## Abstract (arXiv)

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.
