Controlled LLM Training on Spectral Sphere

Summary

The paper introduces the Spectral Sphere Optimizer (SSO), a steepest-descent optimizer that constrains both weights and updates to a spectral sphere of radius $R \propto \sqrt{d_\text{out}/d_\text{in}}$. The conceptual move is to close a well-articulated gap in µP-aligned optimization: existing spectral-norm methods (notably Muon) control the update to the µP scale but leave weights free to drift, degrading feature learning over long horizons. SSO instead solves a Lagrangian tangent-space problem — the update direction is $\mathrm{msign}(\hat{M} + \lambda^{\star}\Theta)$ where $\Theta = u_1 v_1^{\top}$ is the top-singular-value subgradient and $\lambda^{\star}$ is the unique root of $h(\lambda)=\langle\Theta,\mathrm{msign}(G+\lambda\Theta)\rangle$ — followed by a retraction that projects weights back to $\|W\|_2 = R$. The derivation is clean; monotonicity and root-localization of $h(\lambda)$ are proved and empirically corroborated by regenerated figures. The paper also delivers substantial systems work: atomic-module sharding, a ping-pong load balancer for the variable-depth root solver, an adaptive kernel dispatcher, singular-vector caching, and BF16/torch.compile plumbing that recover most of the SSO overhead vs Muon on B200. Empirically the authors report a $1.24\times$ speedup over AdamW on Dense 1.7B / 100B tokens, better µP LR transfer, improved MoE router load balance on an 8B-A1B model, and stabler training on a 200-layer DeepNet. The main conceptual limitation is that the theoretical guarantees (monotone $h$, exact tangent constraint) rely on an exact matrix sign that is only approximated at 8 Polar Express Newton–Schulz steps in practice; the shipped precision check itself shows this approximation fails badly on square matrices (orthogonality residual $\sim 1.18$ at 512×512), so both the µP scope and the theory–implementation gap need bounding on real gradients. The main empirical limitation is that most head-line comparisons are single-seed and the baselines' hyperparameters (LR, weight decay) are not obviously matched to their own optima — the Dense-1.7B setup admits it is tuned for AdamW, but does not document how Muon and SSO's own optimal LR was set at 1.7B scale, and SSO alone omits weight decay via retraction while the baselines keep the inherited $\mathrm{wd}=0.1$.

Strengths

• Conceptual contribution. The framing of Muon as "half-aligned" and SSO as the unique fully µP-aligned steepest descent — with the tangent-projected direction derived by a single Lagrange multiplier and a retraction — is a genuinely clean unification, not a heuristic tweak. It positions SSO as the natural completion of the "old optimizer, new norm" program for spectral norms, and the Discussion's contrast with Stiefel (too rigid) and Hyperball (too loose) is well argued.
• Theoretical scaffolding. Theorems A.2/A.3 establish monotone $h(\lambda)$, root localization in $[-2\|G\|_\ast, 2\|G\|_\ast]$, and the dual-map property, giving a rigorous foundation for the bracket-and-bisect solver rather than treating it as a numerical black box. Re-running figures_in_paper/precision/check_lambda.py regenerates the h(λ) panels at multiple shapes and confirms the monotone-sigmoid behavior with roots on the order of $10^{-5}$ to $10^{-4}$.
• Systems engineering. Section 5 and Table 1 walk through the concrete overhead pipeline (atomic-module sharding, ping-pong placement, adaptive dispatcher with a Triton SYRK path for large matrices, multi-stream execution, BF16 + torch.compile, singular-vector caching). The final $+11.45\%$ per-step overhead vs Muon on B200 (Table 2) is a credible practical trade-off, and the future-work section honestly names remaining pain points (CPU-side bisection sync bubbles, msign-in-BF16). The solver overhead microbenchmark re-runs cleanly and matches the paper's 8.4–12.1× band.
• Ablation coverage. Radius scale $c$ (Figure 5), LR scaler choice (Figure 6a), module granularity (Figure 6b), and Muon weight-decay ablation (Figure 10) each isolate one design axis with matched settings — this is more careful than the norm in optimizer papers.
• Stability results. The tightly bounded activation RMS/AbsMax in Figure 1, the improved MoE MaxVio in Figure 8b (reproduced at the CSV level: AdamW MaxVio mean 0.837 with peaks up to 14.0 vs Spectral Sphere mean 0.157 with peak 0.87), and the smoother DeepNet-200 curves in Figure 9 collectively make a persuasive qualitative case that constraining $\|W\|_2$ is the missing regularizer for long-horizon LLM training.

