Controlled LLM Training on Spectral Sphere

TL;DR

Introduces the Spectral Sphere Optimizer (SSO), an optimizer that constrains both the weight matrices and their updates to lie on a spectral sphere of radius $R = \Theta(\sqrt{d_\text{out}/d_\text{in}})$ , deriving the steepest-descent direction under the spectral norm and projecting back after each step. The motivation is that Muon-style orthogonalized optimizers are only “half µP-aligned” — they control update spectra but let weights drift — whereas SSO closes the loop and gives strict module-wise $\Theta(1)$ activation control end-to-end. The authors implement SSO as a parallel algorithm inside Megatron-LM and report consistent gains over AdamW and Muon across Dense 1.7B, MoE 8B-A1B, and a 200-layer DeepNet, plus practical stability wins (MoE router load-balancing, suppressed outliers, bounded activations).

Key claims

Muon is “half-aligned” with µP: it controls update spectra via Newton–Schulz orthogonalization but places no constraint on weight spectra, so weights drift away from the µP-prescribed scale during training [README §1].
Constraining both weights and updates to a spectral sphere of radius $R = \Theta(\sqrt{d_\text{out}/d_\text{in}})$ realizes a fully µP-aligned optimization process, with tunable radius giving an explicit knob on activation scale [README §1, Abstract].
The SSO update is the steepest-descent direction under the spectral norm, with weights retracted onto the sphere after each step (hard retraction by default, or a step-α dynamic retraction) [README §2, §5.2].
On Dense 1.7B pretraining, SSO outperforms AdamW and Muon by a substantial margin on training loss while maintaining the “healthiest” intrinsic model metrics [README §4, Fig. dense17].
On MoE 8B-A1B, SSO improves on AdamW and Muon and also yields better router load-balancing (lower max-vio) than baseline optimizers [README §4, Fig. moe8B].
On a 200-layer DeepNet, SSO trains stably and outperforms AdamW and Muon, demonstrating that the spectral-sphere constraint scales to extreme depth [README §1, Abstract].
A MuonSphere variant retains equivalent activation control with minimal compute overhead relative to vanilla Muon [README §1].
Spectral-sphere training produces strictly bounded activations and suppressed outliers, which the authors flag as low-precision-training-friendly [README §1, Fig. control_act].
µP learning-rate transfer is preserved/improved under SSO across model widths [README §4, Fig. mup_trans].
System efficiency rests on Atomic Module Sharding, an Adaptive Kernel Dispatcher, and Cached Singular Vectors (power-iteration on the top singular vectors, reused across steps) [README §1, §5.2].

Method

SSO performs steepest descent under the spectral norm, with both weight matrices $W$ and their proposed updates $\Delta W$ constrained to a per-module spectral sphere of radius $R = \Theta(\sqrt{d_\text{out}/d_\text{in}})$ . The update direction is computed by an msign-style Newton–Schulz iteration (8 steps by default) to approximate the matrix sign function — the spectral-norm steepest-descent direction — and the step is then retracted to the sphere either hard (project back to radius $R$ ) or dynamically (move toward the sphere with step $\alpha = 0.05$ ). Module granularity is configurable: QKV can be treated jointly or split per-component / per-group / per-head, SwiGLU gate/up can be split or fused, and MoE experts can be split per-expert. The radius mode (spectral_mup) sets $R$ from µP scaling; the LR scale mode can match Muon, AdamW-RMS-equivalent, or shape-scaling. A Lagrange-multiplier solver (bisection, tol 1e-8, max 20 iters) handles the constrained step.

System-level: the optimizer is implemented as --optimizer spectral_ball_dist inside a Megatron-LM fork, using Atomic Module Sharding (per-module distribution of the spectral computation), an Adaptive Kernel Dispatcher that chooses Newton–Schulz vs. power-iteration vs. direct SVD per shape, and Cached Singular Vectors (power iteration on top singular vectors with state carried across optimizer steps, 20 iters initially).

The paper also defines a set of “model intrinsic health” monitors — per-module update RMS, per-module gradient RMS, MoE max-vio, hidden-state statistics across each transformer sub-block, parameter-spectral norms — that the authors track during training as the empirical signature of “healthy” optimization.

Results

Dense 1.7B pretraining: SSO outperforms AdamW and Muon by a “substantial margin” on training loss, with downstream eval improvements on a standard suite (SciQ, PIQA, WinoGrande, ARC-Easy/Challenge, BoolQ, LogiQA, LAMBADA, HellaSwag) [README §4, Figs. dense17, denseeval].
MoE 8B-A1B: SSO outperforms AdamW and Muon, with improved router load-balancing as a side-effect of the spectral constraint [README §4, Fig. moe8B].
200-layer DeepNet: SSO trains stably where AdamW/Muon struggle, demonstrated as one of the three core experimental settings [Abstract].
µP LR transfer: optimal-LR curves overlap across widths under SSO, i.e. learning-rate transfer holds [Fig. mup_trans].
Activation control: empirical activation scale is tunable via the spectral-sphere radius, with strictly bounded ranges throughout training [Fig. control_act].
WandB runs are public (Baselines, µP search, Radius search) and the Megatron-LM fork is open-sourced.

(Note: the arxiv PDF body was not extractable at filing time; numbers above are from the GitHub README’s published claims and figure callouts. Some quantitative deltas are not pinned to specific values in the fetched material — see linked paper PDF for exact tables.)

Why it’s interesting

This paper takes the opposite stance from How to Set the Learning Rate for Large-Scale Pre-training? on µP’s relevance: rather than declaring µP’s stability motivation stale (because QK-Norm now does that work under SP), SSO argues that µP’s alignment requirement is stronger than current optimizers deliver, and the right response is to constrain both weights and updates spectrally so the µP guarantees actually hold throughout training. Both papers agree that pure parametrization isn’t enough; they disagree on whether the fix lives in the architecture (QK-Norm, RMSNorm, gated rescaling) or in the optimizer (orthogonalized + spectrally retracted updates). For Luma DiTs, this is a directly testable hypothesis: if MMDiT pretraining stability is more about update-spectrum control than about normalization choices, an SSO-style optimizer should be applicable largely unchanged, and the public Megatron fork makes the bake-off cheap.

It is also the first paper on the wiki to put the Muon-vs-µP tension front and center — Muon is now widespread enough that “half-aligned” is a charged claim, and the WandB-public 1.7B/8B-MoE comparisons make it a useful baseline rather than an existence proof.