Muon+通过一步归一化提升大模型预训练效率
MUON+: Towards More Effective Muon via One Additional Normalization Step for LLM Pre-training
- 在极化正交后增加一步归一化，修复梯度更新失衡问题
- 在60M到7B参数模型上，训练和验证困惑度全面优于Muon
- 无需额外状态，适配GPT和LLaMA系列模型的高效预训练
Muon是近期用于大语言模型预训练的强效优化器，通过Newton-Schulz极化迭代实现动量矩阵正交化。直观上看，极化迭代将奇异谱拉平至1，应能消除更新中的行列范数失衡。但实际中，极化步骤反而显著放大了这种失衡。我们称此为极化后更新失衡问题，并证明该失衡会加剧分块下降分析中的二阶项，削弱Muon的每步下降保证。为此，我们提出Muon+，在极化正交后插入一步简单归一化。该方法不引入额外优化器状态。在60M至7B参数的GPT和LLaMA模型上，覆盖计算最优预算及高达约200的词元-参数比扩展实验，Muon+在训练与验证困惑度上均持续优于Muon，带来显著的整体预训练加速。
原文摘要 · Abstract (English)
Muon has recently emerged as a strong optimizer for large language model pre-training, orthogonalizing the momentum matrix via Newton--Schulz polar iterations. A natural intuition is that polar iterations, by flattening the singular spectrum to all ones, should also eliminate column- and row-wise norm imbalance in the update. We show that this is not true in practice: practical polar steps can substantially amplify the imbalance. We term this the post-polar imbalanced update problem, and prove that such imbalance tightens the second-order term in a blockwise descent analysis, weakening Muon's per-step descent guarantee. Motivated by this analysis, we propose Muon+, a one-line fix that inserts a single normalization step after polar orthogonalization. Muon+ adds no optimizer state. Across pre-training experiments on GPT and LLaMA models from 60M to 7B parameters, spanning both compute-optimal budgets and extended token-to-parameter ratios up to approximately 200, Muon+ consistently outperforms Muon in terms of training and validation perplexity, leading to significant overall pre-training speedup.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。