Summary of "MuonGS: Pretraining LLMs with Relaxed Weight-Norm Control"

Summary (Overview)

• Core Contribution: The paper introduces a global perspective on weight-norm control for Muon optimizers in LLM pretraining, using the accumulated effective learning rate $C_T = \sum_{t<T} \eta_t / \|W_t\|_F$ as a measure of exploration capacity over training trajectories.

• Key Theoretical Finding: Vanilla Muon can only guarantee logarithmic accumulation ($C_T = \Theta(\log T)$) under worst-case conditions, while norm-controlled variants (MuonW, MuonH) guarantee at least linear accumulation ($C_T = \Omega(T)$), explaining why norm control improves training.

• New Optimizer Proposal: MuonGS uses Gram–Schmidt projection to remove the radial component of each Muon update without fixing the weight norm, achieving $O(\sqrt{t})$ norm growth and $C_T = \Omega(\sqrt{T})$ accumulation guarantees.

• Empirical Results: MuonGS achieves lower validation perplexity than MuonH on 160M and 400M parameter dense models, with an estimated token-equivalent speedup of up to ~15% at matched validation loss, plus improvements on a 135M active/1B total parameter MoE model.

• Relaxed Norm Control: The paper demonstrates that bounded norms are sufficient but not necessary for good accumulation—sufficiently slow norm growth (e.g., $O(\sqrt{t})$) can sustain exploration throughout training.

Introduction and Theoretical Foundation

Background and Motivation

Weight-norm control is a critical component of successful LLM pretraining optimizers:

• AdamW and MuonW use decoupled weight decay to counter norm growth
• MuonH (Hyperball variant) imposes explicit fixed-norm constraints
• nGPT uses unit-norm weight vectors

The central question: When is norm control needed, why does it help, and how much is sufficient?

Effective Learning Rate

The effective learning rate is defined as the ratio $\eta_t / \|W_t\|_F$ of the nominal learning rate to the current weight norm. For a Muon update $W_{t+1} = W_t - \eta_t O_t$:

$$\frac{\|W_{t+1} - W_t\|_F}{\|W_t\|_F} = \frac{\eta_t}{\|W_t\|_F}\|O_t\|_F$$

Key Theoretical Insight

For bounded update directions ($\|O_t\|_F \leq M$) with constant learning rate $\eta$:

$$\frac{\|W_T - W_0\|_F}{\bar{r}_T} \leq M\eta\frac{T}{\bar{r}_T} \leq M\sum_{t<T}\frac{\eta}{\|W_t\|_F} \tag{1}$$

This shows that sufficient accumulation of effective learning rates is necessary for substantial movement in weight space.

Methodology

Accumulated Effective Learning Rate

The central quantity of the analysis:

$$C_T = \sum_{t=0}^{T-1} \frac{\eta_t}{\|W_t\|_F}$$

Proposition 3.1 establishes that for any interval $[s, T]$:

$$\frac{\|W_T - W_s\|_F}{\max\{r_s, r_T\}} \leq \sqrt{p}(C_T - C_s) \tag{7}$$

where $p = \min\{m, n\}$ for weight matrix $W_t \in \mathbb{R}^{m \times n}$.

For scale-invariant losses $F(W) = f(W/\|W\|_F)$ with $L$-Lipschitz $f$:

$$|F(W_T) - F(W_s)| \leq 2L\sqrt{p}(C_T - C_s) \tag{8}$$

MuonGS Algorithm

The MuonGS update removes the radial component via Gram–Schmidt projection:

$$\widetilde{O}_t = O_t - \frac{\langle W_t, O_t\rangle_F}{r_t^2}W_t, \qquad W_{t+1} = W_t - \eta_t\widetilde{O}_t \tag{12}$$

Algorithm 1: MuonGS for one weight matrix

Require: learning rate η_t, momentum β, initial weights W_0, U_0 = 0
1: for t = 0, 1, 2, ... do
2:   U_{t+1} ← βU_t + ∇f(W_t)                    ▷ gradient momentum
3:   O_t ← Orthogonalize(U_{t+1})                 ▷ Newton–Schulz, as in Muon
4:   Õ_t ← O_t − (⟨W_t, O_t⟩/‖W_t‖²_F) W_t       ▷ remove component along W_t
5:   W_{t+1} ← W_t − η_t Õ_t
6: end for

The projection ensures $\langle W_t, \widetilde{O}_t\rangle_F = 0$, so the squared norm update becomes:

