[R] Controlled LLM Training on Spectral Sphere

[R] Controlled LLM Training on Spectral Sphere

TL;DR: The paper introduces Spectral Sphere Optimizer, which takes steepest descent under spectral norm (Muon) and forces the weights & updates onto a spectral sphere.

Paper: https://www.arxiv.org/pdf/2601.08393

Repo: https://github.com/Unakar/Spectral-Sphere-Optimizer

Abstract:

Scaling large models requires optimization strategies that ensure rapid convergence grounded in stability. Maximal Update Parametrization ( muP) provides a theoretical safeguard for width-invariant theta(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 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 muP-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.

Algorithm:

https://preview.redd.it/f1bvi7yd1cdg1.png?width=1197&format=png&auto=webp&s=88a15a375316f54b092e8101e492a2574dc2ace1

Evals:

https://preview.redd.it/5hefuy7g1cdg1.png?width=1503&format=png&auto=webp&s=8a0864c5279654a1c9a29b7aae57d2a1b160aa4d

https://preview.redd.it/0sy8ih8h1cdg1.png?width=1517&format=png&auto=webp&s=ffd675a60192908ed95652b89540cce8d2110088

https://preview.redd.it/rz6bhc6i1cdg1.png?width=1585&format=png&auto=webp&s=50cd471c7805517d0279877fee235dea3e42954e

https://preview.redd.it/fu5wd7zi1cdg1.png?width=1524&format=png&auto=webp&s=5bfb7668a76ceefa320d7325b6abdb731d985e45

submitted by /u/StartledWatermelon
[link] [comments]

Liked Liked