Hungry Hungry Hippos: Towards Language Modeling with State Space Models

Daniel Y. FuTri DaoKhaled Kamal SaabArmin W. ThomasAtri RudraChristopher Ré

article2023ICLR711 citations

Introduces the H3 state space architecture and FlashConv algorithm, enabling sub-quadratic language models scaled up to 2.7 billion parameters that surpass Transformers on SuperGLUE benchmarks while generating text 2.4 times faster.

Listen

Modern natural language processing and artificial intelligence rely heavily on Transformer models driven by attention mechanisms. However, standard attention scales quadratically with sequence length, creating severe computational and memory bottlenecks when handling long inputs. State space models offer a mathematically efficient alternative that scales near-linearly, but they have historically lagged behind Transformers in language modeling accuracy and suffered from poor hardware utilization on modern accelerators.

The article evaluates the root causes of this expressivity gap and demonstrates new model architectures and computational algorithms that enable state space models to match or exceed Transformer quality while dramatically improving execution speed.

To conduct this evaluation, the authors analyzed state space model failures on synthetic language tasks designed to test key-value associative recall and token comparison capabilities. Building on these diagnostic insights, they introduced a novel architecture named Hungry Hungry Hippo (H3), which stacks shift and diagonal state space layers with multiplicative gating to capture sequential dependencies. To eliminate hardware memory bottlenecks, they designed FlashConv, an algorithm combining fused block matrix-multiplication operations for fast Fourier transform convolutions with a chunked state-passing method for processing long sequences. The authors trained and tested pure H3 and hybrid models—retaining just two standard attention layers—on benchmark corpora including OpenWebText, the 800-gigabyte Pile dataset across scales from 125 million to 2.7 billion parameters, SuperGLUE benchmarks, and non-text modalities such as audio, electroencephalography, and functional magnetic resonance imaging.

The investigation produced four primary findings. First, diagnostic testing revealed that standard state space models struggle to remember tokens appearing after specific events and compare items across a sequence, but the H3 layer fully resolves these synthetic benchmarks. Second, pure H3 closed the language modeling perplexity gap on OpenWebText to within 0.4 points of Transformers, while a hybrid H3-attention model outperformed Transformers by 1.0 perplexity point. Third, when scaled up to 2.7 billion parameters on the Pile, hybrid models achieved lower perplexity and superior zero-shot and few-shot accuracy over Transformer baselines across a majority of SuperGLUE tasks. Fourth, the FlashConv implementation delivered a 2x speedup on long-range sequence benchmarks, accelerated long-sequence training by 4x to 8x over standard attention, and enabled 2.4x higher text generation throughput during inference.

These findings demonstrate that quadratic self-attention is not strictly necessary throughout deep networks to achieve state-of-the-art language understanding. Deploying hybrid architectures drastically lowers the operational compute costs, hardware memory footprints, and latency of running foundation models, unlocking practical real-time inference and the processing of ultra-long contexts such as continuous audio and medical sensor data.

Organizations developing large language and sequence models should explore hybrid state space architectures to reduce serving costs and accelerate data throughput. Decision-makers should consider piloting hybrid models for high-throughput text generation and long-context applications where standard Transformers encounter hardware memory limits. Further exploration is recommended to optimize layer placement and investigate whether similar state space mechanisms can fully replace attention across other multimodal pipelines.

Readers should note that while hybrid models demonstrated robust gains, the authors did not extensively tune training hyperparameters specifically for state space models, relying instead on standard GPT configurations. Additionally, pure H3 models exhibited performance degradation on certain generative zero-shot evaluation formats without prompt examples. Confidence in the reported scalability and inference speedups remains high based on comprehensive evaluations spanning multiple parameter scales and diverse sequential data domains.

Stanford University University at Buffalo computer hardware encoder-decoder few-shot learning state space models Hungry Hungry Hippos language modeling

arXiv: 2212.14052

Paper: Efficiently Modeling Long Sequences with Structured State Spaces, Albert Gu et al. (2022). Introduces the structured state space (S4) framework and continuous-time HiPPO parameterization that H3 diagnoses and adapts specifically for language modeling.
Paper: FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, Tri Dao et al. (2022). Establishes the IO-aware hardware optimization principles and fused kernel designs directly extended by H3's FlashConv algorithm.
Paper: Attention Is All You Need, Ashish Vaswani et al. (2017). Presents the foundational Transformer architecture and quadratic attention mechanism that H3 benchmarks against and partially retains in hybrid configurations.

Paper: Mamba: Linear-Time Sequence Modeling with Selective State Spaces, Albert Gu et al. (2023). Directly succeeds H3 by introducing input-dependent selective state spaces to solve associative recall natively without requiring hybrid attention layers.
Paper: Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality, Tri Dao et al. (2024). Unifies state space models and attention mechanisms through structured matrix duality, advancing the theoretical foundations behind H3 and selective SSMs.
Paper: Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling, Liliang Ren et al. (2025). Continues H3's hybrid sequence modeling strategy by combining selective state space blocks with localized sliding-window attention for long-context language tasks.
Paper: Mamba-3: Improved Sequence Modeling using State Space Principles, Aakash Lahoti et al. (2026). Pushes the boundaries of state space sequence modeling further by introducing refined discretization and complex state updates building upon the H3 and Mamba lineage.

