# How to set AdamW's weight decay as you scale model and dataset size

arXiv: https://arxiv.org/abs/2405.13698 · HTML: https://mlsanity.com/paper/2405.13698/

## Key facts

- Cited by 43 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 50 references that are in the map.
- First posted on arXiv 22 May 2024.

## Builds on (most cited first, 30 of 30)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [OPT: Open Pre-trained Transformer Language Models](https://mlsanity.com/paper/2205.01068/) (2022)
- [Scaling Vision Transformers to 22 Billion Parameters](https://mlsanity.com/paper/2302.05442/) (2023)
- [Symbolic Discovery of Optimization Algorithms](https://mlsanity.com/paper/2302.06675/) (2023)
- [A Downsampled Variant of ImageNet as an Alternative to the CIFAR datasets](https://mlsanity.com/paper/1707.08819/) (2017)
- [S4L: Self-Supervised Semi-Supervised Learning](https://mlsanity.com/paper/1905.03670/) (2019)
- [Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer](https://mlsanity.com/paper/2203.03466/) (2022)
- [An Empirical Model of Large-Batch Training](https://mlsanity.com/paper/1812.06162/) (2018)
- [Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training](https://mlsanity.com/paper/2305.14342/) (2023)
- [Measuring the Effects of Data Parallelism on Neural Network Training](https://mlsanity.com/paper/1811.03600/) (2018)
- [Small-scale proxies for large-scale Transformer training instabilities](https://mlsanity.com/paper/2309.14322/) (2023)
- [Dissecting Adam: The Sign, Magnitude and Variance of Stochastic Gradients](https://mlsanity.com/paper/1705.07774/) (2017)
- [L2 Regularization versus Batch and Weight Normalization](https://mlsanity.com/paper/1706.05350/) (2017)
- [An Exponential Learning Rate Schedule for Deep Learning](https://mlsanity.com/paper/1910.07454/) (2019)
- [Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Model](https://mlsanity.com/paper/1907.04164/) (2019)
- [Three Mechanisms of Weight Decay Regularization](https://mlsanity.com/paper/1810.12281/) (2018)
- [On the SDEs and Scaling Rules for Adaptive Gradient Algorithms](https://mlsanity.com/paper/2205.10287/) (2022)
- [Why Do We Need Weight Decay in Modern Deep Learning?](https://mlsanity.com/paper/2310.04415/) (2023)
- [Depthwise Hyperparameter Transfer in Residual Networks: Dynamics and Scaling Limit](https://mlsanity.com/paper/2309.16620/) (2023)
- [Theoretical Analysis of Auto Rate-Tuning by Batch Normalization](https://mlsanity.com/paper/1812.03981/) (2018)
- [Highly Scalable Deep Learning Training System with Mixed-Precision: Training ImageNet in Four Minutes](https://mlsanity.com/paper/1807.11205/) (2018)
- [Reconciling Modern Deep Learning with Traditional Optimization Analyses: The Intrinsic Learning Rate](https://mlsanity.com/paper/2010.02916/) (2020)
- [Rotational Equilibrium: How Weight Decay Balances Learning Across Neural Networks](https://mlsanity.com/paper/2305.17212/) (2023)
- [Robust Training of Neural Networks Using Scale Invariant Architectures](https://mlsanity.com/paper/2202.00980/) (2022)
- [How to Scale Your EMA](https://mlsanity.com/paper/2307.13813/) (2023)
- [Super Consistency of Neural Network Landscapes and Learning Rate Transfer](https://mlsanity.com/paper/2402.17457/) (2024)
- [Batch size invariant Adam](https://mlsanity.com/paper/2402.18824/) (2024)

## Built on by (most cited first, 30 of 43)

- [Practical Efficiency of Muon for Pretraining](https://mlsanity.com/paper/2505.02222/) (2025)
- [Don't be lazy: CompleteP enables compute-efficient deep transformers](https://mlsanity.com/paper/2505.01618/) (2025)
- [Power Lines: Scaling Laws for Weight Decay and Batch Size in LLM Pre-training](https://mlsanity.com/paper/2505.13738/) (2025)
- [Scaling Optimal LR Across Token Horizons](https://mlsanity.com/paper/2409.19913/) (2024)
- [Straight to Zero: Why Linearly Decaying the Learning Rate to Zero Works Best for LLMs](https://mlsanity.com/paper/2502.15938/) (2025)
- [u-μP: The Unit-Scaled Maximal Update Parametrization](https://mlsanity.com/paper/2407.17465/) (2024)
- [Communication-Efficient Language Model Training Scales Reliably and Robustly: Scaling Laws for DiLoCo](https://mlsanity.com/paper/2503.09799/) (2025)
- [Completed Hyperparameter Transfer across Modules, Width, Depth, Batch and Duration](https://mlsanity.com/paper/2512.22382/) (2025)
- [Rethinking Conventional Wisdom in Machine Learning: From Generalization to Scaling](https://mlsanity.com/paper/2409.15156/) (2024)
- [Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks](https://mlsanity.com/paper/2507.02119/) (2025)
- [Hyperparameter Transfer Enables Consistent Gains of Matrix-Preconditioned Optimizers Across Scales](https://mlsanity.com/paper/2512.05620/) (2025)
- [Weight Decay may matter more than muP for Learning Rate Transfer in Practice](https://mlsanity.com/paper/2510.19093/) (2025)
- [MuLoCo: Muon is a practical inner optimizer for DiLoCo](https://mlsanity.com/paper/2505.23725/) (2025)
- [Fantastic Pretraining Optimizers and Where to Find Them II: Hyperball Optimization](https://mlsanity.com/paper/2606.16899/) (2026)
- [Training Dynamics Impact Post-Training Quantization Robustness](https://mlsanity.com/paper/2510.06213/) (2025)
- [How Much Can We Forget about Data Contamination?](https://mlsanity.com/paper/2410.03249/) (2024)
- [Hyperparameter Transfer with Mixture-of-Expert Layers](https://mlsanity.com/paper/2601.20205/) (2026)
- [Rethinking Language Model Scaling under Transferable Hypersphere Optimization](https://mlsanity.com/paper/2603.28743/) (2026)
- [Decoder-Hybrid-Decoder Architecture for Efficient Reasoning with Long Generation](https://mlsanity.com/paper/2507.06607/) (2025)
- [Scaling with Collapse: Efficient and Predictable Training of LLM Families](https://mlsanity.com/paper/2509.25087/) (2025)
- [FOCUS: First Order Concentrated Updating Scheme](https://mlsanity.com/paper/2501.12243/) (2025)
- [Robust Layerwise Scaling Rules by Proper Weight Decay Tuning](https://mlsanity.com/paper/2510.15262/) (2025)
- [VideoFlexTok: Flexible-Length Coarse-to-Fine Video Tokenization](https://mlsanity.com/paper/2604.12887/) (2026)
- [μnit Scaling: Simple and Scalable FP8 LLM Training](https://mlsanity.com/paper/2502.05967/) (2025)
- [How to Set the Learning Rate for Large-Scale Pre-training?](https://mlsanity.com/paper/2601.05049/) (2026)
- [Efficient Hyperparameter Tuning via Trajectory Invariance Principle](https://mlsanity.com/paper/2509.25049/) (2025)
- [Weight Decay Improves Language Model Plasticity](https://mlsanity.com/paper/2602.11137/) (2026)
- [A Two-Parameter Weibull Framework for Diagnosing Transformer Weight Distributions](https://mlsanity.com/paper/2605.18898/) (2026)
- [Improving Neural Network Training by Decoupling the Magnitude and Direction of Weight Vectors](https://mlsanity.com/paper/2606.25971/) (2026)
- [Predicting Training Re-evaluation Curves Enables Effective Data Curriculums for LLMs](https://mlsanity.com/paper/2509.25380/) (2025)

## Abstract (arXiv)

The scaling of the optimal AdamW weight decay hyperparameter with model and dataset size is critical as we seek to build larger models, but is poorly understood. We show that weights learned by AdamW can be understood as an exponential moving average (EMA) of recent updates. This gives critical insights for how to set the weight decay in AdamW, and how the weight decay should scale with model and dataset size. In particular, the key hyperparameter for an exponential moving average is the EMA timescale. Intuitively, the EMA timescale can be understood as the number of recent iterations the EMA averages over. We find that the optimal timescale, measured in epochs, is roughly constant as we change model and dataset size. Moreover, given a learning rate, there is a one-to-one mapping from the EMA timescale to the weight decay hyperparameter. Thus, if the optimal EMA timescale is constant, that implies that as the dataset size increases, the optimal weight decay should fall and as the model size increases, the optimal weight decay should increase (if we follow the muP recommendation for scaling the learning rate). We validate these scaling rules on ResNet-18 and Vision Transformers trained on CIFAR-10 and ImageNet, and on NanoGPT pre-training on OpenWebText. Finally, we found that as training progresses, muP's learning rate scaling breaks down for AdamW unless weight decay is scaled appropriately.