PaperlayerView exact revision on arXiv

Source-linked AI summary

Hungry Hungry Hippos: Towards Language Modeling with State Space Models

Daniel Y. Fu, Tri Dao, Khaled K. Saab, Armin W. Thomas, Atri Rudra, Christopher Ré

arXiv:2212.14052v3cs.LGcs.CL

TL;DR

SSMs underperform attention in language modeling and remain slower on modern hardware despite favorable sequence-length scaling. The paper uses synthetic tasks to design H3 for recall and comparison, and introduces FlashConv for efficient SSM computation. H3 nearly closes the language-modeling gap, hybrid H3-attention models outperform Transformers on several reported evaluations, and FlashConv supports faster long-sequence training and inference.

  • Problem

    SSMs lag Transformers in language modeling and suffer poor hardware utilization, leaving gaps in expressivity and efficiency.

  • Method

    The paper uses synthetic language tasks to design H3 for recall and comparison, and proposes FlashConv with fused FFT computation and state passing for efficient SSM training.

  • Results

    Hybrid H3-attention models outperform same-size Transformer language models in perplexity and match or outperform them on a majority of SuperGLUE tasks in zero- and few-shot settings.

  • Takeaways & Limitations

    The results suggest that combining targeted SSM expressivity improvements with hardware-aware algorithms is a promising route toward scalable language models.

  • Takeaways & Limitations

    Language-model comparisons in one reported setting are not directly comparable because the models were trained on different data.

Abstract

from arXiv · show

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.

1 Introduction

The paper investigates why SSMs lag attention in language modeling and introduces H3 and FlashConv to address expressivity and hardware-efficiency gaps.

  • Motivation: SSMs achieve state-of-the-art sequence modeling in some domains but often underperform Transformers in language modeling.The paper frames the gap as potentially arising from inductive biases, training resources, or specialized hardware support.
  • Expressivity gap: H3 is a new SSM layer designed to address these capabilities and nearly matches attention in language modeling.It uses synthetic-task insights to guide the layer design.
  • Expressivity gap: Synthetic language tasks reveal that existing SSMs struggle with recalling earlier tokens and comparing tokens across a sequence.These tasks probe capabilities associated with manipulating text and in-context learning.
  • Results: H3 comes within 0.4 perplexity of Transformers on OpenWebText, compared with 3.4 perplexity for existing SSMs.A hybrid H3-attention model retaining two attention layers outperforms Transformers by 1.0 perplexity on OpenWebText.
  • Hardware efficiency: FlashConv improves SSM hardware efficiency through fused FFT-based computation and a state-passing algorithm for longer sequences.The state-passing method processes chunks that fit GPU SRAM while maintaining near-linear compute complexity.

2 Background

SSMs represent sequence transformations through learned state-space dynamics, with equivalent recurrent and convolutional forms that support efficient inference and training.

  • State space models: A continuous-time SSM maps an input signal to an output through a state variable governed by learned matrices A, B, C, and D.The state evolves from the input, while the output combines the state and direct input terms.
  • State space models: A discrete-time SSM maps input tokens to output tokens through a discrete state variable and the relation y_i = Cx_i + Du_i.The state-space representation is used as a layer in a deep learning pipeline.
  • Recurrent view: The recurrent SSM form enables next-step generation from only the current state and permits extrapolation to sequences longer than those seen during training.This avoids requiring the entire input history during inference.
  • Convolutional view: For efficient training, an SSM can express the output sequence as a convolution of the input with a filter, plus initial-state and residual terms.With zero initial state, the output is exactly a linear convolution with a residual connection.
  • Convolutional view: SSM convolution can be computed with FFTs in O(N log N) time instead of naive conventional matrix operations scaling as O(N^2).The FFT method transforms the filter and input, multiplies them pointwise, and applies an inverse FFT.
  • Linear attention: Linear attention assumes a feature-map factorization of similarity, connecting its cumulative updates to recurrent computation.Its outputs depend on incrementally updated state-like quantities.

3 Hungry Hungry Hippos Layer to Model Discrete Sequences

