Memory Determines Learning Direction: A Theory of Gradient-Based Optimization in State Space Models
Not on the map. Why?
Abstract (arXiv)
State space models (SSMs) have gained attention by showing potential to outperform Transformers. However, previous studies have not sufficiently addressed the mechanisms underlying their high performance owing to a lack of theoretical explanation of SSMs' learning dynamics. In this study, we provide such an explanation and propose an improved training strategy. The memory capacity of SSMs can be evaluated by examining how input time series are stored in their current state. Such an examination reveals a tradeoff between memory accuracy and length, as well as the theoretical equivalence between the structured state space sequence model (S4) and a simplified S4 with diagonal recurrent weights. This theoretical foundation allows us to elucidate the learning dynamics, proving the importance of initial parameters. Our analytical results suggest that successful learning requires the initial memory structure to be the longest possible even if memory accuracy may deteriorate or the gradient lose the teacher information. Experiments on tasks requiring long memory confirmed that extending memory is difficult, emphasizing the importance of initialization. Furthermore, we found that fixing recurrent weights can be more advantageous than adapting them because it achieves comparable or even higher performance with faster convergence. Our results provide a new theoretical foundation for SSMs and potentially offer a novel optimization strategy.
Key facts. No citing papers in the map yet (data through 7 Oct 2026). Builds on 14 of its 34 references that are in the map. First posted on arXiv 1 Oct 2025.
Builds on
14 of 34 references.
References · 14
- 2023Mamba: Linear-Time Sequence Modeling with Selective State Spacescited by 4,616
- 2016Gaussian Error Linear Units (GELUs)cited by 3,431
- 2021Efficiently Modeling Long Sequences with Structured State Spacescited by 2,424
- 2024Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Dualitycited by 1,283
- 2020Long Range Arena: A Benchmark for Efficient Transformerscited by 696
- 2016Language Modeling with Gated Convolutional Networkscited by 688
- 2020HiPPO: Recurrent Memory with Optimal Polynomial Projectionscited by 600
- 2022On the Parameterization and Initialization of Diagonal State Space Modelscited by 418
- 2019Overfitting Mechanism and Avoidance in Deep Neural Networkscited by 30
- 2020The interplay between randomness and structure during learning in RNNscited by 20
- 2020Memory and forecasting capacities of nonlinear recurrent networkscited by 11
- 2023Persistent learning signals and working memory without continuous attractorscited by 6
2 more
- 2023Memory of recurrent networks: Do we compute it right?cited by 5
- 2024How noise affects memory in linear recurrent networkscited by 2
Built on by
Not cited yet.
Similar, not cited
Close in wording, no citation either way
- 2024Towards a theory of learning dynamics in deep state space models
- 2024Recurrent neural networks: vanishing and exploding gradients are not the end of the story
- 2023StableSSM: Alleviating the Curse of Memory in State-space Models through Stable Reparameterization
- 2024Mimetic Initialization Helps State Space Models Learn to Recall
- 2024Autocorrelation Matters: Understanding the Role of Initialization Schemes for State Space Models
- 2024HiPPO-Prophecy: State-Space Models can Provably Learn Dynamical Systems in Context
- 2026On the Importance of Gating: Memorization vs. In-Context Learning in State Space Models
- 2024Understanding and Mitigating Bottlenecks of State Space Models through the Lens of Recency and Over-smoothing