Cover for Hungry Hungry Hippos: Towards Language Modeling with State Space Models

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 $\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.

Topics

computer hardware encoder-decoder few-shot learning fft algorithm state space state space model statistical language models

Keywords

state space models Hungry Hungry Hippos language modeling sequence modeling synthetic language tasks linear attention autoregressive language modeling recurrent models

1 Introduction
2 Background
2.1 State Space Models
2.2 Linear attention
3 Hungry Hungry Hippos Layer to Model Discrete Sequences
3.1 Motivation: Synthetic Language Modeling Tasks
3.2 H3 Layer
3.3 Expressivity
4 FlashConv: Efficiently Training SSMs
4.1 Fused Block FFTConv
4.2 State-Passing
5 H3 Evaluation
5.1 Language Modeling
6 FlashConv Evaluation
7 Conclusion
References
A Related Work
B Linear Attention and Time-Varying Systems
C Method details
C.1 Backward Pass
C.2 State-Passing Matrices
D Proofs
D.1 H3 Expressivity
D.1.1 Example Language Λ\Lambda
D.1.2 H3 Model to Solve Λ\Lambda
D.2 Attention Expressivity
D.3 H3 Complexity
D.4 State Passing Correctness
E Experimental Details
E.1 Synthetics
E.2 Model Architecture
E.3 OpenWebText Training
E.4 The Pile Training
E.5 SuperGLUE
E.6 Hardware
F Additional Experiments
F.1 LRA Accuracy
F.2 WikiText103
F.3 PG-19
F.4 Length Extrapolation
F.5 Scaling in Number of Tokens
F.6 H3 Language Model
F.7 Generation Performance
F.8 Non-Text Sequence Modeling

Knowls

Knowl 1 — H3 (Hungry Hungry Hippo) Layer Architecture
model/method
The H3 (Hungry Hungry Hippo) layer is a sequence modeling layer designed to match the expressivity of self-attention in discrete sequence tasks by stacking two linear state-space models (SSMs)—a shift SSM and a diagonal SSM—with multiplicative gating interactions across input projections.

Given an input sequence $u \in \mathbb{R}^{N \times d}$ , where $N$ is the sequence length and $d$ is the hidden dimension, the layer computes linear projections to generate queries $Q = u W_Q$ , keys $K = u W_K$ , and values $V = u W_V$ , with weight matrices $W_Q, W_K, W_V \in \mathbb{R}^{d \times d}$ .

The key projection $K$ is passed through a discrete shift SSM ( $\text{SSM}_{\text{shift}}$ ) with state dimension $m$ , defined by state matrix $A_{\text{shift}} \in \mathbb{R}^{m \times m}$ where $(A_{\text{shift}})_{i,j} = 1$ if $i - 1 = j$ and $0$ otherwise, and input matrix $B_{\text{shift}} = e_1$ . This shifts the input tokens across time, acting as a local delay line.

The representations are partitioned into $H$ attention-like heads, each of dimension $d_h = d / H$ . For each head $h \in \{1, \dots, H\}$ , the outer product of the shifted keys and values is formed at each token step: $K^{(h)} (V^{(h)})^\top \in \mathbb{R}^{N \times d_h \times d_h}$ . This sequence is processed along the sequence length by a diagonal SSM ( $\text{SSM}_{\text{diag}}$ ), whose state transition matrix $A_{\text{diag}}$ is diagonal and initialized via the S4D (HiPPO) framework to maintain long-range memory.

The output of the diagonal SSM is multiplicatively gated by the query projection $Q^{(h)}$ , concatenated across heads, and projected by $W_O \in \mathbb{R}^{d \times d}$ :
Input: Input sequence $u \in \mathbb{R}^{N \times d}$ , weights $W_Q, W_K, W_V, W_O \in \mathbb{R}^{d \times d}$ , shift SSM $\text{SSM}_{\text{shift}}$ , diagonal SSM $\text{SSM}_{\text{diag}}$ , head count $H$ , head dimension $d_h = d / H$
Output: Output sequence $y \in \mathbb{R}^{N \times d}$
Compute $Q = u W_Q$ , $K = u W_K$ , $V = u W_V \in \mathbb{R}^{N \times d}$
Apply shift SSM: $\tilde{K} = \text{SSM}_{\text{shift}}(K) \in \mathbb{R}^{N \times d}$
Split $Q, \tilde{K}, V$ into heads $Q^{(h)}, \tilde{K}^{(h)}, V^{(h)} \in \mathbb{R}^{N \times d_h}$ for $h = 1, \dots, H$
for $h = 1$ to $H$ do
    Form batched outer products: $M^{(h)} = \tilde{K}^{(h)} (V^{(h)})^\top \in \mathbb{R}^{N \times d_h \times d_h}$
    Process via diagonal SSM: $KV^{(h)} = \text{SSM}_{\text{diag}}(M^{(h)}) \in \mathbb{R}^{N \times d_h \times d_h}$
    Multiply by query: $O^{(h)}_i = Q^{(h)}_i KV^{(h)}_i \in \mathbb{R}^{d_h}$ for each timestep $i \in \{1, \dots, N\}$
