ICML 2026 Oral | MuonSSM：用正交化记忆更新稳定状态空间模型

导读

状态空间模型（State Space Models, SSMs）已经成为长序列建模中最重要的注意力替代路线之一。相比 Transformer 的二次复杂度，SSM 可以用线性时间处理长上下文，并通过 scan 机制实现高效并行。但随着序列变长，现有 SSM 往往会遇到一个更底层的问题：记忆状态的更新几何不够稳定。具体来说，很多 SSM 的状态更新本质上是一阶递推。模型不断把新的输入写入状态矩阵，同时用输入相关的转移项传播旧状态。如果每一步写入的方向和尺度没有被良好约束，状态矩阵的奇异值谱会逐渐失衡：少数方向吸收大部分能量，其他方向衰减到接近 0。这会带来梯度退化、长程记忆损失和优化不稳定。 ICML 2026 Oral 论文 MuonSSM: Orthogonalizing State Space Models for Sequence Modeling 提出一个非常明确的解法：不要只盯着 recurrent transition matrix，而是直接调理每一步“写入记忆”的几何。MuonSSM 在 SSM 中加入 momentum pathway，并对低秩输入注入项做轻量 Newton-Schulz 正交化，使每次写入更接近谱条件良好的更新，同时保持 parallel scan 复杂度。这篇论文的价值在于，它把 SSM 的稳定性问题从“状态转移怎么设计”推进到“记忆更新的几何形状怎么被约束”。实验表明，MuonSSM 可以集成到 Mamba、LongHorn、Gated DeltaNet 等多种 SSM backbone 中，在语言、长上下文检索、视觉鲁棒性、COCO/ADE20K 下游任务和时间序列人体动作识别上带来稳定收益。

论文信息

论文标题：MuonSSM: Orthogonalizing State Space Models for Sequence Modeling 作者：Thai-Khanh Nguyen, Ngoc-Bich-Uyen Vo, Thieu N. Vo, Tan M. Nguyen, Cuong Pham 论文链接：https://arxiv.org/abs/2606.30461 PDF：https://arxiv.org/pdf/2606.30461 会议信息：ICML 2026 Oral

背景：SSM 的长序列优势与隐藏瓶颈

SSM 的核心吸引力来自两个方面。第一，它可以在长序列上以近似线性复杂度建模，避免注意力机制随上下文长度平方增长的计算成本。第二，许多现代 SSM 可以被写成 associative scan 形式，因此训练时仍然可以高效并行。但高效并不等于稳定。论文指出，现有 SSM 在长序列中常遇到三类问题。

谱各向异性：状态更新矩阵的奇异值会变得非常不均衡，少数方向主导状态表达。
梯度退化：长时间递推后，早期 token 的梯度容易消失，模型难以保留远距离信息。
记忆干扰：新写入的更新可能覆盖旧信息，使状态矩阵无法形成丰富、可分离的长期记忆。

从这个角度看，MuonSSM 解决的不是“SSM 能不能并行”这种表层问题，而是“每一步写入状态的几何是否健康”。如果写入方向高度相关、尺度失衡，状态容量再大也可能被少数方向占满。

方法：MuonSSM 如何正交化记忆更新

MuonSSM 的方法可以分成两个关键组件：momentum-augmented dynamics 和 Newton-Schulz normalization。

动量增强的状态更新

传统 SSM 在每个时间步会根据输入生成低秩写入项，例如 value 与 key 的外积形式。MuonSSM 在此基础上引入一个额外的 momentum state (M_t)。这个状态并不直接替代原始 memory state (S_t)，而是作为一条辅助通道，把经过几何调理的输入注入项累积起来。直观地说，(S_t) 负责主记忆，(M_t) 负责保留较稳定的更新方向。这样，梯度传播除了穿过原始递推路径之外，还可以通过 momentum pathway 获得额外通道，从而缓解长序列下的梯度衰减。

Newton-Schulz 轻量正交化

