Technical Summary: MuonSSM – Orthogonalizing State Space Models for Sequence Modeling

Problem Statement

State Space Models (SSMs) have emerged as efficient, linear-time alternatives to attention mechanisms for long-sequence modeling. However, existing SSMs (e.g., Mamba, DeltaNet) often suffer from instability and memory degradation over extended horizons. The paper identifies that current SSMs rely on first-order, rank-one updates to their memory states. These updates lead to three primary issues:

1. Spectral Anisotropy: Singular values of the memory matrix become highly non-uniform, causing ill-conditioned optimization landscapes.
2. Gradient Degradation: Repeated contraction through input-dependent transition matrices leads to vanishing gradients over long sequences.
3. Memory Interference: New updates can overwrite previous information due to a lack of temporal inertia and unbalanced update geometry.

While recent variants attempt to mitigate these issues via gating or normalization, they do not explicitly condition the geometry of accumulated updates or introduce temporal inertia to improve long-range signal propagation.

Methodology: MuonSSM

The authors propose MuonSSM, a framework that stabilizes SSM training by explicitly conditioning the geometry of memory updates rather than the recurrent transition matrix. The approach augments standard SSM dynamics with two key components while preserving the affine structure required for efficient parallel scans:

1. Momentum-Augmented Dynamics

MuonSSM introduces an auxiliary momentum matrix $M_t$ that accumulates update directions across timesteps. This mimics momentum-based optimization, providing "temporal inertia" to the state updates.

• The momentum state is updated as: $M_t = \gamma M_{t-1} + \text{NS}(\tau \beta_t v_t k_t^\top)$.
• The memory state $S_t$ is updated by combining the previous state with the momentum term: $S_t = S_{t-1}(\alpha_t(I - \beta_t \eta k_t k_t^\top)) + M_t$.
• This creates an additional channel for gradient flow, mitigating exponential decay in long-range credit assignment.

2. Single-Iteration Newton–Schulz (NS) Normalization

To address spectral imbalance, the input injection (a rank-1 matrix) is normalized using a lightweight, single-step Newton–Schulz iteration before being added to the momentum state.

• Operator: $\text{NS}(X) = (a + b \tilde{X}\tilde{X}^\top + c (\tilde{X}\tilde{X}^\top)^2)\tilde{X}$, where $\tilde{X}$ is the input scaled by its Frobenius norm.
• Function: This operator bounds singular values (keeping them near unity) and preserves the rank of the input without requiring an explicit SVD.
• Geometric Effect: Unlike standard Frobenius normalization, the backward-pass Jacobian of the NS operator amplifies directions orthogonal to the current rank-one write. This encourages "rank enrichment," preventing the memory state from collapsing into a low-dimensional subspace.

Parallelizability

Crucially, the coupled dynamics of memory and momentum can be cast as a block-affine recurrence. This allows the model to utilize parallel associative scans (reducing recurrent depth from $O(L)$ to $O(\log L)$), ensuring that the computational complexity remains linear in the sequence length $L$.

Key Contributions

1. Framework Design: Introduction of MuonSSM, which integrates momentum-based second-order dynamics and implicit spectral conditioning into scan-based recurrence.
2. Theoretical Analysis:
   • Proof that the augmented dynamics preserve parallelizability.
   • Demonstration that the momentum pathway introduces a non-vanishing gradient component, stabilizing long-range propagation.
   • Analysis showing that the NS operator bounds singular values and drives rank enrichment by biasing backward signals toward orthogonal directions.
3. Empirical Validation: Comprehensive experiments across language, vision, and time-series tasks showing consistent improvements in accuracy, robustness, and length generalization.

Experimental Results

The authors evaluated MuonSSM against baselines (Mamba, LongHorn, Gated DeltaNet) across three modalities:

• Language Modeling: On the FineWeb-Edu dataset, MuonSSM variants achieved lower perplexity and superior performance on common-sense reasoning benchmarks. Notably, in "Needle-in-a-Haystack" retrieval tasks, MuonSSM maintained high accuracy up to 8K context lengths, significantly outperforming baselines that degraded rapidly beyond their training horizon (2K tokens).
• Vision: Replacing Mamba mixers with MuonSSM in the MambaVision framework led to consistent gains in ImageNet-1K classification, object detection (COCO), and semantic segmentation (ADE20K). Crucially, MuonSSM showed significantly improved robustness against corruptions (IN-C), renditions (IN-R), and adversarial attacks (IN-A), suggesting a shift from texture bias to shape-invariant feature learning.
• Time-Series: In Human Activity Recognition tasks (MuWiGes, UESTC-MMEA-CL, MMAct), MuonSSM variants outperformed baselines, with the performance gap widening as dataset complexity increased.

Ablation & Analysis:

• Spectral Conditioning: Visualizations confirmed that standard SSMs suffer from spectral collapse (high condition numbers), while MuonSSM maintains a near-isometric spectrum.
• Capacity vs. Geometry: Doubling the state dimension of baselines yielded only marginal gains compared to the substantial improvements from MuonSSM, confirming that the benefits stem from geometric conditioning rather than simple capacity expansion.
• Efficiency: While adding a constant-factor overhead per step, MuonSSM accelerates convergence (approx. 1.3$\times$ faster), resulting in fewer total training steps to reach target performance.

Significance and Claims

The paper posits that geometric conditioning of updates is a principled pathway to stable, scalable sequence modeling. By explicitly controlling the geometry of memory updates via momentum and Newton–Schulz normalization, MuonSSM addresses the fundamental instability of first-order SSMs without sacrificing the linear-time efficiency of parallel scans.

The authors claim that this approach:

• Improves gradient propagation and mitigates spectral amplification.
• Enriches memory representations over long horizons by preventing rank collapse.
• Provides a robust alternative to attention for long-context tasks, demonstrating that stability can be achieved through update dynamics rather than solely through complex transition matrix parameterizations.

The work suggests that future sequence modeling can benefit from decoupling the stability of the update mechanism from the recurrence structure, offering a simple yet effective mechanism for stable training across diverse modalities.