H3 is an SSM layer designed to address the recall and cross-sequence comparison capabilities that existing SSMs lack. It matches attention on synthetic tasks, approaches Transformer language-modeling performance, and supports a strong hybrid model.

  • Motivation: Existing SSMs fail synthetic tasks requiring recall after events and comparison of tokens across a sequence.These failures motivate H3’s design around the missing capabilities.
  • H3 design: H3 combines shift and diagonal SSMs with multiplicative interactions to model discrete sequences.The shift SSM detects relevant events, while multiplicative interactions compare information across positions and gate values into memory.
  • Efficiency: O(d^2N + dN log N) time and O(dN) space characterize H3 when head dimension is O(1).The layer scales as O(N log N) in sequence length, compared with attention’s typical O(N^2d) time and O(N^2) space.
  • Synthetic expressivity: H3 matches attention on the synthetic languages and is expressive enough to solve associative recall.Its mechanism gates a value into the diagonal SSM and repeatedly outputs the stored value.
  • Natural-language evaluation: 0.4 perplexity points separate H3 from Transformers on OpenWebText, versus 3–3.9 points for existing SSM variants.The comparison uses approximately 125M-parameter, 12-layer models trained for 50B tokens.
  • Natural-language evaluation: 1.0 point is the OpenWebText advantage of a hybrid H3-attention model over Transformers.The hybrid retains two self-attention layers while using H3 elsewhere.

4 FlashConv: Efficiently Training SSMs

FlashConv improves SSM efficiency by fusing FFT-based convolution and using block FFTs for shorter sequences, then passing recurrent state between chunks for longer sequences. The state-passing algorithm preserves the exact SSM output while scaling beyond GPU SRAM limits.

  • Short-sequence computation: FlashConv fuses FFT, pointwise multiplication, and inverse FFT, while block FFT uses specialized matrix-multiplication hardware for sequences up to 8K.Kernel fusion reduces memory reads and writes, and block FFT enables use of tensor cores.
  • Short-sequence computation: 2× is the speedup from kernel fusion and block FFT for FFTConv on sequences shorter than 8K.The block FFT may use more FLOPs than a standard FFT but can run faster on specialized matrix-multiplication hardware.
  • Long-sequence computation: Sequences longer than 8K cannot fit the fused computation in A100 GPU SRAM, motivating state passing.The algorithm processes SRAM-sized chunks and carries an additional state vector between them.
  • Long-sequence computation: The state-passing algorithm computes each chunk with FFT convolution and recurrently updates the end-state used to initialize the next chunk.Only FFTs of the chunk size are required, rather than one FFT spanning the full sequence.
  • Correctness: Algorithm 2 returns exactly the same output as evaluating the SSM with a large FFT of size N.This equivalence is stated for input u and SSM matrices A, B, C, and D.

5 H3 Evaluation

Hybrid H3-attention models match or exceed Transformer quality across perplexity and zero-/few-shot evaluations, while generating text substantially faster.

  • Language Modeling: Hybrid H3 models match or exceed Transformer quality in perplexity and zero-/few-shot learning across evaluated language-modeling tasks.Evaluations use models from 125M to 2.7B parameters trained on the Pile.
  • Language Modeling: On the Pile, the 125M hybrid model outperforms GPT-Neo, while hybrid models outperform GPT-Neo and GPT-2 on zero-shot transfer to OpenWebText and WikiText-103.GPT-2 comparisons are not directly comparable because those models were trained on different data.
  • Language Modeling: Hybrid models outperform or match the best Transformer baseline on more than half of SuperGLUE tasks in zero-shot and 3-shot evaluations.The reported 3-shot metric is accuracy using logit scoring.
  • Language Modeling: 2.4× higher throughput is achieved by a 1.3B-parameter hybrid H3 model versus a similarly sized Transformer.The comparison uses A100 80GB inference with batch size 64 and prompt lengths of 512, 1024, or 1536.

6 FlashConv Evaluation

