MuonSSM: Orthogonalizing State Space Models for Sequence Modeling

Summary

MuonSSM augments selective state-space models with two lightweight components: an auxiliary momentum matrix $\mathbf{M}_t$ that maintains an exponentially weighted average of rank-1 input injections $\tau\beta_t\mathbf{v}_t\mathbf{k}_t^\top$, and a single-step Newton–Schulz (NS) normalization applied to each injection before it enters the memory and momentum paths. Section 2 recasts recent SSMs (Mamba, DeltaNet, Gated DeltaNet, LongHorn) as special cases of a common associative-memory recurrence and identifies their shared weakness — repeated rank-1 first-order updates produce spectral anisotropy, gradient decay, and memory interference. Section 3 shows the coupled $(\mathbf{S}_t,\mathbf{M}_t)$ system remains block-affine, so parallel associative scans still apply; that a scalar $\gamma\mathbf{I}_m$ pathway preserves a non-vanishing gradient channel as $\gamma\!\to\!1$; and that the NS backward Jacobian amplifies fully orthogonal writes by $a\!\approx\!3.44$ while contracting mixed directions to $\approx 0.701$, giving a mechanism for momentum-driven rank enrichment. Experiments span 170M-parameter FineWeb-Edu pretraining with S-NIAH retrieval, MambaVision-Tiny for ImageNet/COCO/ADE20K and IN-R/IN-A/IN-C robustness, and three HAR benchmarks, with a $(\gamma,\tau)$ sensitivity sweep, an NS-iteration ablation, a doubled-$d_{\text{state}}$ capacity control, and a Frobenius-vs-NS ablation.

The conceptual move — conditioning the geometry of low-rank memory injections rather than the recurrent transition operator — is well-motivated and composes cleanly with existing selective-scan backbones. Reported improvements are broadly consistent across three modalities and three backbones, and the theory is honest about its linear-recurrence assumptions. The main limitations are that several mechanistic interpretations (spectral "whitening" reducing texture bias, saddle-point escape, near-isometric spectrum) reach beyond what is measured; that several headline-supporting numbers conflict with the prose and the Table 6 F1 columns are arithmetically inconsistent; and that the released code covers only a fraction of the experiments. Under execution, only the two operator-level claims that could run (NS conditioning, timing scaling) reproduced — qualitatively, not in exact magnitude — while every headline performance table was blocked by absent compute, gated datasets, or missing drivers.

Strengths

• Conceptual contribution. Casting Mamba, DeltaNet, Gated DeltaNet, and LongHorn as special cases of a single scalar-gated associative-memory recurrence (Table 1) is a useful unification, and modifying the low-rank injection rather than the recurrent transition operator is a clean design that plugs into existing selective-scan backbones.
• Theoretical scaffolding. The block-affine recurrence (Proposition 3.1) is a tidy parallelizability proof that directly justifies Algorithm 1, and the NS backward-Jacobian eigenvalue analysis (Proposition 3.6: $0$, $a\!+\!b\!+\!c\!\approx\!0.701$, $a\!\approx\!3.44$) is a specific, testable claim distinguishing NS from Frobenius normalization.
• Breadth of evaluation. Three modalities × three backbones × in-distribution + robustness/OOD + long-context retrieval is more evidence than most SSM-modification papers offer, and the $2\times d_{\text{state}}$ capacity control (Section 5.3) is a thoughtful check that gains are not just from doubled state.
• Isolation of NS from momentum. The three-way ablation (Momentum only / +Frobenius / +NS) in Table 12 is the right control for the Remark 3.7 story, and the effective-rank proxy is reasonable.
• Sensitivity discipline. Sweeping $(\gamma,\tau)$ on all three HAR datasets and reporting min/max/% $\ge$ baseline (Table 11) is stronger than a single-cell hyperparameter report.

Weaknesses