$$r_{t+1}^2 = r_t^2 + \eta_t^2\|\widetilde{O}_t\|_F^2 \tag{13}$$

—norm growth is only second-order in the step size.

Empirical Validation / Results

Theoretical Guarantees

Theorem 4.1 (MuonGS growth and accumulation):

$$r_T^2 \leq r_0^2 + p\sum_{t=0}^{T-1}\eta_t^2 \tag{14}$$

For constant $\eta_t = \eta > 0$:

$$C_T \geq \frac{2}{p\eta}\left(\sqrt{r_0^2 + p\eta^2 T} - r_0\right) = \Omega(\sqrt{T}) \tag{15}$$ $$C_{2T} - C_T \geq \frac{\eta T}{\sqrt{r_0^2 + 2p\eta^2 T}} = \Omega(\sqrt{T}) \tag{16}$$

Table 1: Norm and accumulation guarantees at a constant positive learning rate

Proposition 4.2 (Stationarity for exact current-gradient MuonGS): For scale-invariant objectives, every even $T \geq 2$:

$$\min_{T/2 \leq t < T}\|\text{grad}\, f(W_t/r_t)\|_F \leq \frac{K}{\sqrt{T}} \tag{18}$$

Experimental Setup

• Architectures: GPT-style models at 160M and 400M parameters; MoE model with 135M active/1B total parameters
• Data: FineWeb (100B tokens), Ultra-FineWeb for MoE
• Baselines: MuonH (primary), Muon, MuonW
• Training: BF16 mixed-precision, Nesterov momentum 0.95, linear LR decay

Key Results

Dense Models (Figure 5-6):

• MuonGS achieves lower validation perplexity than MuonH at extended training horizons (overtraining, >1 Chinchilla budget)
• Estimated token-equivalent speedup: 7–15% at matched validation loss for longer horizons

MoE Model (Figure 7):

• MuonGS achieves lower best validation perplexity than MuonH at both 10.8B and 21.6B token budgets

Accumulation Analysis (Figure 3): In practical GPT-160M training, vanilla Muon shows much slower $C_T$ growth than MuonW/MuonGS, with the gap widening over later training windows.

Theoretical and Practical Implications

Theoretical Implications

1. Unified Framework: The accumulated effective learning rate $C_T$ provides a unified basis for comparing norm-control strategies, linking norm growth to exploration capacity over entire training trajectories.

2. Necessary vs. Sufficient Conditions: The paper clarifies that accumulation bounds are necessary (not sufficient) conditions for exploration—a large $C_T$ doesn't guarantee good performance, but small accumulation provably limits progress.

3. Relaxed Requirements: Bounded norms ($O(1)$) are sufficient but not necessary for improving on vanilla Muon's logarithmic worst case. Norm growth of $O(t^a)$ with $a < 1$ suffices, opening a broader design space.

4. Later-Window Accumulation: The $C_{2T} - C_T = \Omega(\sqrt{T})$ guarantee shows MuonGS sustains exploration in later training, unlike vanilla Muon where later-window accumulation can be bounded ($\Omega(1)$).

Practical Implications

1. No Weight Decay Tuning: Like MuonH, MuonGS avoids tuning a weight decay coefficient (unlike MuonW), simplifying hyperparameter search.

2. Overtraining Performance: MuonGS's favorable scaling with training horizon is particularly relevant to modern LLM pretraining practice, which commonly extends beyond one Chinchilla budget.

3. Simplicity: The Gram–Schmidt projection is a closed-form, computationally cheap modification to Muon that requires no additional iterative optimization.

4. Broad Applicability: Works for both dense transformers and modern MoE architectures.

Conclusion

The paper reexamines weight-norm control for Muon optimizers through a global trajectory-level perspective. Key takeaways:

1. Accumulated effective learning rate $C_T$ measures exploration capacity and must keep growing for continued progress.

2. Vanilla Muon's logarithmic worst-case accumulation explains its slower empirical progress; norm control (weight decay, fixed norm) guarantees linear accumulation.

3. MuonGS relaxes norm-control requirements while guaranteeing $\Omega(\sqrt{T})$ accumulation (improving on vanilla Muon) without fixing a radius, achieving up to 15% token-equivalent speedup over MuonH.

Open questions:

• Which rate of norm growth works best for a given training budget?
• The accumulation guarantees are worst-case bounds; tighter characterizations for practical settings remain unexplored.
• Whether even weaker norm control (e.g., $O(t^a)$ with $a$ closer to 1) could yield further improvements.