Research map / Architectures and efficiency / RNNs and normalization
Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows that constant learning rate with a short cooldown (and SWA) can replace cosine schedules, matching cosine's scaling behavior while enabling scaling-law experiments with about half the compute by reusing training runs.
Problem
Cosine learning rate schedules require training a separate model from scratch for each training length, making training recipe iteration and scaling-law experiments expensive and inflexible.
Approach
Replace cosine with a constant learning rate followed by a short cooldown that can be started at any checkpoint, plus stochastic weight averaging; this matches cosine's loss at each token count while allowing one run to serve many training lengths, so scaling laws need fewer runs.
Key results
- Constant LR with a 20% cooldown matches best cosine validation loss across training durations on a 210M model, with slightly lower LR sensitivity.
- A 1-sqrt cooldown outperforms linear cooldown for long runs; a 10-sqrt-style cooldown of only 10k steps matched cosine in a 200k-step run.
- Scaling-law experiments with constant LR plus cooldown/SWA reduce compute by about half; estimated Chinchilla suite compute drops from 5.59e23 to 2.36e23 FLOPs.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace cosine decay with constant LR and a cooldown | LLM training · SlimPajama and FineWeb training runs | validation loss | No clear effect |
| use the 1-Sqrt cooldown instead of linear cooldown (ablation) | constant LR + cooldown · SlimPajama training runs | validation loss | Helped |
| extend cooldown to 10–20% of training steps (ablation) | constant LR + cooldown · 124M SlimPajama model | validation loss: 10-20% | Helped |
| use cosine cooldown instead of cosine from the start (ablation) | learning-rate schedule · SlimPajama training runs | validation loss | Mixed |
| vary the exponent in the cooldown function from 0.5 to 0.1 or 0.2 (ablation) | 1-Sqrt cooldown · SlimPajama training runs | validation loss: 0.1 and 0.2 | Hurt |
| anneal cosine LR to zero instead of 10% (ablation) | cosine schedule · FineWeb 1B, 100B tokens | Aggregated Score: 46.26 -> 45.88 | Hurt |
| add SWA during training (ablation) | constant LR and cosine schedules · 210M SlimPajama model | validation loss | Helped |
| replace SWA with EMA (ablation) | weight averaging · LLM training | validation loss | Hurt |
| use schedule-free optimizer with momentum (0.9, 0.95) instead of cooldown (ablation) | 210M model training · SlimPajama long run | training loss | Hurt |
| use schedule-free optimizer with momentum (0.95, 0.99) instead of cooldown (ablation) | 210M model training · SlimPajama long run | training loss | Hurt |
| reduce scaling experiments to one reusable run per model and use cooldowns or SWA for target token counts | scaling-law experiments · SlimPajama scaling experiments | compute and GPU hours: half the time and FLOPs | Helped |
| use 1-Sqrt 20% cooldown instead of cosine to 10% | 1B model · FineWeb 1B, 100B tokens | Aggregated Score: 46.26 -> 46.23 | No clear effect |
Surprising findings
- Constant LR with cooldown matches cosine performance, and a cooldown of only about 5-10% of steps can match cosine for long runs.
- Annealing cosine to zero improves loss but can hurt downstream benchmark metrics due to early saturation.
Limitations
- Experiments mostly on models up to 360M parameters; 1B and 8B runs are short or few, with no full 8B training.
- Downstream benchmark trends are not uniform across tasks; some benchmarks do not benefit from cooldown.
- SFO and SWA do not fully match cooldown performance; SFO is sensitive to momentum hyperparameters.
Experimental setup
Datasets: SlimPajama-6B OpenWebText2 FineWeb FineWeb-Edu
Benchmarks: MMLU ARC OpenBookQA PIQA HellaSwag CommonSenseQA SIQA Winogrande
Metrics: validation loss validation perplexity accuracy (acc_norm) FLOPs GPU hours
Baselines: cosine schedule schedule-free optimizer (SFO) linear cooldown Chinchilla scaling suite
Builds on: Chinchilla scaling laws cosine learning rate schedule stochastic weight averaging schedule-free optimizer WSD schedule
Closest papers by meaning
- A Dynamical Model of Neural Scaling Laws (ICML 2024 · similarity 0.90)
- Power Lines: Scaling laws for weight decay and batch size in LLM pre-training (NeurIPS 2025 · similarity 0.90)
- Navigating Scaling Laws: Compute Optimality in Adaptive Model Training (ICML 2024 · similarity 0.89)
- Resolving Discrepancies in Compute-Optimal Scaling of Language Models (NeurIPS 2024 · similarity 0.88)
- How Does Critical Batch Size Scale in Pre-training? (ICLR 2025 · similarity 0.88)
- A Hitchhiker's Guide to Scaling Law Estimation (ICML 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.
