MuonSSM: Orthogonalizing State Space Models for Sequence Modeling
- URL: http://arxiv.org/abs/2606.30461v1
- Date: Mon, 29 Jun 2026 15:27:19 GMT
- Title: MuonSSM: Orthogonalizing State Space Models for Sequence Modeling
- Abstract summary: State space models (SSMs) have emerged as efficient linear-time alternatives to attention for long-sequence modeling.<n>We introduce MuonSSM, a framework that stabilizes SSM training by explicitly conditioning the geometry of memory updates.<n>We show that MuonSSM improves gradient propagation, mitigates spectral amplification, and enriches memory representations over long horizons.
- Score: 4.4972560829708454
- License: http://creativecommons.org/licenses/by/4.0/
- 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.
Related papers
- DSSMs: State Space Models with Explicit Memory via Delay Differential Equations [9.12852957469023]
State Space Models (SSMs) have emerged as a powerful paradigm for efficient long-sequence modeling.<n>We propose Delay State Space Models (DSSMs), a delay differential equation (DDE)-inspired extension of diagonal SSMs.<n>We show that DSSMs substantially improve targeted delayed-retrieval tasks while outperforming S4D on most standard sequence metrics.
arXiv Detail & Related papers (2026-07-11T10:19:21Z) - Flash PD-SSM: Memory-Optimized Structured Sparse State-Space Models [62.362971723881095]
State-space models (SSMs) face a fundamental trade-off between efficiency and expressivity.<n>We propose Flash PD-SSM, a novel SSM that achieves comparable throughput to widely-used structured SSMs with significantly better expressivity guarantees.<n>We show that Flash PD-SSM defines a new state-of-the-art (SoTA) accuracy among competing SSM methods.
arXiv Detail & Related papers (2026-05-18T22:06:33Z) - Scalable Spatio-Temporal SE(3) Diffusion for Long-Horizon Protein Dynamics [51.85385061275941]
Molecular dynamics (MD) simulations remain the gold standard for studying protein dynamics.<n>Recent generative models have shown promise in accelerating simulations, yet they struggle with long-horizon generation.<n>We present STAR-MD, a scalable diffusion model that generates physically plausible protein trajectories over micro-scale timescales.
arXiv Detail & Related papers (2026-02-02T14:13:28Z) - MS-SSM: A Multi-Scale State Space Model for Efficient Sequence Modeling [60.648359990090846]
State-space models (SSMs) have recently attention as an efficient alternative to computationally expensive attention-based models for sequence modeling.<n>This paper introduces a multi-scale SSM framework that represents sequence dynamics across multiple resolution and processing each resolution with specialized state-space dynamics.
arXiv Detail & Related papers (2025-12-29T19:36:28Z) - The Curious Case of In-Training Compression of State Space Models [49.819321766705514]
State Space Models (SSMs) tackle long sequence modeling tasks efficiently, offer both parallelizable training and fast inference.<n>Key design challenge is striking the right balance between maximizing expressivity and limiting this computational burden.<n>Our approach, textscCompreSSM, applies to Linear Time-Invariant SSMs such as Linear Recurrent Units, but is also extendable to selective models.
arXiv Detail & Related papers (2025-10-03T09:02:33Z) - Quantum-Optimized Selective State Space Model for Efficient Time Series Prediction [39.146761527401424]
We propose a hybrid quantum-optimized approach that integrates state space dynamics with a variational quantum gate.<n>We empirically validate Q-SSM on three widely used benchmarks, i.e., ETT, Traffic, and Exchange Rate.<n>Results show that Q-SSM consistently improves over strong baselines.
arXiv Detail & Related papers (2025-08-29T22:00:48Z) - Long-Context State-Space Video World Models [66.28743632951218]
We propose a novel architecture leveraging state-space models (SSMs) to extend temporal memory without compromising computational efficiency.<n>Central to our design is a block-wise SSM scanning scheme, which strategically trades off spatial consistency for extended temporal memory.<n>Experiments on Memory Maze and Minecraft datasets demonstrate that our approach surpasses baselines in preserving long-range memory.
arXiv Detail & Related papers (2025-05-26T16:12:41Z) - Exemplar-Free Continual Learning for State Space Models [32.73275711666184]
State-Space Models (SSMs) excel at capturing long-range dependencies with structured recurrence.<n>Their evolving internal states pose challenges in adapting them under Continual Learning.<n>We propose Inf-SSM, a novel and simple geometry-aware regularization method.
arXiv Detail & Related papers (2025-05-24T08:59:13Z) - Learning to Dissipate Energy in Oscillatory State-Space Models [51.98491034847041]
State-space models (SSMs) are a class of networks for sequence learning.<n>We show that D-LinOSS consistently outperforms previous LinOSS methods on long-range learning tasks.
arXiv Detail & Related papers (2025-05-17T23:15:17Z) - SMR: State Memory Replay for Long Sequence Modeling [19.755738298836526]
This paper proposes a novel non-recursive non-uniform sample processing strategy to overcome compatibility limitations in parallel convolutional computation.
We introduce State Memory Replay (SMR), which utilizes learnable memories to adjust the current state with multi-step information for generalization at sampling points different from those in the training data.
Experiments on long-range modeling tasks in autoregressive language modeling and Long Range Arena demonstrate the general effectiveness of the SMR mechanism for a series of SSM models.
arXiv Detail & Related papers (2024-05-27T17:53:32Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.