Efficiently Modeling Long Sequences with Structured State Spaces

By Albert Gu, Karan Goel, Christopher R'e. 2021

3,939 citations recorded by the source corpus.

A central goal of sequence modeling is designing a single principled model that can address sequence data across a range of modalities and tasks, particularly on long-range dependencies. Although conventional models including RNNs, CNNs, and Transformers have specialized variants for capturing long dependencies, they still struggle to scale to very long sequences of $10000$ or more steps. A promising recent approach proposed modeling sequences by simulating the fundamental state space model (SSM) $ x'(t) = Ax(t) + Bu(t), y(t) = Cx(t) + Du(t) $, and showed that for appropriate choices of the state matrix $ A $, this system could handle long-range dependencies mathematically and empirically. However, this method has prohibitive computation and memory requirements, rendering it infeasible as a general sequence modeling solution. We propose the Structured State Space sequence model (S4) based on a new parameterization for the SSM, and show that it can be computed much more efficiently than prior approaches while preserving their theoretical strengths. Our technique involves conditioning $ A $ with a low-rank correction, allowing it to be diagonalized stably and reducing the SSM to the well-studied computation of a Cauchy kernel. S4 achieves strong empirical results across a diverse range of established benchmarks, including (i) 91\% accuracy on sequential CIFAR-10 with no data augmentation or auxiliary losses, on par with a larger 2-D ResNet, (ii) substantially closing the gap to Transformers on image and language modeling tasks, while performing generation $60\times$ faster (iii) SoTA on every task from the Long Range Arena benchmark, including solving the challenging Path-X task of length 16k that all prior work fails on, while being as efficient as all competitors.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Model Architectures

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Model Architectures

Citations in the corpus

BLOOM: A 176B-Parameter Open-Access Multilingual Language Model (2022, methodology): “At the time of the design of BLOOM, they consistently underperformed in natural language tasks (Gu et al., 2021).”
FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022, background): “[35] Albert Gu, Karan Goel, and Christopher Ré.”
Efficient Transformers: A Survey (2020, methodology): “A recent promising line of work, based on Structured State Spaces (Gu et al., 2021) also demonstrated very promising results on long range modeling.”
It's Raw! Audio Generation with State-Space Models (2022, background): “[13], we add additional pointwise linear layers after the S4 layer in the style of the feed-forward network in Transformers or the inverted bottleneck layer in CNNs [27].”
General-purpose, long-context autoregressive modeling with Perceiver AR (2022, background): “The recently introduced S4 (Gu et al., 2021) is a very efﬁcient model that avoids attention altogether while producing interesting re-sults, but it is not yet competitive on standard datasets like WikiText-103.”
Gated Delta Networks: Improving Mamba2 with Delta Rule (2024, positive): “The ﬁeld of linear RNNs has rapidly evolved from using data-independent decay mechanisms, as exempliﬁed by models like S4 (Gu et al., 2022), S5 (Smith et al., 2023), LRU (Orvieto et al., 2023), RWKV4/5 (Peng et al., 2023), and RetNet (Sun et al., 2023a), to incorporating data-dependent decay…”
Learning to (Learn at Test Time): RNNs with Expressive Hidden States (2024, neutral): “The general idea of fast weights is to update the parameters of a “fast” model on only the most relevant data, as opposed to the conventional practice of updating a “slow” model on all data [74].”
Mamba-UNet: UNet-Like Pure Visual Mamba for Medical Image Segmentation (2024, background): “Recent developments in State Space Models (SSMs) [6,22,27], especially Structured SSMs (S4) [8], offer a promising solution with their efficient performance in processing long sequences.”
Large Language Models: A Survey (2024, neutral): “While the notion of State Space Models has a long history in machine learning, it should be noted that in the context of language models, SSM is usually used in reference to the newer Structure State Space Model architecture or S4 for short (see Gu et al. [29]).”
Lost in the Middle: How Language Models Use Long Contexts (2023, background): “…task-specific since it depends on the marginal value of the added context and the model’s ability to effectively use long input contexts, but we perform a case study with open-domain question answering on NaturalQuestions-Open to better understand this trade-off in existing language models.”
LocalMamba: Visual State Space Model with Windowed Selective Scan (2024, methodology): “Despite initial challenges in training, owing to their computational and memory intensity, recent advancements [9–11, 16, 43] have significantly ameliorated these issues, positioning deep SSMs as formidable competitors against CNNs and Transformers.”
Gated Linear Attention Transformers with Hardware-Efficient Training (2023, background): “Experiments on language modeling demonstrate that gated linear attention Transformers can perform respectably compared to strong baselines.”
Resurrecting Recurrent Neural Networks for Long Sequences (2023, methodology): “( 𝑗 ) + ) ( ) Note that a similar idea is used in deep SSMs (Gu et al., 2021a) in the context of discretization.”
Parallelizing Linear Transformers with the Delta Rule over Sequence Length (2024, neutral): “An alternative unifying framework might be to view the above sequence transformations as a discretization of a continuous state space model [31, 103, 30], or as a matrix multiplication with a masked structured matrix [73, 85, 45, 18].”
Hyena Hierarchy: Towards Larger Convolutional Language Models (2023, methodology): “• H3: Implicit parametrization using state-space models (SSMs), in particular S4 (Gu et al., 2021; Gupta et al., 2022).”
RS-Mamba for Large Remote Sensing Image Dense Prediction (2024, background): “SSMs have gained significant traction in the field of deep learning in recent years, marking a remarkable evolution in the way long-range dependencies and sequential data are handled [28], [29], [30].”
RWKV: Reinventing RNNs for the Transformer Era (2023, background): “State space models (SSM) like S4 (Gu et al., 2022) and its variants (Dao et al.”
Hungry Hungry Hippos: Towards Language Modeling with State Space Models (2022, methodology): “FlashConv sets state-of-the-art performance on the long range arena benchmark [59] using S4 [25].”
RULER: What's the Real Context Size of Your Long-Context Language Models? (2024, methodology): “Finally, novel architectures (Gu et al., 2022; Fu et al., 2023a; Poli et al., 2023; Fu et al., 2023b; Sun et al., 2023a; Beck et al., 2024; Sun et al., 2024) such as Mamba (Gu & Dao, 2023) and RWKV (Peng et al., 2023) have also been proposed to efficiently handle long-context input.”