arXiv:2606.30461cs.LG2026-06被引 1

通过几何调控更新路径,提升状态空间模型的长期稳定性与性能。

MuonSSM: Orthogonalizing State Space Models for Sequence Modeling

论文配图:MuonSSM: Orthogonalizing State Space Models for Sequence Modeling
图 1 · 摘自论文原文
  • 引入动量路径与低秩输入的牛顿-舒尔变换,规范记忆更新几何结构。
  • 在语言、视觉和时间序列任务中均实现精度与长序列性能的提升。
  • 适合追求高效稳定长序列建模的开发者或研究者参考。

状态空间模型(SSMs)作为长序列建模的线性时间替代方案,正受到广泛关注。然而,现有SSM在长时间跨度下常因一阶更新条件不佳和更新几何失衡而出现训练不稳定与记忆退化问题。本文提出MuonSSM,一种通用框架,通过显式调节记忆更新的几何结构而非递归转移矩阵来稳定训练过程。该方法在SSM中引入基于动量的路径,并对低秩输入注入执行轻量级牛顿-舒尔变换,实现有界且谱条件良好的更新,同时保持并行扫描复杂度。理论分析表明,MuonSSM改善了梯度传播,缓解了谱放大效应,并在长时域内丰富了记忆表征。在语言、视觉与时间序列基准上的大量实验显示,将其集成到多种SSM主干网络中均获得一致的准确率、鲁棒性及长上下文性能提升。这些结果确立了更新几何条件化作为稳定、可扩展序列建模的原理性路径。

原文摘要 · Abstract (English)

State space models (SSMs) have emerged as efficient linear-time alternatives to attention for long-sequence modeling. However, existing SSMs often suffer from instability and memory degradation over extended horizons due to poorly conditioned first-order updates and unbalanced update geometry. We introduce MuonSSM, a general framework that stabilizes SSM training by explicitly conditioning the geometry of memory updates rather than the recurrent transition matrix. MuonSSM augments SSMs with a momentum-based pathway and a lightweight Newton Schulz transformation on low-rank input injections, yielding bounded and spectrally conditioned updates while preserving parallel scan complexity. Theory shows that MuonSSM improves gradient propagation, mitigates spectral amplification, and enriches memory representations over long horizons. Extensive experiments across language, vision, and time-series benchmarks show consistent gains in accuracy, robustness, and long-context performance when integrated into diverse SSM backbones. These results establish geometric conditioning of updates as a principled pathway to stable, scalable sequence modeling.

状态空间模型长序列建模梯度稳定序列架构

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。