Research map / Architectures and efficiency / RNNs and normalization
StableSSM: Alleviating the Curse of Memory in State-space Models through Stable Reparameterization
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proves that state-space models without reparameterization can only stably approximate exponentially decaying memory targets, and proposes stable reparameterizations (plus a gradient-optimal "best" one, f(w)=-1/(w^2+0.5)) that enable stable approximation of any decaying-memory target.
Problem
SSMs without reparameterization inherit the "curse of memory" of RNNs: recurrent weights approach the stability boundary, so they cannot stably approximate targets with polynomial decaying memory, and gradient-based training becomes unstable at scale.
Approach
Define a class of stable reparameterizations of the recurrent eigenvalues (e.g., exponential, softplus) that provably yield stable approximation of any decaying-memory target, and derive a "best" reparameterization from a gradient-over-weight Lipschitz criterion, f(w)=-1/(aw^2+b).
Key results
- On LRA, the 'best' parameterization averages 86.48 vs 86.09 for the exponential (S4) parameterization.
- On MNIST at learning rate 5e-2, Direct and Softplus/Exp diverge or are unstable (NaN for Direct at lr>=5e-3), while 'Best' reaches test loss 0.089073 at lr 5e-2.
- Theory: vanilla SSM memory function is bounded by (d+1)L0 theta_max^2 exp(-beta t), so only exponential-decay targets can be stably approximated.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Best parameterization instead of exponential parameterization | SSM · Long Range Arena average | score: 86.09 -> 86.48 | Helped |
| Use Best parameterization instead of exponential parameterization | SSM · Long Range Arena | score | Mixed |
| Use Best rather than exponential parameterization (ablation) | SSM on MNIST · MNIST, LR 5e-6 | test loss: 2.241486 -> 2.241217 | Helped |
| Use Best rather than exponential parameterization (ablation) | SSM on MNIST · MNIST, LR 5e-5 | test loss: 0.774661 -> 0.765220 | Helped |
| Use Best rather than exponential parameterization (ablation) | SSM on MNIST · MNIST, LR 5e-4 | test loss: 0.093418 -> 0.091924 | Helped |
| Use Best rather than softplus parameterization (ablation) | SSM on MNIST · MNIST, LR 5e-3 | test loss: 0.023795 -> 0.023475 | Helped |
| Use Best rather than softplus parameterization (ablation) | SSM on MNIST · MNIST, LR 5e-2 | test loss: 0.802772 -> 0.089073 | Helped |
| Use Best rather than exponential parameterization (ablation) | SSM on CIFAR10 · CIFAR10, LR 5e-6 | test loss: 1.745816 -> 1.745290 | Helped |
| Use Best rather than exponential parameterization (ablation) | SSM on CIFAR10 · CIFAR10, LR 5e-5 | test loss: 1.218064 -> 1.215510 | Helped |
| Use Best rather than exponential parameterization (ablation) | SSM on CIFAR10 · CIFAR10, LR 5e-4 | test loss: 0.866817 -> 0.870412 | Hurt |
| Use Best rather than exponential parameterization (ablation) | SSM on CIFAR10 · CIFAR10, LR 5e-3 | test loss: 1.567662 -> 1.364697 | Helped |
| Use Best rather than exponential parameterization (ablation) | SSM on CIFAR10 · CIFAR10, LR 5e-2 | test loss: 1.846173 -> 1.713892 | Helped |
Surprising findings
- Direct parameterization yields NaN on CIFAR10 at all tested learning rates, and exponential parameterization diverges at lr 5.0 in image classification.
- Standard S4 exponential parameterization is only slightly outperformed by the proposed 'best' parameterization on LRA, a small average gap (86.48 vs 86.09).
Limitations
- Theory is continuous-time and restricted to real diagonal recurrent weights with time-homogeneous, causal targets; the multi-layer gradient bound is not proven.
- Empirical gains are modest on LRA and mixed across tasks (e.g., Image task favors exponential parameterization).
Experimental setup
Datasets: WikiText-103 MNIST CIFAR10 Long Range Arena synthetic linear functional targets
Benchmarks: WikiText-103 MNIST CIFAR10 Long Range Arena synthetic linear functional approximation
Metrics: test loss gradient-over-weight ratio average accuracy on LRA
Baselines: Direct parameterization Softplus parameterization Exponential parameterization (S4) ReLU parameterization
Builds on: S4 HiPPO Hyena Long Range Arena
Closest papers by meaning
- HOPE for a Robust Parameterization of Long-memory State Space Models (ICLR 2025 · similarity 0.88)
- Recurrent neural networks: vanishing and exploding gradients are not the end of the story (NeurIPS 2024 · similarity 0.87)
- On Feature Learning in Structured State Space Models (NeurIPS 2024 · similarity 0.87)
- Understanding and Mitigating Bottlenecks of State Space Models through the Lens of Recency and Over-smoothing (ICLR 2025 · similarity 0.86)
- On the Importance of Gating: Memorization vs. In-Context Learning in State Space Models (NeurIPS 2026 · similarity 0.86)
- MaRK: Markov-adapted Recurrent Kernels for Dynamic Operator Conditioning in State Space Models (NeurIPS 2026 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.