arXiv:2604.09967cs.LGcs.AI2026-04被引 12

Muon²通过自适应预处理提升优化器效率,减少迭代次数并加快训练。

Muon$^2$: Boosting Muon via Adaptive Second-Moment Preconditioning

论文配图:Muon$^2$: Boosting Muon via Adaptive Second-Moment Preconditioning
图 1 · 摘自论文原文
  • 在正交化前引入自适应二阶矩预处理,改善动量矩阵条件数
  • 减少40% Newton-Schulz迭代次数,相同损失下节省最多1/4训练时间
  • 适用于大模型预训练,尤其适合参数量达130亿的GPT、LLaMA等架构

Muon作为一种新兴优化器,在大规模基础模型预训练中表现优异,其性能依赖于通过迭代正交化实现的更新矩阵结构。然而,正交化质量受Newton-Schulz(NS)迭代次数影响,带来显著计算与通信开销。本文提出Muon²,通过在正交化前引入类似Adam的自适应二阶矩预处理,显著改善动量矩阵的谱特性,从而加速收敛至实用水平的正交化。我们通过方向对齐度量实际正交化质量,证明Muon²在每一步极分解中均显著优于原版Muon。在高达130亿参数的GPT、LLaMA及Mixture-of-Experts预训练实验中,Muon²及其内存高效变体Muon²-F持续超越原有版本,减少40%的NS迭代次数,并在达到相同损失时最多节省1/4训练时间。

原文摘要 · Abstract (English)

Muon has emerged as a promising optimizer for large-scale foundation model pre-training by exploiting the matrix structure of neural network updates through iterative orthogonalization. However, the orthogonalization quality of Muon hinges on the number of Newton--Schulz (NS) iterations performed, which poses efficiency challenges due to its non-trivial computation and communication cost. We propose Muon$^2$, an extension of Muon, to improve both quality and efficiency by applying Adam-style adaptive second-moment preconditioning before orthogonalization. Our key insight is that the core challenge of polar approximation in Muon lies in the ill-conditioned momentum matrix, of which the spectrum is substantially improved by Muon$^2$, leading to faster convergence toward a practically sufficient orthogonalization. We further characterize the practical orthogonalization quality via directional alignment, under which Muon$^2$ demonstrates dramatic improvement over Muon at each polar step. Across GPT, LLaMA, and Mixture-of-Experts pre-training experiments up to 13B parameters, Muon$^2$ (and its memory-efficient variant Muon$^2$-F that preserves most of its benefits) consistently outperforms Muon and its variants while reducing NS iterations by 40%, and saves up to 1/4 training time over Muon when achieving the same loss.

优化器大模型训练正交化自适应

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。