cs.LGJul 26, 2026

Scale Weight Decay and Train Better

Authors: Anuj Apte

Organizations: Global Technology Applied Research, JPMorganChase, New York, NY 10001

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η/η_{\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%\mathbf{30\%} faster at our largest scale across models from 7293072 - 930 million parameters trained at 600\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.

Explore similar work

CardsList
  1. Muown: Row-Norm Control for Muon Optimization

    May 11, 2026Kai Lion, Florian Hübler, Bingcong Li +2Muon OptimizerMuon