Muon Meets Mamba: Spectral Optimization for State Space Models

AI 中文解读

Muon这个新优化器在Transformer上表现亮眼，但到了Mamba这类状态空间模型上效果如何？最新研究给出了出人意料的答案：它只在特定位置起作用。研究者用130M参数的Mamba-2模型做了严格对比，发现把Muon用在输出投影层效果最好，反而比同时用在输入和输出层都强。这就像给汽车换高性能轮胎，只换驱动轮比四个轮子全换提速更明显。更重要的是，这种优势主要体现在训练效率上——用更少的数据就能达到同等效果，而且持续训练也不会消失。有趣的是，原本以为优化器能改善模型“健康指标”（条件数）就能提升性能，但实验推翻了这个猜测：输入投影虽然被优化得更“健康”，却对最终结果没帮助。这项发现对AI训练有实际意义：开发者不必盲目给所有模块都用新优化器，找准关键位置反而省时省力。未来训练大模型时，或许能通过这种“精准用药”的方式，在算力有限的情况下更快训练出高质量模型，让AI应用落地更高效。

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.