MuonSSM: Orthogonalizing State Space Models for Sequence Modeling

2026-06-29

AI Summary

Existing state space models (SSMs) often suffer from training instability and long-range memory degradation in long-sequence modeling due to geometric imbalance in memory updates and ill-conditioned first-order dynamics. This work proposes a novel paradigm that enforces geometric orthogonality not on the transition matrix itself, but directly on the memory update process. By integrating momentum pathways, low-rank input injections, and a lightweight Newton–Schulz iteration, the method ensures bounded and spectrally well-conditioned updates. This approach substantially improves gradient propagation, suppresses spectral amplification, and enhances long-context modeling capabilities. Empirical evaluations across language, vision, and time-series benchmarks demonstrate consistent gains in both accuracy and robustness when applied to diverse SSM backbone architectures.

Abstract

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.

Problem

• state space models
• instability
• memory degradation
• long-sequence modeling
• update geometry

Innovation

• State Space Models
• Geometric Conditioning
• Orthogonalization
• Long-Sequence Modeling
• Newton-Schulz Iteration