FlashConv improves SSM efficiency through hardware-aware convolution algorithms and state passing, maintaining near-linear scaling for long sequences.

  • FlashConv Evaluation: 2× acceleration on Long Range Arena lets FlashConv-powered S4 outperform Transformers by 5.8×.Long Range Arena evaluates long-range sequence modeling.
  • FlashConv Evaluation: 2–3× overall speedup over cuFFT FFTConv is obtained for H3 using block FFT, state passing, and kernel fusion.Kernel fusion helps short sequences, while block FFT and state passing target medium and long sequences.
  • FlashConv Evaluation: FlashConv maintains nearly linear scaling while benchmarking H3 against attention across sequence lengths from 256 to 32K.The benchmark measures forward and backward passes.

7 Conclusion

The paper narrows the modeling and hardware gaps between SSMs and attention with H3 and FlashConv, while identifying hybrid architectures as a promising direction.

  • Conclusion: H3 nearly matches attention in language modeling, and retaining two attention layers produces a hybrid that outperforms Transformers on OpenWebText.The hybrid’s complementary SSM-attention design is presented as a promising scaling direction.
  • Conclusion: FlashConv reduces the hardware barrier by using a BlockFFTConv algorithm based on matrix-multiplication units and the dual recurrent–convolution view of SSMs.The paper frames hardware efficiency as a central dimension of the SSM-attention gap.
  • Conclusion: The paper reports encouraging language-model results up to 1.3B parameters and identifies more sophisticated H3 designs and larger-scale models as future directions.The conclusion also emphasizes combining the complementary strengths of SSMs and attention.

C Method details

The method details describe fused FFT-based convolution gradients and an H3 construction whose shift and diagonal SSMs implement associative recall.

  • C.1 Backward Pass: The backward pass computes convolution gradients by reusing FFT infrastructure, including conjugate FFTs and inverse FFTs.The gradients include du, df, and dD; the direct term contributes Ddy to du.
  • C.1 Backward Pass: The convolution backward pass represents zero-padded signals with Toeplitz structure and derives gradients through matrix multiplication and DFT identities.The construction uses signals extended to length 2L and extracts the relevant first N entries.
  • State Passing: The state-passing update matrix Mux is formed by stacking columns AN′−1B through B, mapping chunk inputs to the accumulated state contribution.Each column AiB is a state-space response vector.
  • H3 Model to Solve Λ: For associative recall, H3 assigns keys to separate heads, shifts key information by one step, and cumulatively sums matching values.A multiplicative interaction suppresses outputs from nonmatching heads, while the output projection decodes the retained value.
  • H3 Model to Solve Λ: Proposition 3 states that the described H3 parameterization solves the associative recall problem for language Λ.The proof sketch identifies the matching head and shows that its output encodes the associated value.

D.2 Attention Expressivity

The attention construction solves associative recall by first retrieving the previous token and then comparing it with the current token to perform a key-value lookup. The section also characterizes H3's complexity and synthetic-task setting.

  • Attention Expressivity: A two-layer attention construction solves associative recall by retrieving the previous token, comparing it with the current token, and outputting the matched value.The first layer shifts token encodings, while the second layer performs the comparison and lookup.
  • Attention Expressivity: Causal softmax attention approximates a shift matrix, placing the previous token encoding in the first half of each output vector.The construction uses positional embeddings for queries and keys, a causal mask, and token values constrained to half the hidden dimension.
  • Attention Expressivity: The second attention layer uses dot-product and softmax comparisons to select earlier positions whose preceding token matches the current token.The resulting output sums the values associated with matching key-value pairs, yielding the token that followed the earlier key.
  • H3 Complexity: O(d2N + dN log N) time and O(dN) space characterize the H3 layer for sequence length N and hidden dimension d.The time combines matrix multiplications with two SSM computations based on FFT convolutions.
  • Attention Expressivity: H3's synthetic tasks model in-context learning through induction-head and associative-recall abilities involving earlier-token retrieval and key-value mappings.These tasks test whether models can learn from examples in the input sequence and use that information to generate the correct output.

E.3 OpenWebText Training

OpenWebText experiments use a Megatron-LM training recipe, a shared validation split, mixed-precision optimization, and sequence length 1024.

  • Training setup: OpenWebText models follow the Megatron-LM training recipe and use the GPT-2 BPE tokenizer.The dataset is split by randomly selecting 0.5% for validation, with all models evaluated on the same validation set.
  • Training setup: All OpenWebText models use sequence length 1024, mixed-precision training, and 100K optimization steps with model-specific learning rates.Training uses AdamW, weight decay 0.1, and an effective batch size of 512 with gradient accumulation.