• Numerical evidence contradicts prose (Tables 3, 4, 6). LongHorn baseline beats +Muon on S-NIAH-PK at 2K (67.9 vs 66.7); Mamba IN-A drops 20.57→20.50 despite the text asserting "higher accuracy on IN-A"; Gated DeltaNet's MMAct gain (+0.22%) is the smallest of its HAR datasets, contradicting the "widening margin with complexity" narrative; and the IN-C reductions are called "significant" with no variance test.
• F1 arithmetic anomaly (Table 6/7). Across multiple rows the reported F1 lies below recall, which is impossible for standard F1 (e.g. MuonMamba MMAct $P=78.16$, $R=74.06$, harmonic mean $\approx 76.05$, reported 74.25). This is systematic, spans backbones, and undermines the headline HAR metric columns; the averaging convention must be specified and the numbers re-derived.
• Mechanistic interpretations exceed the diagnostics. "Spectral whitening" causing IN-R texture-bias reduction, "escape saddle points," and "near-isometric" spectrum are all stated causally but not measured. Notably these conflict with the paper's own theory: Corollary 3.5 gives $\sup\rho\approx 1.2$ (not an isometry) and Proposition 3.6 shows explicitly anisotropic backward geometry.
• "Second-order dynamics" mislabel and "orthogonalizing" overclaim. The augmented dynamics are a first-order block-affine recurrence with a heavy-ball EMA, not curvature-informed second-order updates; and by Corollary 3.5 the single NS step preserves the rank-1 injection (it rescales one singular value), so it does not "orthogonalize" the injection despite the title/abstract framing.
• Proof gaps. The off-diagonal block $\mathbf{B}$ in the proof of Proposition 3.3 — on which the whole non-vanishing-gradient argument hinges — is invoked ("This yields:") but never written in closed form; and the Proposition 3.8 attainability proof for $t>\min(d,m)$ is garbled.
• Presentation and setup ambiguities. Figure 4's NS-iteration ablation is discussed with no numbers and an ungrammatical sentence; the sensitivity range is written as an interval $\gamma\in[0.8,0.99]$ that hides a discrete four-point grid; "identical parameter counts across all experiments" conflicts with the deliberately larger $2\times d_{\text{state}}$ baselines; Appendix B.4 says "Adam" (vs AdamW elsewhere); and the GPU count is "four H100s" in Section 4 but "a single H100" in Appendix B.
• Rank-truncation intervention under-controlled. Only one truncation curve is described; the claim that NS's extra directions are functional needs matched-$k$ curves for the Frobenius/momentum-only arms, not just a generic accuracy drop.

Reproducibility & code

Code was executed. The overall replication score is low (0.0625), but almost entirely because the work needs resources this environment lacks — multi-GPU pretraining, ImageNet/COCO/ADE20K, and three gated HAR datasets — rather than because executed results contradicted the paper.

