Muon相遇Mamba:状态空间模型的谱优化

Muon Meets Mamba: Spectral Optimization for State Space Models

精选理由

Muon优化器之前都在Transformer上测试,这篇论文把它用到Mamba-2 130M上,发现输出投影用Muon最划算,训练效率更高,值得看看。

AI 摘要

Muon是一种通过Newton-Schulz迭代在谱范数下做最速下降的优化器,此前证据几乎全部来自Transformer模型。这篇论文在Mamba-2 130M上对比了Muon与AdamW,只改变哪些权重组用Muon训练。结果显示只有输出投影用Muon时收益最明显,输入投影或两者都用效果更差。优势主要体现在token效率上,在两个语料和两个token预算下都成立,且超过计算最优训练点后依然持续。较低的条件数不能解释增益,因为条件数更低的输入投影反而没有帮助。

原文 · arXiv cs.LG

Muon Meets Mamba: Spectral Optimization for State Space Models

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.