Scale Weight Decay and Train Better
- URL: http://arxiv.org/abs/2607.23777v1
- Date: Sun, 26 Jul 2026 17:55:40 GMT
- Title: Scale Weight Decay and Train Better
- Abstract summary: We propose to scale weight decay by a fraction of the peak learning rate $/_max$.<n>We prove that this preserves the stationarity guarantees of the corresponding unregularized methods for both a gradient descent and a non-Euclidean Muon.<n>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 roughly constant value.
- Score: 0.0
- License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
- 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.
Related papers
- Fantastic Pretraining Optimizers and Where to Find Them II: Hyperball Optimization [67.29294986811314]
Matrix baseds such as Muon can substantially speed up language model pretraining, but their gains over AdamW are observed to shrink when using standard constant weight decay.<n>We propose Hyperball, a simple wrapper that addresses this issue.<n>On Qwen3 style models up to 1.2B parameters, Muon Hyperball achieves 20-30% token equivalent speedup over weight decay baselines.
arXiv Detail & Related papers (2026-06-15T16:09:02Z) - How Neural Reward Models Learn Features for Policy Optimization: A Single-Index Analysis [53.063298916923976]
We study the feedback in a Gaussian single-index model with $r*(x) = *(langle *, xrangle)$ and $x sim N(0, I_d)$.<n>We analyze a two-stage neural reward model that first learns the hidden direction $*$ from reward-weighted samples and then fits the readout layer by weighted ridge regression.
arXiv Detail & Related papers (2026-05-23T22:00:38Z) - Muown: Row-Norm Control for Muon Optimization [41.541871767555484]
Muon is a strong competitor to AdamW for language model pre-training.<n>Muown improves perplexity over Muon, SOAP, AdamW, and Lion.
arXiv Detail & Related papers (2026-05-11T16:26:06Z) - Learning to Forget: Continual Learning with Adaptive Weight Decay [30.558842015806363]
Continual learning agents with finite capacity must balance acquiring new knowledge with retaining the old.<n>Weight decay, viewed as a mechanism for forgetting, can serve this role by gradually discarding information stored in the weights.<n>We introduce Forgetting through Adaptive Decay (FADE), which adapts per- parameter weight decay rates online via approximate meta-gradient descent.
arXiv Detail & Related papers (2026-04-29T18:01:13Z) - PathProx: A Proximal Gradient Algorithm for Weight Decay Regularized
Deep Neural Networks [25.114642281756495]
Weight decay is one of the most widely used forms of regularization in deep learning.
This paper argues that gradient descent may be an inefficient algorithm for this objective.
For neural networks with ReLU activations, solutions to the weight decay objective are equivalent to those of a different objective.
arXiv Detail & Related papers (2022-10-06T17:22:40Z) - Long-Tailed Recognition via Weight Balancing [66.03068252811993]
Naive training produces models that are biased toward common classes in terms of higher accuracy.
We investigate three techniques to balance weights, L2-normalization, weight decay, and MaxNorm.
Our approach achieves the state-of-the-art accuracy on five standard benchmarks.
arXiv Detail & Related papers (2022-03-27T03:26:31Z) - Robust Implicit Networks via Non-Euclidean Contractions [63.91638306025768]
Implicit neural networks show improved accuracy and significant reduction in memory consumption.
They can suffer from ill-posedness and convergence instability.
This paper provides a new framework to design well-posed and robust implicit neural networks.
arXiv Detail & Related papers (2021-06-06T18:05:02Z) - On the Overlooked Pitfalls of Weight Decay and How to Mitigate Them: A Gradient-Norm Perspective [90.39123717733334]
We present the first practical scheduler for weight decay, called the Scheduled Weight Decay (SWD) method.
Our experiments also support that SWD indeed mitigates large gradient norms and often significantly outperforms the conventional constant weight decay strategy for Adaptive Moment Estimation (Adam)
arXiv Detail & Related papers (2020-11-23T00:39:49Z) - Neural networks with late-phase weights [66.72777753269658]
We show that the solutions found by SGD can be further improved by ensembling a subset of the weights in late stages of learning.
At the end of learning, we obtain back a single model by taking a spatial average in weight space.
arXiv Detail & Related papers (2020-07-25T13:23:37Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.