[Paper Review] Hungry Hungry Hippos: Towards Language Modeling with State Space Models

Daniel Y. Fu, Tri Dao+4 more|arXiv (Cornell University)|Dec 28, 2022

Topic ModelingComputer Science116 citations

TL;DR

The paper introduces H3, a new state space model layer designed to bridge the expressivity gap to attention for language modeling, and FlashConv, a hardware-aware FFT-based acceleration enabling efficient training and inference of SSMs on long sequences.

ABSTRACT

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$ imes$ speedup on the long-range arena benchmark and allows hybrid language models to generate text 2.4$ imes$ 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.

Motivation & Objective

• Assess expressivity gaps between state space models (SSMs) and attention for language modeling using synthetic tasks.
• Design an SSM layer (H3) that enables recalling past tokens and comparing tokens across a sequence.
• Improve hardware efficiency of SSMs to match or surpass Transformer performance on standard benchmarks.
• Demonstrate scalable, faster training and generation for large-scale SSM-based models with FlashConv.

Proposed method

• Introduce H3, a two-SSM layer with shift and diagonal matrices and multiplicative input projections to enable token recall and cross-token comparisons.
• Demonstrate H3 matches attention on synthetic language tasks and nearly closes the gap with Transformers on OpenWebText (within 0.4 perplexity).
• Propose a hybrid H3–attention model with two attention layers that outperforms Transformers on OpenWebText by 1.0 perplexity.
• Develop FlashConv, a fused FFT-based convolution with block FFT and a state-passing algorithm to scale SSMs to long sequences and longer models.
• Show that FlashConv yields speedups (e.g., 2× on long sequences) and enables text generation 2.4× faster than Transformers.
• Scale hybrid H3–attention models to up to 2.7B parameters on the Pile and report perplexity and SuperGLUE performance.
• Provide inference throughput comparisons showing Hybrid H3 up to 2.4× faster than Transformer at similar size.

Experimental results

Research questions

• RQ1 Can SSMs exhibit the same expressivity as attention for language modeling on synthetic tasks?
• RQ2 What architectural changes (e.g., shift/diagonal SSMs and multiplicative interactions) enable SSMs to recall past tokens and compare across a sequence?
• RQ3 How can hardware-aware algorithms (like FlashConv) reduce the runtime gap between SSMs and attention on modern accelerators?
• RQ4 Do hybrid H3–attention models achieve competitive perplexity and zero/few-shot performance on standard benchmarks compared to Transformers?
• RQ5 Can FlashConv scale SSM-based models to billions of parameters with feasible training and inference speedups?

Key findings

• H3 matches attention on synthetic language tasks and nearly closes the gap to Transformers on OpenWebText (within 0.4 perplexity).
• A hybrid H3–attention model with two attention layers outperforms Transformers on OpenWebText by 1.0 perplexity.
• Hybrid H3–attention models trained on the Pile with sizes 125M–2.7B outperform Transformer baselines on perplexity and show strong SuperGLUE zero-/few-shot performance.
• FlashConv delivers up to 2× speedup on long sequences and enables SSM-based generation to be 2.4× faster than Transformers.
• SSM-based models scaled to 2.7B parameters on the Pile achieve lower perplexity than Transformers in several settings and show competitive zero- and few-shot results.
• FlashConv achieves state-of-the-art speed on Long Range Arena when compared to Transformer baselines and other SSM methods.