• Mechanism claim partially reproduces. Running a driver over the shipped NS operator on a $d_{\text{state}}=64$ matrix confirmed the Section 5.1 direction (before $\approx 1.9\times10^6$, within ~12% of the paper's $2.2\times10^6$; NS drops it ~1 order of magnitude), but exact scalars differ: $\kappa_{\text{after}}\approx 6.5\times10^4$ vs $1.2\times10^5$ and reduction $\approx 30\times$ vs $18\times$. No shipped script computes this on the recurrent state matrix as the text describes; the quantity was measured on the low-rank injection because $\mathbf{M}_t$ is locked in a compiled kernel.
• NS overshoot confirmed. The shipped Muon quintic overshoots to $\sigma_{\max}\approx 1.19$, matching Corollary 3.5 and contradicting the Appendix B.5 assertion $\sigma_{\max}\le 1$.
• Complexity overhead not the stated "constant factor." A synthetic timing sweep reproduced the $O(L)$ trend, but the +Muon/base ratio grows from $\approx 5.5\times$ at $L{=}512$ to $\approx 27\times$ at $L{=}32\text{k}$, and the measurement is a microbenchmark rather than the paper's MMAct training-time-per-epoch.
• Headline performance tables could not be run. Table 2 (language) needs multi-GPU 10B-token pretraining and ships no working Original-vs-+Muon toggle, data pipeline, or checkpoints; Table 3 (S-NIAH) ships no SFT driver or retrieval harness; Tables 4–5 (vision) ship only IN-R/IN-A/IN-C validators, no ImageNet/COCO/ADE20K training drivers; Tables 6–7 (HAR) are blocked on gated data and the runner wires only the LongHorn backbone with no seed. None of these could be adjudicated.
• Metric columns and ablations not implemented. Precision/recall/F1 are commented out in the HAR scripts (accuracy only), so the Table 6/7 P/R/F1 columns cannot be regenerated; the Frobenius branch and effective-rank routine for Table 12 have no implementation, and the proxy that could run (data-independent) leaves effective rank unchanged under Frobenius, so the reported $12.98\to13.34\to16.62$ ordering is unverified.
• What is covered. The backbone implementations plausibly realize the Table 1 specialization mapping, and the $(\beta,\tau)$ grid in the HAR scripts matches the paper's sweep — the strongest paper–code alignment.

Recommended Changes

Essential

• Reconcile prose with the tables. Rewrite the text around Tables 3, 4, and 6 to acknowledge the LongHorn +Muon loss on S-NIAH-PK at 2K, the Mamba IN-A drop, and Gated DeltaNet's smallest-gain-on-MMAct, and add variance estimates before calling any mCE reduction "significant." (Weaknesses: numerical-evidence contradictions.)
• Fix the Table 6/7 F1 columns. Specify the averaging convention (micro/macro/weighted), recompute every cell where F1 currently falls below recall, and audit the training loop that emitted them. (Weaknesses: F1 anomaly.)
• Release a runnable headline pipeline. Expose the Original-vs-+Muon toggle and all three backbones via a config flag, ship the FineWeb-Edu data/tokenizer pipeline and the Alpaca SFT + S-NIAH harness, and document one working command per row of Tables 2–3. (Reproducibility: language results not reproducible; SFT/S-NIAH absent.)
• Ship the vision and HAR drivers. Add ImageNet-1K, COCO Mask R-CNN, and ADE20K UperNet training drivers with configs; add per-backbone selection, a --seed, and precision/recall/F1 computation to the HAR runner; and provide data-access instructions instead of hardcoded private paths. (Reproducibility: vision drivers absent; HAR runner hard-codes one backbone.)
• Publish the diagnostic code and reconcile its numbers. Release the exact code path and inputs behind Figure 2's condition-number values (stating whether the state matrix or the injection is measured), and the Frobenius branch + effective-rank routine behind Table 12; then align the manuscript with the values those scripts actually produce (~$30\times$ reduction, $\sigma_{\max}\approx 1.19$). (Reproducibility: mechanism magnitudes, effective-rank ablation; Weaknesses: $\sigma_{\max}\le 1$.)

Suggested

• Correct the "modest constant-factor overhead" wording. Report the +Muon/base timing ratio across sequence lengths (which grows to $\approx 27\times$), or provide the MMAct training-time-per-epoch measurement that backs the claim. (Reproducibility: complexity overhead.)
• Soften mechanistic overclaims. Replace "spectral whitening ... prevents ... texture biases," "escape saddle points," "near-isometric," and "throughout training" with hedged wording, or add matching diagnostics (shape-vs-texture metric, curvature statistics, a rank-over-iterations trace). (Weaknesses: mechanistic interpretations.)
• Relabel the mechanism. Use "heavy-ball / two-term recurrence with an auxiliary momentum state" instead of "second-order dynamics," and reconcile the "orthogonalizing" framing with the rank-preserving single NS step. (Weaknesses: second-order mislabel / orthogonalizing overclaim.)
• Complete the proofs. Write out the off-diagonal block $\mathbf{B}$ in closed form (Proposition 3.3) and rewrite the garbled $t>\min(d,m)$ paragraph in the attainability proof. (Weaknesses: proof gaps.)
• Report the NS-iteration ablation numerically and strengthen the truncation study. Add a No-NS/NS×1/NS×5 table (fixing the broken sentence), and overlay matched-$k$ truncation curves for the three Table 12 arms. (Weaknesses: Figure 4 without numbers; rank-truncation under-controlled.)
• Fix presentation nits. Use set notation for the $(\gamma,\tau)$ grid, reconcile "identical parameter counts" with the $2\times d_{\text{state}}$ control, state Adam vs AdamW, resolve the four-vs-one GPU discrepancy, and align the MMAct channel count and num_classes between code and paper. (Weaknesses: presentation/setup ambiguities; Reproducibility: MMAct data spec.)