Research map / Architectures and efficiency / RNNs and normalization
Power Lines: Scaling laws for weight decay and batch size in LLM pre-training
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 the optimal AdamW timescale tau = B/(eta*lambda*D) follows a power law in tokens-per-parameter, so optimal weight decay can be predicted for any N, D, B, and that optimal and critical batch sizes scale as power laws in dataset size D.
Problem
Choosing learning rate and weight decay (and batch size) for LLM pre-training at scales where hyperparameter tuning is infeasible; prior scaling laws for batch size used compute or loss and held weight decay fixed.
Approach
Treats the AdamW timescale tau = B/(eta*lambda*D) as the scaling variable, fits its optimum as a power law in tokens-per-parameter (D/N) using muP-trained GPT2-like models, and fits optimal batch size and critical batch size as power laws in D rather than compute or loss.
Key results
- Optimal tau decreases as a power law in TPP (R^2=0.975), from about 1.0 at 1 TPP to 0.01 at 1000 TPP, with exponent around -0.5, and generalizes to larger-scale runs spanning about 3 orders of magnitude in compute.
- Optimal batch size scales as a power law in D (R^2=0.984) and critical batch size as a power law in D_min (R^2=0.940), with fitted CBS exponent 0.47 close to prior work's 0.462.
- Tuning weight decay rather than learning rate as batch size grows was strictly superior in 6 of 8 settings at the 610M 20TPP scale.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Tune λ instead of η while holding η at the μP proxy-tuned value | AdamW pre-training · 610M 20TPP, B=63–8064 | validation loss: Tuning λ was strictly superior in 6 of 8 cases | Helped |
| Tune λ instead of η while holding η at the μP proxy-tuned value | AdamW pre-training · 111M 200TPP, varying D | validation loss: 2.810 -> 2.805 | Helped |
| Scale optimal λ linearly with batch size below the critical batch size | AdamW pre-training · 610M 20TPP, B in [63, 2016] | optimal τ: 0.21 | No clear effect |
| Scale optimal τ with tokens per parameter rather than keeping it constant | AdamW pre-training · LLM pre-training across TPP | R²: R²=0.975 | Helped |
| Tune λ for each batch size rather than hold it fixed | Batch-size scaling-law fits · LLM pre-training across D | fit quality (R²) | Helped |
| Use D (or dmin) rather than compute C or loss L as the scaling variable | Optimal and critical batch size · LLM pre-training across model scales | power-law fit R²: R²=0.984; R²=0.940 | Helped |
| Use larger batch sizes beyond bopt | LLM pre-training · B > cbs | training efficiency | Hurt |
| Train overtrained rather than undertrained models under the FLOPs/B time model | Time–compute Pareto configurations · Target loss L̂=2.6 | Pareto optimality | Helped |
| Start large-batch training from a 20 TPP+ base setting rather than target 20 actual TPP | Time–compute Pareto configurations · Large B | Pareto optimality | Helped |
Surprising findings
- Optimal learning rate fails to scale with batch size up to CBS and training becomes unstable above a maximum learning rate, so weight decay is the more flexible knob.
- Loss degrades at very small batch sizes rather than approaching the asymptotic minimum data predicted by theory.
Limitations
- Single dataset (SlimPajama), vocabulary, context length, and LR schedule; results may not transfer.
- Full cross-product of learning rate and weight decay not swept at scale; learning rate fixed via muP proxy.
- Performance degrades at very small batch sizes, and the EMA perspective is only a conceptual approximation.
Experimental setup
Datasets: SlimPajama
Benchmarks: held-out SlimPajama validation loss (1.1B tokens)
Metrics: validation cross-entropy loss R^2 bootstrapped parameter percentiles
Baselines: DeepSeek LLM batch size law Zhang et al. critical batch size law fixed weight decay 0.1 default eta-tuning at fixed lambda
Builds on: maximal update parameterization AdamW Scaling laws for neural language models Batch size scaling laws for LLMs (critical batch size) EMA view of AdamW (timescale)
Closest papers by meaning
- Scaling Optimal LR Across Token Horizons (ICLR 2025 · similarity 0.91)
- Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations (NeurIPS 2024 · similarity 0.90)
- How Does Critical Batch Size Scale in Pre-training? (ICLR 2025 · similarity 0.90)
- A Multi-Power Law for Loss Curve Prediction Across Learning Rate Schedules (ICLR 2025 · similarity 0.89)
- Resolving Discrepancies in Compute-Optimal Scaling of Language Models (NeurIPS 2024 · similarity 0.89)
- Scaling Laws for Sparsely-Connected Foundation Models (ICLR 2024 · 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.