提出谱球优化器，让大模型训练更稳定高效
Controlled LLM Training on Spectral Sphere
- 在权重与更新上施加谱约束，实现理论对齐的优化
- 在多种大模型上优于AdamW和Muon，激活值严格受控
- 适合追求训练稳定性和高性能的大模型研究者
扩大模型规模需要确保快速收敛且稳定的优化策略。最大更新参数化（μP）提供了宽度无关的Θ(1)激活控制理论保障，而新兴优化器如Muon仅部分符合这些约束：控制更新但允许权重漂移。为解决此问题，我们提出谱球优化器（SSO），在模块层面严格施加权重与更新的谱约束。通过推导谱球上的最速下降方向，SSO实现了完全μP对齐的优化过程。为支持大规模训练，我们在Megatron中实现了高效的并行算法。在包括Dense 1.7B、MoE 8B-A1B和200层DeepNet在内的多种架构上进行广泛预训练，SSO始终优于AdamW和Muon。此外，我们观察到显著的实际稳定性提升，包括改善MoE路由负载均衡、抑制异常值，并严格控制激活值。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。