MuonSSM 的另一个关键是对低秩输入注入项做单步 Newton-Schulz 变换。Newton-Schulz 迭代常用于近似矩阵正交化或谱归一化。论文采用轻量单步版本，不做显式 SVD，因此额外成本较小。这个操作的作用是把每一步写入的奇异值控制在更健康的范围内，避免某些方向过度放大、某些方向迅速坍缩。作者强调，MuonSSM 并不是去正交化整个 recurrent transition matrix，而是调理 input-dependent memory updates。这一点很重要，因为它保留了原 SSM backbone 的 scan-compatible 结构。

理论分析：为什么它仍然可并行、也更稳定

论文从四个角度分析 MuonSSM 的性质。

并行性

作者将原始 memory state 和 momentum state 拼接为扩展状态，并证明其耦合动态仍然可以写成 block-affine recurrence。这意味着 MuonSSM 不会破坏 SSM 的 parallel associative scan 结构，理论上仍能保持 (O(\log L)) 的并行深度和 (O(L)) 的总工作量。换句话说，MuonSSM 加了动量通道和 Newton-Schulz 调理，但没有把 SSM 退化成必须串行展开的 RNN。

梯度稳定性

Momentum pathway 为梯度提供了一条额外传播路径。相比单一路径反复乘输入相关转移矩阵，MuonSSM 的扩展状态可以缓解长期递推中梯度不断收缩的问题。论文的附录实验也用梯度热图显示，MuonMamba 相比 Vanilla Mamba 在长上下文中具有更均匀的梯度传播。

谱条件化

Newton-Schulz 归一化直接改善写入项的奇异值谱。实验证据显示，标准 Mamba 的 recurrent state matrix 出现严重谱坍缩，条件数约为 (2.2 \times 10^6)；加入 MuonSSM 后，条件数降到约 (1.2 \times 10^5)，约改善 18 倍。更平坦的谱意味着状态空间中更多方向被有效使用，优化地形也更稳定。

记忆表示丰富化

论文还讨论了 rank enrichment：如果每次写入都集中在相近方向，状态矩阵的有效秩会偏低；Newton-Schulz 的反向几何会鼓励学习信号朝当前 rank-one 写入不那么共线的方向传播。与 momentum 累积结合后，模型可以形成更高有效秩的记忆表示。

实验：跨语言、视觉与时间序列的稳定收益

MuonSSM 的实验覆盖多种任务和 backbone，重点不是证明一个单一模型最强，而是证明该几何调理机制可以作为通用插件改善多类 SSM。

语言建模与常识推理

语言实验中，作者在 FineWeb-Edu 10B tokens 上预训练，并比较 Mamba、LongHorn、Gated DeltaNet 的原始版本与 Muon 版本。结果显示，Muon 集成后在困惑度和多个常识推理任务上普遍提升。例如，Mamba 的平均准确率从 42.13 提升到 44.23；LongHorn 从 41.23 提升到 43.27；Gated DeltaNet 从 41.32 提升到 43.10。更重要的是，困惑度也下降明显，例如 LongHorn 的 LMB perplexity 从 96.80 降到 80.98，Gated DeltaNet 从 97.92 降到 83.47。

长上下文检索

在 Needle-in-a-Haystack 类任务中，论文测试 PassKey、Number、UUID 三种检索变体，并覆盖 2K、4K、8K 上下文长度。MuonSSM 对长上下文尤其有帮助。例如 Mamba 在 S-NIAH-PK 8K 上从 8.8 提升到 15.8；LongHorn 在 8K 上从 20.0 提升到 39.3；Gated DeltaNet 从 25.7 提升到 44.5。这类结果与方法动机高度一致：如果更新几何更稳定，模型在长距离传播中更不容易丢失早期写入的信息。

视觉鲁棒性与下游任务

