Controlled LLM Training on Spectral Sphere

TL;DR

The paper presents the Spectral Sphere Optimizer (SSO), a training method for large language models that enforces strict spectral constraints on both weights and updates to achieve stable, width-aware learning aligned with Maximal Update Parametrization (muP). By solving a constrained steepest-descent problem on the spectral sphere and applying a retraction to keep the weight norm fixed, SSO achieves faster convergence and stronger stability than AdamW and Muon, while enabling robust activation bounding and MoE router balance. The authors provide a Megatron-LM implementation with architectural and infrastructural guidelines (atomic modularity, load balancing, adaptive kernels, and caching) to scale SSO to hundreds of billions of parameters. Empirical results across Dense 1.7B, MoE 8B-A1B, and DeepNet 200-Layer models show improved convergence speed, bounded activations, reduced outliers, and improved MoE routing stability, indicating strong practical impact for scalable LLM pretraining. The work also outlines future GPU-native solvers, kernel optimizations, and low-precision pathways to further enhance efficiency and applicability.

Abstract

Scaling large models requires optimization strategies that ensure rapid convergence grounded in stability. Maximal Update Parametrization ($\boldsymbolμ$P) provides a theoretical safeguard for width-invariant $Θ(1)$ activation control, whereas emerging optimizers like Muon are only ``half-aligned'' with these constraints: they control updates but allow weights to drift. To address this limitation, we introduce the \textbf{Spectral Sphere Optimizer (SSO)}, which enforces strict module-wise spectral constraints on both weights and their updates. By deriving the steepest descent direction on the spectral sphere, SSO realizes a fully $\boldsymbolμ$P-aligned optimization process. To enable large-scale training, we implement SSO as an efficient parallel algorithm within Megatron. Through extensive pretraining on diverse architectures, including Dense 1.7B, MoE 8B-A1B, and 200-layer DeepNet models, SSO consistently outperforms AdamW and Muon. Furthermore, we observe significant practical stability benefits, including improved MoE router load balancing, suppressed outliers, and strictly bounded activations.

Figures (12)

• Figure 1: Training dynamics of Dense-1.7B activations (log-scaled cross-layer averages). Our Spectral Sphere maintains constant activation magnitudes throughout training because its $\boldsymbol{\mu}$P-metrized constraints on the spectral manifold ensure that the activation RMS remains at $\Theta(1)$ scale. Muon activations show a mild drift due to learning rate decay and weight decay. By contrast, AdamW proves the most unstable, generating significantly larger activations, with attention AbsMax and FFN RMS reaching $\sim\!100\times$ magnitude compared to those spectral optimizers.
• Figure 2: $\boldsymbol{\mu}$P width scaling across 25$\times$ model size (70M to 1.8B). Although $\boldsymbol{\mu}$P aims for width-invariant scaling, Muon still exhibits notable optimal learning-rate drift. In contrast, our Spectral Sphere achieves stable LR transfer, while also obtaining lower optimal loss thanMuon. More details and related experiments are provided in \ref{['app:mup_width']}.
• Figure 3: Empirical curves of$h(\lambda) = \langle {\bm{\Theta}}, \operatorname{msign}({\bm{G}} + \lambda {\bm{\Theta}}) \rangle$for random matrices. $h(\lambda)$ is monotonic non-decreasing in $\lambda$, and its root $\lambda^\star$ lies close to zero (proof in \ref{['app:localization']}). Here, each curve is obtained by averaging over 5 repeats, and each matrix is initialized from $\mathcal{N}(0,0.02^2)$.
• Figure 4: Geometry of Steepest Descent Update Directions. The left solid arc denotes the ${\bm{W}}$ sphere, while the right dotted arc denotes the $\Delta {\bm{W}}$sphere (unit ${\bm{\Phi}}$ scaled by $\eta$). The shaded region represents the feasible set within the tangent space of the ${\bm{W}}$ sphere at step ${\bm{W}}_i$. Under weight constraint, projecting ${\bm{G}}$ onto the tangent space (Spectral Sphere) yields the largest update angle.
• Figure 5: Ablation of radius scaling on optimization and activation. (a) Final loss vs. learning rate for varying radius scales $c$. A moderate scale (e.g. $c=2.0$) achieves the best performance. (b) AbsMax and (c) RMS of FFN activations during training. AbsMax monotonically follows the radius scale, whereas RMS follows a clear power-law scaling with $c$.

Theorems & Definitions (6)

• Theorem A.1
• proof
• Theorem A.2
• proof
• Theorem A.3
• proof