Hungry Hungry Hippos: Towards Language Modeling with State Space Models

By Tri Dao, Daniel Y. Fu, Khaled Kamal Saab, Armin W. Thomas, Atri Rudra, Christopher R'e. 2022

677 citations recorded by the source corpus.

State space models (SSMs) have demonstrated state-of-the-art sequence modeling performance in some modalities, but underperform attention in language modeling. Moreover, despite scaling nearly linearly in sequence length instead of quadratically, SSMs are still slower than Transformers due to poor hardware utilization. In this paper, we make progress on understanding the expressivity gap between SSMs and attention in language modeling, and on reducing the hardware barrier between SSMs and attention. First, we use synthetic language modeling tasks to understand the gap between SSMs and attention. We find that existing SSMs struggle with two capabilities: recalling earlier tokens in the sequence and comparing tokens across the sequence. To understand the impact on language modeling, we propose a new SSM layer, H3, that is explicitly designed for these abilities. H3 matches attention on the synthetic languages and comes within 0.4 PPL of Transformers on OpenWebText. Furthermore, a hybrid 125M-parameter H3-attention model that retains two attention layers surprisingly outperforms Transformers on OpenWebText by 1.0 PPL. Next, to improve the efficiency of training SSMs on modern hardware, we propose FlashConv. FlashConv uses a fused block FFT algorithm to improve efficiency on sequences up to 8K, and introduces a novel state passing algorithm that exploits the recurrent properties of SSMs to scale to longer sequences. FlashConv yields 2$\times$ speedup on the long-range arena benchmark and allows hybrid language models to generate text 2.4$\times$ faster than Transformers. Using FlashConv, we scale hybrid H3-attention language models up to 2.7B parameters on the Pile and find promising initial results, achieving lower perplexity than Transformers and outperforming Transformers in zero- and few-shot learning on a majority of tasks in the SuperGLUE benchmark.

Subfields: Artificial Intelligence, Large Language Models, 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 Large Language Models

In Model Architectures

Citations in the corpus

BLOOM: A 176B-Parameter Open-Access Multilingual Language Model (2022, background): “Both of these approaches are promising, and have now demonstrated competitive results–at large scales for MoE (Fedus et al., 2022; Srivastava et al., 2022), and at smaller scale for state-space models with H3 (Anonymous, 2023).”
Parallelizing Linear Transformers with the Delta Rule over Sequence Length (2024, neutral): “We adapt F LASH L INEAR A T - TENTION [116] to implement Eq.”
Jamba: A Hybrid Transformer-Mamba Language Model (2024, background): “Closest are perhaps H3 [14], a specially designed SSM that enables induction capabilities, and a generalization called Hyena [35].”
VMamba: Visual State Space Model (2024, positive): “These models are then integrated into large representation models [39, 37, 13].”
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.”
The MiniPile Challenge for Data-Efficient Language Models (2023, background): “structurally different architectures [16, 63, 32], pre-training schemes [50, 41], optimizers schemes [3, 26, 24], differential privacy [2, 5], mechanistic interpretability [15, 46, 6], etc.”
U-KAN Makes Strong Backbone for Medical Image Segmentation and Generation (2024, background): “In contrast, structured state-space sequence models (SSMs) [16, 18, 67] have recently shown high efficiency and effectiveness in long-sequence modeling, and are promising solutions for long-term dependency modeling in visual tasks.”
Gated Linear Attention Transformers with Hardware-Efficient Training (2023, positive): “The synthetic MQAR task (Arora et al., 2023a) is a more challenging multi-query version of the induction head task (Fu et al., 2023b) in which a model has to recall the token following a query token multiple times.”
Retentive Network: A Successor to Transformer for Large Language Models (2023, methodology): “During training, RetNet also achieves 25-50% memory saving and 7 × acceleration than standard Transformer and an advantage towards highly-optimized FlashAttention (Dao et al., 2022a).”
MambaHSI: Spatial–Spectral Mamba for Hyperspectral Image Classification (2025, background): “Fu et al. [51] filled the performance gap between SSMs and Transformers in language modeling.”
ChangeMamba: Remote Sensing Change Detection With Spatiotemporal State Space Model (2024, methodology): “After that, the H3 model [47] further advanced these concepts, achieving performance on par with Transformers in language modeling tasks.”
Mechanistic Interpretability for AI Safety - A Review (2024, background): “This focus on cherry-picked results from toy models is a form of streetlight interpretability (Casper, 2023), examining AI systems under only ideal conditions of maximal interpretability.”
VideoMamba: State Space Model for Efficient Video Understanding (2024, positive): “Based on it, various models have been developed ( e.g. , S5 [67], H3 [21], and GSS [57]).”
FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision (2024, neutral): “RWKV [42], H3 [17], MEGA [35], Retnet [54] enhance the expressivity of the simple cumulative sum in linear attention with more sophisticated recurrences.”
HyenaDNA: Long-Range Genomic Sequence Modeling at Single Nucleotide Resolution (2023, background): “On language, the H -family of implicit convolution language models, H3 and Hyena , [12, 37] used long convolutions and gating to match Transformer performance in O ( L log 2 L ) time, notably lower than the O ( L 2 ) of attention-based models.”
Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model (2024, neutral): “[18] designs a new SSM layer, H3, that nearly fills the performance gap be-tween SSMs and Transformer attention in language modeling.”
RWKV: Reinventing RNNs for the Transformer Era (2023, methodology): “, 2022) and its variants (Dao et al., 2022b; Poli et al., 2023) are also proposed.”