How Many Heads Make an SSM? A Unified Framework for Attention and State Space Models
Not on the map. Why?
Abstract (arXiv)
Sequence modeling has produced diverse architectures -- from classical recurrent neural networks to modern Transformers and state space models (SSMs) -- yet a unified theoretical understanding of expressivity and trainability trade-offs remains limited. We introduce a unified framework that represents a broad class of sequence maps via an input-dependent effective interaction operator Wij(X), making explicit two recurring construction patterns: (i) the Unified Factorized Framework (Explicit) (attention-style mixing), in which Wij(X) varies through scalar coefficients applied to shared value maps, and (ii) Structured Dynamics (Implicit) (state-space recurrences), in which Wij is induced by a latent dynamical system. Using this framework, we derive three theoretical results. First, we establish the Interaction Rank Gap: models in the Unified Factorized Framework, such as single-head attention, are constrained to a low-dimensional operator span and cannot represent certain structured dynamical maps. Second, we prove an Equivalence (Head-Count) Theorem showing that, within our multi-head factorized class, representing a linear SSM whose lag operators span a k-dimensional subspace on length-n sequences requires and is achievable with H=k heads. Third, we prove a Gradient Highway Result, showing that attention layers admit inputs with distance-independent gradient paths, whereas stable linear dynamics exhibit distance-dependent gradient attenuation. Together, these results formalize a fundamental trade-off between algebraic expressivity (interaction/operator span) and long-range gradient propagation, providing theoretical grounding for modern sequence architecture design.
Key facts. Cited by 1 paper in the map (data through 14 Oct 2026). Builds on 8 of its 16 references that are in the map. First posted on arXiv 17 Dec 2025.
Builds on
8 of 16 references.
References · 8
- 2017Attention Is All You Needcited by 60,361
- 2023Mamba: Linear-Time Sequence Modeling with Selective State Spacescited by 4,718
- 2021Efficiently Modeling Long Sequences with Structured State Spacescited by 2,465
- 2020Transformers are RNNs: Fast Autoregressive Transformers with Linear Attentioncited by 2,022
- 2020Rethinking Attention with Performerscited by 1,601
- 2024KAN: Kolmogorov-Arnold Networkscited by 649
- 2020HiPPO: Recurrent Memory with Optimal Polynomial Projectionscited by 605
- 2024Jamba: A Hybrid Transformer-Mamba Language Modelcited by 401
Built on by
Citing papers per year: 2026: 1.
1 paper.
Similar, not cited
Close in wording, no citation either way
- 2025Advancing Intelligent Sequence Modeling: Evolution, Trade-offs, and Applications of State-Space Architectures from S4 to Mamba
- 2024Understanding the differences in Foundation Models: Attention, State Space Models, and Recurrent Neural Networks
- 2026State Rank Dynamics in Linear Attention LLMs
- 2025Eigenvalues as a Metric for Memory Dynamics in Sequence Models
- 2025On Structured State-Space Duality
- 2022Simplifying and Understanding State Space Models with Diagonal Linear RNNs
- 2026Sessa: Selective State Space Attention
- 2024Theoretical Foundations of Deep Selective State-Space Models