Hungry Hungry Hippos: Towards Language Modeling with State Space Models
Authors: Daniel Y Fu
Tri Dao
Khaled Kamal Saab
Armin W Thomas
Atri Rudra
Christopher Re
Venue: ICLR 2023 notable top 25% (2023)
Keywords: language modeling
state space models
efficiency

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.

Published as a conference paper at ICLR 2023
HUNGRY HUNGRY HIPPOS: TOWARDS LANGUAGE MODELING WITH STATE SPACE MODELS
Daniel Y. Fu"i, Tri Dao"i, Khaled K. Saab', Armin W. Thomas Atri Rudra', Christopher Ré
i Stanford University, University at Buffalo, SUNY {danfu, tridao aesidaaeecssstaffferr.. stanford. edu, (ksaab, at hms (ataas)ssaofrd..co atrisbuffalo. edu, chri stanford, edu
ABSTRACT
State space models (SSMs) have demonstrated state-of-the-an 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 of Transformers on WebTee. Furthermore, hybrid 125M-parameter H3-attention model that retains surprisingly outperforms Transformers on OpenWeeTeeee by 1.0 PPI.. Next, to improve the efficiency of training SSMs on modern hardware, propose FLASHHONV. FLASHCONV uses a fused block FFT algorithm to improve efficiency on sequences up 8K, and introduces a novel state passing algorithm that exploits the recurrent properties of SSMs to scale to longer sequences. FLASHCONV yields 2 speedup the long-range arena benchmark and allows hybrid language models to generate text 2.4x faster than Transformers. Using FLASHCONV, we scale hybrid H3-attention language models up to 2.7B parameters the Pile and find promising initial results, achieving lower perplexity than Transformers and outperforming Transformers in zero- and few-shot learning on majority of tasks in the SuperGLUE benchmark.
attention layers
INTRODUCTION
State space models (SSMs) have achieved -tate-of-the-art sequence modeling performance in domains ranging from time series analysis (Gu et al., 2022a) to audio generation (Goel et al., 2022). However, they have yet to match the performance of Transformers on language modeling, often underperforming Transformers by multiple points in perplexity (Gu et al., 2022a). An natural question is whether this gap in performance is due to inherent inductive biases and capabilities in attention (Edelman al., 2022; Olsson et al., 2022), or whether it is a function of the significant organizational resources that have been spent training and tuning large attention- based language models (Chowdhery al., 2022; Hoffmann et al., 2022; Zhang et al., 2022), as well as specialized hardware support for attention, ranging from tensor cores (NVIDIA, 2017) to transformer chips (NVIDIA, 2022h; Kao et al., 2021).
We take first steps towards answering these questions in this paper. First, we use synthetic language modeling tasks to show that there is an expressivity gap between SSMs and attention. Using our insights. we design new SSM layer that nearly matches attention language modeling. Second, we propose better hardware-awart algorithms for SSMs that allow them to take advantage of modern accelerators-and run faster than attention.
Understanding the Expressivity Gap. To understand the gap between SSMs and attention, we draw on synthetic language modeling tasks that have been proposed as mechanistic basis for in-context learning in Transformers (Olsson et al., 2022) These synthetic languages focus on the ability to manipulate text-recalling tokens from earlier time steps, or comparing tokens from different points in sequence. We find that existing SSMs struggle model these synthetic languages. To probe how important these
"Equa Contribution, Order determined by coin flip.

---

Published as a conference paper at ICLR 2023
nulf
Yeu
nulf
3
Out
null
Diag SSM
Diag
Diag
Diag
(x--a)
Il (a)
Fused Block
SSM State
Fused Block FFTConv
SSM State
Shift SSM
Shift
Shift
Shift
FFTConv
NXgmal
K
Input
,,,
X
Store key Store val
Recall val
8K Chunk
H3 Layer
H3 for Associative Recall
FlashConv
Figure I: Left: stacks two discrete SSMs with shift and diagonal matrices and uses multiplicative interactions between input projections and their outputs to model comparisons between points in sequence. Middle: H3 can perform associative recall- scall-which is easy for attention, but not existing SSMs. Right: FLASHCONV uses a new state-passing algorithm over fused block FFTConv to increase hardware efficiency of SSMs, allowing H3 to scale to
models.
skills are for language modeling, we propose H3 (Hungry Hungry Hippo), a new SSM-based layer designed to solve these language modeling tasks. H3 stacks two SSMs, with multiplicative interactions between their outputs and input projections. The SSMs allow H3 keep a log of tokens (to recall them later), while the multiplicative interactions allow for comparisons across the sequence.
H3 matches attention on the synthetic languages and almost closes the gap with Transformers on language modeling-coming within 0.4 perplexity Transformers on Open WebText (compared to 3.4 ppl for ex- isting SSMs-even those explicitly designed for language modeling (Mehta et al., 2022)). Furthermore, a simple hybrid H3-attention model that retains two attention layers surprisingly oulperforms Transformers on Open by 1.0 perplexity. To further evaluate H3 on language modeling, we train 125M-, 355M-, 1.3B-and 2. 2.7a-parrmeter hybrid H3-attention language models on the Pile (Gao et al., 2020), using hyperparameters from GPT-3 (Brown et al., 2020). These hybrid models outperform Transformer-base language models of the same size in perplexity, and match or outperform them on a majority of tasks the SuperGLUE benchmark in zero- and few-shot learning. Since the SSM layers in these hybrid models admit a recurrent view, they can also perform 2.4x faster inference than Transformers.
Scaling SSMs. Next, we improve the efficiency of SSMs on modern hardware, to reduce the hardware barrier between attention and SSMs. SSMs scale nearly linearly in sequence length instead of quadratically like attention, but still run slower on modern hardware due to poor hardware utilization. To close this gap, we propose FLASHCONV, hierarchical algorithm for computing SSMs, inspired by IO-Aware attention (Dao et al., 2022b). The technical challenge that SSMs require FFT-based convolution over the input sequence, which requires an FFT, pointwise multiply, and inverse FFT. When implemented in cuFFT (NVIDIA, 2022a), this operation incurs expensive GPU memory reads/writes, and cannot utilize the specialized matrix multiply units available on modern hardware To use specialized matrix multiply units, we appeal classical techniques that split the FFT into blocks and compute using series of matrix multiplications. Combined with kernel fusion, this "block" FFT solution increases hardware efficiency, but only as long as the sequence length can fit into GPU SRAM (on-chip memory, analogous to L1 cache on the CPU) up to sequence length 8K on modern A 1000.
To scale to sequences longer than 8K, we propose a state passing algorithm (Figure right), specialized to SSMs. The key insight is that we can use the recurrent properties of SSMs to process the input in chunks-aa as long we keep track of an additional state vector. The state passing algorithm splits the input into the largest chunks that can fit into GPU SRAM, efficiently computes the FFT- FFT-based convolution using block FFT, and updates an intermediate state to start the next chunk. Using this state- satte-passig algorithm, FLASHCONV can scale SSMs to any sequence length- even longer than can fit on GPU SRAM at once once-while while maintaining near linear compute complexity. FLASHCONV sets state-of of-the-art speed on long range arena using S4 (Gu et al., 2022a), outperforming Transformers by 5.8x and previous S4 models by 2x. FLASHCONV trains H3 4-8 x times faster than attention for long sequences, and is a critical component for scaling billion parimmee models?.
'n A 1000 GPU has a maximum of 312 TFLOPs/s of FP16 with tensor cores, but only 20 TFLOPs/s of FP32 (and 40 TFLOPs/s of FP16) without tensor cores (NVIDIA, 2020). This trend started with the V100 GPUs (NVIDIA, 2017) and has continued with the H100 GPUs (NVIDIA, 2022h). Code for H3 is available at https:: tuub. com wtt/ Haw.///azyeeeeeeee raceesearchh// /H.

---

Published as a conference paper at ICLR 2023
2
BACKGROUND
present some background on state space models and linear attention, which inspired our H3 layer.
2.1 STATE SPACE MODELS
A continuous-ting state-space representation (Brogan, 1974) defines a linear mapping from an input
signal u(t) E (as function of time to an output signal y(t) E IR through state-variable x=t) () Rm with the following differential equation, for some matrices A E Immm B E R C E R R
D P fxt)= = Ar(())+ Btt(). vtt C+(()) +Du(t).
Similarly, a discrete-time state-space representation defines linear mapping from a discrete input signal u; (for i = ..). to a discrete output signal ys though a state-variable R
22 = -== Bu,
y = Co, Dus-
state-space model (SSM) uses these representations as a layer in deep learning pipeline, where the matrices A,B, C,D are learned from data (e.g.. with gradient-based optimization). One often has d of these SSMs parallel, each corresponding to one hidden dimension. To preserve the sequence history, HiPPO (Gu et al., 2020) projects the history having SSMs whose A,B matrices are initialized to some special matrices.
basis of orthogonal polynomials, which translates
This recurrent form of SSMs allows efficient inference (i.e., generation): to generate the output of the next time-step, one only needs the state of the current time-step, not the entire input history. Furthermore, SSMs can freely extrapolate to sequences longer than seen during training.
SSMs as Convolution. For efficient training, given the entire sequence of the input u1.-.,uN,, the output sequence .NN can also be written the convolution of the input with the filter (Gu al., 2021):
f [CB, CAB, C.., BB).
That is, from initial condition Zo. we have y; = CA'Bzo + -u]; +Du, where (f*u) denotes linear convolution between and u. If set the initial condition zg to be zero, then y exactly a linear convolution of with residual connection Du. More generally, any linear time-invariant system (of which SSMs are special case) can be written as a convolution.
Given a ID input sequence u RR N of length N, we denote the 1D output sequence y E E R.N of an SSM parameterized by matrices A,B,C,D
y= SSMMA,B..D
To simplify notation, we omit the reference to C,B.C,D and write y = SSM(u) if they are clear from context. When is multidimensional of dimension d, we stack d of these SSMs together that defines a mapping from ue RR N'xx to y E IR using the same notation y= SS.(Mu
To construct the filter / from A,B,C B efficiently, A is often constrained to be diagonal (Gupta et al., 2022; Gu et al., 2022b), diagonal plus low-rank (Gu et al., 2022a).
SSM through FFTss Computing the convolution naively through conventional matrix operations is expensive for long kernels, scaling as O (22). Instead, we can use FFTs:: take the FFT of f and u, multiply them together pointwise, and then take the inverse FFT. This yields O(Nlog algorithm.
2.2 LINEAR ATTENTION
We describe linear attention (Katharopoulos et al., 2020) and its connection to RNNs, which inspired model design (Section 3).
In standard attention (Vaswani et al., 2017), have N query/key/value tokens K,, V, Rd for i= 1.....NN N, where N is the sequence length and dl is the head dimension. For some similarity metric Sim: Rd R + R. want compute the output:
Σ - Si,
O,
Sin_ (Q,K,) E ERd
For standard softmax attention, Sim(q.k) E (often the dot product is scaled by 1/Vd). Linear atten- tion makes the assumption that Sim has the form Sim(g.k)) = )(q) ( ((k), for some (nonlinear) function ø.
The output then O, x(Q+) ((K)) V
Let S, -
gadd 24 oj=, gd
x(Q,)
d, = ((Q.) 2, E Then O, = TQQ.) This connects linear attention to RNNs: the output O, is function of S, and 2,, both of which are incrementally updated (as cumulative sums).
d.

---

Published as a conference paper at ICLR 2023
3 HUNGRY HUNGRY HIPPOS LAYER TO MODEL DISCRETE SEQUENCES
To understand the gap between SSMs and attention on language modeling, we examine two synthetic language modeling tasks. These tasks motivate our H3 layer to add a discrete SSM (hased on shift matrix) and multiplicative interaction to effectively model discrete sequences. We then show that the H3 layer is expressive enough to solve these synthetic tasks, and that this understanding leads to better performance on a real language modeling benchmark.
3.1 MOTIVATION: SYNTUETIC LANGUAGE MODELING TASKS
We describe two coosely-related synthetic tasks, summarized in Table 1. Olsson et al. (Olsson et al., 2022) argue that the ability to solve (variants of) these tasks accounts for the majority of the in-context learning capability of Transformers, and more intuition is given in Appendix E.
Table 1: Synthetic language modeling tasks.
Task
Input fbc xyz
Output
Sequence Length
Vocab Size
Induction Head
/
30
20
Associative Recall
a2c4b3dla
20
10
The Induction Head task tests how well a model can recall content after a special token (e.g., f in Table I). At the end of the sequence, the model must recall the token that appeared immediately after the special token earlier in the sequence. Associative Recall (Ba et al., 201 6) is similar to the induction head task, but requires the model remember multiple key-value pairs. At the end of the sequence, the model must recall a specific value belonging to a specific key.
Table 2: Evaluation of 2-layer models on synthetic language tasks.
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
Table 2 (for two-layer models) shows that S4D (Gu et al., 2022b) and Gated State Spaces (Mehta et al., 2022) both fail to model these synthetic languages, which suggests they may not have the expressivity for general language. We argue that these failures suggest two missing capabilities: (i) to remember tokens that appear after a particular event (e.g.. .. the special token in the induction head task), and (ii) to compare tokens across the sequence (c.g., comparing keys to decide which value to recall). Attention has both these capabilities: can compare tokens by constructing the quadratic attention matrix QK" and it can recall tokens by direct copying (multiplying softmax QK with V). In Section 3.2, we design our new layer H3 enable these capabilities in SSMs, narrowing the expressivity gap between SSMs and attention.
3.2 H3 LAYER
H3 uses SSMs with shift and diagonal matrices, along with multiplicative operations against projections of the input to capture the missing capabilities identified by the synthetics.
High-level Intuition. (i) To remember tokens from the past, we want the state 24; to copy from the input
Wi and then pass that information to the next state I/+1- As Fi-1 relates to br Ar, we use a discrete SSM with shift matrix A (described formally below) that shifts the elements of state vector (e.g., mapping bo, - (0.o,b). (ii) To compare tokens across the sequence, we use multiplicative interaction: the output of an SSM, containing information from previous time steps, is multiplied with the input at the current time steps, thus measuring similarity between tokens.
H3 is loosely inspired by linear attention (Section 2): we project the input u to get three signals Q.K, V. Then we replace the non-linearity @ (K) with an SSM where A is a shift matrix (SSMahit), and we replace the summation S; with SSM with diagonal A (SSMaiag). The output, for the case of head dimension d; - is:
Q (SSMahin (K)&V),
where @ denotes pointwise multiplication. We can view this form as stacking two SSMs with multiplicative interaction (each is a "hungry hippo", hence the name of our layer). A more formal connection between linear attention, time-varying systems, and H3 can be found in Appendix B.
Shift and Diagonal SSMs. The shift and diagonal SSMs designed address the capability to log tokens after particular events. In the shift SSM, we constrain A E Rm to be
Remembering Key Tokens:
fori-1=j
a shift matrix A 0 otherwise The action of this matrix on the hidden state x; is to shift each coor-
dinate down by one thereby creating a "memory" of the previous states. For example, if B =1, the first

---

Published as a conference paper at ICLR 2023
basis vector, then 2; [,,, contains the inputs from the previous time steps. We learn B and (B can also be fixed 61 for simplicity, in which case the output is a ID conv. with kereel size m).
The diagonal SSM constrains A to be diagonal and initializes it from the diagonal version of HiPPO (S4D (Gu et al., 2022b)). This parameterization allows the model to remember state over the entire sequence. The shift SSM can detect when a particular event occurs, and the diagonal SSM can remember token afterwards for the rest of the sequence.
Multiplicative Interaction for Comparison. We take the multiplicative interactions from linear attention, but they provide another missing capability when combined with a shift matrix: comparing tokens across the sequence. The multiplicative interactions between the output of the shift SSM and the V projection mimics local multiplicative interactions in linear attention (depending on the size of the hidden state). Similarly, multiplicative interactions with the projection and the output of the diagonal SSM allows comparisons between tokens over the entire sequence.
H3 Layer. The overall layer is given in Algorithm and shown schematically in Figure (left). We use the H3 layer construct model in the same style Transformers by interleaving it with MLPs, connected by residual connection and layer norm (i.e., pre-norm architecture (Baevski & Auli, 2018)). We will also consider hybrid H3 -1tttteenooon model (two attention layers while the rest are H3, Sections 3.3 and 5).
Algorithm 1
H3 Layer
Require: Input sequence R N from the previous layer, weight matrices Wa, WK, Wy, Wo WoeRexd ashift SSM a diagonal SSM SSMaiag head dimension da.
Compute Q u W K W K, V uWv Nxd
2. Pass K through the shift SSM: K = SSMabin (K)RR*
E Nxd
Split Q.K. Vin into H "heads" (Q
(h) ( forh= H), each sequence of N vectors size dy =d/ H.
for 1 H do
5: Take the batched outer product K (h)
(V()) ER
Nxd, x (batched in the N-dimension) and pass it through (V(h)) ERR Xd,,
diagonal SSM: K (h) = SSMaing K
6: Batch-multiply by Q: O (h) = [q, KV O
(M)
KV
KV (batched the N-dimension).
7: end for
8: Concatenate the output O of cach head, and multiply by the output projection matrix Wo ERdxd
Efficiency We show that H3 scales as ONN log N) with sequence lenghh a-asymptoticllly more efficient than attention, which typically requires O (N time and O(N2) space" (proof in Appendix D.3). Proposition 1. Let N be the sequence length, d be the hidden dimension, and assume that the head dimen- sion dy of order O(1). Then the H3 layer takes Odd N +dN +dNlogN N) time and (dN) space to compute.
3.3 EXPRESSIVITY
We show that H3 can model our synthetic languages, as well as natural language on OpenWeb- Text (Gokaslan et al., 2019). We also present a hybrid H3-attention extension that outperforms Transformers on OpenWeeTexx
Mechanism for Solving Associative Recall with H3. H3 is expressive enough to solve our synthetic language modeling tasks, shown in Table 2. Figure/ 1 (middle) shows mechanism for single H3 layer to solve the associative recall task for a particular key-value pair (a,3). The shift SSM and following multiplicative interaction act gate whether to let value through to the diagonal SSM, based on whether the previous token was key a. The diagonal SSM stores the value 3 in memory, and continually outputs it. The final multiplicative interaction gates whether to let the diagonal SSM's 's' output through-bast on whether the current input token is the key a. We formally construct the weights of an H3 layer to solve this task in Appendix D.I.
Table 3: Perplexity of SSM variants compared to Transformers on Open WebText. All models have 12 layers, with size around 125M, and are trained with the same hyperpameters, for 50B tokens.
113
H1 Hybrid Attn)
S4D
GSS GSS Hybrid (2 Atln)
Transformer
21.0
19.6
24.9
24.0
19.8
20.6
Better Synthetic Language Modeling Translates to Better Natural Language Modeling. We validate that when H3 can solve these synthetic tasks, it also improves the modeling capability on natural language
'There are several memory-efficient algorithms for attention (Rabe & Staats, 2021; Dao et al., 2022b), though their time complexity is still quadratic in N, which is a lower-hound for attention (Keles et al., 2022).

---

Published as a conference paper at ICLR 2023
(e.g., the OpenWebText dataset). As shown in Table H3 comes within 0.4 perplexity points of Transformers when trained for 50B tokens on OpenWebText, and performs much better than existing variants (S4D, GSS), by 3.9 points.
Extension: H3-attention Hybrid Model. A simple hybrid H3 language model surprssingyy outperforms Transformers on Text by 1.0 point. Our hybrid model simply retains two self-attention layers: one in the second layer, and one in the middle (layer 2 2+-N/2 2 for an N-layer model, N even). The H3-attention hybrid also outperforms the GSS-attention hybrid (Mehta et al., 2022).
4 FLASHCONV: EFFICIENTLY TRAINING SSMS
To improve the efficiency of SSMs on modern hardware, we propose FLASHICONV. FLASHCONV fuses the FFT, pointwise multplly, and inverse FFT to reduce memory reads/writes. It also uses a block FFT algorithm to make use of specialized matrix multiply units (e.g., tensor cores on A A00) for sequence lengths up to 8K. For sequences longer than 8K, the computation no longer fits in GPU SRAM', so we propose a novel state- plate-passing algorithm that splits the sequence into chunks to compute the FFT convolution one chunk a time. FLASHCONV can speed up any SSMs (not just H3).
4.1 FUSED BLOCK FFTCONV
We deploy two techniques to speed up the FFT-hased convolution for sequences shorter than 8K: kemne fusion and block FFT. Kernel fusion addresses IO bottlenecks due reading and writing of intermediate results, while block FFT allows the FFT-based convolution to utilize specialized matrix multiplication units. These techniques allow us speed up FFTConv by 2x (Section 6) for sequences shorter than 8k.
Naive implementations of FFTConv using standard libraries such as cuFFT are IO-bound due to repeated reading and writing of intermediate results. The FFT convolution in an SSM with input and filter f has the form iFFT (F (T @ FFT()) (where O denotes pointwise multiplication). It requires reading and writing intermediate results to GPU memory-whieh can dominate the runtime. Following FLAS W...SATTTETTON (Dao et al., 2022b), we first fuse the entire FFTCCnv into a single kernel and compute it in SRAM to avoid this overhead.
Kernel Fusion.
Block FFT. To further speed up the computation of FFT-based convolution, we exploit specialized matrix multiplication hardware on modern GPUs (e.g.. Tensor Cores on Nvidia GPUs perform fast 16 x 16 matrix multiplication). We appeal to classical results that show that the FFT can be written as a series of block-diagonal matrix multiplications interleaved with permutation. We note that such algorithms are not new, but our setting (fused FFTConv on GPU) introduces new bottlenecks- by removing the I0 bottlenecks, compute becomes the bottleneck (note that single FFT on GPU is usually IO bound).
Suppose that we want to perform an --point FFT, which is equivalent to multiply by the DFT matrix F N.. Suppose that N = N,,, N2 for some integers N2, By the Cooley- Sukey decomposition of the DFT (Cooley & Tukey, 1965; Bailey, 1990) (also known as the four-step FFT algorithm), we can write FN= P[IN, Fxx, D @ F N₂ P. where P denotes a fixed permutation that reshapes the input as N₁ x N₂ array and then transpose it, & denotes Kroneker product, D is a N x N diagonal matrix (called the twiddle factors) (Dao et al., 2022a), and IN, and F N, are the identity and DFT matrix of size N N₁. As NN2 & N and IN, F N2 are just block-diagonal matrices, we can make matmul units to perform these multiplications. Similarly, if N = N₁ N3 then we can decompose the N-point FFT into series of (block) FFT of size N1, N2. and N3. interleaved by permutation.
of specialized
The block FFT algorithm incurs ONVrlog /lgr)) FLOPs for a sequence length N, if N can be written as rP for two integers r.p. This incurs more FLOPs than standard FFT (O(NlogNN) but can run faster when we using specialized matrix multiplication hardware.
4.2 STATE-PASSING
However, the fused kernel cannot run if the sequence is too long to fit into GPU SRAM (longer than 8K on A A100). We show how to exploit the particular form of the FFT in SSM to speed it up for long sequences.
The recurrent nature of SSMs allows us split the FFTConv of a length- NN sequence into chunks of size N' each (N' is the longest FFT we can fit into SRAM), assuming N is a multiple of N'). We use FFTConv to compute each chunk and use a recurrence to connect the chunks. In particular, we split the inputs u into C= C/N chunks a(e) N for c= 1,...,C. C. Similarly, split the states x into xx()) e R xm
and the output y into y(t) e N for -ori=1a I....,C. We will only need the end-state N of each chunk c.
(c)
'SRAM, or ch-p memory, is much faster than off- GPU memory, but usually much smaller, on the order of around 100KB for each streaming processor.

---

Published as a conference paper at ICLR 2023
Table Perplexity (lower is better) of models the Pile, Texx and WikiText- 1033. GPT-Neo and hybrid H3 trained the Pile, while GPT2 trained WehText. All models the GPT2 tokenizer. We report the perplexity of GPT-2 models they were trained on different data.
Pile (*) context, though the performance not directly comparable since
Model
Pile
Open Weh
Tex
GPT-2 small (125l (125M)
19.0*
22.6
29.9
GPT-Neo- 125M
9.4
22.6
26.3
Hybrid H3-125M
8.8
20.9
23.7
GPT-2 medium medhum (355M)
13.9
17.0
21.8
Hybrid H3-355M
7.1
15.9
16.9
GPT-2 XL (L.5B)
12.4*
12.9
17.0
GPT-Neo-1.3. 1.3B
6.2
13.1
13.3
Hybrid H3-1.3B
6.0
12,4
12.5
GPT-Neo-2.7B
11.7
11.5
Hybrid H3-2.7B
5.4
11.0
10.6
Let f (CB/=CCBBCAB.CABBB B...B be the SSM filter. Recall from Section that for each chunk
CAN'
Vi (0) CA'BE (c-1) + ( (c) since Tiy" (c-1) the end-state of the previous chunk (e 1)
(c-1) + f (D is the initial condition for the current chunk In vector notation, (()) = M
for some matrix May E NR xm, Additionally need update the end-state of each chunk with xx = N' (c-1), +Maz (()) for some matrix MmxX" (derivation in Appendix C.2). In essence, we can compute the output for each chunk with FFT-based convolution long remember the end-state of the previous chunk, and the end-state of each chunk can be updated recurrently. This yields a state-passing algorithm for long sequences, where only compute FFT of length N', and update some hidden state each iteration.
N
Let BLOCKFFTCONY refer to our fused block FFTConv kernel. Then, the state-passing algorithm for ID input given by Algorithm 2. For inputs of dimension d where stack d SSMs, simply batch Algorithm 2 along the d-dimension.
Algorithm 2 State Passing Algorithm
Require: Input E R SSM parameterized by matrices A E R mm size N' where N is multiple of N',
BERTXI CERIXII 1mm
DERIXI chunk
Precompute A N' E Rm M.. BA.."'
Mc.,IIC,CA,,,.CAA CA CA
2: Split the imputs W1.N into N/N chunks (e) "i: for==1,....C.
Let the initial state be "xx' ""... for 15c5 C do
(0)
5. Compute (() = =M (c-1)
BLOCKFFTCON(( "
Update state: TN A
(c)
N' (c-1)
M (0
end for
Return y (1.....))
prove that Algorithm yields the same output as if one has computed the SSM using large FFT of size N (proof in Appendix D.4):
Proposition For input u ERN and matrices A,B,C,D, the output E N returned by Algorithm 2 is the same the output defined by the SSM parameterized by A.B.C.D.
H3 EVALUATION
To understand how well capturing the synthetics in Section 3.1 translates to language modeling, we train two hybrid hybrid H3 -tttteenioon language models sizes 125M, 355M, 1.3B, and 2.7B, and we evaluate their performance against Transformers. The hybrid models match or exceed the quality of Transformers perplexity and zero/few-she learning. We also validate that H3 models retain strong performance on non-text sequence modeling. Appendix F contains additional experiments on more datasets, length extrapolation, and scaling with data.
5.1 LANGUAGE MODELING
We compare hybrid H H3-attention language models against Transformer-basee language models. We evaluate language modeling performance using perplexity, zero-shot learning. and few-shot learning performance. Hybrid H3 models outperform Transformers, which suggests that closing the gap between

---

Published as a conference paper at ICLR 2023
Table 5: Zero-shot on SuperGLUE with logit scoring. Best results in bold, second best underline.
Model
WSC
ReCoRD
BoolQ
COPA
Average
OPT-125M
39.4
52.0
48.7
57.4
58.9
44.9
59.6
60.0
50.1
GPT-Nen- 125M
36.5
53.6
53.1
41.1
59.9
39.6
62.2
50.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 medium (355M)
50.0
52.0
51.3
28.6
59.5
53.3
61.0
65.0
52.6
OPT-350M
53.5
53.4
35.7
60.0
Hybrid H13-55MM
37.5
51.7
55.2
41,1
59.5
62.3
61.5
69.0
54.7
OPT-1.3B
36.5
49.5
53.4
39.3
58.3
61.8
55.0
69.0
52,9
GPT.Neo 1.3B
41.3
50.0
52.3
33.9
57.9
55.5
59.9
66.0
52.1
Hybrid H.3-L.3B
52.9
30.3
53.4
33.9
58.2
67.8
61.7
74.0
56.5
OPT-2.70
GPT.Neo-2
37.5
50.0
52.3
50.0
59.1
60.0
61.1
67.0
54.6
Hybrid H3-2.7B
36.5
Table 6: 3-shot acc. on SuperGLUE with logit scoring. Best results bold, second best underline.
Model
WSC
WIC
RTE
CB
MultiCC
ReCoRD
BoolQ
COPA
Average
OPI-125M
36.3
50.2
47.3
57.9
41.9
60.0
47.9
GPT-Neo-112 125M
38.5
50.0
53.1
17.9
56.3
39.6
62.1
60.0
47.2
Hybrid H3-125M
43.3
49.1
48.9
55.0
67.0
53.7
GPT-2 medium (355M)
36.5
50.5
45.0
8.9
43.5
53.3
58.8
65.0
45.6
OPT-350M
37.5
50.0
45.8
44.6
49.8
51.4
61.7
60.0
50.1
Hybrid H3-355M
42.3
47.5
50.5
28.6
59.7
62.3
60.5
69.0
52.6
OPT-L33
44.2
31.1
53.4
16.1
59,9
62.1
38.3
70.0
49.4
GPI-Neo-1_3B
35.6
67.0
Hybrid H3-1.33
36.5
49.2
55.2
23.2
59.3
67.6
56.9
76.0
53.0
OPT-2.75
442
71.0
GPT-Nco-2.7B
49.0
31.9
51.6
21.4
57.0
60.0
56.0
68.0
51.9
Hybrid H3-2,7B
36.5
45.6
47.3
46.4
59.4
71.1
60.6
77.0
55.5
Table Inference throughput on A100 80GB, 1.3B models. Batch size 64, prompt length 512, 1024, or 1536, and generating 128 tokens per sequence in batch (i.e., 64 tokens a batch). Hybrid H3 is up to 2.4x faster than a Transformer of similar size in inference. The difference is larger for longer sequences.
Tokens/s Transformer- Hybrid H3-1.3B
Prompt length 512
Prompt length 1024
Prompt length 1536
1340
770
520
1980
1580
1240
SSMs and attention on the synthetic languages translates real language modeling capabilities. We also report the generation speed of hybrid H3 models compared to Transformers; since SSMs are recurrent models, they can generate tokens 2.4 faster than Transformers. Appendix F shows performance of pure H3 language models on these same evaluation metrics.
Setup We train hybrid models sizes 125M, 355M, 1.3B, and 2.7B the Pile (Gao al., 2020) for 400B tokens. We compare against checkpoints of equivalent sizes from Open-Al (Radford et al., 2019) and GPT-Neo (Black al., 2021), from HuggingFace (Wolf al., 2020).
Perplexity Table 4 shows perplexity on the Pile (Gao et al., 2020), WebnWeeTexxt (Gokaslan et al., 2019), and WikiText- 103 (Merity et al., 2016). On the Pile, 125M hybrid model outperforms GPT-Neo, which was also trained on the Pile. Our hybrid models also outperform GPT-Neo models and GPT-2 models zero-shot transfer Open WebText and WikiText113. We report the PPL of GPT-2 models for context, though the performance is not directly comparable since they were trained on different data.
Zero- and Few-shot Performance We compare the zero- and few-shot performance of hybrid H3 language models against OPT (Zhang et al., 2022), GPT-Neo, and GPT-2 models, where public checkpoints available. We report performance with rank classification on the logits of the possible choices (see Appendix E.7 for raw generation). Table 5 reports zero-shot performance on the SuperGLUE benchmark, and Table 6 reports the 3-shot performance. Following the perplexity results, the hybrid language models outperform or match the best the Transformer baseline on more than half the tasks.
Language Modeling Inference Finally, since SSMs are recurrent models, they admit faster text generation than Transformers. Table shows inference throughput of a 1.3B -3B-maraeettr hybrid model compared to a Transformer. The hybrid model has up to 2.4x higher throughput.
FLASHCONY EVALUATION
We evaluate how well FLASHCONY speeds up SSMs. FLASHCONY sets state-of-the-art performance on the long range arena benchmark (Tay et al., 2020) using S4 (Gu et al., 2022a). We report performance
There is no pretrained GPT-Neo at the 350M size.

---

Published as a conference paper at ICLR 2023
Table 8: Speedup the LRA benchmark.
Models
Speedup
Transformer
1x
FlashAttention (Dao et., 2022b)
2,4x
Block-sparse FlashAttention (Dao et al., 2022b)
2.8x
54 (Gu et al., 2022c)
2.9x
S4 with FLASHCONV
5.8x
Benchmarking FlashConv
Inset (256-4K)
35x
30
Time (ms)
Time (ms)
2x
20
10
1K
2K Sequence Length
4K
1K
16K
Sequence Length
FlashAttention
cuFFT Conv
Fused Block FFT Conv
FlashConv
Fused Conv
Figure 2: We compare the speed of different algorithms to perform FFT- FFs-based convolution. along with FlashAtten- tion (Dao al., 2022b) (the fastest attention implementation we know of). We use batch size 8, hidden dimension 1024, and varying sequence length from 256 to 32k, and mcasure on an A 1000SMMX4400GB GPU, We see that kernel fusion gives up 3.4x speedup over naive FFTCon for short sequences (up to 512), block FFT gives to 2 speedup for medium length sequences (1k- 8k), and state- passing allows 2.3 x faster FFTConv for long sequences (16k and above).
of training H3 module with FLASHCONV compared to attention at various sequence lengths, from 256 to 32K and demonstrate nearly linear scaling.
Long Range Arena The Long Range Arena (LRA) benchmark (Tay et al., 2020) is benchmark for long-range sequence modeling. The state-of-the-are approach, $4 (Gu al., 2022c), is SSM. Table 8 shows that FLASHHONV accelerates S4 by 2x, outperforming Transformers by 5.8x.
Benchmark H3 Against Attention We benchmark the time to run the forward and backward pass of H3 with FLASHCONY against attention. FLASHCONY maintains nearly linear scaling, even to very long sequence lengths. Fig. 2 shows overall 2-3x speedup over FFTConv with cuFFT using our techniques (block FFT, state-passing). Simple kernel fusion (even without block FFT) can yield speedup over cuFFT for short sequences, since memory reads/writes the bottleneck for short sequences. For long sequences, SSMs using state passing can be dozens of times faster than even the fastest attention implementation.
7
CONCLUSION
Our main goal is to understand and narrow the gap between attention and SSMs in language modeling in terms of modeling capabilities and hardware efficiency. Our exploration based synthetic language tasks motivated us to design the H3 layer, which is surprisingly competitive with attention. Our BLoCKFFTCom
algorithm exploits matrix multiplication units and the dual recurrent conoooluoooooiioon view SSMs to substantially speed up SSMs, reducing the hardware barrier between attention and SSMs. We are excited about several future directions. Our H3 layer is simple combination of two SSMs. and more sophisticated designs could more expressive. Our encouraging results on language models up to 1.3B parameters suggests that scaling SSMs to larger sizes promising avenue. Since simply adding two attention layers to H3 models already outperforms both the pure H3 model and Transformers, we optimistic about combining the complementary strengths of SSMs and attention in the future.
Reproducibility Statement. To facilitate the reproducibility of our algorithms and results, (i) we include a link to downloadable source code in supplementary materials, (ii) for our theoretical statements and results, include clear explanations of any assumptions and complete proof of the claims in Appendix D; for any datasets used in the experiments, complete description of the data processing steps in Appendix E. We will also release model checkpoints for all our models.
Ethics Statement. Our work seeks understand the fundamental capabilities and limitations of newly-emerging model architectures. As the amount of data and model size grows, we also week to understand how to make training these models more efficient- and run inference more efficiently. This potentially connects to energy savings during model development and deployment. We also note that

---

Published as a conference paper at ICLR 2023
the relative underutilization of tensor cores in the FFT convolutions of state space models (even with our block FFT) suggests that consumer GPUs may be able to train models at a cheaper price point.
However, as with any language model training, developing new techniques may impact a wide range of upplications, each with potential benefits and harms. For example, making language model training cheaper and making inference more efficient make it cheaper to spread disinformation. Similarly, improving the efficiency of model training may not reduce the overall environmental footprint of training, since the same resources may be used to train more models, or train the same models for longer. While our work makes partial progress the fronts of efficiency and understanding, it does not explicitly address the issues of fairness and bias in language models.
ACKNOWLEDGMENTS
We thank Albert Gu for helpful discussion regarding the model architecture, and more importantly for sending us daily hippo videos. We thank Together Computer for providing portions of the compute used to train models in this paper. We gratefully acknowledge the support of NIH under No. U54EB02040S (Mobilize), NSF under Nos. CCF1763315 (Beyond Sparsity), CCF1563078 (Volume to Velocity), and 1937301 (RTML); US DEVCOM ARL under No. W91 (Interactive Au Teaming); ONR under No. N000141712266 (Unifying Weak Supervision); ONR
Understanding and Applying Non-Fuclidean Geometry in Machine Learning: 22012225 (NEPTUNE); NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASF, Accenture, Ericsson, Qualcomm, Analog Devices, Google Cloud, Salesforce, Total, the HAI-GCP Cloud Credits for Research program, the Stanford Data Science Initiative (SDSI), Department of Defense (DoD) through the National Defense Science and Engineering Graduate Fellowship (NDSEG) Program, Wu Tsai Neuroscience Stanford Interdisciplinary Graduate Fellowship. and members of the Stanford DAWN project: Facebook, Google, and VMWare. The U.S. Government is authorized reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed or implied, of DARPA, NIH, ONR, or the U.S. Government. Atri Rudra's 's research is supported by NSF grant CCF-1763481.
REFERENCES
Jimmy Ba, Geoffrey E Hinton, Volodymyr Mnih, Joel 7. Leibo, and Catalin Ionescu. Using fast weights to attend to the recent past. Advances in neural information processing systems, 29, 2016.
Alexei Baevski and Michael Auli. Adaptive input representations for neural language modeling. International Conference on Learning Representations, 2018.
David H Bailey. FFTss in external or hierarchical memory. The journal of Superccoppuing 4(1):23-35, 1990.
Sid Black, Leo Gao, Phil Wang, Connor Leahy, and Stella Biderman. GPT-Neo: Large Scale Autoregressive Language Modeling with Tensorfwww,, March 2021. URL attps://dd /doi. orq/10. 5281 /2t.//222222..// 5297715. If you use this software, please cite it using these metadata.
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.
Willian Brogan, Modern control theory, 1974.
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. 2020.
Kyunghyun Cho, Bart Van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder-decoder approaches. arXiv preprinl a4Xi /259, 2014.
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.
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: 2022.
10

---

Published as a conference paper at ICLR 2023
James W. Cooley and John W. Tukey. An algorithm for the machine caleulation of complex fourier series. Mathematicc of Computation, 19(90):297-30,, 1965. ISSN 00255718, 10886842. URL
http://wwit www.wjstooee .esttororggee stable hat.aa/eeee0035--. 2003354.
Kamalaker Dadi, Gaël Varoquaux, Antonia Machlouzarides- Shalit, Krzysztof J Gorgolewski, Demian Wassermann, Bertrand Thirion, and Arthur Mensch. Fine-graie atlases of functional modes for fmri analysis. Neurolmage, 221:117126,, 2020.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime G Carbonell, Quoe Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. In Proceedings of the 57th Annual Meeting of the Association Compulaaion Linguistics, pp. 2978-2988, 2019.
Tri Dao, Beidi Chen, Nimit Sohoni, Arjun Desai, Michael Poli, Jessica Grogan, Alexander Liu, Aniruddh Rao, Atri Rudra, and Christopher Ré. Monarch: Expressive structured matrices for efficient and accurate training. In International Conference on Machine Learning (ICML), 2022a.
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. In Advances in Neural Information Processing Systems, 2022b.
Giannis Daras, Nikita Kitaev, Augustus Odena, and Alexandros G Dimakis. Smyrr-efficientaattennni attentioon using asymmetric clustering. Advances in Neural Information Processing Systems, 33:6476-6444 6489,, 2020.
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, pp. 5793-5831. PMLR, 2022.
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. html.
Bruce Fischl. Freesurfer. Neuroimage, 6222:774-71,,,
Robert Fisher, Carlos Acevedo, Alexis Arzimanoglou, Alicia Bogacz, 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, 4554):475-442, 2014
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
2020.
Karan Goel, Albert Gu, Chris Donahue, and Christopher Ré. It's raw! audio generation with state-space models. arXiv preprint arXiv:2202.09779, 2022.
Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanie Tellex. Openwebtext corpus, 2019.
Albert Gu, Tri Dao. Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo:
Recurrent memory with opti- mal polynomial projections. Advances in Neural Information Processing Systems, 33: 1487 2020.
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-tiin time models with linear state-space layers. Advances in neural information processing systems, 34, 2021.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In The International Conference on Learning Representations (ICLR), 2022a.
Albert Gu, Ankit Gupta, Karan Goel, and Christopher Ré. On the parameterization and initialization of diagonal state space models. In Advances in Neural Information Processing Systems, 2022b.
Albert Gu, Isys Johnson, Aman Timalsina, Atri Rudra, and Christopher Ré. How to train your hippo: State space models with generalized orthogonal basis projections. arXiv preprint arXiv:2206. /2037, 2022c.

---

Published as a conference paper at ICLR 2023
Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are as effective structured state spaces. In Advances in Neural Information Processing Systems, 2022.
Curtis Hawthorne, Andrew Jaegle, Cätälina 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.
Sepp Hochreiter and Jürgen Schmidhuber. Lstm solve hard long time lag problems. Advances in neural information processing systems, 9, 1996.
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.
Sara Hooker. The hardware lottery. Communications of the ACM, ((112):58-65200)
Sheng- Sheng-Chun Kao, Suvinay Subramanian, Gaurav Agrawal, and Tushar Krishna. An optimized dataflow for mitigating attention performance bottlenecks. arXiv preprint arXiv:2107.06419, 2021.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are RNNs: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, pp. 5p.51.. 5156-5165. PMLR, 2020.
Feyza Duman Keles, Pruthuvi Mahesakya Wijewardena, and Chinmay Hegde. On the computational complexity of self-attention. arXiv preprint arXiv:2209, 228904888,,,, 2022.
Michael Patrick Kerr. The impact of epilepsy patients' lives. Acta Neurologica Scandinavica, 126: 2012.
Maedbh King, Carlos R Hernandez- Casttllo,,,,, Russell Poldrack, Richard B Ivry, and Jöm Diedrichsen. Functional boundaries in the human cerebellum revealed by multi-domain task battery, Nature neuroscience, 22(8):1371- 1377-1388 2019.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In The International Conference on Machine Learning (ICML), 2020.
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, 14:22795-22207, 2021.
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.
Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces, arXiv preprint arXiv:2206. 13947, 2022.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016.
Eric Nguyen, Karan Goel, Albert Gu, Gordon Downs, Preey Shah, Tri Dao, Stephen Baccus, and Christopher Ré. $4nd: Modeling images and videos as multidimensional signals with state spaces. In Advances in Neural Information Processing Systems, 2022.
NVIDIA. Nvidia Tesla V100 GPU architecture, 2017.
NVIDIA. Nvidia A100 tensor core GPU architecture, 2020.
NVIDIA. cufft v1 1.7.1 documentation, 2022a. http:///dc...viiaacccocc/uddldtt
NVIDIA. Nvidia H100 tensor core GPU architecture, 2022b.
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 Hatfiel-Doodss. 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.
12

---

Published as a conference paper at ICLR 2023
Alan V Oppenheeim. Applications of digital signal processing. Englewood Cliffs, 1978.
Alan V Oppenheem, John R Buck, and Ronald W Schafer. Discrete-time signal processing. Vol. 2. Upper Saddle River, NJ: Prentice Hall, 2001.
Markus N Rabe and Charles Staats. Self-attention does not need O(n²) memory. arXiv preprint arXiv:2/ 12.05682, 2021.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, lla Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog, 18))9, 2019.
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.
Khaled Saab, Jared Dunnmon, Christopher Ré, Daniel Rubin, and Christopher Lee-Messer. Weak supervision an efficient approach for automated seizure detection in electroencephaalograpyy NPJ digital medicine, 3(1):1-12, 2020.
Vinit Shah, Eva Von Weltin, Silvia Lopez, James Riley McHugh, Lillian Veloso, Meysam Golmohammadi, lyad Obeid, and Joseph Picone. The temple university hospital seizure detection corpus. Frontiers in neuroinformatiis, 12:83, 2018.
Mohammad Khubeb Siddiqui, Ruben Morales- -Meeende, Xiaodi Huang, and Nasir Hussain. A review of epileptic seizure detection using machine learning classifiers. Brain informatics, 7(1):1-18, 2020.
Siyi Tang, Jared Dunnmon, Khaled Kamal Saab, Xuan Zhang, Qianying Huang. Florian Dubost, Daniel Rubin, and Christopher bes-er Self-super graph neural networks for improved electroen- cephalographic seizure analysis. In International Conference on Learning Representations, 2021.
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.
Armin W Thomas, Christopher Ré, and Russell A Poldrack. Self-supervised learning of brain dynamics from broad neuroimaging data, arXiv preprint arXiv: 2220..1141,,0
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.
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.
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768, 2020.
Pete Warden. Speech commands: A dataset for limited-vocabulary speech recognition. arXiv preprint 2018.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi 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-arr natural language processing. Empirical Methods in Natural Language Processing: System Demonstrations, pp. 38-48 48- Online, October 2020. Association for Computational Linguistics. URL https:t/p//.ccccc. / www.//wwwww.ooo.w.ooo.... ac web. org/ aat/antthoooooooo oggy/ 2000 emn emnlp-demos. 6.
In Proceedings of the 2020 Conference
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.0m 0/068 2022.
13

---

Published as a conference paper at ICLR 2023
A RELATED WORK
State space models have shown promise in modeling sequential data, including time series data (Gu et al., 2022a), audio (Goel et al., 2022), and visual data (Nguyen et al., 2022). Our model builds off work on simplifying and parameterizing diagonal versions of S4 (Gu et al., 2022b; Gupta et al., 2022; Gu et al., 2022c). Gated state spaces (Mehta et al., 2022) also aim to adapt SSMs to language modeling, but our results suggest that the GSS model does not perform as well as H3 (or even as well as earlier SSMs like S4D). The idea to combine SSMs with attention in hybrid models is not new; Mehta et al. (Mehta et al., 2022) also showed that interleaving attention with their GSS layer can improve performance, which we also validate on our Open OpenWebTex, Teb experiments. These positive results suggest that attention and SSMs are complementary, and that hybrid models may be a promising direction for future work.
Large language foundation models (Bommasani et al., 2021) have demonstrated the power of scaling attention-based networks to billions of parameters and training them on trillions of tokens (Hoffmaann et al., 2022). Understanding the mechanistic basis (Elhage et al., 2021) behind these models may yield insights into better design choices for future models. These and similar explorations have informed the design of H3 and our selection of synthetic languages. A number of recent works have also explored how to address the shortcomings of attention by approximating the attention computation (Wang et al., 2020; Katharopoulos al., 2020; Choromanski et al., 2020; Tay et al., 2020; Kitaev et al., 2020; Daras et al., 2020). We believe these efforts are complementary to SSMs, and we are excited to see how they can be combined in future work.
Linear attention (Katharopoulos et al., 2020) and classical sequence models like RNNs serve as inspiration for H3. Appendix B draws direct connection between linear attention and LTI systems. Luo et al. (Luo et al., 2021) also propose variant of linear attention that can achieve Onlogg)) scaling in sequence length. Appendix F evaluates linear attention on language modeling, and finds that it underperforms exact attention, whereas H3 outperforms attention. The multiplicative interactions H3 are reminiscent of gating mechanisms in LSTMs (Hochreiter & Schmidhuber, 1996) and GRUs (Cho et al., 2014), which suggests that architectural lessons from these sequence models may be useful for adapting SSMs to language modeling. A number of algorithms for scaling attention to longer sequences have also been proposed, such as Transformer-XL. (Dai et al., 2019), Reformer (Kitaev et al., 2020), Performer (Choromanski et al., 2020), and Perceiver AR (Hawthorne et al., 2022). Some of these approaches underperform exact attention on language modeling, and may be slower in wall-clock speed (Dao et al., 2022b). A thorough comparison of these alternatives to exact attention and how well they scale in model size and amount of training data is fruitful future work.
FFT algorithms are used in a wide variety of applications, including signal processing (Oppenheim, 1978), control theory (Brogan, 1974), and more. Various algorithms for computing the FFT have existed for decades (Oppenheim et al., 2001). We hope our work on appealing to these classic algorithms to accelerate new applications such as learned SSMs will inspire future algorithmic exploration, even if hardware is not designed for them (Hooker, 2021).
B LINEAR ATTENTION AND TIME- VARYING SYSTEMS
We draw some connections from linear attention to LTI systems and SSMs.
We first present linear attention linear -vrriig system, and show how converting it to a linear lime-invarian system matches H3.
Linear time- vime-varying system and linear attention In general, a layer in a sequence model takes in a sequence and outputs a sequence. Many of these take the form of a linear time-varying system (thanks to the Sitard-Linndlol theorem, nonlinear systems can be approximated by series of linear system):
T; A B,
3/1 C,r;+DDr0
This has the same form as SSMs (Section 2), except that the matrices can depend on the timestep.
Recall the form of linear attention from Section 2. For the purpose of approximation, we ignore the denomi-
nator in linear attention Section (i.e., assuming d; = =1). We see that S, is just cumulative sum, satisfying
the recurrence S = =5 + K5+1 VV Similarly, O satisfies the recurrence O:41 = (6(Q+11) S This is a linear time-varying system of the form xi+1 A Bu4+1 and yi+1 = =C4110141 (with A = T, B= T, w4= o VT', -C That is, A and B are constant, but C' is time-variant.
14

---

Published as a conference paper at ICLR 2023
To convert this into linear time-invarian version, treat the time-variam C, post- step. We instead of a fixed for the LTI. This yields an LTI:
Ii-1 = Ax, Bo(K)V
1/-1 Cais
for some matrices A, A,B, that are learned. We then apply post- cost-processing by multiply Vs+1 with o(Q.)² Replacing (G)) with a shift SSM yields an analogue to H3.
C METHOD DETAILS
Since we have described the forward pass in Section 3, we describe here the backward pass in details.
C.I BACKWARD PASS
We show how to compute the backward pass in fused kernel.
y= Du. In case, have f and u have the same length, they are symmetric as far the convolution is concerned.
Suppose we are given dy = at (where I some loss function). We wish to compute du, df, and dD by (which are 3ü 87 and aD respectively).
The most challenging part is computing the gradient through the convolution operator but we see that we can re-use our FFT infrastructure for it. The rest of the operations are straightforward; we have dD= dyu".
Gradient of the Convolution Here, we'll discuss how to compute df by integrating w.r.t. the convolution operator *, As an immediate consequence, we ll able to compute du as well.
Since f and u are the same length L, f and u*] have the same result. Thus, we'll start from f here.
For some notation, O=u+f. Then, dO dy. Recall that 0li]= b [|]
We'll start by extending and f with zeros, to give them length 2T. Let = (u [0] [1] ........
(T. -T 1).0,....,0
and f' extended similarly. Let O' f',, and 0= O'|'|N]] Assume that we have all the values of dO' (we only have them for the first half, but we'll ll see that it doesn't matter in the end).
Let's construct Toeplitz matrix H such that u'
H.
u' [0)
0
[1]
u'
H
w/ -2 " (22 -
u [0]
Since, we have -0 for i 2 L, we can actually fill in the zeros the above matrix as well:
0'O u'
w'(2L-1 u u'(0)
n [2]
H
w w2 1]
u 1|22-1 2]
n [0]
Then, we can use the matrix multiplication chain rule to find that:
2 [0]
u/
w/
2
w/ 122
w [0]
[2L
df' HH'==",d''==
[2]
u d'
w [0
w/
(2L
1)]
w/
[0]
1 22l- 1)] w' w21 2)]
[0]
where we use u'[-i to mean w 22 - i]. Notice that this matrix has the same format as Hw! Let
u/, [[=, -2N 1)]]. Then:
df' (u,, #dO').
So how we compute ul, efficiently? Naively, we might incur some nasty memory access issues. But nice property about the DFT saves us!
Let Ui]] be the i-th element of the DFT of signal Note that U[i is complex. We have:
U = -[]-
15

---

Published as a conference paper at ICLR 2023
where here the represents the complex complexeoojuuaatee We can use this property to compute df' efficiently; df' =V, =i,= =f==d'== =4= where FFT denotes taking the complex conjugate of the FFT, and dy' denotes dy, padded with zeros.
(FFT (dy (:
Computing du We can use this same trick to compute du, except we need to add in the contribution from the original Du term. We end up with:
du=du' [: N)+
C.2 STATE-PASSING MATRICES
We show how to derive M., for the state update in our state- state-passing algorithm.
We wish to construct matrix v.Ra.. E mR"NN' such that M Σn A - Note that
A'B E [dx1 is column vector. We can simply stack these column vectors to form matrix: N' -BB........ B.....B
Mar A N' A
D PROOFS
We show parameterizations of H3 and attention that solves the associative recall task. We prove Propo- sition and Proposition
D.I H3 EXPRESSIVITY
This section formally describes parameterization of H3 that solves the associative recall task.
D.1.1 EXAMPLE LANGUAGE
Consider simple language with 4 keys and 4 values. For concreteness, will use the keys {k,,2, k2, k5, = Lx and the values {11,02, 03, sequence of key-value pairs with one key at the end, want model to generate the value associated with the key at the end. Assume that the key at the end appeared in the sequence.
02, 24] = Lv, i.e. our language I= Lx U Lv. Given a
More formally, let +11 be the length of the sequence, N even. The language A consists of sequences LN+1, Each sequence has an associated mapping Je LK - Lv. For each sequence, the odd indices are randomly sampled from Lx. for ;., N The even indices are defined by sz: T2- = (1 for 1 < N/2. The last item in the sequence IN+1. is randomly drawn from the keys that have appeared in a already, i.e. xN++1 E [U The goal of this language modeling task is produce fa +1) the end of the sequence.
D.1.2 H3 MODEL TO SOLVE A
We describe toy H3 model that can solve A.
Consider model consisting of embedding layer, an H3 model, and an output projection with softmax. Recall that d is the dimension of the H3 model, m is the dimension of its hidden states, and H is the number of heads. Let d= (0 8.m H. =4. 4. Let the embedding layer map each key k; to the e, basis vector, and map each value to the 64-; basis vector.
Let Bahi/t and Cshift denote the parameters of the shift SSM, and Adiag. Baiag, and Cittgg denote the parameters of the diagonal SSM (let D be zero for both). Let Bahift = BBiag = Cdiag =1 Let Cshift [01]. Let Ading be diagonal matrix with Is along its diagonal for each H3.
Remark. The action of diagonal SSM parameterized by Adiag, B Bdiag. and Cdiag is to act as a cumulative sum over all its imput. The action of shift SSM parameterized by Bahaft and Cahift shift its input by one time step.
Recall that the H3 layer maps its input to Q, K, and V by applying a W a a w K. and Wy. Let Wo and W K be the following:
1
0
0
0
0
0
07
1
0
0
0
0
1
1
0
0
0
0
We = W
0
0
0
0
0
0
0
0
16

---

Published as a conference paper at ICLR 2023
Recall that Q and K are split into H heads (Q K for (1,2,3,4)), each of which is sent to independent H3.
Remark. The action of W Q and W k "assign" each key different H3 head, i.e., Q is only
non-zero when 2t kg Similarly, K, is only non-zero when zt-1 k2 (since K, K due to the time delay of the shift SSM).
Let Wy be the following:
0
0
07
0 0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
W
0
0
0
0
0
0
1
1
1 0
0
1
1
1
1
Remark. The action of this matrix is to encode the input value (as "binary" ") and send it to all H3 heads.
E.g.. v = -v(2) V -v for all i, and ( [0,0] +00)=2) +, U1,
V = [0.1] + 22 U2 etc.
We claim that for #N+1 k; O (1)
N+1 will be multiple of the binary encoding of fx (k:)), and all the other
(j] heads of the output O N $ijj will be zero. Let the output projection Wa be such that, with a non-linearity afterwards, inverts the binary encoding to produce the embedding of the desired output fz (k;).. We will assume such projection exists, proof left to the reader.
Proposition 3. The model described above solves the associative recall problem for the language A.
Proof. Proof sketch. WLOG, let ÆN+1 hş- Then Q [1.1], but Q =0,00] for j * i. Thus. 0 [0,0] for j i due to the multiplicative interaction.
Since (+)) = [1,1], (+) is the output of the diag SSMs in the H3 head corresponding to k; (recall that each head has two independent shift SSMs and two independent diag SSMs). The output of the diag SSMs are the cumulative sum of all the inputs they have seen in the sequence.
For one of the diag SSMs to see non-zero input, its preceding shift SSM must have non-zero output. The only times this can happen the sequence are when fz-1 1 = k2. But then x; = fa (kz). Thus, the input to the diag SSMs are precisely the binary encoding of fa (k:). Then the output O is multiple of the binary encoding of (k;), Wo decodes this output to the embedding of fz (k:).
D.2 ATTENTION EXPRESSIVITY
We provide an informal sketch of two-layer attention model that can solve the associative recall task, inspired by the construction of (Olsson et al., 2022). The first layer of the attention model outputs the embedding of the previous token the sequence, and concatenates it with the current token in the sequence. The second layer compares the current token to the previous token embeddings, and outputs the paired embedding when there a match-whice is exactly the key-value lookup.
The construction proceeds as follows:
-e-g.Pi--1 In the first layer, let Q: be mapped to the positional embedding of token x;- (e.g.. Pi-1 if p; denotes the positional embedding of token 2;)), and K, be mapped to the positional embedding of token P;.
The attention matrix A is computed as QK", with causal mask (i.c., =if=
Then, softmar(A (A) approximates the shift matrix (see Section 3).
Let V; be an encoding of token I;, constrained to the first half of the hidden dimension.
Then, for output = =00ftmazz(((t V, the first half of the vector O is the encoding of token Fi-1-
In the second layer, assume that you have a skip connection, that maps the encoding of the input token v; the second half of the vector O,.
Then, the input to the second layer encodes both .;-1 and Zi.
In the second layer, let Q; extract the encoding of x;, and let K; extract the encoding of 2;-1-
17

---

Published as a conference paper at ICLR 2023
Apply causal mask on QK" Then, the value of softmar(QK")) )ag is large if = 22-1, and j-j
Let V, extract the encoding of z.--
Then, output O; is the sum of values 2; such as z = 2;. But then O is exactly a lookup of the token that came after z; when appeared previously in the sequence- which exactly solves associative recall.
We note that the above construction requires the ability for the positional encodings to select the previous to- ken based on the dot product and softmax, and for token comparisons through the dot product and softmax.
D.3 H3 COMPLEXITY
We prove Proposition which states that the H3 layer takes for sequence length N and hidden dimension
time and O(dN) space
Proof. We first analyze the time complexity. Consider the matrix multiplies in H3, where the input R. xd is multiplied by three weight matrices of size d x These take time O). The output is also multiplied with ouiput projection weight matrix of size d x d, also taking time OdP Therefore the matrix multiplies take time Odd
Now consider the two SSMs in H3. The first SSM involves a convolution of K E R Rxd (in the N-dimension) with a kernel of size N d. This reduces an FFT, a pointwise multiply, and inverse FFT (in the N N-dimension). This takes time O(dNlog.N) The second SSM involves H convolutions, inputs of size N d, d,, along the N N-dimension. This takes time: ((HH,NNoog.
(log, OddNlog,,,
where the fact that da d/ H and that dn O(). Therefore the two SSMs take total time O(dNlogN). As a result, the H3 layer takes time:
Old N +dNlogN).
Now we analyze the space complexity. The matrix multiplies all take space O(dN). The FFTs, pointwise multiplies, and inverse FFTs of the two SSMs takes (ddN space and O(Hd = space. Therefore the overall space complexity (ddN.
D.4 STATE PASSING CORRECTNESS
We prove Proposition 2. We assume that the BLOCKFFTCONY algorithm is correct, i.e., the output BBB..CCFFFCCNN (J.u) is equal the output of an SSM with convolution kernel f and input u.
Proof. Proof by induction on C.
Base case: 1= WTS y - (1(1)), M 10)
M May =IN.
In this case, note that N = N', Then ( M May (L BWLC////III/DLCCCFFTTOOWWWWW//
But u = u1, y0 =p0")))))(()))1 Additionally, by the recursive definition of state space,
xN A N
20
Bu;
N A
A
N'
Zo
Bu (1)
i=1
M Mz
+ A
B.A N'
B B,..., ().
M May No
(0)
+M (1)
18

---

Published as a conference paper at ICLR 2023
Inductive step: C>1. Assume that (v y....
and UN (C-1) -I WTS
that -C NCC). and M MazuN N
(C-1) +M -IN. Let denote 1C))-
For i (C-CCCCN 1)N2, have:
yi = "AA''Bzz; "
= CA'' + (f*u )i
= "CA'"'Z:: BLOCKEFTTOOC )-NN'
= (LOKK
(May (C-1) BLOCKFFICONN((11
-I-N"
Thus, yy) - Similarly,
.N NN= - N'
A N' BButt
5(C-1)N'
N
= A' AIN'
(0-1)
N' - Bu
(C)
IN.'
1=1
-M N"
(C-1)
[A
B,A N' -2 BB....,BBuu (C)
(C-1)
M (O).
MaZ N'
E EXPERIMENTAL DETAILS
E.1 SYNTHETICS
Our synthetic tasks, inspired by (Olsson et al., 2022), are designed to mimic the in in-context learning capability of large language models-the the ability learn from examples in the imput sequence, and use information from the input to generate the right answer for the output. For example, the induction head task requires memorizing the token that appears after the special H token in the input sequence, and the associative recall task requires learning the mapping from keys to tokens from the input sequence.
We evaluate synthetics by training two-layer versions of our GPT models, with different modules replacing attention. We train models with inner dimension 32, and MLP dimension 128. For all the synthetics, we use a learning rate of 5e-4 and a weight decay of 0.1. We sample 5000 training examples and 500 test examples from the same distribution, and we train for 200 epochs. Again, we use embedding dropout of 0.1 and residual dropout of O.O.
E.2 MODEL ARCHITECTURE
For our 125M models, we use 12 layers, with hidden dimension 1024, and MLP dimension 4096. For our 355M models, we use 24 layers. with the same hidden dimension and MLP dimension. The 1.3B models have 24 layers, with hidden dimension 2048, and MLP dimension 8 8192. The 2.7B models have 32 layers, hidden dimension 2560, and MLP dimension 10240. The hybrid models have 12, 16, 16, and 20 heads for the 125M, 355M 1.3B, and 2.7B models, respectively. The 125M hybrid model has an attention layers at layers 1 and 7, the 355M and 1.3B hybrid models have attention layers at layers I and 13. and the 2.7B hybrid model has attention layers layers 10 and 21. For both our hybrid models and our H3 models, we use SSM state size 64. Our hybrid model uses head dimension 1 for H3, while our pure H3 model uses head dimension 8. We run models with mixed-preciste
training, with 6 for the MLP's and attention. When training language models, we use fp32 for the FFTConv.
E.3 OPENWEBTEXT TRAINING
For the 125M models trained Open WebTex, we follow the training recipe of the Megatron-LM repo.
We effective batch size of 512, and use gradient accumulation fit into available GPU memory. We use the Adam W optimizer, with learning rate 6e-4 for GPT-2 small and 1.5e-4 for GPT-2 medium, and decay of 0.1. All models are trained with the same hyperparameters for 100K steps. We run all im- plementations with mixed-precision training (Py Torch AMP). We train models with sequence length 1024.
19

---

Published as a conference paper at ICLR 2023
We use the Openwebtext dataset, with the GPT-2 BPE tokenizer. We randomly select 0.5% of the dataset as the validation set, with the rest being used as training set. This random selection validation set is done once, and all models are evaluated on the same validation set.
E.4 THE PILE TRAINING
For the 125M and 355M models trained on the Pile, we follow the training recipe GPT-3. We use batch size 256, with sequence length 2048. We train our models for 800K steps. We use residual dropout 0.0 and embedding dropout 0.1. We use the Adam optimizer, with learning rate 6e-4 for the 125M model and 3e-4 for the 355M model, and weight decay of 0.1. We use cosine schedule with 8000 steps for linear warmup, and decay the learning rate to 10% by 300B tokens, then continue training 10% learning rate for another 100B tokens. We suspect that there exist better hyperparameters for H3 language models, but we did have the resources tune them.
For the 1.3B models, we double the batch size to 512 (with sequence length 2048), again following the train- ing recipe of GPT-3. The number of training steps are halved so that we train the same number of tokens.
For the Pile dataset, we again use the GPT-2 BPE tokenizer, similar to GPT-3 and GPT-Neo.
E.5 SUPERGLUE
We follow the prompts used in the GPT-3 paper (Brown et al., 2020). For rank classification on the binary classification tasks, we use yes/no for WSC, WIC, MultiRC, and BoolQ, and we use true/false for RTE. For CB, we use /reether as the three choices. For COPA and ReCoRD, we use the continuations provided by the task.
E.6 HARDWARE
All models were trained on either a single 16x 1000400-0 node a cluster of 1008000 nodes.
F
ADDITIONAL EXPERIMENTS
F.1 LRA ACCURACY
We evaluate the accuracy of H3 on LRA. We compare accuracy to S4D (Gu al., 2022b), since H3 uses an S4D kernel as component in its layer. We use the same hyperparameters as S4D, and make the layer bidirectional by making two copies and running them in opposite directions.
Table 9: LRA performance of H3 compared to S4D.
Model
ListOps
Text
Retrieval
Image
Pathfinder
Path X
Avg
S4D (Gu et al., 2022b)
58.3
87.3
90.7
87.5
93.6
92.3
85.0
H3
57.5
88.2
91.0
87.3
93.0
91.8
84.8
Table 9 shows that H3 performs well on the LRA benchmark, even thought it was designed for autoregressive language modeling. H3 outperforms S4D on two of the LRA tasks, and comes within 1 point on the others.
F.2 WIKITEXT103
We train 125M-sized models on WikiText103 (Merity al., 2016) and compare their PPL transform- ers, as well as other variants of efficient or long-range attention. We use the same hyperparameters and setup as training Open WebText. We also provide results from Transformer-XL and Perceiver AR for context, though the results may not be directly comparable due to differences in model size and tokenizer.
Table 10: Test PPL on WikiText103.
Models
PPL
Transformer (125M)
18.6
Hybrid H3 (125M)
18.5
Performer (125M) (Choromanski al., 2020)
26.8
Reformer (125M) (Kitaev et al., 2020)
26.0
Linear Attention (125M) (Katharopoulos et al., 2020)
25.6
Perceiver AR (358M) (Hawthorne et al., 2022)
18.4
Transformer-XI.c (285M) (Dai al., 2019)
18.4
Table 10 shows that the Hybrid H3 model is competitive with Transformers of the same size, as well as larger models such as the 358M Perceiver AR and 285M Transformer- XL. The hybrid H3 model also significantly outperforms transformers with performer, reformer, and linear attention.
We note that the Transformer-Xi XI and Perceiver AR PPI numbers are from the original papers, and may not be directly comparable to our results. In particular, they use a tokenizer with a different vocab size,
20

---

Published as a conference paper at ICLR 2023
which means that the PPLs are not directly comparable. In addition, the larger vocab size necessitates a change in the model (adaptive softmax) that may affect performance. The top five numbers in Table trained with the same setup and directly comparable each other.
F.3
PG-19
We evaluate models trained on the PG-19 dataset (Rae et al., 2019), a natural language dataset comprised of texts from books. We compare the performance of Hybrid H3 compared against Transformers and linear attention. We use the same setup as evaluating on WebTet.
Table 11: Test PPL on PG-19.
Models
PPL
Transformer (T25M)
17.0
Hybrid H3 (125M)
16.2
Linear Attention (125M) (Katharopoulos al., 2020)
19.1
Table 11 shows that Hybrid H3 outperforms transformers and linear attention.
F.4 LENGTH EXTRAPOLATION
One property of SSMs is that they can naturally extrapolate sequence lengths longer than those seen dur- ing training. We use the synthetic associative recall task to demonstrate that H3 maintains this capability. To do so, train two-layer H3 model sequences of length 20 drawn from the associative recall syn- thetic language. Then, we evaluate accuracy of the last token prediction on sequences of length 20 and 40.
Table 12: Accuracy of H3 model trained for associative recall sequences of length 20, evaluated sequences length and 40.
Models
Acc, seqlen 20
Acc, seqlen 40
H3
99.8
98.4
Table 12 shows that H3 maintains accuracy on sequences of length 40, which twice the length of the training sequences.
F.5 SCALING IN NUMBER OF TOKENS
We evaluate how well a Hybrid H3 model scales with the number of tokens seen during training, compared Transformer. For these experiments, we train 125M Hybrid H3 model and 125M Transformer the Pile for 5B, 1OB, and 15B tokens. We use a learning rate of 6e-4, adjusting the warmup to be 1% of the total training time, and adjusting the decay rate to decay the learning rate 6e-5 by the end of training.
Table 13: Test PPL on the Pile for models trained with fewer tokens.
Train Tokens
Hybrid H3 (T25M)
Transtormer (125M)
5B
I1.8
12.7
10B
10.7
11.3
15B
10.2
10.7
Table 13 shows the results. Both the Hybrid H3 model and Transformer model improve as the number of training tokens increases.
F.6 H3 LANGUAGE MODEI.
Table 14: Zero-shot performance SuperGLUE with rank classification. Best results for each model size in bold.
Model
WSC
WIC
RTE
CB
MultiRC
ReCoRD
BoolQ
COPA
Average
OPT-T25M
39.4
52.0
48.7
37.4
58.9
44.9
59.6
60.0
50.1
HPT-NNe-125M
36.5
53.6
53.1
41.1
59.9
39,6
62.2
60.0
50.8
H3-125M
61.5
50.0
53.1
41.1
4.6
15.8
46.4
51.0
40.4
Hybrid H3-125M
39.4
51.4
59.2
48.2
51.4
55.0
59.6
67.0
53.9
Table 3-shot performance on SuperGL. with rank classification. Best results each size in bold, second
best underline.
Model
WSC
WIC
RTE
CB
MultiRC
ReCoRD
BoolQ
COPA
Average
OPT-125M
36.5
50.2
47.3
44.6
57.9
44.9
41.9
60.0
47.9
GPT-Neo-1255
38.5
50.0
53.1
17.9
56.3
39.6
62.1
60.0
47.2
H3-125M
63.5
50.0
52.3
48.2
32.6
15.8
37.8
51.0
43.9
Hybrid H3-125M
43.3
49.1
58.1
51.8
48.9
55.0
56.1
67.0
53.7
We report the results of pure H3 language model evaluaattt.. We train 1 125M model the Pile
for 400B tokens. Tables and 15 show zero-shot and few-shot performance on SuperGI.UE respectively.
21

---

Published as a conference paper at ICLR 2023
F.7
GENERATION PERFORMANCE
Table 16: Zero-shot performance on SuperGLUE with generation. Best results for each size in bold, second best underline.
Model
WSC
WTC
RTE
CB
MultiCC
ReCoRD
BookQ
COPA
Average
OPT-125M
36.5
48.4
49.8
8.9
39.1
44.9
45.9
50.0
41.7
GPT.Nen- 125M
27.9
Hybrid H3-125M
0.0
0.0
47.3
8.9
4,4
55.0
47.6
67.0
28.8
GPT-2 medium (355M)
50.0
50.2
16.2
21.4
10.5
53.3
38.4
65.0
38.1
3PT-350M
41.3
34.8
49.5
16.1
23.6
51.4
39.7
60.0
39.6
Hybrid H3-355M
22.1
21.5
47.3
8.9
17.1
62.3
44.4
69.0
36.6
Table 17: 3-shot perforance on SuperGLUEE with generation. Best results for each size in bold, second best
Model
WSC
WTC
RTE
CB
MuliKCC
ReCoRD
BeniQ
COPA
Average
OPT-125M
36.5
49.1
47.3
33.9
35.5
44,8
38.5
60.0
43.2
GPT-Neo-1250 125M
38.5
50,0
53.1
42.9
22.5
39.7
61.2
68.0
47.0
H3-1125M
0.0
0.0
47.3
8.9
0.0
15.4
37.8
53.0
20.3
Hybrid H3-125M
43.3
49.1
40.3
55.2
49.5
67.0
50.5
GPT medium (355M)
36.5
50.5
47.3
28.6
35.3
53.1
37.8
63.0
44.0
OPT-350M
37.5
59.0
Hybrid H3-355M
42.3
47.5
50.5
37.5
57.5
61.4
45,4
73.0
51.9
We report results on SuperGLUE for generation. Instead taking rank classification, we instead let the model generate response, and we search for the gold label (i.e., "yes" or "no" for the yes/no questions) in the output. Tables 16 and 17 report the results. The trends for few-shot learning match with the logit results, but the hybrid and H3 models perform very poorly in zero-shot performance on some tasks. In these cases, the models tend to generate long text responses that are not relevant to the answer. The few-shot learning examples help the models generate answers in parsable format.
F.8 NON-TEXT SEQUENCE MODELING
We show that H3 outperforms Transformers classification and seizure classification over raw EEG signals. H3 sets state-of-the-art performance on seizure classification and matches S4 speech classification whichh suggests that H3, one of its hybrids, may be a strong candidate for multimodal foundation model. Appendix E gives experimental details, and Appendix F gives additional experiment brain fMRI data.
non-text sequence modeling tasks: speech
Seizure Classification from EEG Seizures, which are characterized by uncontrolled brain activity, are some of the most common neurological disorders (Fisher et al., 2014). Chronic seizures, or epilepsy, cause range of psychiatric and psycho-social disorders and impact the lives roughly one percent of the global population (Kerr, 2012). The first step treating epilepsy is manual analysis of scalp EEG by board-certified neurologists. However, the vast amount of EEG data produced by each patient (which can be up to days of data) makes manual EEG analysis costly and time-consuming process.
To mitigate the costs associated with EEG monitoring. recent deep learning techniques have began show promise in flagging abnormal EEG segments for potential scizure events (Siddiqui et al., 2020). A challenge with classifying EEG data is the trade-off between increasing input sequence length, where more context has been shown to improve seizure classification performance Saab et al. (2020), with the increased difficulty of training deep learning models long sequences (e.g., EEG signal sampled at 200Hz produces 12,000 time steps per minute). As a result, many techniques involve domain-speciaalize models and pre-processing steps, such as FFT transforms and graphical representations Tang et al. (2021).
We use the largest publicly available EEG corpus, TUSZ v v1.5.2 (Shah et al., 2018), which includes 5,612 EEG signals from 636 patients, with 3,050 annotated seizures. Signals are segmented into 60 6e-sec clips, and split into train/val/test by patient. The train set contains 39765 clips, the val set contains 4351 clips, and the test set contains 10001 clips.
Table 18: Performance (AUROC) on 60s seizure classification from raw REG (sequence length 12000).
H3
Transformer Dense-CNN CNN-LSTM LSTM
ID-CNN
83.2
x
78.0
68.6
69.3
69.7
We evaluate binary seizure classification of 60-sec EEG clips. sampled at 200Hz, with 19 electrodes: E RI 2|2,000x19 and y (0.1) on the TUSZ v1.5.2 (Shah et al., 2018) corpus. Transformers cannot process the long sequence length of EEG signals without running out of GPU memory, whereas H3 can-and sets state-of-the-arr performance.
Raw Speech Classification The SC10 speech commands task (Warden, 2018) contains raw audio signals one second in length, sampled at 16kH Similarly EEG signals, Transformers cannot process the long sequence length. Table shows that H3 comes within half a point of S4, the state-of-the-arr method.
22

---

Published as a conference paper at ICLR 2023
Table 19: SC 10- classificatin audio (sequence length 16000).
H3
54
WaveCinnDD Transformer
Performer
CKConv
97.04
97.50
96.25
30.77
71.66
Functional Magnetic Resonance Imaging Data Functional Magnetic Resonance Imaging (IMRI) data are typically represented in four dimensions, indicating the measured (BOLD) signal in temporal sequences S = (Vi.... Vi} of 3-dimensional volumes V E R each indicating the BOLD signal for all spatial locations of the brain (as defined by three spatial dimensions a, y, and z). A key challenge for the analysis of fMRI data is the high dimensionality and low sample size of its datasets, which typically contain many hundred thousand dimensions (i.e., voxels) for each of several hundred volumes V in each of tens to hundreds of sequences S. In this setting, where the number of features exceed the number of samples, standard machine learning approaches are prone to overfitting.
In spite of the low sample size of individual datasets, neuroimaging research can be considered recently entering a big data era because researchers more frequently share their collected datasets publicly (Markiewicz al., 2021). The availability of these data open up the opportunity for pre-training in neuroimaging at scale, as recently demonstrated by (Thomas et al., 2022), enabling models to utilize the knowledge that can be learned from public functional neuroimaging data for the analysis of individual datasets. Specifically, (Thomas et al., 2022) evaluate the performance of several self-supervised learning frameworks for functional neuroimaging data by first pre-training models on broad fMRI dataset spanning 11,980 fMRI runs from 1,726 individuals across 34 datasets and subsequently adapting the pre-trained models to two downstream mental state decoding datasets (namely, the HCP (Van Essen et al., 2013) and MDTB (King et al., 2019) datasets). In mental state decoding, predictive models are tasked with identifying (i.e., decoding) some set of mental states (e.g.. answering questions about prose story or math problem) from measured brain activity. The authors find that a GPT- model, pre-trained in causal learning framework, performs best in decoding the 20 (HCP) and 26 (MDTB) mental states of the two downstream datasets.
To evaluate the performance of H3 on fMRI data, we replicate this analysis, using the up- and downstream fMRI datasets that were published by (Thomas et al., 2022), treating H3 as drop-in replacement for the GPT model. To alleviate the high dimensionality challenge of fMRI data, and due to the generally high spatial correlation of brain activity, the original authors have reduced the volumetric time series S to cerooonlll-ineeeddeeee
a set θ E f ,.. of 1.024 of functional modes (DiFuMo) brain atlas (Dadi et al., 2020)), each describing the BOLD signal for some subset of voxels U/. E V, such that the resulting sequences X < describe the activity pattern of each brain network n for time points t.
brain networks @ (as defined by the dictionaries
In line with (Thomas et al., 2022), we first pre-train models f(() to predict the distribution of brain activity for the next time point j of an input sequence X, using mean absolute error (Tree)) training objective given the model's output X E 1R:x";
Liree = 1 - [X,, - Xjali Xx, = bn + I,,(E' (EE We. n
= + Eros + ba + -n Xt..Wn.e. Here, E E Re and EPos ER. represent learnable embeddings for each possible time point and position of an input sequence (for details, see (Thomas et al., 2022)). As the sampling frequency of fMRI is variable, the same position of an input sequence can correspond to different time points. Note that f(-) processes the input in a oover-dimensional embedding representation E E RR* (with e 768 dimensions).
We evaluate two model architectures for f(-), namely, the GPT variant used in (Thomas et al., 2022), with hidden layers and 12 attention heads, and a corresponding H3 variant with hidden layers (with H = 64 and m 1; see section 3). For both models, the sequence of hidden-states outputs of the last model layer are used to determine X.
Just as (Thomas al., 2022), we randomly divide the upstream data into distinct training and validation datasets by randomly designating 5% of the fMRI runs of each of the 34 upstream datasets as validation data (at minimum of 2 runs per dataset) and using the rest of the runs for training. During upstream learning, we then randomly sample sequences of 10 to 100 time points from the fMRI runs and train models with the ADAM optimizer (with 31 0.9, 3g = 0.999, and Ie - ) for 5, 000 steps a mini-batch size of 512, and a learning rate of 5e -4 We apply a linear learning rate decay schedule (with warm-up up phase of 1 % of the total number of training steps), gradient norm clipping at 1.0, and L2- .2-regularisation (weighted by 0.1). We also apply dropout at a rate of 0.1 for the Gaseb model (based on (Thomas al., 2022)) and evaluate three dropout rates for H3: 0.1. 0.2, and 0.3.
We find that the H3 variant trained with 0.2 2 dropout performs on par with the GPT model, in terms of mean absolute emrrr (Fig. 3), and therefore continue all further analyses with this model variant. We also find that
23

---

Published as a conference paper at ICLR 2023
Training
Evaluation
0.80
0.80
GPT
H3 (dropout: 10%)
0.78
0.78
(dropout: 20%)
H3 (dropout: 30%)
0.76
0.76
Lrec
Lrec
0.74
0.74
0.72
0.72
0.70
0.70
0
2000
4000
0
2000
4000
Training steps
Training steps
Figure 3:
Upstream absolute (Lree) in training and evaluation datasets the course of model training.
both models exhibit almost identify Lree error distributions throughout the brain, with relatively higher errors in the posterior parietal, occipital, and cingulate cortices as well parts of the limbic system (Fig. 4).
GPT
H3
L
R
L
R
lateral
ventral medial
0.000
0.001
0.002
0.003
0.000 0.001
0.002
0.003
Figure 4:
Mean absolute error (Lree) the final pre-trained models for each voxel the brain projected onto the inflated cortical surface of the FsAverage template (Fischl, 2012).
To adapt the pre-trained models mental state decoding, add learnable classification embedding gels E " to the end of input sequences X and forward the model's prediction f(E to a decoding head p(-), composed of dense hidden layer with model units (one for each embedding dimension, with tanh activation) as well as a sof ftmar output layer with one model unit for each considered mental state in the data. Accordingly,
adapt models by optimizing standard cross entropy loss objective: Less = E,y,oog pf (E ))x, where y indicates a binary variable that 1 is the correct mental state and otherwise.
During downstream adaptation, begin training with the respective pre-trained model parameters and then allow all parameters to change freely. Similar to (Thomas et al., 2022), we randomly split each of the two downstream datasets into distinct training and test datasets, each comprising 40 (HCP) 10 (MDTB) distinct individuals. We adapt models for 750 steps at mini-batch size of 256 and a learning rate 5e
24

---

Published as a conference paper at ICLR 2023
(otherwise using the same learning parameters for upstream training). Importantly, repeat each down- stream training run 20 times using different random seeds, leading to different random splits of the data and variability in other non non-deterministic training factors (such as random initialization and data shuffling).
As for the upstream data, the H3 and GPT-based models generally perform par in their mental state decoding performances in the two left-out test datasets (Table 20).
Table 20: Downstream adaptation performance models pre-trained on fMRI data, averaged over 20 training runs
with varying random seeds. Fl Pl-scores are macro-averaged.
Datasct
Model
Acc. (+95%CCI
HCP
GPT
(
(H0.39)
H3
88.75 (+:0.33)
87.16 (4:0.37)
MDTB
GPT
89.47 (#0.44)
(H0.54))
H3
88.25 (+:0.15)
(1.066))
25


Reviews:
[{"review_id": "zjzG3bctoc", "reviewer": "Reviewer_EkLp", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This work aims to improve the state space models (SSMs), which achieve state-of-the-art performance on the LRA benchmark but perform worse on language modeling and large-scale language pre-training. The authors identified two problems related to the model design through two synthetic tasks. Based on the observation, the authors further proposed some modifications which lead to strong LM performance. At the same time, the authors also leverage Flash to improve the efficiency of the inference.", "summary_of_the_review": "The authors solved an important problem for state-space models, and the empirical study is solid. I hope the authors can address my concerns, and currently, I give a weak acceptance.", "strength_and_weaknesses": "Pros:\n\n1. It is interesting to use synthetic tasks to investigate the drawback of the state-space model.\n2. It is important to improve S4 on large-scale language modelling tasks. This work make an essential step towards this goal\n3. The empirical results is adquante to support the advantage of the model\n\nCros:\n\nMajor questions:\n\n1.The authors proposed two synthetic tasks, Induction Head and Associative Recall. However, the relationship between these two tasks and the language model performance remains unknown. Why does the poor performance of the language model have a strong connection with these two tasks? This connection should be discussed clearly since it motivates the modification, which (as the authors stated) leads to better language model performance. \n\n2.The change improves the language modelling performance, but will the new model hurt the LRA performance. If I didn't miss anything, there is only an efficiency study on LRA without performance comparison. This work will be problematic if the model fails to work on LRA but only work on LMs.\n\nMinors:\n\n1. The authors claim that standard attention has the ability to complete Induction Head and Associative Recall. Is this a rigorous statement or not? If yes, I recommend the authors give formal proof of the statement. I am not that sure whether attention can succeed #for any length#\n\n2. The authors were motivated by linear attention and proposed an O(nlogn) architecture with FFT. The work \"Stable, Fast and Accurate: Kernelized Attention with Relative Positional Encoding\" achieves the same efficiency with FFT based on linear attention. I believe the relationship should be discussed as a reference ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "xWn5efDnERm", "reviewer": "Reviewer_UjQC", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposed the H3 architecture to narrow the gap between state space models and attention models on language modeling. Basically, H3 uses a shift SSM and a diagonal SSM to replace the non-linear function in linear attention. The shift SSM is designed to remember tokens in the input sequence while the multiplication between the two SSM outputs for token comparison. \nTo further improve H3 efficiency on hardware the authors proposed FlashFFTConv, which uses fused block FFT algorithm. \n\nExperiments were conducted on training multiple variants of H3 on the PILE dataset, and compared with corresponding GPT2 checkpoints, on OpenWebText and WikiText103. Moreover, the authors also evaluate H3 on zero-shot and few-shot classification.\nAt last, they evaluate efficiency of the proposed FlashFFTConv algorithm on the LRA benchmark, achieving about 2 times speed up.", "summary_of_the_review": "To sum up, I am worried that the paper needs to re-organize its experiments and analysis before publication.", "strength_and_weaknesses": "Strengths: The design of H3 is novel and well-motivated and the paper writing is clear and easy to follow. The authors also clearly claimed their contributions.\n\nWeaknesses: The main concerns are from experiments. I found that the experimental setup on language models is not standard, which makes the comparisons unfair. It is difficult to position H3 among other neural models in language modeling. I elaborated my concerns in the following questions.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "pr5qJlPfla", "reviewer": "Reviewer_PnYM", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper presents a new state space model (SSM) that achieves on par with transformer based language models while inferring up to 2 times faster. The authors first highlight the drawback of previous SSMs compared to attention -- SSMs can't recall earlier tokens in the sequence and can't compare tokens across the sequence. The authors propose an improvement that captures these features by using `shift` and `diagonal` SSM layers. Shift layer shifts dimensions in an input (or steps if the input is a sequence) and diagonal SSM is convenient for computing kernel in the convolution. By borrowing the structure of linear attention, the authors propose a new SSM layer, called H3, that utilizes recurrence structure and can scale to longer sequences better. They show that H3 improves previous SSMs and achieves on par with transformers on synthetic as well as real OpenWebText tasks. They further propose a mechanism, called FastFFTConv, for efficient estimation of convolution where the convolution is replaced with `iFFT (FFT(u) * FFT(f))`. They divide input into chunks -- each chunk fits into the GPU memory -- and keep end-state of each chunk for a more efficient convolution which results in up to 3x speed up against attention.", "summary_of_the_review": "H3 is a simple yet novel extension of previous SSMs. It presents improvements compared to Transformers and can infer 2 times faster. FastFFTConv is useful for efficient training and can be applied to other SSM works.", "strength_and_weaknesses": "**Strengths** I found the paper interesting and useful for thinking beyond typical transformer architectures. The paper proposes a simple SSM layer that can replace any self-attention module and broadly applicable. Experimental results are also promising and efficiency improvements should help other work in state space models.\n\n**Weaknesses**\n1. You mention that SSMs are capable of extrapolating to sequences longer than observed during training. I found no empirical results to support this claim in the paper.\n\n2. There are missing references in long-term dependency line of work, especially *Perceiver* style models. Given that one of the things that the paper claims is capturing long term dependencies, I think these need to be cited and compared against.\n\n3. How would you compare to other models in Table 2 is H3 is trained from data rather than by construction? Please also clarify how other models, excluding attention, are obtained.\n\n4. One thing that Transformers have shown is they scale with data and model sizes. While from 125M to 335M, there is a clear improvement in Table 4, I would like to see if there is a similar pattern for data size as in Table 6.\n\nGeneral-purpose, long-context autoregressive modeling with Perceiver AR. Curtis Hawthorne, Andrew Jaegle, Cătălina Cangea, Sebastian Borgeaud, Charlie Nash, Mateusz Malinowski, Sander Dieleman, Oriol Vinyals, Matthew Botvinick, Ian Simon, Hannah Sheahan, Neil Zeghidour, Jean-Baptiste Alayrac, João Carreira, Jesse Engel.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%