视觉实验基于 MambaVision 风格的层级架构，并评估 ImageNet-1K、ImageNet-R、ImageNet-A、ImageNet-C，以及 COCO2017 检测/分割和 ADE20K 语义分割。整体看，Muon 版本在大多数设置中取得小幅但稳定的提升。例如 LongHorn 在 ImageNet-1K Top-1 上从 81.63 提升到 82.01，在 ImageNet-R 上从 45.44 提升到 46.28，在 ImageNet-A 上从 23.76 提升到 25.27。COCO/ADE20K 下游任务中，Muon 版本也在 APbox、APmask、mIoU 等指标上通常优于原始 backbone。

时间序列人体动作识别

作者还在 MuWiGes、UESTC-MMEA-CL、MMAct 三个人体动作识别数据集上测试。MuonSSM 在这些时间序列任务中提升更明显，说明它不仅适用于语言和视觉，也适用于传感器/动作序列这类模态。例如在 UESTC-MMEA-CL 上，Mamba 从 87.74% accuracy 提升到 91.62%；LongHorn 从 89.06% 提升到 91.56%；在 MMAct 上，MuonMamba 从 71.46% 提升到 74.65%，MuonLongHorn 从 72.47% 提升到 74.40%。

经验分析：收益来自几何调理，而不是简单扩容

论文专门做了 capacity vs. geometric conditioning 的对比。作者把 LongHorn 和 Mamba 的状态维度翻倍，得到 2×dstate baseline，然后与 Muon 版本比较。结果显示，单纯扩容只带来有限提升，而 Muon 版本明显更强。在 MMAct 上，LongHorn 从 72.47 提升到 72.88，而 MuonLongHorn 达到 74.40；Mamba 从 71.47 提升到 72.52，而 MuonMamba 达到 74.65。这说明性能提升主要来自更新几何被调理，而不是状态容量变大。

收敛与复杂度

论文还展示了训练动态：MuonSSM 相比 LongHorn baseline 收敛更快，验证 loss 和 perplexity 更低，并在图中标注约 1.3× faster。复杂度方面，Muon 引入 Newton-Schulz 矩阵乘会带来常数因子开销，但曲线趋势仍保持与序列长度相关的线性可扩展性，不改变 scan-based SSM 的渐近性质。

Newton-Schulz 迭代次数

消融实验比较了 Momentum Only、单步 Newton-Schulz 和五步 Newton-Schulz。结果显示，不加 NS 的动量通道改善有限；单步 NS 效果最好；五步 NS 反而可能过度刚性，削弱模型捕捉必要非正交相关性的能力。作者因此采用单步 Newton-Schulz 作为默认设置，在性能和计算效率之间取得平衡。

鲁棒性与归纳偏置

论文用 GradCAM 分析 MambaVision-Tiny 在 ImageNet-R 这类域外样本上的关注区域。标准 MambaVision 更容易受纹理误导，例如把艺术化金鱼或鲨鱼样本误判为纹理相近类别；MuonSSM 则更关注对象结构形状，如金鱼身体和锤头鲨头部轮廓。这说明谱正交化不只是数值稳定技巧，也可能改变模型学到的归纳偏置：从表面高频纹理转向更稳健的形状和结构特征。

总结

MuonSSM 的核心观点可以概括为一句话：长序列 SSM 的稳定性，不只取决于状态转移如何设计，也取决于每一步写入状态的几何是否良好。通过 momentum pathway 和单步 Newton-Schulz 正交化，MuonSSM 让低秩输入注入项更稳定、更均衡、更不容易谱坍缩，同时保留 SSM 的 parallel scan 优势。从实验看，MuonSSM 在语言、长上下文检索、视觉鲁棒性、下游视觉任务和时间序列识别中都带来稳定提升，并且收益不是简单来自扩容，而是来自更新几何被改善。作为 ICML 2026 Oral，这篇论文给 SSM 研究提供了一条清晰路线：未来的长序列模型不应只追求更复杂的状态结构，也要认真设计记忆写入本身的谱性质、秩结构和梯度路径。

成为VIP会员查看完整内容