Research map / Architectures and efficiency / RNNs and normalization
How to set AdamW's weight decay as you scale model and dataset size
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 AdamW's weights are an EMA of recent updates with timescale 1/(eta*lambda), finds the optimal timescale in epochs is roughly constant across dataset and model sizes, and derives weight decay scaling rules (decrease with dataset size, increase with width under muP).
Problem
How the optimal AdamW weight decay should scale with model width and dataset size, which is poorly understood, and why muP learning-rate transfer breaks down under AdamW with fixed weight decay.
Approach
Interprets AdamW's decoupled-weight-decay weights as an exponential moving average of recent updates, so the EMA timescale tau = 1/(eta*lambda) (measured in epochs) is the key hyperparameter; holding the epoch-timescale fixed gives scaling rules for weight decay with dataset size and model width (lambda ∝ s under muP).
Key results
- Optimal epoch-timescale tau_epoch is far more stable than optimal weight decay across ImageNet subset sizes (80 to 1.28M samples) and NanoGPT dataset fractions.
- Under muP with fixed weight decay the optimal base learning rate shifts dramatically with ResNet-18 width; scaling lambda ∝ s with eta ∝ 1/s restores stable optimal learning rate.
- Weight magnitudes in AdamW scale approximately as 1/lambda over 1000-epoch ResNet-18 CIFAR-10 runs.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Increase weight decay as dataset size increases to keep the EMA timescale in epochs approximately constant | AdamW · ImageNet subsets; OpenWebText subsets | optimal weight decay | Helped |
| Scale weight decay proportionally to model width while scaling the learning rate with 1/s | ResNet-18 on ImageNet 32x32 · 320K ImageNet subset | stability of optimal weight decay across widths | Helped |
| Scale weight decay proportionally to model width while scaling the learning rate with 1/s | 8-layer GPT · OpenWebText pre-training | optimal weight decay consistency across widths | Helped |
| Scale weight decay with model size while applying μP learning-rate scaling | ResNet-18 · CIFAR-10 and 320K ImageNet subset | stability of optimal base learning rate across model sizes | Helped |
| Scale weight decay with model size while applying μP learning-rate scaling | ViT · ImageNet 32x32 subset | stability of optimal base learning rate across model sizes | Helped |
| Add output batch normalization (ablation) | ResNet-18 · CIFAR-10 | performance dependence on learning rate at fixed timescale | Mixed |
| Use learning-rate-dependent initialization (ablation) | ResNet-18 · CIFAR-10 | performance dependence on learning rate at fixed timescale | Helped |
| Decouple the learning rate for normalization parameters from the learning rate for other parameters (ablation) | ResNet-18 · CIFAR-10 | performance dependence on learning rate at fixed timescale | Helped |
| Add layer normalization after all linear layers (ablation) | ViT · CIFAR-10 | performance dependence on learning rate at fixed timescale | Helped |
| Keep weight decay fixed while using μP learning-rate scaling (ablation) | ResNet-18 · CIFAR-10 and 320K ImageNet subset | stability of optimal base learning rate across model sizes | Hurt |
| Keep weight decay fixed while scaling model width (ablation) | ResNet-18 · 320K ImageNet subset | stability of optimal weight decay across widths | Hurt |
| Keep weight decay fixed while scaling model width (ablation) | 8-layer GPT · OpenWebText pre-training | optimal weight decay across widths | Hurt |
Surprising findings
- Under fixed weight decay, muP learning-rate transfer breaks down for AdamW, and increasing weight decay with width restores it.
- The optimal weight decay increases with model width under muP while it falls with dataset size.
Limitations
- Scaling experiments are small-scale (ResNet-18, ViT on 32x32 ImageNet, 124M NanoGPT) and LLM runs use a single seed.
- EMA interpretation is approximate because AdamW updates are correlated with the current weights.
- Super-linear relationship between width and optimal lambda not established; larger-scale experiments needed.
Experimental setup
Datasets: CIFAR-10 ImageNet 32x32 OpenWebText
Benchmarks: CIFAR-10 ImageNet 32x32 OpenWebText validation loss
Metrics: test accuracy test loss validation loss
Baselines: standard muP scaling with fixed weight decay
Builds on: muP (maximal update parametrization) AdamW (decoupled weight decay)
Closest papers by meaning
- Power Lines: Scaling laws for weight decay and batch size in LLM pre-training (NeurIPS 2025 · similarity 0.88)
- Straight to Zero: Why Linearly Decaying the Learning Rate to Zero Works Best for LLMs (ICLR 2025 · similarity 0.86)
- Implicit Bias of AdamW: $\ell_\infty$-Norm Constrained Optimization (ICML 2024 · similarity 0.86)
- Scaling Exponents Across Parameterizations and Optimizers (ICML 2024 · similarity 0.86)
- Rethinking Weight Decay for Robust Fine-Tuning of Foundation Models (NeurIPS 2024 · similarity 0.86)
- Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations (NeurIPS 2024 · similarity 0.85)
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.