Research map / Learning theory and algorithms / Sample complexity
Resolving Discrepancies in Compute-Optimal Scaling of Language Models
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.
Reproduces the Kaplan compute-optimal scaling law on OpenWebText2 and RefinedWeb, identifies last-layer FLOP counting, warmup length, and scale-dependent optimizer tuning as causes of its discrepancy with the Chinchilla law, and shows constant learning rate suffices to recover the Chinchilla exponent.
Problem
Kaplan et al. and Hoffmann et al. give substantially different compute-optimal scaling laws for model size versus compute budget, and the reason for the discrepancy is not well understood.
Approach
Systematically reproduces the Kaplan setup and then corrects one factor at a time (head FLOP counting, warmup duration proportional to model size, per-size tuning of learning rate, batch size, and AdamW beta2), showing these account for the gap to the Chinchilla law without needing learning rate decay.
Key results
- Reproducing the Kaplan setup gives exponent a=0.864 on OpenWebText2 and 0.835 on RefinedWeb, close to Kaplan's 0.88.
- Correcting head FLOPs, warmup, and cosine decay lowers the exponent to about 0.57-0.60, and tuning learning rate, batch size and beta2 with a constant learning rate yields 0.518 (OWT2) and 0.497 (RefinedWeb), matching Chinchilla's 0.5 within 0.6%.
- With tuned hyperparameters and constant learning rate, the predicted optimal model size at Chinchilla compute is within 15% of Chinchilla's size.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Count last-layer FLOPs instead of excluding the head from model size (ablation) | Kaplan-reproduction setup · OpenWebText2 and RefinedWeb | compute-optimal exponent a: 0.864 -> 0.699; 0.835 -> 0.706 | Helped |
| Set warmup tokens equal to model size rather than using a long fixed warmup (ablation) | Kaplan-reproduction setup with last-layer FLOPs counted · OpenWebText2 and RefinedWeb | compute-optimal exponent a: 0.699 -> 0.603; 0.706 -> 0.602 | Helped |
| Use per-model, per-compute cosine learning-rate decay (ablation) | Corrected-warmup setup · OpenWebText2 and RefinedWeb | compute-optimal exponent a: 0.603 -> 0.574; 0.602 -> 0.571 | Helped |
| Use per-model optimizer tuning for learning rate, batch size, and β₂ with a constant learning-rate schedule (ablation) | Corrected-warmup setup · OpenWebText2 and RefinedWeb | compute-optimal exponent a: 0.574 -> 0.518; 0.571 -> 0.497 | Helped |
| Tune AdamW β₂ above 0.95 for smaller batch sizes (ablation) | Small-batch AdamW training · Small batch sizes (128 and below) | performance | Helped |
| Use cosine learning-rate decay instead of a constant learning rate (ablation) | Tuned optimizer training · Compute-optimal loss across compute scales | loss | Helped |
| Shorten warmup duration (ablation) | Compute-optimal training · Low compute scales | loss: up to 0.5 nat per token | Helped |
| Tune optimizer hyperparameters (ablation) | Compute-optimal training · Low compute scales | loss: up to 0.5 nat per token | Helped |
| Use a constant learning-rate schedule instead of cosine decay (ablation) | Tuned optimizer training · Compute-optimal loss | loss | Hurt |
Surprising findings
- Careful learning rate decay is not essential for the Chinchilla scaling law; a constant learning rate reproduces it.
- Setting AdamW beta2 to 0.95 is suboptimal at small batch sizes; higher beta2 (0.99 or 0.999) is needed.
Limitations
- Limited compute scale (up to about 8e19 FLOPs) compared with Kaplan and Chinchilla experiments, so scaling effects may differ at larger scale.
- Hyperparameter sweep targeted only smaller models and trained each for 20 tokens per parameter, possibly favoring that token-to-parameter ratio.
- Compute-optimality concerns only pretraining loss, not downstream capabilities.
Experimental setup
Datasets: OpenWebText2 RefinedWeb
Benchmarks: held-out validation loss on OpenWebText2 and RefinedWeb
Metrics: compute-optimal model size exponent a R^2 of power-law fit predicted optimal model size at Chinchilla compute validation log loss
Baselines: Kaplan et al. scaling law Hoffmann et al. (Chinchilla) scaling law Besiroglu et al. adjusted scaling law
Builds on: Kaplan scaling laws Chinchilla scaling law OpenLM Llama
Closest papers by meaning
- Language models scale reliably with over-training and on downstream tasks (ICLR 2025 · similarity 0.93)
- Scaling Laws for Precision (ICLR 2025 · similarity 0.93)
- A Hitchhiker's Guide to Scaling Law Estimation (ICML 2025 · similarity 0.93)
- Scaling Inference-Efficient Language Models (ICML 2025 · similarity 0.92)
- Functional Scaling Laws in Kernel Regression: Loss Dynamics and Learning Rate Schedules (NeurIPS 2025 · similarity 0.91)
- Observational Scaling Laws and the Predictability of Langauge Model Performance (NeurIPS 2024 · similarity 0.91)
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.