[论文] MuonGS: Pretraining LLMs with Relaxed Weight-Norm Control

论文概要

研究领域: ML / 优化器 作者: Athanasios Glentis, Dawei Li, Mingyi Hong 发布时间: 2026-10-27 arXiv: 2610.37647

中文摘要

权重衰减和显式权重范数约束可改善语言模型预训练中的优化器性能，但关于为何有效以及需要多少控制的统一理解仍不完整。本文用累积有效学习率C_T = Σ η_t/||W_t||_F分析Muon整个训练轨迹上的范数控制。对于 bounded update方向，任意训练区间上的累积可界定相对权重位移。研究表明，vanilla Muon仅能对数级累积，可能限制探索能力，解释了其较慢的实证进展；而MuonW和MuonH保证至少线性累积。更进一步，足够慢的范数增长保证比对数最坏情况更快的累积。基于这一洞察，本文提出MuonGS，用Gram-Schmidt投影移除每个Muon更新的径向分量但不固定权重范数。在有界更新方向和恒定名义学习率下，权重范数t步后最多以√t增长，因此MuonGS保证C_T = Ω(√T)。在160M和400M稠密模型及135M激活/1B总参数MoE模型的预训练实验中，MuonGS在长训练周期上实现比MuonH更低的验证困惑度，稠密模型的token等效加速最高约15%。

原文摘要

Weight decay and explicit weight-norm constraints improve optimizer performance in language model pretraining, yet a unified understanding of why such control helps and how much is needed remains incomplete. We analyze norm control in Muon over the entire training trajectory using the \emph{accumulated effective learning rate}, \(C_T=\sum_{t<T}η_t/\|W_t\|_F\), where \(η_t\) is the nominal learning rate and \(η_t/\|W_t\|_F\) is the effective learning rate relative to the weight norm. For bounded update directions, accumulation over any training interval bounds relative weight displacement. We show that vanilla Muon can accumulate only logarithmically, potentially limiting exploration and helping explain its slower empirical progress, whereas MuonW and MuonH guarantee at least linear accumulation under the stated conditions. Going beyond bounded-norm control, we show that sufficiently slow norm growth guarantees faster accumulation than the logarithmic worst case, opening up less restrictive optimizer designs. Based on this insight, we propose MuonGS, which uses Gram--Schmidt (GS) projection to remove the radial component of each Muon update without fixing the weight norm. With bounded update directions and a constant nominal learning rate, we prove that the weight norm grows at most proportionally to \(\sqrt{t}\) after \(t\) steps. Consequently, MuonGS guarantees \(C_T=Ω(\sqrt{T})\), with accumulation sustained even in later training. Further, for smooth scale-invariant objectives we establish its \(O(T^{-1/2})\) stationarity guarantee. In pretraining experiments with 160M- and 400M-parameter dense models as well as a 135M active- and 1B total- parameter MoE model, MuonGS achieves lower validation perplexity than MuonH at long training horizons. In particular, for the dense model it achieves an estimated token-equivalent speedup up to about \(15\%\) at matched validation loss to MuonH.

---

*自动采集于 2026-10-29*

#论文 #arXiv #AI #小凯