Scale Weight Decay and Train Better

Author: Anuj Apte (Global Technology Applied Research, JPMorganChase, New York)

---

Summary (Overview)

• Core proposal: Scale the decoupled weight decay coefficient by the learning-rate fraction $\eta_t / \eta_{\text{max}}$, so the shrink term enters at order $\mathcal{O}(\eta_t^2)$ instead of $\mathcal{O}(\eta_t)$, satisfying Robbins–Monro summability conditions.
• Key empirical finding: Muon with scaled weight decay (Muon-SW) trains mixture-of-experts (MoE) transformers 21.7%–29.4% faster (in optimizer steps) than standard Muon across model sizes 72M–932M parameters, with the largest speedup at width 1024 (~30% fewer steps to reach the same validation loss).
• Theoretical contribution: Proves that scaled weight decay preserves the asymptotic stationarity guarantees of unregularized SGD and Muon, while constant weight decay provably converges to a non-minimizer on simple shifted quadratic objectives (Propositions 3.1 and 3.2).
• Mechanistic insight: A steady-state analysis shows constant weight decay drives the weight norm to a value roughly 14× smaller than scaled weight decay at the same hyperparameters, explaining the norm collapse seen in standard training.
• Practical impact: Requires only a one-line modification to existing training pipelines with no additional computational overhead; the same principle applies to other matrix-valued optimizers (Shampoo, SOAP, Aurora).

---

Introduction and Theoretical Foundation

Background and Motivation

Scaling laws have driven training of neural networks on ever-increasing data volumes. Modern pre-training uses decoupled weight decay:

$$W_{t+1} = (1 - \eta_{t} \lambda) W_{t} - \eta_{t} U_{t}, \tag{1}$$

where $W_t$, $\eta_t$, $U_t$ are the weight, learning rate, and pre-conditioned gradient update at step $t$, and $\lambda$ is the weight decay coefficient.

Key observation: With constant weight decay, if training converges, it converges to a point where $W_t = -U_t/\lambda$, which may be far from an optimum of the original loss. The convergence guarantees apply to modified optimization dynamics, not the original objective.

The Proposed Fix

The author proposes scaling the decay by the learning-rate fraction:

$$W_{t+1} = \left(1 - \frac{\eta_{t}^{2}}{\eta_{\mathrm{max}}} \lambda\right) W_{t} - \eta_{t} U_{t}. \tag{2}$$

This is motivated by the Robbins–Monro conditions:

$$\sum_{t} \eta_{t} = \infty \quad \text{and} \quad \sum_{t} \eta_{t}^{2} < \infty. \tag{6}$$

• An $\mathcal{O}(\eta_t)$ perturbation (constant decay) can move the stationary set.
• An $\mathcal{O}(\eta_t^2)$ perturbation (scaled decay) is summable and preserves the stationarity guarantee.

Muon-SW Algorithm

Muon orthogonalizes the momentum buffer via the polar factor: $O_t = \mathrm{polar}(M_t) = UV^\top$ (from SVD $M_t = U\Sigma V^\top$). The standard Muon update with constant decay is:

$$W_{t+1} = (1 - \eta_{t} \lambda) W_{t} - \widehat{\eta}_{t} O_{t}, \tag{10}$$

Muon-SW replaces this with:

$$W_{t+1} = \left(1 - \frac{\eta_{t}^{2}}{\eta_{\mathrm{max}}} \lambda\right) W_{t} - \widehat{\eta}_{t} O_{t}, \tag{11}$$

where $\widehat{\eta}_t = 0.2 \eta_t \sqrt{\max(m,n)}$ rescales the learning rate for a matrix of shape $m \times n$.

---

Methodology

Theoretical Analysis Framework

The paper establishes convergence guarantees under the following setup:

• Objective: $f(W) = \mathbb{E}_{z \sim \mathcal{D}}[\ell(W; z)]$ (expected loss over data distribution)
• Gradient estimate: $G_t = \nabla f(W_t) + \xi_t$ with zero-mean noise of bounded variance
• Momentum buffer: $M_t = \beta M_{t-1} + G_t$

Key Results

Theorem 3.1 (SGD with scaled weight decay): Under L-smoothness, bounded iterates, and Robbins–Monro conditions, the iterates from

$$x_{t+1} = x_{t} - \eta_{t} g_{t} - q \eta_{t}^{2} x_{t}, \quad q = \frac{\lambda}{\eta_{\mathrm{max}}} \tag{23}$$

satisfy:

$$\frac{1}{S_T} \sum_{t=0}^{T-1} \eta_{t} \mathbb{E}\|\nabla F(x_{t})\|_2^2 \rightarrow 0, \quad S_T = \sum_{t=0}^{T-1} \eta_{t} \tag{24}$$

Theorem 3.2 (Strongly convex case): For a $\mu$-strongly convex objective, $\mathbb{E}\|x_t - x^\star\|_2^2 \rightarrow 0$ with rate $O(1/t)$ for the harmonic schedule.

Theorem 3.3 (Muon with scaled weight decay): Muon-SW satisfies:

$$\limsup_{T \to \infty} \frac{1}{S_T} \sum_{t=0}^{T-1} \eta_{t} \mathbb{E}\|\nabla f(W_t)\|_* \leq 2\sigma \sqrt{\frac{(1-\beta) r}{B}}, \tag{27}$$

where $\|\cdot\|_*$ is the nuclear norm, $r = \min(m,n)$, $B$ is the batch size, and $\sigma$ is the noise standard deviation.

Negative Results for Constant Weight Decay

Proposition 3.1 (SGD): On the shifted quadratic $f(w) = \frac{a}{2}(w-c)^2$, gradient descent with constant weight decay converges to:

