Muon Meets Mamba: Spectral Optimization for State Space Models

📄 arXiv: 2608.03941v1 📥 PDF

作者: Arslan Battalov, Karim Kramin, Alexander Markotenko, Sofia Sinitsina

分类: cs.LG

发布日期: 2026-08-04

备注: 15 pages, 25 figures, 8 tables

💡 一句话要点

提出Muon优化器以提升状态空间模型的训练效率

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 优化器 状态空间模型 牛顿-舒尔茨 令牌效率 深度学习 模型训练 计算资源 自然语言处理

📋 核心要点

  1. 现有优化方法在状态空间模型中的应用效果尚不明确，尤其是在不同权重组的训练中表现不一。
  2. 论文提出Muon优化器，通过正交化权重更新，提升状态空间模型的训练效率，尤其是在输出投影上。
  3. 实验结果显示，Muon在输出投影的训练中显著优于输入投影，提升了令牌效率，并在不同条件下保持优势。

📝 摘要（中文）

Muon是一种新型优化器，通过牛顿-舒尔茨迭代正交化每个权重矩阵的更新，能够在谱范数下执行最陡下降。尽管其在Transformer模型上的表现已有所报道，但在状态空间模型上的应用尚未得到充分探讨。本文将Muon与AdamW在Mamba-2 130M模型上进行比较，结果表明Muon在输出投影上的表现优于输入投影或同时训练两者，主要体现在令牌效率的提升上。该优势在两个语料库和两种令牌预算下均得以保持，并且在训练超出计算最优点时仍然存在。条件数的降低并未能解释这一增益，训练的更好条件输入投影并未带来相应的帮助。

🔬 方法详解

问题定义：本文旨在解决现有优化器在状态空间模型训练中的效率不足，尤其是在不同权重组的训练效果不均衡的问题。

核心思路：Muon优化器通过牛顿-舒尔茨迭代正交化权重更新，旨在提升模型的训练效率，尤其是在输出投影的训练中。

技术框架：整体架构包括对权重矩阵的正交化更新过程，分为输入投影和输出投影的训练阶段，比较不同权重组的训练效果。

关键创新：Muon优化器的核心创新在于其正交化更新机制，显著降低了训练过程中的条件数，与传统优化器如AdamW相比，能够更有效地利用计算资源。

关键设计：在实验中，Muon的参数设置经过精细调整，损失函数设计为适应谱范数的最陡下降，确保在不同训练阶段的稳定性和效率。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，Muon在输出投影的训练中表现优于输入投影，令牌效率显著提升。具体而言，Muon在两个语料库和两种令牌预算下均取得了优于AdamW的结果，且在训练超出计算最优点时仍保持优势。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、时间序列预测和控制系统等，能够有效提升模型训练的效率和性能，尤其是在资源受限的环境中。未来，Muon优化器有望在更多类型的模型中推广应用，推动相关领域的研究进展。

📄 摘要（原文）

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.