arXiv:2605.10797cs.LG2026-05被引 4

Muown改进了Muon优化器,解决大模型训练中的谱范数漂移问题。

Muown: Row-Norm Control for Muon Optimization

  • 将权重矩阵行幅值显式建模为优化变量,沿l∞几何更新。
  • 在124M到2.7B参数模型上,比Muon等优化器更低困惑度。
  • 显著降低对权重衰减的敏感性,适合大规模语言模型训练。

Muon已成为语言模型预训练中媲美AdamW的优化器,但在大规模训练中对权重衰减敏感。近期研究发现,无解耦权重衰减的Muon会导致权重矩阵的谱范数随训练上升。通过将谱范数分解为行幅值因子与行相干因子,我们识别出行幅值是实际驱动因素,而行相干保持稳定。受此诊断启发,我们提出Muown,作为Muon的即插即用替代方案:将行幅值向量作为显式优化变量,基于分解诱导的ℓ∞几何进行更新,其余方向仍沿用原有Muon策略。理论证明,Muown在确定性和随机环境下均达到最优非凸收敛率,且随机噪声系数在训练全程低于Muon。在124M至2.7B参数的GPT类模型上,于FineWeb-Edu数据集上的预训练实验表明,Muown在困惑度上优于Muon、SOAP、AdamW和Lion;同时拓宽了近优学习率的平台区间,降低对权重衰减的敏感性,并在适当分片下以可忽略的额外计算开销避免谱范数漂移。

原文摘要 · Abstract (English)

Muon has emerged as a strong competitor to AdamW for language model pre-training, yet its behavior at scale is sensitive to weight decay. Recent work has observed that, for Muon without decoupled weight decay, the spectral norm of weight matrices drifts upward over training. Through a decomposition of the spectral norm into a row-magnitude factor and a row-coherence factor, we identify the former as the empirical driver of this drift under Muon, while the latter remains well-behaved along the trajectory. Motivated by this diagnosis, we introduce Muown, a drop-in replacement for Muon that treats the row-magnitude vector as an explicit optimizer variable, updating it under the $\ell_\infty$ geometry induced by the decomposition, while applying Muon unchanged to the remaining direction component. We prove that Muown attains the optimal non-convex rates in both deterministic and stochastic regimes under a dual norm aligned with the underlying geometries and with a stochastic noise coefficient that empirically remains below that of Muon throughout training. Across GPT-style pre-training on FineWeb-Edu with model sizes from 124M up to 2.7B parameters, Muown improves perplexity over Muon, SOAP, AdamW, and Lion. It also widens the plateau of near-optimal learning rates across model scales, reduces sensitivity to weight decay, and avoids the spectral norm drift at negligible step-time overhead when appropriately sharded.

优化器大模型谱范数训练稳定

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