Research map / Architectures and efficiency / RNNs and normalization
Scaling Law with Learning Rate Annealing
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.
Proposes a scaling law L(s)=L0+A·S1^(-α)−C·S2 that models full training loss curves under arbitrary learning rate schedules using forward LR area and an LR-annealing momentum area, fitted from one or two curves.
Problem
Classical scaling laws fit only final-step losses from many runs, ignoring intermediate training dynamics and the loss drop caused by learning rate annealing, which makes them costly and unable to predict full loss curves under different learning rate schedules.
Approach
Adds an LR-annealing term to the power-law scaling law: loss depends on the forward area S1 (sum of LR) and a momentum-weighted annealing area S2 (decayed sum of LR drops), with only one extra parameter C, so full loss curves for unseen schedulers can be predicted from a single or two fitted curves.
Key results
- Fitting one constant+cosine pair (20K steps) and predicting unseen schedulers at 60K steps yields mean prediction error about 0.2% and fit R^2 > 0.999.
- The N-extended law (with S2 scaled by N^gamma) fits and predicts a 1.2B-size run with R^2 > 0.998.
- The fitted formulation uses less than 1% of the compute of Chinchilla-style fitting for 100 points.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add the LR annealing area term to the scaling law | Language-model validation-loss curve prediction across unseen schedulers · Unseen learning-rate schedulers | mean prediction error: 0.2% | Helped |
| Use the model-size-extended scaling law | Loss-curve prediction for an unseen larger model · Unseen model size | R²: $R^2>0.998$ | Helped |
| Use full loss curves instead of final-loss-only fitting | Scaling-law fitting · Pretraining across learning-rate schedulers | computational cost: 100% -> $<$1.0% | Helped |
| Set cosine cycle length to total steps and minimum LR to 0 | Cosine learning-rate schedule · Final training step | validation loss | Helped |
| Use constant LR instead of cosine LR | Language-model pretraining · Small total-step runs | final validation loss | Helped |
| Use cosine LR instead of constant LR | Language-model pretraining · Large total-step runs | final validation loss | Helped |
| Use WSD or multi-step cosine instead of cosine LR | Language-model pretraining · Pretraining loss | validation loss | Helped |
| Use 1-sqrt rather than cosine annealing | WSD schedule · 10% annealing ratio | final loss | Helped |
| Use cosine rather than 1-sqrt annealing | WSD schedule · 50% annealing ratio | final loss | Helped |
| Use a moderate WSD annealing ratio, typically 10% to 20% | WSD schedule · Different total training steps | predicted final loss | Helped |
| Increase the re-warmup maximum learning rate | Continual pretraining · Re-warmup phase | initial peak loss | Hurt |
| Increase re-warmup steps | Continual pretraining · Final validation loss | validation loss | No clear effect |
Surprising findings
- 1-sqrt annealing beats cosine annealing only at small annealing ratios (10%); at 50% annealing cosine wins, contradicting prior claims that 1-sqrt is consistently best.
- Adam beta2 does not affect the LR-annealing delay phenomenon, so Adam momentum is not its root cause.
Limitations
- Lambda (decay factor) is fixed to 0.999 as a hyperparameter rather than fitted; choice affects predictions.
- Assumes no data distribution shift; continual pre-training with distribution shift is not fully covered.
- Root cause of the delay phenomenon is not identified.
Experimental setup
Datasets: FineWeb RedPajama-CC multilingual mixture (Setting C) BLOOM-176B loss curves OLMo-1B loss curves
Benchmarks: validation loss on RedPajama-CC BLOOM-176B loss curve OLMo-1B loss curve
Metrics: R^2 mean prediction error Huber loss
Baselines: Chinchilla scaling law WSD scheduler cosine learning rate schedule 1-sqrt annealing multi-step cosine
Builds on: Chinchilla scaling law WSD learning rate scheduler Kaplan scaling laws
Closest papers by meaning
- Universal One-third Time Scaling in Learning Peaked Distributions (ICML 2026 · similarity 0.92)
- Functional Scaling Laws in Kernel Regression: Loss Dynamics and Learning Rate Schedules (NeurIPS 2025 · similarity 0.92)
- A Hitchhiker's Guide to Scaling Law Estimation (ICML 2025 · similarity 0.91)
- Tokens-per-Parameter Coverage Is Critical for Robust LLM Scaling Law Extrapolation (NeurIPS 2026 · similarity 0.91)
- Resolving Discrepancies in Compute-Optimal Scaling of Language Models (NeurIPS 2024 · similarity 0.90)
- A Multi-Power Law for Loss Curve Prediction Across Learning Rate Schedules (ICLR 2025 · similarity 0.89)
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.