Scale Weight Decay and Train Better
Not on the map. Why?
Abstract (arXiv)
The discovery of scaling laws has motivated training neural networks on ever increasing quantities of data. This is typically done with a constant decoupled weight decay which causes the network weights to shrink steadily over the course of training. Taking inspiration from the Robbins--Monro conditions, we propose to scale weight decay by the fraction of the peak learning rate η/ηmax. We prove that this scaled weight decay preserves the asymptotic stationarity guarantees of the corresponding unregularized methods for both stochastic gradient descent and the non-Euclidean spectral optimizer Muon, thereby avoiding the additional asymptotic bias introduced by constant decoupled weight decay. This retains the stability benefits of weight decay without changing the asymptotic optimization target. Using a steady-state analysis, we explain why under standard weight decay the weight norm shrinks steadily as training proceeds, whereas under scaled weight decay it settles to a roughly constant value. When applied to the training of mixture-of-experts models, Muon with scaled weight decay (Muon-SW) consistently outpaces Muon with identical hyperparameters, reaching the same validation loss 30% faster at our largest scale across models from 72 - 930 million parameters trained at ~ 600 tokens per active parameter. If this trend continues to hold, the method promises to substantially accelerate the pre-training of frontier models while requiring only a few lines of code to implement.
Key facts. Cited by 2 papers in the map (data through 14 Oct 2026). Builds on 58 of its 70 references that are in the map. First posted on arXiv 26 Jul 2026.
Builds on
58 of 70 references.
References · 58
- 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 30,186
- 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,305
- 2022Training language models to follow instructions with human feedbackcited by 15,971
- 2023Llama 2: Open Foundation and Fine-Tuned Chat Modelscited by 13,028
- 2024DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Modelscited by 9,379
- 2023Direct Preference Optimization: Your Language Model is Secretly a Reward Modelcited by 8,065
- 2020Scaling Laws for Neural Language Modelscited by 6,586
- 2025DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learningcited by 6,330
- 2016SGDR: Stochastic Gradient Descent with Warm Restartscited by 5,045
- 2023Mamba: Linear-Time Sequence Modeling with Selective State Spacescited by 4,718
- 2021RoFormer: Enhanced Transformer with Rotary Position Embeddingcited by 4,703
- 2017Deep reinforcement learning from human preferencescited by 3,912
38 more
- 2017Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layercited by 3,509
- 2021Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsitycited by 3,073
- 2022Training Compute-Optimal Large Language Modelscited by 3,019
- 2020GLU Variants Improve Transformercited by 1,804
- 2019Root Mean Square Layer Normalizationcited by 1,480
- 2024Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Dualitycited by 1,336
- 2024The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scalecited by 1,160
- 2016Optimization Methods for Large-Scale Machine Learningcited by 1,150
- 2026DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligencecited by 1,040
- 2021Scaling Vision Transformerscited by 999
- 2024Gated Delta Networks: Improving Mamba2 with Delta Rulecited by 532
- 2025GLM-4.5: Agentic, Reasoning, and Coding (ARC) Foundation Modelscited by 467
- 2025Muon is Scalable for LLM Trainingcited by 428
- 2018Shampoo: Preconditioned Stochastic Tensor Optimizationcited by 426
- 2021Explaining Neural Scaling Lawscited by 363
- 2022Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfercited by 288
- 2024SOAP: Improving and Stabilizing Shampoo using Adamcited by 192
- 2024Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modelingcited by 174
- 2024Scaling Laws for Fine-Grained Mixture of Expertscited by 173
- 2017L2 Regularization versus Batch and Weight Normalizationcited by 124
- 2025The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithmcited by 110
- 2023A Spectral Condition for Feature Learningcited by 109
- 2023Tensor Programs VI: Feature Learning in Infinite-Depth Neural Networkscited by 97
- 2025Dion: Distributed Orthonormalized Updatescited by 68
- 2025Benchmarking Optimizers for Large Language Model Pretrainingcited by 68
- 2025ASGO: Adaptive Structured Gradient Optimizationcited by 57
- 2023Rotational Equilibrium: How Weight Decay Balances Learning Across Neural Networkscited by 53
- 2024Why Warmup the Learning Rate? Underlying Mechanisms and Improvementscited by 46
- 2025Towards Greater Leverage: Scaling Laws for Efficient Mixture-of-Experts Language Modelscited by 45
- 2024Scaling Optimal LR Across Token Horizonscited by 40
- 2025MuonBP: Faster Muon via Block-Periodic Orthogonalizationcited by 35
- 2025On the Convergence of Muon and Beyondcited by 31
- 2025Convergence Bound and Critical Batch Size of Muon Optimizercited by 28
- 2025Why Gradients Rapidly Increase Near the End of Trainingcited by 22
- 2026Muon2: Boosting Muon via Adaptive Second-Moment Preconditioningcited by 17
- 2026Fantastic Pretraining Optimizers and Where to Find Them II: Hyperball Optimizationcited by 16
- 2025Cautious Weight Decaycited by 13
- 2025Dion2: A Simple Method to Shrink Matrix in Muoncited by 13
50 most cited of 58 All 58 (JSON)
Built on by
Citing papers per year: 2026: 2.
2 papers.
- 2026Convergence of Practical Muoncited by 0
- 2026Second-Moment Stochastic Approximation Methodscited by 0
- 2026Second-Moment Stochastic Approximation Methods29 Sep 2026
- 2026Convergence of Practical Muon27 Sep 2026
Similar, not cited
Close in wording, no citation either way
- 2025Muon: Training and Trade-offs with Latent Attention and MoE
- 2026Spectral Weight Decay: Inducing Low-Rank Structure in Neural Network Weights
- 2026An Isotropy-Preserving Spectral Cap for Muon: Theory and Three Case Studies
- 2026Training-Loss Guarantees for Muon with Finite-Step Newton--Schulz Orthogonalization
- 2026Improving Neural Network Training by Decoupling the Magnitude and Direction of Weight Vectors
- 2026Muon Meets Mamba: Spectral Optimization for State Space Models
- 2025LiMuon: Light and Fast Muon Optimizer for Large Models
- 2024OvSW: Overcoming Silent Weights for Accurate Binary Neural Networks