Source-linked AI summary
Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations
Alexander Hägele, Elie Bakouch, Atli Kosson, Loubna Ben Allal, Leandro Von Werra, Martin Jaggi
TL;DR
Large-model scaling research is costly and is complicated by cosine schedules whose cycle length must be fixed to the training duration. The paper evaluates constant learning rates with cooldowns and stochastic weight averaging, finding comparable or strong performance while enabling reusable runs and cheaper scaling experiments.
Problem
Large-model training is expensive, while cosine scheduling requires matching the cycle length to the training duration and complicates reusable scaling experiments.
Method
The paper studies constant learning rates with cooldowns, stochastic weight averaging, and schedule-free optimization as alternatives to cosine scheduling.
Results
Constant learning rate with cooldown matches cosine performance, while stochastic weight averaging provides strong performance during training without additional training overhead.
Takeaways & Limitations
Reusable checkpoints enable scaling-law experiments with fewer training runs and substantially reduced compute and GPU hours.
Takeaways & Limitations
The experiments cover models up to 8B parameters, and training behavior can be more brittle at modern scales and extremely long training.
Abstract
from arXiv · showhide
Scale has become a main ingredient in obtaining strong machine learning models. As a result, understanding a model's scaling properties is key to effectively designing both the right training setup as well as future generations of architectures. In this work, we argue that scale and training research has been needlessly complex due to reliance on the cosine schedule, which prevents training across different lengths for the same model size. We investigate the training behavior of a direct alternative -- constant learning rate and cooldowns -- and find that it scales predictably and reliably similar to cosine. Additionally, we show that stochastic weight averaging yields improved performance along the training trajectory, without additional training costs, across different scales. Importantly, with these findings we demonstrate that scaling experiments can be performed with significantly reduced compute and GPU hours by utilizing fewer but reusable training runs. Our code is available at \url{https://github.com/epfml/schedules-and-scaling/}.
1 Introduction
The paper argues that cosine scheduling makes scaling research unnecessarily expensive and proposes reusable training runs with constant learning rates, cooldowns, and weight averaging. These alternatives match cosine performance while enabling cheaper scaling experiments and more flexible training.
- 1 Introduction: Large-model training is costly and algorithmically complex, so researchers rely on small experiments before extrapolating to larger scales.This motivates reducing the number of independent training runs needed for scaling research.
- 1 Introduction: Cosine scheduling complicates experiments because its cycle length must match the intended training duration, preventing convenient stopping before or continuation beyond the cycle.This constraint affects both intermediate evaluation and final training decisions.
- 1 Introduction: Constant learning rate with cooldown matches cosine performance and scales reliably across decay forms and lengths, outperforming cosine for sufficiently long cooldowns.The authors present this as a direct alternative to cosine for large-model training.
- 1 Introduction: Stochastic weight averaging provides strong performance throughout training without additional training, while schedule-free optimization offers another strong but non-optimal alternative.These methods can replace learning-rate decay when a performance gap is acceptable.
- 1 Introduction: Scaling-law experiments can use one sufficiently long run per model size, followed by cooldown or averaging from reusable checkpoints, reducing compute and GPU hours.This replaces repeated training from scratch across different training lengths.
2 Background: Cosine Learning Rate Schedule for LLMs
Cosine scheduling reaches its best reported loss only when its cycle matches the training duration, making the planned endpoint a critical design choice. This creates problems for evaluating models during training and for continuing training afterward.
- Revisiting the optimality of the cosine schedule: Cosine is widely used because it combines a high learning rate after warm-up with gradual cooling toward a low final rate.The schedule typically decreases to 10% of the maximum learning rate.
- The pitfalls of cosine: Cosine achieves the best perplexity at specific step counts when its schedule length matches the training duration.Thus, the intended token count must be known in advance to obtain the best model under this schedule.
- The pitfalls of cosine: Cosine is suboptimal during training because it underestimates model performance at the same token count.The schedule's final performance does not represent the best attainable intermediate performance.
- The pitfalls of cosine: After cosine reaches the end of its cycle, its learning rate is generally too low for continued progress, while rewarming causes loss spikes and slow recovery.These effects make continuation beyond the planned cycle brittle.
3 A Different Route: Constant Learning Rate with Cooldown
Constant learning rate with cooldown matches well-tuned cosine performance without requiring the total training duration in advance, while enabling reusable runs and flexible training decisions. Cooldown shape and duration affect results, with (1-sqrt) often outperforming linear decay and downstream performance generally matching cosine.
- Experimental comparison: An almost perfect match between the best cosine and cooldown schedules holds across different training durations, with cooldown slightly less sensitive to learning-rate variation.
- Different cooldown schedules: The (1-sqrt) cooldown consistently outperforms linear decay, and both can outperform an untuned cosine schedule for sufficiently long training.
- Constant LR with cooldown provides similar performance to well-tuned cosine without requiring the number of training steps in advance.The cooldown can be initiated retrospectively from a checkpoint, supporting flexible stopping decisions and cheaper scaling studies.
- Cooldown duration: Cooldown benefits plateau around 20% of training, while long runs can match cosine with less than 20% when the cooldown has enough absolute steps.A 200k-step run found that 10k cooldown steps almost perfectly matched cosine performance.
- Optimization behavior: The cooldown produces a smooth loss transition into a connected basin, while independently sweeping maximum and final learning rates is important because final-LR choices affect outcomes.
- Scaling and downstream evaluation: Downstream benchmark performance generally matches between schedules, including at 1B-scale training, although not all benchmarks show the same cooldown-related improvement.The paper also reports matching loss values and no instability in a short 8B investigation.
4 Do We Even Need to Cooldown?
The paper evaluates whether cooldowns can be replaced by weight averaging or a schedule-free optimizer. SWA improves performance throughout training without extra overhead, but explicit cooldowns remain stronger, while SFO is sensitive to momentum choices.
- 4.1 Stochastic Weight Averaging (SWA): SWA averages parameters within fixed training windows, retaining one additional model copy and enabling later evaluation of longer averaging windows.The experiments use 500-step windows and save window averages as checkpoints.
- 4 Do We Even Need to Cooldown?: The methods target strong models at intermediate training points, potentially replacing shorter separately trained runs when the remaining performance gap is acceptable.SWA is especially useful when avoiding a separate cooldown is preferred.
- 4.1 Stochastic Weight Averaging (SWA): SWA improves performance along the training trajectory without additional overhead, but does not match explicit cooldowns.The improvement appears with both constant-LR and cosine training, reducing the gap without a separate decay phase.
- 4.2 Schedule-Free Optimizer: SFO does not consistently outperform cooldown because its performance depends sensitively on momentum parameters and training length.The (0.9, 0.95) setting performs significantly worse and increases loss near the end of training.
5 The Implications for Scaling Law Research
The paper argues that cooldowns and weight averaging make scaling-law experiments reusable across training lengths instead of requiring separate cosine runs. This reduces the number of training runs and substantially lowers compute and GPU-hour costs.
- Importance of scaling laws: Scaling laws model loss as a function of parameter count and training tokens, supporting compute allocation, data-quality trade-offs, and architecture comparisons.The paper presents the power-law form L(N, D) = A N^α + B D^β + E.
- Implications for scaling law research: A single sufficiently long run per model size, followed by cooldown or averaging from checkpoints, reduces scaling-law experiments to the model-scaling axis.The approach effectively divides the required number of training runs by one order of magnitude while supporting flexible continual training.
- Experimental setup: The experimental setup compares three cosine runs per model with one longest run for averaging or cooldown across 33M–360M models and 0.3B–10B tokens.All models use the same SlimPajama 6B dataset, with token counts spaced around D/N=20.
- Results: Cosine and cooldown reach nearly identical losses at matched token counts, while SWA shows a similar relationship in the cross-model comparison.Figure 12 places cosine and cooldown performance almost perfectly on the diagonal.
- Compute savings: The experiments save half the time and FLOPs when runs are spaced at token ratios of 10, 20, and 30.The reported savings cover both compute and GPU hours.
- Takeaway: Scaling experiments can use fewer reusable constant-LR runs with ad-hoc cooldowns, substantially reducing compute and GPU hours.The paper also estimates savings for the Chinchilla model suite under 10% cooldowns after one run per model size.
6 Limitations
The experiments cover models up to 8B parameters and long runs, but the authors caution that training can become more brittle at modern scales and extremely long durations.
- 6 Limitations: The study evaluates models up to 8B parameters, so behavior at larger scales is not directly established.The authors report consistent trends across tested scales but note a boundary for extrapolation.
- 6 Limitations: Training behavior can become more brittle at modern scales and under extremely long training.The authors note that high-learning-rate instabilities can be alleviated, but do not remove the stated scope boundary.
7 Related work
The paper situates its contribution among work on cosine schedules, weight averaging, and scaling-law methodology. It emphasizes that prior scaling-law studies established training procedures but did not address efficient reusable experiments.
- Cosine Schedules and Alternatives for Transformers: Cosine decay became the de-facto language-model schedule, while earlier vision work used stepwise or cyclic learning rates for exploration across epochs.The paper contrasts the schedule’s origins with its later use as a single-cycle LLM schedule.
- Weight Averaging: Prior work connects weight averaging with improved convergence or generalization and reports equivalence between weight averaging and decaying learning-rate schedules.These results motivate testing SWA as an alternative in LLM training.
- Scaling Law Experiments for Neural Language Models: Scaling-law research progressed from Kaplan’s laws to Chinchilla’s multi-duration cosine procedure and later studies of data quality, repetition, and token allocation.The paper identifies efficient scaling-law experiments as the missing focus of these prior works.
- Scaling Law Experiments for Neural Language Models: Concurrent work finds that constant learning rates can recover the optimal token-to-parameter ratio, but cooldown remains necessary for properly estimating performance.The paper specifically highlights downstream-task performance as a setting where cooldown matters.
8 Conclusion
The paper finds that constant learning rate with cooldown can reliably replace cosine, while SWA provides strong intermediate estimates without extra training overhead. These methods support continual training and substantially reduce scaling-law research costs, without claiming a universally best schedule.
- Constant learning rate with cooldown reliably matches cosine, while stochastic weight averaging provides strong but non-optimal intermediate estimates without additional training overhead.
- The authors do not claim to have established the best learning-rate schedule, but demonstrate a simple recipe with compelling practical advantages.
- The proposed methods enable more flexible continual training and reduce the cost of scaling-law research, especially for models trained beyond compute-optimality.
A Experimental Details
The appendix documents model, optimizer, dataset, evaluation, infrastructure, and FLOPs-estimation details. It also provides configuration tables and figures illustrating loss behavior under cosine and cooldown schedules.
- Architecture and training: Experiments use decoder-only Transformer models trained with AdamW, short warmups, weight decay, gradient clipping, and predominantly 512-token sequences.The standard setup includes beta parameters (0.9, 0.95), weight decay 0.1, and gradient clipping at 1.0.
- Dataset and evaluation: The main experiments use roughly 6B SlimPajama tokens and evaluate on fixed validation batches, while larger runs use FineWeb-based configurations and specialized GPU setups.
- Experimental documentation: The appendix supplies model and training-parameter tables plus code and figures documenting configurations, FLOPs computation, and cosine-versus-cooldown loss curves.
- Evaluation: Downstream evaluation covers common language-model benchmarks including MMLU, ARC, OpenBookQA, PIQA, HellaSwag, CommonSenseQA, SIQA, and Winogrande.
- Compute accounting: FLOPs are computed from explicit embedding, attention, and MLP operations rather than the common 6=ND approximation.
B.1 More Results on Cooldown
Additional cooldown experiments examine functional form, cooldown length, exponent choice, and learning-rate transfer. The results favor the 1-sqrt form but indicate that schedule rankings can depend on cooldown length and that cosine’s final learning rate should be tuned independently.
- Cooldown functions: The 1-sqrt cooldown outperforms linear decay, and this improvement persists across fewer decay steps, learning rates, and decay timestamps.
- Scope: Cooldown-function rankings may change for substantially different cooldown lengths, so the reported comparison focuses on practically relevant 10% and 20% cooldowns.
- Exponent sweep: Exponents 0.1 and 0.2 perform noticeably worse because the learning rate remains too low for many steps, while a=0.5 remains best among the other tested exponents.
- Learning-rate sensitivity: The optimal learning rate transfers across different cooldown lengths in the reported experiments.
- Cosine comparison: Cosine annealing should use a final learning rate below 10% of its maximum, while maximum and final rates should ideally be swept independently.
B.2 Additional Results and Compute Savings
The additional experiments show robust cooldown and SWA behavior across schedule variants, model sizes, training lengths, and OpenWebText2. Reusable runs reduce scaling-study compute and GPU-hour requirements, with savings increasing for larger models and longer runs.
- Cooldown robustness: The 1-sqrt cooldown remains the strongest tested cooldown function, with the same ordering broadly preserved across learning rates and 10% versus 20% cooldown lengths.
- Compute savings: Scaling experiments save a factor of 12 in FLOPs across models, with even larger savings expected when using more runs across training lengths.
- Compute savings: GPU-hour savings become especially prominent for bigger models and longer training runs, although the 360M SWA runtime is slightly affected by cluster congestion.
- Transfer: Cooldown and SWA findings transfer to OpenWebText2, where cooldown matches cosine and SWA improves performance during training.
- Long training runs: On 100B- and 460B-token FineWeb runs, cosine and cooldown match, while cosine-to-zero saturates before training ends and hurts overall performance.
B.5 Full Results of Large Model Runs
The 1B model’s downstream performance is comparable between cosine and cooldown schedules at both 100B and 460B tokens. Cooldowns improve some benchmarks during training, but benefits vary by metric and longer cooldowns are not consistently better.
- The detailed downstream benchmark results cover 1B-model runs evaluated throughout training and at the final checkpoints.
- Cooldowns produce an uptick in MMLU and HellaSwag performance, whereas OpenBookQA benefits less clearly.The uptick occurs when the cooldown starts at 80B tokens in the 100B-token run.
- Cosine and cooldown schedules have comparable final evaluation results after both 100B and 460B tokens.The comparison is reported in Table 4 for 100B tokens and Table 5 for 460B tokens.
- Longer cooldowns do not necessarily improve metrics, including when increasing cooldown length from 5% to 20%.