MuonSSM: Orthogonalizing State Space Models for Sequence Modeling

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.

Modern machine learning models often need to process long sequences, such as long documents, videos, sensor signals, or time-series data. Transformers are powerful for these tasks, but they become expensive as the sequence gets longer. State space models offer a more efficient alternative, yet their internal memory can become unstable or gradually lose useful information over long contexts. In this work, we introduce MuonSSM, a method that improves how state space models update and preserve their memory. Instead of only changing the main recurrent transition, MuonSSM adds a momentum-like memory pathway and a lightweight normalization step that keeps each update more balanced. This helps information and gradients travel more reliably across long sequences while preserving the efficient parallel computation that makes state space models attractive. We test MuonSSM on language modeling, long-context retrieval, vision tasks, and human activity recognition. Across these settings, MuonSSM consistently improves accuracy, robustness, and long-context performance over several state space model baselines. These results suggest that better-shaped memory updates can be a simple and effective way to build stable and scalable sequence models.