스펙트럼 스피어 기반의 제어된 LLM 훈련

Controlled LLM Training on Spectral Sphere

대규모 모델의 확장은 안정성을 기반으로 빠른 수렴을 보장하는 최적화 전략을 필요로 합니다. Maximal Update Parametrization (µP)은 너비에 무관하게 Θ(1)의 활성화 제어를 위한 이론적 안전장치를 제공하지만, Muon과 같은 새로운 최적화 알고리즘은 이러한 제약 조건과 "부분적으로만" 일치합니다. 즉, 업데이트를 제어하지만 가중치가 드리프트될 수 있습니다. 이러한 제한 사항을 해결하기 위해, 본 논문에서는 가중치와 업데이트 모두에 대해 엄격한 모듈별 스펙트럴 제약을 적용하는 **스펙트럼 스피어 최적화기 (SSO)**를 소개합니다. SSO는 스펙트럴 스피어 상의 가장 가파른 하강 방향을 유도하여, 완전한 µP-정렬 최적화 프로세스를 구현합니다. 대규모 훈련을 가능하게 하기 위해, SSO를 Megatron 내의 효율적인 병렬 알고리즘으로 구현했습니다. Dense 1.7B, MoE 8B-A1B 및 200-layer DeepNet 모델을 포함한 다양한 아키텍처에서 광범위한 사전 훈련을 통해, SSO는 AdamW 및 Muon보다 일관되게 우수한 성능을 보였습니다. 또한, MoE 라우터 부하 균형 개선, 이상치 억제 및 엄격하게 제한된 활성화와 같은 상당한 실질적인 안정성 향상을 관찰했습니다.

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.