end for
Concatenate heads: $O = [O^{(1)}, \dots, O^{(H)}] \in \mathbb{R}^{N \times d}$
Compute final projection: $y = O W_O \in \mathbb{R}^{N \times d}$
return $y$
Knowl 2 — State-Passing Algorithm for Scaling SSMs to Long Sequences
algorithm
The state-passing algorithm enables exact evaluation of state-space model (SSM) convolutions on sequence lengths $N$ that exceed the capacity of GPU fast on-chip SRAM (typically $N > 8192$ on NVIDIA A100 GPUs) by chunking the sequence and recurrently propagating the boundary hidden state between chunks.

An SSM defines a mapping $x_i = A x_{i-1} + B u_i$ , $y_i = C x_i + D u_i$ for state $x_i \in \mathbb{R}^m$ , input $u_i \in \mathbb{R}$ , and parameter matrices $A \in \mathbb{R}^{m \times m}$ , $B \in \mathbb{R}^{m \times 1}$ , $C \in \mathbb{R}^{1 \times m}$ , $D \in \mathbb{R}^{1 \times 1}$ . Given chunk size $N'$ (the largest sequence fitting in SRAM) such that $N = C \cdot N'$ , the input $u$ is partitioned into chunks $u^{(c)} \in \mathbb{R}^{N'}$ for $c = 1, \dots, C$ .

The output for chunk $c$ and its final state $x^{(c)}_{N'}$ are computed using the precomputed matrices $M_{ux} = [A^{N'-1} B, A^{N'-2} B, \dots, B] \in \mathbb{R}^{m \times N'}$ and $M_{xy} = [C^\top, (CA)^\top, \dots, (CA^{N'-1})^\top]^\top \in \mathbb{R}^{N' \times m}$ :

$y^{(c)} = M_{xy} x^{(c-1)}_{N'} + \text{BlockFFTConv}(f, u^{(c)}) + D u^{(c)}$ $x^{(c)}_{N'} = A^{N'} x^{(c-1)}_{N'} + M_{ux} u^{(c)}$

where $f = [CB, CAB, \dots, CA^{N'-1}B]$ is the length- $N'$ convolution kernel, and $\text{BlockFFTConv}$ denotes a fused block FFT convolution kernel executed entirely within SRAM.
Input: Input sequence $u \in \mathbb{R}^N$ , SSM parameters $(A, B, C, D)$ , chunk length $N'$ where $N$ is a multiple of $N'$ , filter $f = [CB, \dots, CA^{N'-1}B]$
Output: Output sequence $y \in \mathbb{R}^N$
Precompute $A^{N'} \in \mathbb{R}^{m \times m}$ , $M_{ux} = [A^{N'-1}B, \dots, B] \in \mathbb{R}^{m \times N'}$ , $M_{xy} = [C^\top, \dots, (CA^{N'-1})^\top]^\top \in \mathbb{R}^{N' \times m}$
Split $u$ into $C = N / N'$ chunks: $u^{(1)}, u^{(2)}, \dots, u^{(C)} \in \mathbb{R}^{N'}$
Initialize boundary state $x^{(0)}_{N'} = 0 \in \mathbb{R}^m$
for $c = 1$ to $C$ do
Compute chunk output: $y^{(c)} = M_{xy} x^{(c-1)}_{N'} + \text{BlockFFTConv}(f, u^{(c)}) + D u^{(c)}$
Update chunk end state: $x^{(c)}_{N'} = A^{N'} x^{(c-1)}_{N'} + M_{ux} u^{(c)}$
end for
Concatenate outputs: $y = [y^{(1)}, y^{(2)}, \dots, y^{(C)}] \in \mathbb{R}^N$
return $y$
Knowl 3 — Fused Block FFTConv in FlashConv
model/method
Fused Block FFTConv is a hardware-aware convolution method that accelerates FFT-based state space model convolutions on modern GPUs (such as NVIDIA A100/H100) by addressing both memory-bandwidth (IO) bottlenecks and compute efficiency for sequences up to 8K tokens.

Standard SSM training computes the linear convolution $y = \text{iFFT}(\text{FFT}(u) \odot \text{FFT}(f)) + D u$ . In standard libraries (such as cuFFT), each FFT and inverse FFT performs separate global memory reads and writes, making the operation bandwidth-bound. Fused Block FFTConv resolves this via two mechanisms:
1. Kernel Fusion: The entire FFT convolution sequence—forward FFT of input $u$ , pointwise multiplication with filter transform $\text{FFT}(f)$ , inverse FFT, and the residual term $Du$ —is fused into a single GPU kernel executed in fast on-chip SRAM.
2. Block FFT via Matrix Multiplications: To utilize specialized Tensor Core matrix multiplication hardware instead of scalar FP32 units, an $N$ -point Discrete Fourier Transform (DFT matrix $F_N$ ) with $N = N_1 N_2$ is decomposed using the Cooley-Tukey four-step FFT algorithm:
$F_N = P (I_{N_2} \otimes F_{N_1}) P^\top D (I_{N_1} \otimes F_{N_2}) P$

where $P$ denotes the permutation matrix that reshapes and transposes an $N_1 \times N_2$ array, $\otimes$ is the Kronecker product, $D$ is the $N \times N$ diagonal twiddle factor matrix, and $I_{N_i}$ is the $N_i \times N_i$ identity matrix. Because $(I_{N_2} \otimes F_{N_1})$ and $(I_{N_1} \otimes F_{N_2})$ are block-diagonal matrices, the FFT is computed as small matrix multiplications mapped directly to GPU Tensor Cores. For sequence lengths $N = r^p$ , the algorithm executes $O(N r \log N / \log r)$ FLOPs, achieving higher compute throughput on Tensor Cores than standard $O(N \log N)$ FFT implementations.
Knowl 4 — Hybrid H3-Attention Model Architecture
model/method
The Hybrid H3-Attention model is an autoregressive neural network architecture that combines H3 layers with a minimal number of standard multi-head self-attention layers to achieve superior language modeling performance while maintaining the fast generation capabilities of recurrent SSMs.

In an $N$ -layer transformer-style architecture (using pre-layer normalization and interleaved multi-layer perceptrons (MLPs)), all but two self-attention layers are replaced with H3 layers:
- First Attention Layer: Positioned at layer 2 (0-indexed layer 1).
- Second Attention Layer: Positioned in the middle of the network at layer $2 + N/2$ for an $N$ -layer model with $N$ even (e.g., layer 7 for a 12-layer model, layer 13 for a 24-layer model, or layers 10 and 21 for a 32-layer model).
Because $N - 2$ layers are H3 layers, which can be evaluated during autoregressive generation as constant-time $O(1)$ state updates without maintaining a key-value cache, the hybrid model achieves up to $2.4\times$ higher generation throughput than standard full-attention Transformers.
Knowl 5 — Computational Complexity of the H3 Layer
theoretical result
Let $N$ be the input sequence length, $d$ be the hidden dimension, and assume that the head dimension $d_h = d / H$ is of order $O(1)$ , where $H$ is the number of heads.

The time complexity of computing a forward pass of the H3 layer is:

$\mathcal{O}(d^2 N + d N \log N)$

and the space complexity is:

$\mathcal{O}(d N)$

Conditions and Scope: This holds when linear projections require $\mathcal{O}(d^2 N)$ time, the first 1D shift SSM convolution takes $\mathcal{O}(d N \log N)$ time using FFTs, and the second multi-head diagonal SSM takes $H \cdot \mathcal{O}(d_h^2 N \log N) = \mathcal{O}(d d_h N \log N) = \mathcal{O}(d N \log N)$ time. This represents an asymptotic improvement over standard self-attention, which requires $\mathcal{O}(d N^2)$ time and $\mathcal{O}(N^2)$ space.
Knowl 6 — Exactness of Chunked State-Passing SSM Convolution
theoretical result
For any 1D input sequence $u \in \mathbb{R}^N$ and state-space model parameters $(A, B, C, D)$ with $A \in \mathbb{R}^{m \times m}$ , $B \in \mathbb{R}^{m \times 1}$ , $C \in \mathbb{R}^{1 \times m}$ , and $D \in \mathbb{R}^{1 \times 1}$ , the output sequence $y \in \mathbb{R}^N$ produced by the state-passing algorithm (using chunk size $N'$ dividing $N$ ) is mathematically identical to the output defined by the full-length linear time-invariant SSM convolution:

$y_k = C A^k x_0 + \sum_{j=1}^k C A^{k-j} B u_j + D u_k$

with initial state $x_0 = 0 \in \mathbb{R}^m$ .

Preconditions: The convolution within each chunk is computed with the true SSM impulse response filter $f = [CB, CAB, \dots, CA^{N'-1}B]$ , and the boundary states between chunks are updated via $x^{(c)}_{N'} = A^{N'} x^{(c-1)}_{N'} + \sum_{i=1}^{N'} A^{N'-i} B u^{(c)}_i$ .

Knowl 7 — Synthetic Language Modeling Performance on Induction and Recall Tasks

empirical result

On synthetic language benchmarks that evaluate core in-context learning primitives—the Induction Head task (recalling the token following a special character earlier in the sequence) and the Associative Recall task (retrieving the value paired with a specific key across multiple stored key-value pairs)—2-layer models show that previous SSMs fail while H3 matches full self-attention.

All models were evaluated with 2 layers, inner dimension 32, MLP dimension 128, trained for 200 epochs on 5,000 examples. S4D and Gated State Spaces fail because their linear filters cannot simultaneously shift/log tokens after specific events and compare distinct token representations across the sequence. H3 solves both tasks via its shift SSM and multiplicative query-key-value gating.

Knowl 8 — Language Modeling Perplexity on OpenWebText, the Pile, and WikiText-103

empirical result

When evaluated on autoregressive language modeling, pure H3 substantially closes the gap between SSMs and Transformers, and Hybrid H3-Attention models outperform standard Transformer baselines across model scales from 125M to 2.7B parameters.

On 12-layer 125M models trained on OpenWebText for 50B tokens with identical hyperparameters:

S4D: 24.9 PPL
Gated State Spaces (GSS): 24.0 PPL
Pure H3: 21.0 PPL (within 0.4 PPL of Transformer)
Transformer: 20.6 PPL
GSS Hybrid (2 attention layers): 19.8 PPL
Hybrid H3 (2 attention layers): 19.6 PPL

Across models trained on the Pile for 400B tokens (evaluated on the Pile test set, zero-shot OpenWebText, and zero-shot WikiText-103):

Model	Pile PPL	OpenWebText PPL	WikiText103 PPL
GPT-Neo-125M	9.4	22.6	26.3
Hybrid H3-125M	8.8	20.9	23.7
GPT-2 medium (355M)	13.9*	17.0	21.8
Hybrid H3-355M	7.1	15.9	16.9
GPT-Neo-1.3B	6.2	13.1	13.3
Hybrid H3-1.3B	6.0	12.4	12.5
GPT-Neo-2.7B	5.7	11.7	11.5
Hybrid H3-2.7B	5.4	11.0	10.6

Note: GPT-2 was trained on WebText, so Pile PPL is for context.

Knowl 9 — SuperGLUE Downstream Zero-Shot and Few-Shot Evaluation

empirical result

Hybrid H3-Attention language models match or outperform similarly sized Transformer baselines (OPT and GPT-Neo) across zero-shot and 3-shot evaluation on the SuperGLUE benchmark using logit rank scoring.

Zero-shot SuperGLUE accuracy across benchmark sizes:

Model	WSC	WIC	RTE	CB	MultiRC	ReCoRD	BoolQ	COPA	Average
OPT-125M	39.4	52.0	48.7	37.4	58.9	44.9	59.6	60.0	50.1
GPT-Neo-125M	36.5	53.6	53.1	41.1	59.9	39.6	62.2	60.0	50.8
Hybrid H3-125M	39.4	51.4	59.2	48.2	51.4	55.0	59.6	67.0	53.9
GPT-Neo-1.3B	41.3	50.0	52.3	33.9	57.9	55.5	59.9	66.0	52.1
OPT-1.3B	36.5	49.5	53.4	39.3	58.3	61.8	55.0	69.0	52.9
Hybrid H3-1.3B	52.9	50.3	53.4	33.9	58.2	67.8	61.7	74.0	56.5
GPT-Neo-2.7B	37.5	50.0	52.3	50.0	59.1	60.0	61.1	67.0	54.6
OPT-2.7B	51.0	50.8	50.5	41.1	57.4	65.9	60.9	66.0	55.5
Hybrid H3-2.7B	36.5	51.3	57.0	37.5	58.7	71.3	61.1	81.0	56.8

In 3-shot evaluation, Hybrid H3-125M attains 53.7% average (vs. OPT-125M at 47.9% and GPT-Neo-125M at 47.2%), Hybrid H3-1.3B attains 53.0% (vs. OPT-1.3B at 49.4% and GPT-Neo-1.3B at 51.2%), and Hybrid H3-2.7B attains 55.5% (vs. OPT-2.7B at 53.0% and GPT-Neo-2.7B at 51.9%).

Knowl 10 — Execution Speedups and Inference Throughput with FlashConv

empirical result

FlashConv provides substantial training and inference speedups by accelerating FFT-based convolutions on modern accelerators and utilizing the recurrent formulation of SSMs during text generation.

Long Range Arena (LRA) Benchmark Speedup: Evaluating S4 with FlashConv on LRA yields a $2.0\times$ speedup over the previous S4 implementation and a $5.8\times$ speedup over standard Transformer:

Transformer: $1.0\times$
FlashAttention: $2.4\times$
Block-sparse FlashAttention: $2.8\times$
S4: $2.9\times$
S4 with FlashConv: $5.8\times$

Long Sequence Training Speed: On an A100-40GB GPU (batch size 8, hidden dimension 1024), kernel fusion accelerates short sequence FFTConv by up to $3.4\times$ (sequence lengths $\le 512$ ), block FFT yields up to $2.0\times$ speedup for sequence lengths 1K–8K, and state-passing allows $2.3\times$ faster computation for sequence lengths 16K–32K, running up to $35\times$ faster than FlashAttention at 32K length.
Autoregressive Inference Throughput: For 1.3B-parameter models on an 80GB A100 GPU (batch size 64, generating 128 tokens per prompt):

Hybrid H3-1.3B generates up to $2.4\times$ faster than the standard Transformer baseline at longer prompt lengths.

Knowl 11 — Non-Text Sequence Modeling Performance: Audio, EEG, and fMRI
empirical result
H3 models generalize beyond text to raw audio, EEG, and fMRI sequence modeling tasks that feature long sequence lengths:
1. Seizure Classification on Raw EEG: On the TUSZ v1.5.2 corpus (60-second clips sampled at 200 Hz, sequence length 12,000 with 19 electrodes), Transformers fail due to out-of-memory errors, while H3 achieves state-of-the-art performance with 83.2 AUROC (compared to Dense-CNN at 78.0, CNN-LSTM at 68.6, LSTM at 69.3, and 1D-CNN at 69.7).
2. Raw Audio Classification: On Speech Commands SC10 (1-second raw audio at 16 kHz, sequence length 16,000), Transformers run out of memory, while H3 achieves 97.04% accuracy, competitive with S4 at 97.50% and outperforming WaveGan-D (96.25%), Performer (30.77%), and CKConv (71.66%).
3. fMRI Mental State Decoding: Pre-training H3 on 11,980 fMRI runs from 34 datasets and adapting downstream to mental state decoding yields performance on par with a causal GPT baseline: HCP test accuracy of 88.75% (H3) vs. 88.44% (GPT), and MDTB test accuracy of 88.25% (H3) vs. 89.47% (GPT).

Coverage note — Omitted mathematical derivations and proofs (such as the gradient derivation of the convolution operator in Appendix C.1 and the inductive proofs of Propositions 2 and 3) in accordance with the requirement to state results without derivations or intermediate proofs.

References

1.Jimmy Ba, Geoffrey E Hinton, Volodymyr Mnih, Joel Z Leibo, and Catalin Ionescu. Using fast weights to attend to the recent past. Advances in neural information processing systems, 29, 2016.
2.Alexei Baevski and Michael Auli. Adaptive input representations for neural language modeling. In International Conference on Learning Representations, 2018.
3.David H Bailey. FFTs in external or hierarchical memory. The journal of Supercomputing, 4(1):23–35, 1990.
4.Sid Black, Leo Gao, Phil Wang, Connor Leahy, and Stella Biderman. GPT-Neo: Large Scale Autoregressive Language Modeling with Mesh-Tensorflow, March 2021. If you use this software, please cite it using these metadata.
5.Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint arXiv:2108.07258, 2021.
6.Willian L Brogan. Modern control theory, 1974.
7.Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901, 2020.
8.Kyunghyun Cho, Bart Van Merri¨enboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder-decoder approaches. arXiv preprint arXiv:1409.1259, 2014.
9.Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. In International Conference on Learning Representations (ICLR), 2020.
10.Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311, 2022.
11.James W. Cooley and John W. Tukey. An algorithm for the machine calculation of complex fourier series. Mathematics of Computation, 19(90):297–301, 1965.
12.Kamalaker Dadi, Ga¨el Varoquaux, Antonia Machlouzarides-Shalit, Krzysztof J Gorgolewski, Demian Wassermann, Bertrand Thirion, and Arthur Mensch. Fine-grain atlases of functional modes for fmri analysis. NeuroImage, 221:117126, 2020.
13.Zihang Dai, Zhilin Yang, Yiming Yang, Jaime G Carbonell, Quoc Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 2978–2988, 2019.
14.Tri Dao, Beidi Chen, Nimit Sohoni, Arjun Desai, Michael Poli, Jessica Grogan, Alexander Liu, Aniruddh Rao, Atri Rudra, and Christopher R´e. Monarch: Expressive structured matrices for efficient and accurate training. In International Conference on Machine Learning (ICML), 2022.
15.Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher R´e. Flashattention: Fast and memory-efficient exact attention with io-awareness. In Advances in Neural Information Processing Systems, 2022.
16.Giannis Daras, Nikita Kitaev, Augustus Odena, and Alexandros G Dimakis. Smyrf-efficient attention using asymmetric clustering. Advances in Neural Information Processing Systems, 33:6476–6489, 2020.
17.Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. In International Conference on Machine Learning, pages 5793–5831. PMLR, 2022.
18.Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. https://transformer-circuits.pub/2021/framework/index.html.
19.Bruce Fischl. Freesurfer. Neuroimage, 62(2):774–781, 2012.
20.Robert S Fisher, Carlos Acevedo, Alexis Arzimanoglou, Alicia Bogacz, J Helen Cross, Christian E Elger, Jerome Engel Jr, Lars Forsgren, Jacqueline A French, Mike Glynn, et al. Ilae official report: a practical clinical definition of epilepsy. Epilepsia, 55(4):475–482, 2014.
21.Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027, 2020.
22.Karan Goel, Albert Gu, Chris Donahue, and Christopher R´e. It’s raw! audio generation with state-space models. arXiv preprint arXiv:2202.09729, 2022.
23.Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanie Tellex. Openwebtext corpus, 2019.
24.Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher R´e. Hippo: Recurrent memory with optimal polynomial projections. Advances in Neural Information Processing Systems, 33:1474–1487, 2020.
25.Albert Gu, Karan Goel, and Christopher R´e. Efficiently modeling long sequences with structured state spaces. In The International Conference on Learning Representations (ICLR), 2022.
26.Albert Gu, Ankit Gupta, Karan Goel, and Christopher R´e. On the parameterization and initialization of diagonal state space models. In Advances in Neural Information Processing Systems, 2022.
27.Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher R´e. Combining recurrent, convolutional, and continuous-time models with linear state-space layers. Advances in neural information processing systems, 34, 2021.
28.Albert Gu, Isys Johnson, Aman Timalsina, Atri Rudra, and Christopher R´e. How to train your hippo: State space models with generalized orthogonal basis projections. arXiv preprint arXiv:2206.12037, 2022.
29.Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are as effective as structured state spaces. In Advances in Neural Information Processing Systems, 2022.
30.Curtis Hawthorne, Andrew Jaegle, C˘at˘alina Cangea, Sebastian Borgeaud, Charlie Nash, Mateusz Malinowski, Sander Dieleman, Oriol Vinyals, Matthew Botvinick, Ian Simon, et al. General-purpose, long-context autoregressive modeling with perceiver ar. arXiv preprint arXiv:2202.07765, 2022.
31.Sepp Hochreiter and J¨urgen Schmidhuber. Lstm can solve hard long time lag problems. Advances in neural information processing systems, 9, 1996.
32.Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556, 2022.
33.Sara Hooker. The hardware lottery. Communications of the ACM, 64(12):58–65, 2021.
34.Sheng-Chun Kao, Suvinay Subramanian, Gaurav Agrawal, and Tushar Krishna. An optimized dataflow for mitigating attention performance bottlenecks. arXiv preprint arXiv:2107.06419, 2021.
35.Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and Fran¸cois Fleuret. Transformers are RNNs: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, pages 5156–5165. PMLR, 2020.
36.Feyza Duman Keles, Pruthuvi Mahesakya Wijewardena, and Chinmay Hegde. On the computational complexity of self-attention. arXiv preprint arXiv:2209.04881, 2022.
37.Michael Patrick Kerr. The impact of epilepsy on patients’ lives. Acta Neurologica Scandinavica, 126:1–9, 2012.
38.Maedbh King, Carlos R Hernandez-Castillo, Russell A Poldrack, Richard B Ivry, and J¨orn Diedrichsen. Functional boundaries in the human cerebellum revealed by a multi-domain task battery. Nature neuroscience, 22(8):1371–1378, 2019.
39.Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In The International Conference on Machine Learning (ICML), 2020.
40.Shengjie Luo, Shanda Li, Tianle Cai, Di He, Dinglan Peng, Shuxin Zheng, Guolin Ke, Liwei Wang, and Tie-Yan Liu. Stable, fast and accurate: Kernelized attention with relative positional encoding. Advances in Neural Information Processing Systems, 34:22795–22807, 2021.
41.Christopher J Markiewicz, Krzysztof J Gorgolewski, Franklin Feingold, Ross Blair, Yaroslav O Halchenko, Eric Miller, Nell Hardcastle, Joe Wexler, Oscar Esteban, Mathias Goncavles, et al. The openneuro resource for sharing of neuroscience data. Elife, 10:e71774, 2021.
42.Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces. arXiv preprint arXiv:2206.13947, 2022.
43.Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016.
44.Eric Nguyen, Karan Goel, Albert Gu, Gordon Downs, Preey Shah, Tri Dao, Stephen Baccus, and Christopher R´e. S4nd: Modeling images and videos as multidimensional signals with state spaces. In Advances in Neural Information Processing Systems, 2022.
45.NVIDIA. Nvidia Tesla V100 GPU architecture, 2017.
46.NVIDIA. Nvidia A100 tensor core GPU architecture, 2020.
47.NVIDIA. cufft v11.7.1 documentation, 2022. https://docs.nvidia.com/cuda/cufft/index.html.
48.NVIDIA. Nvidia H100 tensor core GPU architecture, 2022.
49.Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread, 2022. https://transformer-circuits.pub/2022/in-context-learning-and-induction-heads/index.html.
50.Alan V Oppenheim. Applications of digital signal processing. Englewood Cliffs, 1978.
51.Alan V Oppenheim, John R Buck, and Ronald W Schafer. Discrete-time signal processing. Vol. 2. Upper Saddle River, NJ: Prentice Hall, 2001.
52.Markus N Rabe and Charles Staats. Self-attention does not need O(n 2 ) memory. arXiv preprint arXiv:2112.05682, 2021.
53.Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019.
54.Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, Chloe Hillier, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling. In International Conference on Learning Representations, 2019.
55.Khaled Saab, Jared Dunnmon, Christopher R´e, Daniel Rubin, and Christopher Lee-Messer. Weak supervision as an efficient approach for automated seizure detection in electroencephalography. NPJ digital medicine, 3(1):1–12, 2020.
56.Vinit Shah, Eva Von Weltin, Silvia Lopez, James Riley McHugh, Lillian Veloso, Meysam Golmohammadi, Iyad Obeid, and Joseph Picone. The temple university hospital seizure detection corpus. Frontiers in neuroinformatics, 12:83, 2018.
57.Mohammad Khubeb Siddiqui, Ruben Morales-Menendez, Xiaodi Huang, and Nasir Hussain. A review of epileptic seizure detection using machine learning classifiers. Brain informatics, 7(1):1–18, 2020.
58.Siyi Tang, Jared Dunnmon, Khaled Kamal Saab, Xuan Zhang, Qianying Huang, Florian Dubost, Daniel Rubin, and Christopher Lee-Messer. Self-supervised graph neural networks for improved electroencephalographic seizure analysis. In International Conference on Learning Representations, 2021.
59.Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers. In International Conference on Learning Representations, 2020.
60.Armin W Thomas, Christopher R´e, and Russell A Poldrack. Self-supervised learning of brain dynamics from broad neuroimaging data. arXiv preprint arXiv:2206.11417, 2022.
61.David C Van Essen, Stephen M Smith, Deanna M Barch, Timothy EJ Behrens, Essa Yacoub, Kamil Ugurbil, Wu-Minn HCP Consortium, et al. The wu-minn human connectome project: an overview. Neuroimage, 80:62–79, 2013.
62.Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
63.Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768, 2020.
64.Pete Warden. Speech commands: A dataset for limited-vocabulary speech recognition. arXiv preprint arXiv:1804.03209, 2018.
65.Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, R´emi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-the-art natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pages 38–45, Online, October 2020. Association for Computational Linguistics.
66.Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. OPT: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068, 2022.

Citation

MLA

Fu, D. Y., et al. “Hungry Hungry Hippos: Towards Language Modeling with State Space Models”. arXiv, 2022, http://arxiv.org/abs/2212.14052v3.

APA

Fu, D. Y., Dao, T., Saab, K. K., Thomas, A. W., Rudra, A., & Ré, C. (2022). Hungry Hungry Hippos: Towards Language Modeling with State Space Models. arXiv. http://arxiv.org/abs/2212.14052v3

Chicago

Fu, D. Y., T. Dao, K. K. Saab, A. W. Thomas, A. Rudra, and C. Ré. 2022. “Hungry Hungry Hippos: Towards Language Modeling with State Space Models”. arXiv. http://arxiv.org/abs/2212.14052v3.

Harvard

Fu, D.Y. et al. (2022) “Hungry Hungry Hippos: Towards Language Modeling with State Space Models”, arXiv [Preprint]. Available at: http://arxiv.org/abs/2212.14052v3.

Vancouver

1. Fu DY, Dao T, Saab KK, Thomas AW, Rudra A, Ré C (2022) Hungry Hungry Hippos: Towards Language Modeling with State Space Models. arXiv

BibTeX

@article{fu2022hungry,
  title = {Hungry Hungry Hippos: Towards Language Modeling with State Space Models},
  author = {Fu, Daniel Y. and Dao, Tri and Saab, Khaled K. and Thomas, Armin W. and Rudra, Atri and Ré, Christopher},
  year = {2022},
  journal = {arXiv},
  url = {http://arxiv.org/abs/2212.14052v3},
  eprint = {2212.14052}
}

Metadata:arXiv

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: Authors

Model	Prompt len 512	Prompt len 1024	Prompt len 1536
Transformer-1.3B	1340 tokens/s	770 tokens/s	520 tokens/s
Hybrid H3-1.3B	1980 tokens/s	1580 tokens/s	1240 tokens/s

Task	Random	S4D	Gated State Spaces	H3	Attention
Induction Head	5.0	35.6	6.8	100.0	100.0
Associative Recall	25.0	86.0	78.0	99.8	100.0

Hungry Hungry Hippos: Towards Language Modeling with State Space Models

Abstract

Topics

Keywords

Table of Contents

Knowls

Knowl 1 — H3 (Hungry Hungry Hippo) Layer Architecture

Knowl 2 — State-Passing Algorithm for Scaling SSMs to Long Sequences

Knowl 3 — Fused Block FFTConv in FlashConv

Knowl 4 — Hybrid H3-Attention Model Architecture

Knowl 5 — Computational Complexity of the H3 Layer

Knowl 6 — Exactness of Chunked State-Passing SSM Convolution

Knowl 7 — Synthetic Language Modeling Performance on Induction and Recall Tasks

Knowl 8 — Language Modeling Perplexity on OpenWebText, the Pile, and WikiText-103

Knowl 9 — SuperGLUE Downstream Zero-Shot and Few-Shot Evaluation

Knowl 10 — Execution Speedups and Inference Throughput with FlashConv

Knowl 11 — Non-Text Sequence Modeling Performance: Audio, EEG, and fMRI

References

Citation

BibTeX

Access the Paper