Research map / Architectures and efficiency / Attention and transformers
State Space Models are Provably Comparable to Transformers in Dynamic Token Selection
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.
Shows theoretically that state space model layers interleaved with feedforward networks can perform input-dependent token selection comparably to Transformers, on input copying, associative recall, and piecewise smooth nonparametric regression.
Problem
Single SSM layers use input-independent filters and so lack dynamic token selection, which Transformers have via attention; it is unclear whether SSMs combined with nonlinear FNN layers can match Transformers on such tasks.
Approach
SSM layers sandwiched between FNN layers can emulate attention's key-query inner-product token selection, so SSM+FNN stacks match Transformer expressivity for copying, associative recall, and piecewise gamma-smooth regression rates.
Key results
- Two-layer SSM+FNN achieves copying error bounds with poly-log parameters, versus linear memory needed by a single SSM for error at most 1/2.
- Associative recall solved with O(poly log |W|) parameters, improving on Hyena's O(sqrt(|W|) log^2 |W|) bound.
- Piecewise gamma-smooth regression estimation rate n^{-2a/(2a+1)} (up to log factors) matches Transformer rates.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add a second SSM layer with FNNs before and after it (ablation) | SSM + FNN · Input copying | copying performance | Helped |
| Add a second SSM layer with FNNs before and after it (ablation) | SSM + FNN · Associative recall | recall performance | Helped |
| Add a second SSM layer with FNNs before and after it (ablation) | SSM + FNN · Nonparametric regression | loss | Helped |
| Use two SSM layers with FNNs instead of a single SSM layer | SSM · Input copying | copying performance | No clear effect |
| Use two SSM layers with FNNs instead of a single SSM layer | SSM · Associative recall | recall performance | No clear effect |
| Mask tokens judged unimportant for the target sequence | Pre-trained Hyena fine-tuned on DNA base sequences · Genomic Benchmark Dataset | classification accuracy | No clear effect |
| Mask the same token positions across other samples as were unimportant in the target sequence | Pre-trained Hyena fine-tuned on DNA base sequences · Genomic Benchmark Dataset | classification accuracy | Hurt |
Surprising findings
- Pretrained Hyena on DNA classification remains accurate even after masking most input tokens, indicating sparse important tokens.
Limitations
- Does not analyze whether SSMs can be optimized efficiently by gradient methods.
- Does not cover input-dependent SSMs such as Mamba or linear attention.
- Uses block diagonal A rather than the diagonal parameterization used in practice.
Experimental setup
Datasets: Genomic Benchmark Dataset synthetic input copying, associative recall, induction heads, selective copying, nonparametric regression data
Benchmarks: Genomic Benchmark Dataset
Metrics: error probability for copying accuracy MSE estimation rate
Baselines: single-layer SSM (SSM+FNN) Transformer Hyena
Builds on: Transformers as copying machines (Jelassi et al.) Hyena Mamba-like SSM theory Schmidt-Hieber FNN estimation theory
Closest papers by meaning
- Longhorn: State Space Models are Amortized Online Learners (ICLR 2025 · similarity 0.89)
- Demystifying the Token Dynamics of Deep Selective State Space Models (ICLR 2025 · similarity 0.89)
- From Layers to States: A State Space Model Perspective to Deep Neural Network Layer Dynamics (ICLR 2025 · similarity 0.88)
- Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality (ICML 2024 · similarity 0.88)
- Theoretical Foundations of Deep Selective State-Space Models (NeurIPS 2024 · similarity 0.88)
- Limits of Deep Learning: Sequence Modeling through the Lens of Complexity Theory (ICLR 2025 · similarity 0.87)
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.