Scale Weight Decay and Train Better

Standard neural network training typically employs decoupled weight decay with a constant coefficient. This practice causes network weights to shrink steadily throughout training, which can introduce asymptotic bias and prevent the model from converging to the true minimizer of the loss function. The author asks whether weight decay can be applied in a more principled manner that improves training speed while preserving the stationarity guarantees of unregularized optimization methods.

Inspired by the Robbins-Monro conditions for stochastic approximation, the author proposes scaling the weight decay coefficient by the ratio of the current learning rate to the peak learning rate (η/η_max). This modification ensures that the weight decay term enters the optimization dynamics at a higher order (O(η^2)) as the learning rate decreases, effectively reducing the "shrinkage" pressure late in training. The author integrates this into the Muon optimizer—a spectral-norm-constrained optimizer for matrix-valued weights—to create Muon-SW. The method is evaluated on mixture-of-experts (MoE) models ranging from 72 to 930 million parameters.

• Improved Convergence: Muon-SW consistently outperforms standard Muon with identical hyperparameters. At the largest scale tested, it reaches the same validation loss 30% faster.
• Weight Stability: Unlike standard weight decay, which causes the root mean squared (RMS) norm of weights to decline continuously, scaled weight decay allows the weight norm to settle to a stable, roughly constant value.
• Theoretical Guarantees: The author proves that scaled weight decay preserves the asymptotic stationarity guarantees of unregularized methods. Conversely, constant weight decay is shown to converge to a biased point that is not the minimizer of the original objective, particularly when the target weights have large spectral norms.