E.4 The Pile Training

Pile experiments train models with GPT-3-style settings and evaluate H3 on long-range accuracy. H3 outperforms S4D on two LRA tasks and comes within one point on the others.

  • The Pile training: Pile models use GPT-3-style training with sequence length 2048, AdamW optimization, cosine learning-rate decay, and model-size-specific batch sizes.The 125M and 355M models train for 800K steps, while 1.3B models use doubled batch size and half as many steps for the same token count.
  • The Pile training: The Pile experiments use the GPT-2 BPE tokenizer and run models on either one 16xA100-40GB node or eight 8xA100-80GB nodes.The benchmark prompts follow the GPT-3 paper, with task-specific label formats for SuperGLUE evaluation.
  • LRA evaluation: H3 outperforms S4D on two LRA tasks and comes within 1 point on the others.The evaluation compares H3 with S4D using the same hyperparameters and bidirectional copies of the layer.

F.2 WikiText103

WikiText103 and PG-19 evaluations find hybrid H3 competitive with or superior to several Transformer baselines, while longer-context and generation analyses reveal both extrapolation strengths and zero-shot weaknesses.

  • WikiText103: Hybrid H3 is competitive with same-size Transformers and larger Perceiver AR and Transformer-XL models on WikiText103.It also significantly outperforms Transformers using performer, reformer, and linear attention.
  • WikiText103: Transformer-XL and Perceiver AR perplexities are not directly comparable because their original results use different tokenizers and model configurations.The top five WikiText103 results were trained with the same setup and are directly comparable.
  • PG-19: Hybrid H3 outperforms Transformers and linear attention on PG-19.The evaluation uses the same setup as OpenWebText on this natural-language book dataset.
  • Length extrapolation: H3 maintains associative-recall accuracy at sequence length 40 after training on sequences of length 20.The evaluation tests last-token prediction at both the training length and twice that length.
  • Scaling with training tokens: Hybrid H3 and Transformer models both improve as the number of Pile training tokens increases.The comparison trains 125M models on 5B, 10B, and 15B tokens.
  • SuperGLUE generation: Few-shot generation trends match logit-based results, but hybrid and H3 models perform poorly zero-shot on some tasks because outputs can be irrelevant and unparsable.Few-shot examples help models generate responses in a parsable format.

F.8 Non-Text Sequence Modeling

H3 extends sequence modeling beyond language to raw EEG, speech, and fMRI. It achieves strong performance on long non-text sequences while matching GPT-based models on fMRI pretraining and decoding.

  • Raw speech and EEG: H3 outperforms Transformers on raw speech classification and achieves state-of-the-art seizure classification over raw EEG signals.These evaluations cover raw audio and EEG rather than text sequences.
  • Raw speech and EEG: Transformers cannot process the 12,000-step EEG clips without running out of GPU memory, whereas H3 can.The clips are 60 seconds long, sampled at 200Hz, and recorded with 19 electrodes.
  • Raw speech and EEG: H3 comes within half a point of S4 on the one-second, 16kHz SC10 raw speech commands task.The task contains raw audio sequences that Transformers cannot process at the stated sequence length.
  • Functional magnetic resonance imaging: The fMRI pipeline represents activity as sequences of 1,024 functionally independent brain-network signals and trains models to predict the next time point.Models are pretrained on sequences sampled from upstream fMRI runs before downstream mental-state decoding.
  • Functional magnetic resonance imaging: On fMRI pretraining, H3 with 0.2 dropout performs on par with GPT in mean absolute error and shows nearly identical brain-wide error distributions.Both architectures use four hidden layers; the GPT variant has 12 attention heads, while the H3 variant uses H = 64 and m = 1.
  • Functional magnetic resonance imaging: H3 and GPT-based models generally perform on par when adapted for mental-state decoding on two held-out fMRI test datasets.The models are initialized from their respective pretrained parameters and then adapted for classification.
Loading 2212.14052v3…