Summary (Overview)

• Hyperball is an optimizer wrapper that constrains weight matrices and their updates to fixed Frobenius norms, transforming any base optimizer (Adam, Muon, KL-SOAP) into a "Hyperball variant" (AdamH, MuonH, KL-SOAP-H)
• Key finding: MuonH achieves 20–30% token equivalent speedup over weight decay baselines on 1.2B parameter Qwen3-style models, compared to only ~10% for MuonW at the same scale
• Improved hyperparameter transfer: Hyperball reduces optimal learning rate drift to ~1.4× across depths (L ∈ [4, 512]) and widths (d ∈ [128, 2048]), versus 2–4× for AdamW/MuonW baselines
• Theoretical contribution: Proves that weight decay in scale-invariant layers acts as an indirect controller of angular optimization speed rather than classical regularization, with explicit formulas for equilibrium weight norms and angular step sizes
• Versatility: Works with any base optimizer—KL-SOAP-H reaches the fastest loss-at-step result on the modded-nanogpt Track 3 benchmark (3.2780 in 3125 steps)

Introduction and Theoretical Foundation

Motivation

Previous work observed that matrix-based optimizers like Muon provide ~30% speedup over AdamW at small scales, but this shrinks to ~10% as model size and data scale grow. The paper asks: can we preserve optimizer speedups at higher compute?

The Weight Decay Puzzle

In modern Transformers with normalization layers, many weight matrices are scale invariant: for scalar c > 0, rescaling a matrix leaves the loss unchanged, $L(cW) = L(W)$. Classical ℓ₂ regularization intuition fails here—if the loss is unchanged by weight scale, penalizing $\|W\|_F$ cannot be the primary benefit.

Key Theoretical Insight

The paper shows weight decay acts through norm dynamics:

1. Weight decay establishes an equilibrium weight norm $R_\star$
2. This radius then determines the angular learning rate—how fast the weight matrix direction changes
3. The angular step size depends on learning rate and weight decay mainly through the product $\eta\lambda$

Hyperball Update Rule

For a parameter matrix $W_t$, base update $u_t$, learning rate $\eta_t$, and fixed radius $R = \|W_0\|_F$:

$$W_{t+1} = R \cdot \text{Normalize}\left(W_t - \eta_t R \cdot \text{Normalize}(u_t)\right)$$

where Normalize$(X) := X/\|X\|_F$. This constrains the trajectory to a hypersphere of radius R, keeping weight norms and update norms constant.

Methodology

Where Constraints Are Applied

• Hyperball constraint: Applied to attention and MLP weight matrices in prenorm Transformers
• Standard Adam: Used for embeddings, normalization gains, and scalar parameters (where norm carries semantic information)

Design Choice: Frobenius vs. Spectral Norm

The paper chooses Frobenius norm over spectral (operator) norm for two reasons:

1. Computational cost: Frobenius projection is $O(N^2)$ per matrix vs. $O(N^3)$ for SVD-based spectral projection
2. Theoretical motivation: The stable rank ratio $\mathcal{R}(W) := \|W\|_F^2 / \|W\|_{op}^2$ indicates when Frobenius control approximates spectral control well

Experimental Setup

• Architecture: Qwen3-style decoder-only with QK-Norm
• Data: DCLM-baseline, StarCoder, ProofPile 2, FineWeb-Edu
• Scale: 1.2B parameter models, Chinchilla ratios 1×–8×
• Metrics: Token-equivalent speedup (ratio of tokens AdamW needs to match a method's loss), learning rate drift across scales

Empirical Validation / Results

End-to-End Speedup

• MuonH sustains 20–30% speedup over AdamW scaling law, growing with training duration
• MuonW alone yields only ~10% at 1.2B scale
• On 8B model (Marin Ferries), MuonH achieves 0.04 loss improvement over AdamW baseline

modded-nanogpt Track 3 Benchmark Results

Hyperparameter Transfer

• Depth scaling (d=128, L ∈ [4, 512]): Drift ≈1.4× for AdamH/MuonH vs. ≈3× for AdamW, ≈4× for MuonW
• Width scaling (L=4, d ∈ [128, 2048]): Same ≈1.4× drift for both Hyperball variants

Overtrained Setting (130M model)

• MuonH attains lower best C4 validation loss across token budgets 1B–128B
• Fitted asymptotes: $L_\infty = 3.065$ (MuonH) vs. 3.079 (MuonW)

Theoretical and Practical Implications

Key Theoretical Results

Update norm and autocorrelation (Lemma 4.3): For idealized AdamW, SVD Muon, and Moonlight-scaled Muon, the update second moments satisfy:

$$\mathbb{E}\|u_t\|_F^2 = U^2, \quad \mathbb{E}\langle u_t, u_{t-h}\rangle = U^2 c_h$$

with $U = \sqrt{\frac{1-\beta_1}{1+\beta_1}}\sqrt{pq}$ (AdamW), $\sqrt{p}$ (Muon), $0.2\sqrt{pq}$ (Moonlight).

Equilibrium weight norm (Theorem 4.6):

$$R_\star = \eta U \sqrt{\frac{1 + 2\alpha C_\alpha}{1-\alpha^2}}$$

where $\alpha = 1 - \eta\lambda$ and $C_\alpha = \sum_{h=1}^\infty \alpha^{h-1}c_h$ is the autocorrelation sum.

Angular step size:

$$(\eta^{\text{ang}})^2 = \frac{2(1-\alpha)(1-(1-\alpha)C_\alpha)}{1+2\alpha C_\alpha}$$

For AdamW, this simplifies to: $\eta^{\text{ang}} = \sqrt{\frac{2\eta\lambda(1-\beta_1)}{1+(1-\eta\lambda)\beta_1}}$

Inverse gradient scaling (Lemma 4.7): For scale-invariant losses, $\nabla_W L(cW) = \frac{1}{c}\nabla_W L(W)$, implying $\|\nabla_W L(W_t)\|_F \propto 1/R_t$.

Empirically Validated Phenomena

1. Weight norms track LR schedule: Rise during warmup, shrink during decay
2. Gradient norms increase through training: Scale as 1/R_t as weight norms shrink
3. Fixed ηλ product: Nearly identical loss with weight norms proportional to η
4. WD vs. no-WD crossover: WD starts worse but converges lower due to larger effective step size
5. Transfer sensitivity: Less sensitive to initial weight scale than to weight decay scaling

Practical Implications

• Hyperball makes the "angular learning rate" an explicit design choice rather than an indirect consequence of weight decay calibration
• The method is a drop-in wrapper compatible with any base optimizer
• The Frobenius constraint is computationally cheap ($O(N^2)$) while capturing the essential mechanism

Conclusion

Hyperball replaces the implicit norm control of weight decay with an explicit constraint on matrix norms and update norms, improving scaling behavior of matrix-based optimizers and making learning rate transfer more reliable across model widths, depths, and training budgets.

Future Directions

1. Alternative constraints: Spectral, row-wise, column-wise, hybrid, or architecture-dependent constraints may better match some models
2. Sharper theory: Develop better understanding of weight-normalized training, including Weight Normalization-style parameterizations and explicit norm constraints
3. Broader application: The wrapper's compatibility with any base optimizer (demonstrated with KL-SOAP) suggests wide applicability across optimization algorithms

The core message: weight decay in scale-invariant layers is fundamentally an angular learning rate controller, and making this explicit through norm constraints yields both practical speedups and more predictable hyperparameter transfer.