Weaknesses

• LR-matching asymmetry between optimizers. Section 6.1 only reports an LR sweep for AdamW at 1.7B ([$10^{-3}, 10^{-2}$], optimum $5\mathrm{e}{-3}$). Figure 7's caption then states Muon and SSO use "the higher optimal rate ($1\mathrm{e}{-2}$)" but never says how that value was picked at 1.7B — inferring it from the Appendix A.5 width-scaling sweep would rely on µP transfer, but Figure 2 itself claims Muon exhibits LR drift with width. The headline $1.24\times$ speedup rests on whichever LR Muon actually ran at, and this choice is not justified. (The CSV-level speedup reproduces cleanly — 19.4% for Spectral Sphere, 11.9% for Muon — so the arithmetic is not in doubt; the LR choice under which those numbers were obtained is.)
• Weight-decay confound. SSO removes weight decay on hidden 2D weights and replaces it with retraction; Muon and AdamW keep the inherited $\mathrm{wd}=0.1$. No configuration matches implicit regularization across optimizers (e.g. Muon with the same target-radius retraction added, or SSO with $\mathrm{wd}=0.1$ on top of retraction). Some of SSO's edge could therefore be a better regularizer rather than a better direction.
• Theory–implementation gap on msign. Theorems 3.2, A.2, A.3 assume the exact matrix sign; the released solver uses 8-step Polar Express Newton–Schulz in FP32. The shipped precision check confirms this is accurate to $\sim 10^{-7}$ on rectangular shapes but fails outright on square 512×512 matrices (orthogonality residual $\sim 1.18$, 5-vs-8 iteration gap $\sim 3\mathrm{e}{-3}$, above the paper's own 1e-3 threshold). Transformer models contain square projection matrices, so this is not purely academic; the paper does not empirically bound the resulting tangent-constraint residual on real gradients.
• Causal claim on MoE routing overstated. "Weight spectral normalization effectively promotes balanced routing" is drawn from a four-optimizer correlational comparison in which optimization geometry, WD policy, retraction, and LR scaler change together. All four presumably share the auxiliary-loss-free bias-update mechanism, which is the routing-balance mechanism by design. A one-factor toggle (spectral retraction on/off, everything else fixed) is needed to justify a causal reading.
• "Consistently outperforms" is single-seed. No error bars, per-seed variance, or repeats are reported for any main result. Table 3 gaps (Avg $56.35$ vs $56.19$ for MuonSphere; PIQA $75.95$ vs $75.35$) are within plausible seed noise for downstream benchmarks at 1.7B.
• Restricted comparator set. The framing repeatedly invokes the broader landscape of metrized / spectral optimizers (Shampoo, Scion, Hyperball, Stiefel), but the empirical evaluation stops at AdamW, Muon, and MuonSphere (a µP-normalized Muon, i.e. a close SSO variant). "Uniquely preserves stable µP LR transfer" and "consistently outperforms" would be more defensible with at least one non-Muon-family spectral baseline.
• Under-specified MoE and DeepNet configurations. "Largely follows DeepSeek-V3" is misleading — DeepSeek-V3 uses 256 routed experts + top-8; here it is 64 + top-4 with unstated expert-FFN width, so the "8B-A1B" label is not reconstructible from the text. Similarly the DeepNet section says "extended the baseline's 28 layers to 200", but the appendix launcher scripts imply a much narrower width (hidden 512, FFN 1536, 8 heads / 4 KV heads) not mentioned in the main text — a narrow-and-deep stress test is a different regime from a matched-width one.
• Termination / fallback of the $\lambda$-solver. Section 6.1 sets max iterations to 20 and tolerance to $2\mathrm{e}{-4}$, but does not specify (i) whether the cap counts bracketing steps too, (ii) whether the tolerance is on $\lambda$ or on $|h(\lambda)|$, and (iii) what update is used if neither triggers. Given that Figure 3 shows roots as small as $3\mathrm{e}{-5}$, the interpretation matters — a $\lambda$-tolerance looser than the true root effectively degrades SSO to MuonSphere.
• Retraction variants introduced but never compared. Appendix A.3 presents two variants (hard projection vs. dynamic soft decay) with different implicit-regularization semantics, and ties the dynamic one to AdamW's WD coefficient — but no ablation shows which variant is used in the main experiments or how they trade off.
• Degenerate top singular value not empirically checked mid-training. The uniqueness argument for $\sigma_1$ is invoked from random-matrix theory (measure-zero coincidence at init). During training, matrices develop low-rank structure and the gap $\sigma_1 - \sigma_2$ can shrink; the paper does not measure this, and both the retraction and the tangent projector $\Theta = u_1 v_1^\top$ depend on the gap being non-trivial.
• µP monotonicity is scoped only to zero-mean Gaussians. Figure 3 and its proofs are for zero-mean random matrices. If mid-training weights drift to a non-zero mean, both empirical monotonicity of $h(\lambda)$ and the "root close to zero" property can fail (root shifts by orders of magnitude, monotonicity breaks). The scope is stated implicitly and should be flagged.
• Presentation issues. Figure 6a's caption says "validation loss" while the plotted y-axis says "Train Loss"; only two of the three scalers are compared in the caption (Spectral Kaiming silently omitted). Table 1's SYRK-kernel row is a $+9.8\%$ regression against the prior stage yet retained in the cumulative pipeline without explanation. Table 3's blanket "the two Sphere optimizers outperform the baselines on the common-sense reasoning tasks" is not literally true task-by-task.

Reproducibility & code

Full-repo replication was possible only for the auxiliary numeric harness — everything at 1.7B scale requires an external Megatron-LM fork (Unakar/Megatron-LM, branch SSO_main) and multi-node training, and Table 2 additionally requires B200. Of the 17 tracked claims, 4 reproduced (headline 1.24× speedup at CSV level, MoE loss/MaxVio ordering, h(λ) monotonicity, solver-overhead microbenchmark), 5 came back partial with genuine internal inconsistencies, and 8 were unreachable in the environment. The following points concern artifacts that are shipped:

• Headline aggregates match at the CSV level, but not from scratch. Interpolating standard_baseline.csv reproduces $1.241\times$ (Spectral Sphere) and $1.135\times$ (Muon) speedups vs the paper's $1.24\times/1.12\times$ — well within tolerance. The MoE loss ordering (Spectral Sphere 2.4197 < MuonSphere 2.4348 < Muon 2.4399 < AdamW 2.5032) also reproduces from moe_lmloss.csv. Both are internal-consistency checks on committed W&B exports, not re-runs of the training itself.
• Optimizer pipeline sources are not in the released repo. The launcher scripts invoke --optimizer spectral_ball_dist, but the SSO/retraction/adaptive-dispatcher/Triton SYRK implementations live in an external fork. The entire Python surface of the release is sso.py, numeric_tests/, figures_in_paper/, plus shell launchers. As a result Table 1's five-stage optimization ablation and Table 2's four-optimizer per-step latency cannot be reproduced on any hardware from the shipped tree.
• MoE routing-factor estimator (Listing 1) is paper-only. No estimate_scaling_factor routine or MoE simulation code appears anywhere in the repo; the M≈2.0 factor cannot be re-derived from the release, even though the simulation is cheap and self-contained.
• Radius ablation is internally inconsistent. raw_data/radius/radius_search.csv puts the minimum at $c=2$ (loss 2.5887) as the paper claims, but the designated figure radius_ablation_loss.png is generated from raw_data/details/radius_v2.csv, in which $c=1$ finishes lowest (2.5478) and $c=2$ (2.5647) is essentially tied with $c=0.5$ (2.5645). The two shipped radius artifacts disagree about the optimum, and the designated figure contradicts the "$c\approx 2.0$ best" statement.
• µP width sweep is scaled-down relative to the paper's framing. The shipped µP CSVs cover widths {256, 512, 1024, 2048} — an 8× span, not the "25× (70M to 1.8B)" described. The reproduced per-width minima (Muon 2.6608, Spectral Sphere 2.6481 at width 2048) match the star annotations in mup_comparison.png but sit ~7% above the numbers referenced in the paper's discussion (2.4835 / 2.4644), which do not appear anywhere in the release. Qualitative ordering (Spectral Sphere below Muon at every width) does hold.
• LR-scaler figure crops away the region where Align-Adam-RMS wins. lrscaler.csv has Align-Adam-RMS strictly lower for the first ~4000 steps; the plot is cropped to 5000–13500. The plotted quantity is also labeled "Train Loss" despite the caption's "validation loss".
• Module-granularity figure lacks the no-split baseline the caption references. The shipped module.csv has only three arms (all-split, split-qkv-only, split-ffn-only); the claim that split-FFN "produces nearly the same loss curve as no-split" cannot be verified from the release.
• msign precision harness is confounded and has a documented failure mode. numeric_tests/check_msign.py varies iteration count and coefficient set at the same time (quintic-5 vs polar_express-8), so it does not implement the controlled 5-vs-8 sweep the footnote asserts. On the square 512×512 shape, 8-step polar_express has orthogonality residual ~1.18 (essentially not orthogonal) and a 5-vs-8 gap ~$3\mathrm{e}{-3}$, exceeding the paper's 1e-3 threshold.
• DeepNet-200, Figure 8b (MuonSphere/Muon MaxVio), Table 3 per-task numbers, and Table 1/2 B200 timings are not backed by shipped data. The DeepNet launchers exist but no loss CSV; the merged MoE MaxVio CSV contains only AdamW and Spectral Sphere; the eval scripts run a superset of Table 3's tasks with no per-task JSON shipped; the only timing artifact is a $\lambda$-solver microbenchmark, not end-to-end 4M-token latencies.
• Bracket-count wording overstates the authors' own benchmark. Re-running the shipped benchmark_lambda_solver.py at tol=1e-4 gives bracket mean 2.02 across 9 transformer-like shapes, matching the authors' committed reference exactly. The Appendix's "<2 steps" is a rounding-level overstatement of that data.
• Repo hygiene. No requirements.txt/setup.py; check_msign.py's absl-py dependency is undeclared; check_lambda.py runs BF16 GEMMs on CPU (projected ~40h on a workstation) and contains a hardcoded developer save path. Only one of the twelve output PDFs the script should produce is committed — consistent with scratch code that was not run to completion by the authors.
• Data seed / index files not shipped. The 100B/1B split is described as deterministic via offline indices, but no seed, script, or index file is included.

Recommended Changes

Essential

• Justify each optimizer's LR at 1.7B. In Section 6.1, state the LR each of AdamW / Muon / MuonSphere / Spectral Sphere actually used in Figure 7, and either report a per-optimizer sweep at 1.7B or an explicit µP-transfer justification for the "higher optimal rate ($1\mathrm{e}{-2}$)", addressing the LR-matching asymmetry weakness that underpins the $1.24\times$ headline.
• Add a regularization-matched Muon ablation. Report at least one row of Table 3 where Muon uses the same target-radius retraction (or SSO uses $\mathrm{wd}=0.1$ on top of retraction), per the weight-decay confound weakness, so the reader can separate direction from regularizer.
• Report per-seed variance for one central result. Rerun the Dense 1.7B val-loss trajectory and downstream Avg for at least 3 seeds and report the standard deviation, per the single-seed weakness. Without this the "consistently outperforms" wording should be softened.
• Quantify the approximate-msign residual and its scope. Report $|h(\lambda^\star)|$ and $|\langle\Theta,\Phi\rangle|$ (the achieved tangent violation) as a function of training step for 8-step Polar Express msign on real 1.7B gradients, per the theory–implementation gap weakness. Explicitly address the square-shape failure documented in check_msign.py (orthogonality residual ~1.18 at 512×512): either restrict the guarantees to rectangular shapes, or bound the tangent violation on square projection matrices in-model.
• Isolate the causal mechanism for MoE MaxVio. Add a Muon-plus-target-radius-retraction run (all other MoE settings fixed) to Figure 8, per the MoE causal claim weakness. If retraction alone recovers the improvement, say so; if not, weaken the causal wording.
• Ship the SSO optimizer sources or pin the fork. Vendor spectral_ball_dist, retraction, adaptive dispatcher, Triton SYRK kernel, and the load-balance / all-gather harness into the release, or commit-pin the external Megatron-LM fork with the exact hash, per the Table 1 pipeline not shipped and optimizer sources not shipped reproducibility points.
• Ship the MoE routing-factor simulation. Add Listing 1's estimate_scaling_factor as an executable script in the repo, per the routing-scale-factor reproducibility point.
• Reconcile the radius ablation. Either re-run both radius artifacts on a matched grid so radius_search.csv and radius_v2.csv agree, or explicitly state which experiment supports the "$c\approx 2.0$ best" claim and re-plot radius_ablation_loss.png from the supporting data, per the radius ablation internally inconsistent reproducibility point.
• Extend or re-scope the µP sweep. Either release the missing widths reaching 1.8B (with a width→parameter-count table) so the numbers annotated in the paper's Figure 2 discussion (~2.4835 / ~2.4644) are recoverable, or narrow the framing to the actually-shipped 8× width span, per the µP width sweep scaled-down reproducibility point.
• Ship raw eval outputs for Table 3. Release the per-task JSON/CSV for the four trained checkpoints and document exactly which subset of tasks the "Avg." column averages (the eval scripts run sciq/logiqa/… which are not in Table 3), per the Table 3 unbacked reproducibility point.
• Ship raw B200 timing traces (or a benchmark script) for Table 1/2. Release the per-step timing logs for the five ablation snapshots and the four optimizers, or ship a benchmark harness that reproduces them, per the B200 latencies unbacked reproducibility point.
• Ship the DeepNet-200 loss traces. Release the LM-loss CSVs and plot script for Figure 9, per the Figure 9 unbacked reproducibility point — the extreme-depth stability claim is currently only visual.

Suggested

• State the $\lambda$-solver termination policy. In Section 5 or 6.1, clarify whether the 20-iteration cap counts bracketing plus bisection or bisection only, whether the $2\mathrm{e}{-4}$ tolerance is on $\lambda$ or on $|h(\lambda)|$, and what update is emitted if neither triggers, per the termination/fallback weakness.
• Report which retraction variant is used and add an ablation. In Section 6.1, specify hard vs. dynamic retraction per experiment and add a small hard-vs-dynamic loss comparison, per the retraction variants weakness.
• Broaden the comparator set. Add at least one head-to-head against Shampoo/SOAP or Scion at a matched proxy scale, per the restricted comparator set weakness.
• Add a parameter-budget table for MoE 8B-A1B and DeepNet-200. Report total params, active params per token, hidden dim, expert FFN dim, heads, and shared-expert size for the MoE run, and hidden dim / FFN width / heads for DeepNet-200, per the under-specified configurations weakness.
• Add an empirical $\sigma_1 - \sigma_2$ trace. Report the top-singular-value gap over training on at least one Dense 1.7B layer type, per the degenerate top singular value weakness.
• Scope the h(λ) monotonicity statement. In Section 3.2 or Appendix A.2, explicitly note that monotonicity and near-zero root are proven for zero-mean weight matrices and add a check on mid-training weight means, per the µP monotonicity scope weakness.
• Reconcile Figure 6a caption. Fix "validation loss" vs. "Train Loss" y-axis, name the Spectral Kaiming outcome in the caption, and either plot the full x-range in lrscaler_ablation_loss.png or label the crop, per the presentation issues and LR-scaler figure crops points.
• Fix the module-granularity CSV. Add the no-split baseline column to module.csv so the "nearly the same loss curve as no-split" claim is auditable, per the module-granularity figure reproducibility point.
• Complete Figure 8b's MaxVio traces. Add Muon and MuonSphere MaxVio columns to the merged CSV, per the Figure 8b MaxVio incomplete reproducibility point.
• Publish the data-sampling seed / index files. Ship the seed or offline index files used to draw the 100B/1B train/val split from OLMo-Mix-1124, per the data/seed reproducibility point.
• Annotate ablation model sizes. Add the underlying model size to the radius, LR-scaler, and module-granularity ablation CSVs and captions, per the ablation model size not annotated reproducibility point.
• Update the '<2 steps' bracket wording. Change to "≈ 2 steps" to match both the authors' own committed benchmark and the reproduced one (2.02), per the bracket count point.
• Add a requirements.txt/pyproject.toml and remove the hardcoded path in check_lambda.py. Declare absl-py and make the device selection configurable, per the repo hygiene point.
• Soften Table 3 wording. Replace "outperform the baselines on the common-sense reasoning tasks" with "achieve the highest average" and enumerate per-task exceptions (MuonSphere wins LAMBADA-PPL and BoolQ, etc.), per the Table 3 sweeping claim weakness.
• Justify keeping SYRK in the cumulative pipeline. Add one sentence to Table 1 or Section 5.1 clarifying that SYRK is a net regression at the isolated stage but wins under the adaptive dispatcher, per the SYRK kernel weakness.