$$\bar{w} = \frac{ac}{a+\lambda} \neq c, \tag{14}$$

with a strictly positive loss gap:

$$f(\bar{w}) - f(c) = \frac{a}{2}\left(\frac{\lambda}{a+\lambda}\right)^2 c^2 > 0. \tag{15}$$

Proposition 3.2 (Muon): On the shifted matrix quadratic $f(W) = \frac{1}{2}\|W - A\|_F^2$, constant-decay Muon confines every singular value of $W_t$ to a ball of radius $1/\lambda$. If $\sigma_{\max}(A) > 1/\lambda$, the iterates never converge to the minimizer.

Steady-State Analysis of Weight Norm

For a single matrix weight, the RMS norm $r_t = \|W_t\|_F/\sqrt{mn}$ evolves as:

$$r_{t+1}^2 \approx (1 - 2\varphi_t) r_t^2 - 2(0.2\eta_t) a_t r_t + (0.2\eta_t)^2, \tag{34}$$

where $a_t$ is the weight–update alignment (cosine) and $\varphi_t = \lambda\eta_t$ (constant) or $\varphi_t = \lambda\eta_t^2/\eta_{\text{max}}$ (scaled).

The quasi-steady target is:

$$r_t^\star = \frac{0.2\eta_t}{2\varphi_t}\left(|a_t| + \sqrt{a_t^2 + 2\varphi_t}\right). \tag{36}$$

Key limits as $\eta_t \to 0$:

• Scaled decay: $r_t^\star \rightarrow \frac{0.2\eta_{\text{max}}}{2\lambda}(k + \sqrt{k^2 + 2\lambda/\eta_{\text{max}}}) \approx 0.20$ (plateau)
• Constant decay: $r_t^\star \rightarrow \frac{0.2a_0}{\lambda} \approx 0.015$ (collapse)

The constant-decay limit is ~14× smaller than the scaled-decay limit at the same hyperparameters.

---

Empirical Validation / Results

Experimental Setup

• Architecture: LLaMA-style MoE decoder with 8 experts, top-2 routing
• Dataset: FineWeb, tokenized with GPT-2 tokenizer
• Batch: 256 sequences × 2,048 tokens = 524,288 tokens/step
• Training: 600–650 tokens per active parameter
• Hyperparameters: $\lambda = 0.1$, cosine LR decay to 1/10 of peak, muP-transferred peak learning rates

Main Results

Table 1: Muon vs. Muon-SW for MoE training

Key observations:

• Muon-SW improves the best validation loss at every width
• Speed-up (reduction in steps to reach Muon's best loss) rises from 21.7% at width 256 to 29.4% at width 1024
• The trend suggests larger benefits at greater model sizes

Weight Norm Behavior

• Constant decay: Weight norm peaks early, then steadily decays (~60% drop from peak by end of training)
• Scaled decay: Weight norm settles to a roughly constant value
• The steady-state analysis accurately predicts the separation between the two regimes (Figure 3)

---

Theoretical and Practical Implications

Theoretical Significance

1. Preserves stationarity guarantees: Scaled weight decay retains the asymptotic convergence properties of unregularized optimizers—the first principled weight decay variant with this property for Muon.

2. Exposes flaws in prior work: Proposition 3.2 shows that the convergence guarantee for Muon with constant weight decay claimed by Sato et al. [36] is incorrect; their Lemma C.1 bounds $\|W_t\|_{\text{op}}$ by $1/\lambda$ rather than 1.

3. Unified framework: The Robbins–Monro summability argument provides a principled basis for the $\eta_t/\eta_{\text{max}}$ scaling rule, which Defazio [22] introduced for a different motivation (counteracting gradient norm growth).

Practical Implications

• Training efficiency: At frontier scale, a 30% reduction in training steps translates to substantial cost savings.
• Implementation simplicity: Requires only a one-line change to existing training pipelines with zero computational overhead.
• General applicability: The principle extends beyond Muon to Shampoo, SOAP, Aurora, and other matrix-valued optimizers.
• Industry validation: Thinking Machines Lab's Inkling model independently uses Muon with decay $\propto \lambda\eta_t^2$, reporting stable weight norms consistent with the paper's findings.

Limitations and Caveats

• The speed-up is not strictly monotonic across widths (768 slightly lower than 512)
• The largest run (932M parameters) is orders of magnitude below true frontier scale
• The adiabatic tracking analysis breaks down in the final decay tail (near 300k steps for width 1024)

---

Conclusion

This paper proposes scaled weight decay, which multiplies the decoupled decay coefficient by $\eta_t/\eta_{\text{max}}$, and demonstrates its effectiveness both theoretically and empirically:

1. Theory: Scaled weight decay preserves the asymptotic stationarity guarantees of unregularized SGD and Muon, while constant weight decay provably converges to non-minimizers on shifted quadratics.

2. Empirics: Muon-SW achieves the same validation loss 30% faster than standard Muon at the largest scale tested (932M parameters), with speed-ups generally increasing with model size.

3. Mechanism: A steady-state analysis explains why constant weight decay collapses the weight norm (~14× smaller limit) while scaled weight decay maintains a stable plateau.

Future directions:

• Test on other architectures (Mamba, Samba, Gated DeltaNet, Parallax)
• Apply to other matrix-valued optimizers (Shampoo, SOAP, Aurora)
• Validate across domains (vision, audio)
• Study post-training impact (SFT, RLHF, DPO, RL for reasoning)

The method's simplicity—a one-line modification—combined with its substantial empirical gains and clean theoretical foundation makes it a compelling drop-in improvement for frontier model pre-training.