Scale Weight Decay and Train Better

2026-07-26

AI Summary

This work addresses a critical limitation in conventional model training: fixed decoupled weight decay continually compresses network weights, introducing asymptotic bias and deviating from the original optimization objective. To resolve this, the authors propose Scaled Weight Decay (SW), which dynamically adjusts decay strength according to the peak learning rate and, for the first time, incorporates the Robbins–Monro condition into weight decay design to achieve asymptotically unbiased regularization. When integrated with stochastic gradient descent and the non-Euclidean optimizer Muon, Muon-SW accelerates convergence by 30% compared to standard Muon—reaching the same validation loss across Mixture-of-Experts (MoE) models ranging from 72M to 930M parameters—while maintaining optimization stability.

Abstract

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 $\mathbf{30\%}$ faster at our largest scale across models from $72 - 930$ million parameters trained at $\sim 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.

Problem

• weight decay
• scaling laws
• asymptotic bias
• neural network training
• optimization target

Innovation

• scaled weight decay
• asymptotic stationarity
• learning rate scaling
• mixture-of-experts
• optimization bias