Balancing learning-rate warmup duration for longer model training
Balancing Early Performance Sacrifices with Long-Term Gains: Scaling Learning-Rate Warmup Duration Across Training Horizons
Machine Learning
Summary
Training large language models involves gradually increasing the learning rate to avoid problems early on, a process called warmup. How long warmup should last isn’t clear though—it can be either a fixed time or grow with training length, which makes a difference as training gets longer. The authors explain this by modeling training directions that respond differently to peak learning rates, showing warmup helps some but slows others. Their findings offer a simple formula to pick warmup duration depending on training length and learning rate, helping to balance early slow progress with better long-term learning.
What this means in practice
- •For machine learning engineers: Optimize warmup schedules in large-scale language model training for improved long-term performance using the proposed horizon-aware scaling law.
- •For ai infrastructure teams: Predict warmup durations for extended training runs from shorter preliminary runs to better allocate computational resources in model development.
Authors
Kristi Topollai, Anna Choromanska
Abstract
Learning-rate warmup is a standard technique in language-model training, yet its duration remains largely heuristic. Common approaches use either a fixed number of updates or a fixed fraction of the training horizon, two choices that imply very different scaling as training gets longer. When should warmup stay fixed, and when should it grow with the horizon? We address this question with a quadratic model whose modes respond differently to the peak learning rate. Warmup slows progress in directions that already contract well at the peak rate, but can remove persistent error in directions near the stability edge, with higher peak rates shifting the balance toward longer warmup durations. This yields a compact horizon scaling law that captures regimes ranging from essentially no warmup, through fixed-duration warmup, to durations that grow with the training horizon, and explains how the preferred regime changes with peak learning rate. Because the law captures the tradeoff between giving up early progress and improving the trajectory that follows, it can be fit using shorter runs and used to predict warmup at substantially longer horizons. Together, our results explain several familiar properties of warmup through a single tradeoff and suggest treating warmup duration as a horizon-dependent hyperparameter rather than a fixed training heuristic.