# Scale Weight Decay and Train Better

> Scaling weight decay by the learning-rate fraction preserves convergence guarantees and speeds up MoE transformer training by up to 29.4%.

- **Source:** [arXiv](https://arxiv.org/abs/2607.23777)
- **Published:** 2026-09-26
- **Permalink:** https://picx.dev/p/qpVyh1
- **Whiteboard:** https://picx.dev/p/qpVyh1/image

## Summary

# 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**

| Width | Total Params | Active Params | Steps | Tokens | Tok./Act. | Peak LR | Loss (Muon / SW) | Speed-up |
|-------|-------------|---------------|-------|--------|-----------|---------|------------------|----------|
| 256 | 72.7M | 30.2M | 35k | 18.4B | ~610 | $1.0\times10^{-2}$ | 3.350 / 3.297 | 21.7% |
| 512 | 264.9M | 95.0M | 110k | 57.7B | ~610 | $7.0\times10^{-3}$ | 2.971 / 2.915 | 27.5% |
| 768 | 520.0M | 180.3M | 223k | 116.9B | ~650 | $5.6\times10^{-3}$ | 2.799 / 2.754 | 26.5% |
| 1024 | 932.4M | 309.6M | 376k | 197.1B | ~640 | $4.8\times10^{-3}$ | 2.675 / 2.630 | 29.4% |

**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.

---

_Markdown view of https://picx.dev/p/qpVyh1, served by PicX — AI-generated visual whiteboard summaries of research papers._