Papers

MuonGS:用 Gram-Schmidt 投影放松权重范数控制,预训练提速约 15%

MuonGS: Pretraining LLMs with Relaxed Weight-Norm Control

Why we picked it

改进 Muon 优化器的论文:不硬压权重范数,改成投影掉径向分量,实测比 MuonH 省 15% token,训练 LLM 的可以看看。

论文提出 MuonGS 优化器,用 Gram-Schmidt 投影去掉 Muon 每步更新的径向分量,不固定权重范数。作者用累积有效学习率 C_T 分析证明,vanilla Muon 只能对数级累积,而 MuonGS 可达到 Ω(√T)。在 160M、400M 参数稠密模型和 135M 激活/1B 总参数的 MoE 模型上,MuonGS 的验证困惑度低于 MuonH。稠密模型在相同验证损失下估计有最高约 15% 的 token 等效加速。

Original · arXiv cs.AI

MuonGS: Pretraining LLMs with Relaxed Weight-Norm Control

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.