谱球面上的受控LLM训练
机器学习
2026-03-06 v3 人工智能
摘要
扩展大模型需要以确保稳定性为前提的快速收敛优化策略。最大更新参数化（P）为宽度不变的激活控制提供了理论保障，而诸如Muon等新兴优化器仅与这些约束“半对齐”：它们控制更新但允许权重漂移。为解决这一局限性，我们引入了谱球面优化器（SSO），对权重及其更新强制执行严格的逐模块谱约束。通过推导谱球面上的最速下降方向，SSO实现了完全P对齐的优化过程。为实现大规模训练，我们在Megatron中将SSO实现为一个高效的并行算法。通过对多种架构进行广泛预训练，包括Dense 1.7B、MoE 8B-A1B和200层DeepNet模型，SSO始终优于AdamW和Muon。此外，我们观察到显著的实用稳定性优势，包括改善的MoE路由器负载均衡、被抑制的异常值以及严格有界的激活值。
引用
@article{arxiv.2601.08393,
title = {Controlled LLM Training on Spectral Sphere},
author = {Tian Xie and Haoming Luo and Haoyu Tang and Yiwen Hu and Jason Klein Liu and Qingnan Ren and Yang Wang and Wayne Xin Zhao and Rui Yan and Bing Su and Chong Luo and Baining Guo},
journal= {arXiv preprint arXiv:2601.08393},
year = {2026}
}