Muon优化器在Mamba模型中仅优化输出投影层时效果最佳,提升令牌效率。
Muon Meets Mamba: Spectral Optimization for State Space Models
- 仅对输出投影层使用Muon优化,优于全层或输入投影层优化
- 在两种语料和预算下均提升令牌效率,且持续训练仍有效
- 优势来自降低条件数,但输入投影的改善无实际帮助
Muon是一种近期提出的优化器,通过Newton-Schulz迭代使权重更新正交化,在谱范数下实现最陡下降。现有证据几乎全部来自Transformer模型,其在状态空间模型中的表现尚不明确。本文在受控实验中对比了Muon与AdamW在Mamba-2 130M上的表现,仅改变哪些权重组使用Muon。结果表明,收益具有局部性:仅对输出投影使用Muon优于对输入投影或两者同时使用。优势主要体现在令牌效率上,在两个语料库和两个令牌预算下均成立,且在超过计算最优点后仍持续有效。条件数降低虽普遍发生,但更好的输入投影条件并未带来性能提升。
原文摘要 · Abstract (English)
Muon is a recent optimizer that orthogonalizes the update to each weight matrix with a Newton-Schulz iteration, which performs steepest descent under the spectral norm. Almost all the evidence for it comes from Transformer models, and its behavior on state-space models is largely unreported. We compare Muon with AdamW on Mamba-2 130M under a controlled protocol that varies only which weight groups are trained with Muon. The benefit is localized. Muon on the output projection alone beats Muon on the input projection or on both. The advantage is mainly one of token efficiency. It holds on two corpora and two token budgets, and persists when training continues well past the compute-optimal point. Conditioning does not explain the gain. Muon lowers the condition number of whichever projection it trains, but the better-conditioned input projection is not the one that helps.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。