H UNGRY H UNGRY
H IPPOS :
T OWARDS
L ANGUAGE
M ODELING WITH S TATE S PACE M ODELS
Anonymous authors
Paper under double-blind review
A BSTRACT
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 ﬁnd 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 H3-attention model that retains two attention layers surprisingly outperforms Transformers on
OpenWebText by 1.0 PPL.
When trained on the Pile at small/medium scale (125M and 355M parameters), hybrid H3-attention language models display 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.
Next, to improve the efﬁciency of training SSMs on modern hardware, we propose F LASH FFTC ONV .
F LASH FFTC
ONV uses a fused block FFT algorithm to improve efﬁciency on sequences up to 8K, and introduces a novel state passing algorithm that exploits the recurrent properties of SSMs to scale to longer sequences.
F LASH FFTC
ONV yields 2 × speedup on the long-range arena benchmark and allows hybrid language models to generate text 1.6 × faster than Transformers.
I NTRODUCTION
State space models (SSMs) have achieved state-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 et al., 2022;
Olsson et al., 2022), or whether it is a function of the signiﬁcant organizational resources that have been spent training and tuning large attention-based language models (Chowdhery et 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, 2022b; Kao et al., 2021).
We take ﬁrst 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 a new SSM layer that nearly matches attention in language modeling.
Second, we propose better hardware-aware 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 a 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 a sequence.
We ﬁnd that existing SSMs struggle to model these synthetic languages.
To probe how important these 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
ShiftSSM
DiagSSM
X
K Q
V
Y
H3 Layer H3 for Associative Recall
null null
Diag
I(x t-1 =a)
I(x t+1 =a) I ( x t + 1 = a )
I(x N-1 =a) null
I ( x t = a ) a
Shift
I(x N =a) a null y t null y t+
Shift Shift a
I(x t =a) 3 x t x t+
Diag
I(x t =a)  a x N
Diag
I ( x N = a ) 3 y N
[Figure]
FlashFFTConv between their outputs and input projections.
The SSMs allow H3 to keep a log of tokens (to recall them later), while the
H3 matches attention on the synthetic languages and almost closes the gap with Transformers on language modeling—coming within 0.4 perplexity of Transformers on OpenWebText (compared to 3.4 ppl for existing 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 outperforms Transformers on OpenWebText by 1.0 perplexity.
To further evaluate H3 on language modeling, we train 125M- and 355M-parameter 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-based language models of the same size in perplexity, and match or outperform them on a majority of tasks in 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 1.6 × faster inference than
Transformers.
Furthermore, H3 maintains quality on non-text sequence modeling, matching the S model (Gu et al., 2022a) on raw speech classiﬁcation and setting state-of-the-art performance on seizure classiﬁcation over raw EEG signals.
These results suggest that H3, or other SSM-attention hybrids, may be a promising direction for future language models or multimodal foundation models, especially with more resources towards ﬁnding optimal hyperparameters and training schedules for SSMs.
Scaling SSMs.
Next, we improve the efﬁciency 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 F LASH FFTC ONV , a hierarchical algorithm for computing SSMs, inspired by IO-Aware attention (Dao et al., 2022b).
The technical challenge is that SSMs require a 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 1 .
To use specialized matrix multiply units, we appeal to classical techniques that split the FFT into blocks and computes it using a series of matrix multiplications.
Combined with kernel fusion, this “block” FFT solution increases hardware efﬁciency, but only as long as the sequence length can ﬁt into GPU SRAM (on-chip memory, analogous to L1 cache on the CPU)—up to sequence length 8K on modern A100.
To scale to sequences longer than 8K, we propose a state passing algorithm
(Figure 1 right), specialized to SSMs.
The key insight is that we can use the recurrent properties of SSMs to process the input in chunks—as long as we keep track of an additional state vector.
The state passing algorithm splits the input into the largest chunks that can ﬁt into GPU SRAM, efﬁciently computes the FFT-based convolution using block FFT, and updates an intermediate state to start the next chunk.
Using this state-passing algorithm,
F LASH FFTC
ONV can scale SSMs to any sequence length—even longer than can ﬁt on GPU SRAM at once—while maintaining a near linear compute complexity.
F LASH FFTC
ONV sets state-of-the-art speed on long range arena using S4 (Gu et al., 2022a), outperforming Transformers by 5.8 × and previous
S4 models by 2 × .
F LASH FFTC
ONV trains H3 4-8 × times faster than attention for long sequences.
B ACKGROUND
We present some background on state space models and linear attention, which inspired our H3 layer.
2. S TATE S PACE M ODELS
A continuous-time state-space representation (Brogan, 1974) deﬁnes a linear mapping from an input signal u ( t ) ∈ R (as a function of time t ) to an output signal y ( t )
∈ R through a state-variable x ( t ) ∈
R m , with the following differential equation, for some matrices A ∈ R m × m , B ∈ R m × 1 , C ∈ R 1 × m ,
D ∈ R 1 × 1 : ˙ x ( t )=
A x ( t )+ B u ( t ) , y ( t )=
C x ( t )+ D u ( t ) .
Similarly, a discrete-time state-space representation deﬁnes a linear mapping from a discrete input signal u i (for i =1 , 2 ,... ) to a discrete output signal
y i though a state-variable x i ∈ R m : x i =
A x i − 1 + B u
i y i
=
C x i + D u i .
A state-space model (SSM) uses these representations as a layer in a 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 in parallel, each corresponding to one hidden dimension.
To preserve the sequence history,
HiPPO (Gu et al., 2020) projects the history on a basis of orthogonal polynomials, which translates to having SSMs whose A , B matrices are initialized to some special matrices.
This recurrent form of SSMs allows efﬁcient 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 efﬁcient training, given the entire sequence of the input u 1 ,...,u N , the output sequence y 1 ,...,y N can also be written as the convolution of the input with the ﬁlter (Gu et al., 2021):
[Equation]
That is, from an initial condition x 0 , we have y i =
CA
i
B x 0 +( f ∗ u )
i + D u i , where ( f ∗ u ) denotes a linear convolution between f and u .
If we set the initial condition x 0 to be zero, then y is exactly a linear convolution of u , with a residual connection D u .
More generally, any linear time-invariant system (of which SSMs are a special case) can be written as a convolution.
Given a 1D input sequence u ∈ R N of length N , we denote the 1D output sequence y ∈ R N of an SSM parameterized by matrices A , B , C , D as
[Equation]
To simplify notation, we omit the reference to A , B , C , D and write y =
SSM( u ) if they are clear from context.
When u is multidimensional of dimension d , we stack d of these SSMs together that deﬁnes a mapping from u ∈ R N × d to y ∈ R N × d , using the same notation y =
SSM( u ) .
To construct the ﬁlter f from A , B , C efﬁciently, A is often constrained to be diagonal (Gupta et al., 2022;
Gu et al., 2022b), or diagonal plus low-rank (Gu et al., 2022a).
SSM through FFTs.
Computing the convolution naively through conventional matrix operations is expensive for long kernels, scaling as O ( N 2 ) .
Instead, we can use FFTs: take the FFT of f and u , multiply them together pointwise, and then take the inverse FFT.
This yields an O ( N log N ) algorithm.
2. L INEAR ATTENTION
We describe linear attention (Katharopoulos et al., 2020) and its connection to RNNs, which inspired our model design (Section 3).
In standard attention (Vaswani et al., 2017), we have N query/key/value tokens Q i ,K i ,
V i ∈ R d for i = 1 ,...,N , where N is the sequence length and d is the head dimension.
For some similarity metric
Sim: R d × R d → R , we want to compute the output:
[Equation]
R
[Equation]
For standard softmax attention, Sim( q,k )=
e q (cid:62) k
(often the dot product is scaled by 1 / √ d ).
Linear attention makes the assumption that Sim has the form Sim( q,k )=
φ ( q ) (cid:62) φ ( k ) , for some (nonlinear) function φ .
The output is then O i =
φ ( Q i ) (cid:62) (cid:80) ij =1 φ ( K j ) V (cid:62)
j φ ( Q i ) (cid:62) (cid:80) ij =1 φ ( K j ) .
Let S i
= (cid:80) ij =
1 φ ( K j ) V (cid:62) j ∈ R d × d , z i = (cid:80) ij =1 φ ( K j ) ∈ R d , d i = φ ( Q i ) (cid:62) z i ∈ R .
Then O i =
φ ( Q i ) (cid:62) S i
d i .
This connects linear attention to RNNs: the output O i is a function of S i and z i , both of which are incrementally updated (as cumulative sums).
H UNGRY H UNGRY H IPPOS L AYER TO M ODEL D ISCRETE S EQUENCES
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 (based 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. M OTIVATION : S YNTHETIC L ANGUAGE M ODELING T ASKS
We describe two closely-related synthetic tasks, summarized in Table 1.
Olsson et al.
(2022) argues 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]
The Induction Head task tests how well a model can recall content after a special token (e.g., (cid:96) in Table 1).
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., 2016) is similar to the induction head task, but requires the model to remember multiple key-value pairs.
At the end of the sequence, the model must recall a speciﬁc value belonging to a speciﬁc key.
[Table]
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 (e.g., comparing keys to decide which value to recall).
Attention has both these capabilities: it can compare tokens by constructing the quadratic attention matrix QK (cid:62) , and it can recall tokens by direct copying (multiplying softmax( QK (cid:62) ) with V ).
In Section 3.2, we design our new layer H3 to enable these capabilities in SSMs, narrowing the expressivity gap between SSMs and attention.
3. H3 L AYER
H3 uses SSMs with shift and diagonal matrices, along with multiplicative operations against projections of the input to capture the missing capabilities identiﬁed by the synthetics.
High-level Intuition.
(i)
To remember tokens from the past, we want the state x i to copy from the input u i , and then pass that information to the next state x
i +1 .
As x i +1 relates to x i by A x i , we use a discrete
SSM with a shift matrix A (described formally below) that shifts the elements of a state vector (e.g., mapping [ a,b,c ] → [0 ,a,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 ( SSM shift ), and we replace the summation S i with a SSM with diagonal A ( SSM diag ).
The output, for the case of head dimension d h =1 , is:
[Equation]
where (cid:12) 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.
Remembering Key Tokens: Shift and Diagonal SSMs.
The shift and diagonal SSMs are designed to address the capability to log tokens after particular events.
In the shift SSM, we constrain A ∈ R m × m to be a shift matrix A i,j = (cid:26) 1 for i − 1= j 0 otherwise .
The action of this matrix on the hidden state x i is to shift each coordinate down by one—thereby creating a “memory” of the previous states.
For example, if B = e 1 , the ﬁrst basis vector, then x i =[
u i ,u i − 1 ,...,u i − m +1 ] contains the inputs from the previous m time steps.
We learn
B and C ( B can also be ﬁxed to e 1 for simplicity, in which case the output is a 1D conv.
with kernel 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 a 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 Q 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 1 and shown schematically in Figure 1 (left).
We use the
H3 layer to construct a model in the same style as 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 a hybrid H3-attention model (two attention layers while the rest are H3, Sections 3.3 and 5).
Require:
Input sequence u ∈ R N × d from the previous layer, weight matrices W Q , W K , W V , W O ∈ R
d × d , a shift
SSM SSM shift , a diagonal SSM SSM diag , head dimension d h .
1: Compute Q = u W Q , K = u W K , V = u W V ∈ R N × d .
2: Pass K through the shift SSM: K =SSM shift (
K ) ∈ R N × d .
Efﬁciency We show that H3 scales as O ( N log N ) with sequence length N —asymptotically more efﬁcient than attention, which typically requires O ( N 2 d ) time and O ( N 2 ) space 2 (proof in Appendix D.3).
Proposition 1.
Let N be the sequence length, d be the hidden dimension, and assume that the head dimension d h is of order O (1) .
Then the H3 layer takes O ( d 2 N + dN log N ) time and O ( dN ) space to compute.
E XPRESSIVITY
We show that H3 can model our synthetic languages, as well as natural language on OpenWebText (Gokaslan et al., 2019).
We also present a hybrid H3-attention extension that outperforms
Transformers on OpenWebText.
Mechanism for Solving Associative Recall with H3.
H3 is expressive enough to solve our synthetic language modeling tasks, as shown in Table 2.
Figure 1 (middle) shows a mechanism for a 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 as a gate on whether to let a 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 ﬁnal multiplicative interaction gates whether to let the diagonal SSM’s output through—based 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.1.
[Table]
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
(e.g., on the OpenWebText dataset).
As shown in Table 3, H3 comes within 0.4 perplexity points of
Transformers when trained for 50B tokens on OpenWebText, and performs much better than existing
SSM variants (S4D, GSS), by 3 − 3 .
9 points.
Extension: H3-attention Hybrid Model.
A simple hybrid H3-attention language model surprisingly outperforms Transformers on OpenWebText 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+ N/ 2 for an N -layer model,
N even).
The H3-attention hybrid also outperforms the GSS-attention hybrid (Mehta et al., 2022).
F LASH FFTC ONV : E FFICIENTLY T RAINING SSM S
To improve the efﬁciency of SSMs on modern hardware, we propose F LASH FFTC ONV .
F LASH
FFT
C ONV fuses the FFT, pointwise multiply, 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 A100) for sequence lengths up to 8K. For sequences longer than 8K, the computation no longer ﬁts in GPU SRAM 3 , so we propose a novel state-passing algorithm that splits the sequence into chunks to compute the FFT convolution one chunk at a time.
F LASH FFTC ONV can speed up any SSMs (not just H3).
4. F USED B LOCK FFTC ONV
We deploy two techniques to speed up the FFT-based convolution for sequences shorter than 8K: kernel fusion and block FFT.
Kernel fusion addresses IO bottlenecks due to reading and writing of intermediate results, while block FFT allows the FFT-based convolution to utilize specialized matrix multiplication units.
These techniques allow us to speed up FFTConv by 2 × (Section 6) for sequences shorter than 8k.
Kernel Fusion.
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 u and ﬁlter f has the form iFFT ( FFT ( u ) (cid:12) FFT ( f ))
(where (cid:12) denotes pointwise multiplication).
It requires reading and writing intermediate results to GPU memory—which can dominate the runtime.
Following F LASH A TTENTION (Dao et al., 2022b), we ﬁrst fuse the entire FFTConv into a single kernel and compute it in SRAM to avoid this overhead.
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 ×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
IO bottlenecks, compute becomes the bottleneck (note that a single FFT on GPU is usually IO bound).
Suppose that we want to perform an N -point FFT, which is equivalent to multiply by the DFT matrix
F N .
Suppose that N = N 1 N 2 for some integers N 1 ,N 2 .
By the Cooley-Tukey decomposition of the
DFT (Cooley & Tukey, 1965; Bailey, 1990) (also known as the four-step FFT algorithm), we can write
F N = P ( I N 2 ⊗ F N 1 ) P (cid:62) D
( I N 1 ⊗ F N 2 ) P , where P denotes a ﬁxed permutation that reshapes the input as a N 1 × N 2 array and then transpose it, ⊗ denotes Kroneker product, D is a N × N diagonal matrix (called the twiddle factors)
(Dao et al., 2022a), and I N i and F N i are the identity and DFT matrix of size
N i × N i .
As I N 2 ⊗ F N 1 and I N 1 ⊗ F N 2 are just block-diagonal matrices, we can make use of specialized matmul units to perform these multiplications.
Similarly, if N = N 1 N 2 N 3 then we can decompose the
N -point FFT into a series of (block) FFT of size N 1 , N 2 , and N 3 , interleaved by permutation.
The block FFT algorithm incurs O
( Nr log N/ log r ) FLOPs for a sequence length N , if N can be written as r p for two integers r,p .
This incurs more FLOPs than standard FFT ( O ( N log N )) , but can run faster when we using specialized matrix multiplication hardware.
4. S TATE -P ASSING
However, the fused kernel cannot run if the sequence is too long to ﬁt into GPU SRAM (longer than 8K on
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 to split the FFTConv of a length- N sequence into chunks of size N (cid:48) each ( N (cid:48) is the longest FFT we can ﬁt into SRAM), assuming N is a multiple of N (cid:48) ).
We use
FFTConv to compute each chunk and use a recurrence to connect the chunks.
In particular, we split the inputs u into C = N/N (cid:48) chunks u ( c ) ∈ R N (cid:48) for c
=1 ,...,C .
Similarly, split the states x into x ( c ) ∈ R N (cid:48)
× m and the output y into y ( c ) ∈ R N (cid:48) for i =1 ,...,C .
We will only need the end-state x ( c ) N (cid:48) of each chunk c .
Let f =[ CB , CAB , CA 2 B ,..., CA N (cid:48) − 1 B ] be the SSM ﬁlter.
Recall from Section 2 that for each chunk c , y ( c ) i =
CA
i B x ( c − 1) N (cid:48)
+( f ∗ u ( c ) )
i + D u ( c ) i , since x ( c − 1) N (cid:48) , the end-state of the previous chunk ( c − 1) is the initial condition for the current chunk c .
In vector notation, y ( c ) =
M xy x ( c − 1) N (cid:48)
+ f ∗ u ( c ) + D u ( c ) for some matrix M xy ∈ R N (cid:48) × m .
Additionally we need to update the end-state of each chunk with x cN (cid:48)
=
A N (cid:48)
x ( c − 1) N (cid:48)
+ M ux u ( c ) for some matrix M m × N (cid:48) ux (derivation in Appendix C.2).
In essence, we can compute the output for each chunk with FFT-based convolution as long as we 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 we only compute FFT of length N (cid:48) , and update some hidden state each iteration.
Let B LOCK FFTC
ONV refer to our fused block FFTConv kernel.
Then, the state-passing algorithm for 1D input is given by Algorithm 2.
For inputs of dimension d where we stack d SSMs, we simply batch Algorithm 2 along the d -dimension.
Require: Input u ∈ R N , SSM parameterized by matrices A ∈ R m × m , B ∈ R m × 1 , C ∈ R 1 × m , D ∈ R 1 × 1 , chunk size N (cid:48) where N is a multiple of N (cid:48) .
1: Precompute A N (cid:48) ∈
R m × m , M ux =[
A N (cid:48) − 1 B ,..., B ] ∈ R m × N (cid:48) , M xy =[ C , CA ,..., CA N (cid:48) − 1 ] ∈ R N (cid:48)
× m .
2: Split the inputs u 1: N into C = N/N (cid:48) chunks u ( c ) 1: N (cid:48) for c
=1 ,...,C .
3: Let the initial state be x (0) N (cid:48)
=0 ∈ R m .
4: for 1 ≤ c ≤ C do 5:
Compute y
( c ) =
M xy x ( c − 1) N (cid:48)
+ B LOCK FFTC ONV ( f , u j )
+ D u ( c ) ∈
R N (cid:48) .
Update state: x ( c ) N (cid:48)
=
A N (cid:48)
x ( c − 1) N (cid:48)
+ M ux u ( c ) .
7: end for 8: Return y =[ y (1) ...
y ( C ) ] .
We prove that Algorithm 2 yields the same output as if one has computed the SSM using a large FFT of size N (proof in Appendix D.4):
Proposition 2.
For input u ∈ R N and matrices A , B , C , D , the output y ∈ R N returned by Algorithmis the same as the output deﬁned by the SSM parameterized by A , B , C , D .
H3 E VALUATION
To understand how well capturing the synthetics in Section 3.1 translates to language modeling, we train two hybrid hybrid H3-attention language models at sizes 125M and 355M and evaluate their performance against Transformers.
The hybrid models match or exceed the quality of Transformers in perplexity and zero/few-shot 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. L ANGUAGE M ODELING
We compare hybrid H3-attention language models against Transformer-based 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
SSMs and attention on the synthetic languages translates to 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 1.6 × faster than Transformers.
Appendix F shows performance of pure H3 language models on these same evaluation metrics.
Setup We train hybrid models at sizes 125M and 355M on the PILE (Gao et al., 2020) for 300B tokens, on a single node with 16 × A100-40GB GPUs.
We compare against checkpoints of equivalent sizes from
Open-AI (Radford et al., 2019) and GPT-Neo 4 (Black et al., 2021), from HuggingFace (Wolf et al., 2020).
Perplexity Table 4 shows perplexity on the Pile (Gao et al., 2020), OpenWebText (Gokaslan et al., 2019), and WikiText-103 (Merity et al., 2016).
On the Pile, our 125M hybrid model outperforms GPT-Neo, which was also trained on the Pile.
Our hybrid models also outperform GPT-Neo models and GPTmodels on zero-shot transfer to OpenWebText and WikiText103. We report the PPL of GPT-2 models for context, though the performance is not directly comparable since they were trained on different data. [Table] [Table]
[Table] 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 [Table] and GPT-2 models, where public checkpoints are available.
We report performance with rank classiﬁcation on the logits of the possible choices (see Appendix F.7 for raw generation).
Table 5 reports zero-shot performance on the SuperGLUE benchmark [Table] 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
[Table]
[Table] Language Modeling Inference Finally, since SSMs are recurrent models, they admit faster text generation than Transformers.
Table 7 shows inference throughput of a 125M-parameter hybrid model compared to a Transformer.
The hybrid model has 1.6 × higher throughput.
5. N ON -T EXT S EQUENCE M ODELING
We show that H3 outperforms Transformers on two non-text sequence modeling tasks: raw speech classiﬁcation and seizure classiﬁcation over raw EEG signals.
H3 sets state-of-the-art performance on seizure classiﬁcation and matches S4 on speech classiﬁcation—which suggests that H3, or one of its hybrids, may be a strong candidate for a multimodal foundation model.
Appendix E gives experimental details
[Table]
and Appendix F gives an additional experiment on brain fMRI data
[Table] [Table]
Seizure Classiﬁcation from EEG We evaluate binary seizure classiﬁcation of 60 -sec EEG clips, sampled at 200 Hz, with 19 electrodes: x ∈ R 12 , 000 × 19 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-art performance.
[Table] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
Raw Speech Classiﬁcation The SC10 speech commands task (Warden, 2018) contains raw audio signals one second in length, sampled at 16kHz.
Similarly to EEG signals, Transformers cannot process the long sequence length.
Table 9 shows that H3 comes within half a point of S4, the state-of-the-art method.
F LASH FFTC ONV E VALUATION
We evaluate how well F LASH FFTC
ONV speeds up SSMs.
F LASH FFTC
ONV 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 of training H3 module with F LASH FFTC ONV 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 a benchmark for long-range sequence modeling.
The state-of-the-art approach, S4 (Gu et al., 2022c), is an SSM.
Tableshows that F LASH FFTC ONV accelerates S4 by 2 × , outperforming Transformers by 5.8 × .
Benchmark H3 Against Attention
We benchmark the time to run the forward and backward pass of H with F LASH FFTC ONV against attention.
F LASH FFTC
ONV maintains nearly linear scaling, even to very long sequence lengths.
Fig.
2 shows overall 2-3 × 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 are the bottleneck for short sequences.
For long sequences,
SSMs using state passing can be dozens of times faster than even the fastest attention implementation.
C ONCLUSION
Our main goal is to understand and narrow the gap between attention and SSMs in language modeling in terms of modeling capabilities and hardware efﬁciency.
Our exploration based on synthetic language tasks motivated us to design the H3 layer, which is surprisingly competitive with attention.
Our
B LOCK FFTC
ONV algorithm exploits matrix multiplication units and the dual recurrent–convolution view of 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 a simple combination of two SSMs, and more sophisticated designs could be more expressive.
Our encouraging results on small/medium language models suggests that scaling SSMs to larger sizes is a promising avenue.
Since simply adding two attention layers to H3 models already outperforms both the pure H3 model and Transformers, we are 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, we include clear explanations of any assumptions and a complete proof of the claims in Appendix D; for any datasets used in the experiments, a complete description of the data processing steps is in
Appendix E. We will also release model checkpoints for all our models.
Ethics Statement.
Our work seeks to 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 efﬁcient—and run inference more efﬁciently.
This potentially connects to energy savings during model development and deployment.
We also note that 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 applications, each with potential beneﬁts and harms.
For example, making language model training cheaper and making inference more efﬁcient make it cheaper to spread disinformation.
Similarly, improving the efﬁciency 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 on the fronts of efﬁciency and understanding, it does not explicitly address the issues of fairness and bias in language models.
A R ELATED W ORK
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.
(2022) also showed that interleaving attention with their GSS layer can improve performance, which we also validate on our OpenWebText 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 (Hoffmann 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 et 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 a direct connection between linear attention and LTI systems. Luo et al.
(2021) also proposes a variant of linear attention that can achieve O ( n log n ) scaling in sequence length. Appendix F evaluates linear attention on language modeling, and ﬁnds that it underperforms exact attention, whereas H3 outperforms attention. The multiplicative interactions in 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 L INEAR A TTENTION AND T IME -V ARYING S YSTEMS
We draw some connections from linear attention to LTI systems and SSMs.
We ﬁrst present linear attention as a linear time-varying system, and show how converting it to a linear time-invariant system matches H3.
our H3 layer.
Linear time-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 Picard-Lindelof theorem, nonlinear systems can be approximated by a series of linear system):
[Equation]
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 denominator in linear attention Section 2 (i.e., assuming d i =1 ).
We see that S
i is just a cumulative sum, satisfying the recurrence S
i +
= S i + φ ( K i +1 )
V Ti +1 .
Similarly, O i satisﬁes the recurrence
O i +
=
φ ( Q i +1 ) T S
i +1 .
This is a linear time-varying system of the form x i +1 =
A x i + B u
i +
and y i +
=
C i +1 x
i +1 (with A
= I ,
B =
I , u i =
φ ( K i )
V Ti , C i =
φ ( Q i ) T ).
That is, A and B are constant, but C is time-variant.
To convert this into a linear time-invariant version, we treat the time-variant C i as a post-processing step.
We instead of a ﬁxed C for the LTI.
This yields an LTI:
[Equation]
for some matrices A , B , C that are learned.
We then apply post-processing by multiply y i +1 with φ ( Q i ) T .
Replacing φ ( K i ) with a shift SSM yields an analogue to H3.
C M ETHOD DETAILS
Since we have described the forward pass in Section 3, we describe here the backward pass in details.
[Equation]
We show how to compute the backward pass in a fused kernel.
Let y =
f ∗ u + D u .
In our case, we have f and u have the same length, so they are symmetric as far as the convolution is concerned.
Suppose we are given dy = ∂l∂y (where l is some loss function).
We wish to compute du , df , and dD (which are ∂l∂u , ∂l∂f , and ∂l∂ D , respectively).
Themostchallengingpartiscomputingthegradientthroughtheconvolutionoperator-butwe’llseethatwe can re-use our FFT infrastructure for it.
The rest of the operations are straightforward; we have d D = dyu T .
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 be able to compute du as well.
Since f and u are the same length L , f ∗ u and u ∗ f have the same result.
Thus, we’ll start from u ∗ f here.
For some notation, let O =
u ∗ f .
Then, dO = dy .
Recall that O
[ i ]= (cid:80) i − 1 j =
u
[ i − j ] f [ j ] .
We’ll start by extending u and f with zeros, to give them length 2 L .
Let u (cid:48) =[
u [0] ,u [1] ,...,u
[ L − 1] , 0 ,..., 0] , and f (cid:48) extended similarly.
Let O (cid:48)
=
u (cid:48) ∗ f (cid:48) , and O = O (cid:48)
[: N ] .
Assume that we have all the values of dO (cid:48)
(we only have them for the ﬁrst half, but we’ll see that it doesn’t matter in the end).
Let’s construct a Toeplitz matrix H u (cid:48) such that u (cid:48) ∗ f (cid:48)
=
H u (cid:48) f (cid:48) :
H u (cid:48)
=
  u (cid:48)
[0] ... u (cid:48)
[1] u (cid:48)
[0] ... u (cid:48)
[2 L − 1] u (cid:48)
[2 L − 2] ... u (cid:48)
[0]  
Since, we have u (cid:48)
[ i ]= f (cid:48)
[ i ]=0 for i ≥ L , we can actually ﬁll in the zeros of the above matrix as well:
H u (cid:48)
=
  u (cid:48)
[0] u (cid:48)
[2 L − 1] ... u (cid:48)
[1] u (cid:48) [1] u (cid:48)
[0] ... u (cid:48)
[2] ... u (cid:48)
[2 L − 1] u (cid:48)
[2 L − 2] ... u (cid:48)
[0]  
Then, we can use the matrix multiplication chain rule to ﬁnd that:
[Equation] dO (cid:48)
=
  u (cid:48)
[0] u (cid:48) [1] ... u (cid:48)
[2 L − 1] u (cid:48)
[2 L − 1] u (cid:48) [0] ... u (cid:48)
[2 L − 2] ... u (cid:48)
[1] u (cid:48)
[2] ... u (cid:48)
[0]   =   u (cid:48)
[0] u (cid:48) [ − (2 L − 1)] ... u (cid:48) [ − 1] u (cid:48)
[ − 1] u (cid:48)
[0] ... u (cid:48)
[ − 2] ... u (cid:48)
[ − (2 L − 1)] u (cid:48)
[ − (2 L − 2)] ... u (cid:48) [0]   , where we use u (cid:48)
[ − i ] to mean u (cid:48)
[2 L − i ] .
Notice that this matrix has the same format as H u (cid:48) !
Let u (cid:48)∗ =[ u (cid:48) [0] ,u (cid:48)
[ − 1] ,...,u (cid:48)
[ − (2 N − 1)]] .
Then: df (cid:48) =(
u (cid:48)∗ ∗ dO (cid:48) ) .
So how do we compute u (cid:48)∗ efﬁciently?
Naively, we might incur some nasty memory access issues.
But a nice property about the DFT saves us!
Let U [ i ] be the i -th element of the DFT of a signal u .
Note that U
[ i ] is complex.
We have:
U ∗ [ i ]= U [ − i ] , where here the ∗ represents the complex conjugate.
We can use this property to compute df (cid:48) efﬁciently: df (cid:48)
[Equation]
u (cid:48)∗ ∗ dO (cid:48) =
iFFT (
[Equation]
( u (cid:48) )
[Equation]
( dO (cid:48) ))
[Equation]
df
[Equation]
df (cid:48)
[: N ]
[Equation]
iFFT (
[Equation]
[Equation]
(cid:48) ) [Equation] ( dy (cid:48) )) [: N ] , where FFT ∗ denotes taking the complex conjugate of the FFT, and dy (cid:48) denotes dy , padded with zeros.
Computing du We can use this same trick to compute du , except we need to add in the contribution from the original D u term.
We end up with:
[Equation]
C. S TATE -P ASSING M ATRICES
We show how to derive M ux for the state update in our state-passing algorithm.
We wish to construct a matrix vM ux ∈ R m × N (cid:48) such that M ux u
= (cid:80) N (cid:48)
i =
1 A N (cid:48) − 1 B u i .
Note that
A i B ∈ R
d × 1 is a column vector.
We can simply stack these column vectors to form a matrix:
M ux =[
A N (cid:48) − 1 B , A N (cid:48)
− 2 B ,..., B ] .
D P ROOFS
We show parameterizations of H3 and attention that solves the associative recall task.
We prove Proposition 1 and Proposition 2.
D. H3 E XPRESSIVITY
This section formally describes a parameterization of H3 that solves the associative recall task
. D.1. E XAMPLE L ANGUAGE Λ
Consider a simple language with 4 keys and 4 values.
For concreteness, we will use the keys { k 1 ,k 2 ,k 3 ,k 4 } =
L K and the values { v 1 ,v 2 ,v 3 ,v 4 } =
L V , i.e. our language L =
L K ∪ L V .
Given a sequence of key-value pairs with one key at the end, we want a model to generate the value associated with the key at the end.
Assume that the key at the end appeared in the sequence.
More formally, let N +1 be the length of the sequence, N even.
The language Λ consists of sequences x ∈ L N +1 .
Each sequence has an associated mapping f x : L K → L V .
For each sequence, the odd indices are randomly sampled from L K , for x 1 ,x 3 ,...,x N − 1 .
The even indices are deﬁned by f x : x 2 ∗ i =
f x ( x 2 ∗ i − 1 ) , for 1 ≤ i ≤ N/ 2 .
The last item in the sequence x N +1 , is randomly drawn from the keys that have appeared in x already, i.e. x N +1 ∈∪{ x 1 ,x 3 ,...,x N − 1 } .
The goal of this language modeling task is to produce f x ( x N +1 ) at the end of the sequence.
D.1. H3 M ODEL TO S OLVE Λ
We describe a toy H3 model that can solve Λ .
Consider a model consisting of an 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 =8 ,m =2 ,H =4 .
Let the embedding layer map each key k
i to the e i basis vector, and map each value v i to the e
+ i basis vector.
Let B shift and C shift denote the parameters of the shift SSM, and A diag , B diag , and C diag denote the parameters of the diagonal SSM (let D be zero for both).
Let B shift =
B diag =
C diag =
e 1 .
Let C shift =
Let A diag be a diagonal matrix with 1 s along its diagonal for each H3.
Remark.
The action of a diagonal SSM parameterized by A diag , B diag , and C diag is to act as a cumulative sum over all its input.
The action of shift SSM parameterized by B shift and C shift is to shift its input by one time step.
Recall that the H3 layer maps its input to Q , K , and V by applying u W Q , u W K , and u W V .
Let W Q and W K be the following:
[Equation]
 1 1 0 0 0 0 00 0 1 1 0 0 00 0 0 0 1 1 00 0 0 0 0 0 10 0 0 0 0 0 00 0 0 0 0 0 00 0 0 0 0 0 00 0 0 0 0 0 0 
Recall that Q and K are split into H heads ( Q ( i ) , K ( i ) for i ∈ { 1 , 2 , 3 , 4 } ), each of which is sent to an independent H3.
Remark.
The action of W Q and W K are to “assign” each key to a different H3 head, i.e., Q ( i ) t is only non-zero when x t = k i .
Similarly, K ( i ) t is only non-zero when x t − 1 = k
i (since K t =
K t − 1 due to the time delay of the shift SSM).
Let W V be the following:
W [Equation] =   0 0 0 0 0 0 00 0 0 0 0 0 00 0 0 0 0 0 00 0 0 0 0 0 00 0 0 0 0 0 00 1 0 1 0 1 01 0 1 0 1 0 11 1 1 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 (1) t =
V (2) t
=
V (3) t =
V (4) t for all i , and V ( i ) t =[0 , 0]
⇔ x t = v 1 , V ( i ) t =[0 , 1] ⇔ x t = v 2 , etc.
We claim that for x N +1 =
k i , O ( i ) N +1 will be a multiple of the binary encoding of f x ( k i ) , and all the other heads of the output O ( j ) N +1 , 1 ≤ j ≤ 4 ,j (cid:54) =
i , will be zero.
Let the output projection W O be such that, with a non-linearity afterwards, it inverts the binary encoding to produce the embedding of the desired output
f x ( k i ) .
We will assume such a projection exists, proof left to the reader.
Proposition 3.
The model described above solves the associative recall problem for the language Λ .
Proof.
Proof sketch.
WLOG, let x N +
= k i .
Then Q ( i ) =
[1 , 1] , but Q ( j ) =
[0 , 0] for j (cid:54) = i .
Thus,
O ( j ) =[0 , 0] for j (cid:54) =
i due to the multiplicative interaction.
Since Q ( i ) =[1 , 1] , O ( i ) is the output of the diag SSMs in the H3 head corresponding to k i (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 a non-zero input, its preceding shift SSM must have a non-zero output.
The only times t this can happen in the sequence are when x t − 1 = k i .
But then x t = f x ( k i ) .
Thus, the input to the diag SSMs are precisely the binary encoding of f x ( k i ) .
Then the output O ( i ) is a multiple of the binary encoding of f x ( k i ) , W O decodes this output to the embedding of f x ( k i ) . D. A TTENTION E XPRESSIVITY
We provide an informal sketch of a two-layer attention model that can solve the associative recall task, inspired by the construction of (Olsson et al., 2022).
The ﬁrst layer of the attention model outputs the embedding of the previous token in 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 is a match—which is exactly the key-value lookup. The construction proceeds as follows: • In the ﬁrst layer, let Q i be mapped to the positional embedding of token x i − 1 (e.g., p i − 1 if p i denotes the positional embedding of token x i ), and K i be mapped to the positional embedding of token x i .
{+• The attention matrix A is computed as QK T , with a causal mask (i.e., A i,j =
+}
if j >i ).
• Then, softmax ( A ) approximates the shift matrix (see Section 3).
• Let V i be an encoding of token x i , constrained to the ﬁrst half of the hidden dimension.
• Then, for output
O = softmax ( QK T ) V , the ﬁrst half of the vector
O i is the encoding of token x i − 1 .
In the second layer, assume that you have a skip connection, that maps the encoding of the input token x i to the second half of the vector
O i .
Then, the input to the second layer encodes both x i − 1 and x i .
In the second layer, let Q i extract the encoding of x i , and let K i extract the encoding of x i − 1 .
• Apply a causal mask on QK T .
Then, the value of softmax ( QK T ) i,j is large if x i =
x
j − 1 , and i>j − 1 .
• Let V i extract the encoding of x i .
{+Then, output O i is the sum of values x j such as x j − 
+}
= x i .
But then O i is exactly a lookup of the token that came after x i when it 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 token based on the dot product and softmax, and for token comparisons through the dot product and softmax.
We prove Proposition 1, which states that the H3 layer takes O ( d 2 N + dN log N ) time and O ( dN ) space for sequence length N and hidden dimension d .
Proof.
We ﬁrst analyze the time complexity.
Consider the matrix multiplies in H3, where the input u ∈ R N × d is multiplied by three weight matrices of size d × d .
These take time O
( d 2 N ) .
The output O is also multiplied with an output projection weight matrix of size d × d , also taking time O
( d 2 N ) .
Therefore the matrix multiplies take time O
( d 2 N ) .
Now consider the two SSMs in H3.
The ﬁrst SSM involves a convolution of K ∈ R N × d (in the
N -dimension) with a kernel of size N × d .
This reduces to an FFT, a pointwise multiply, and an inverse
FFT (in the N -dimension).
This takes time O ( dN log N ) .
The second SSM involves H convolutions, inputs of size N × d h × d h , along the N -dimension.
This takes time:
[Equation]
where we use the fact that d h = d/H and that d h
= O (1) .
Therefore the two SSMs take total time
O ( dN log N ) .
As a result, the H3 layer takes time:
O ( d 2 N + dN log N ) .
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 O ( dN ) space and O ( Hd 2 h N )
=
O ( dd h N )=
O ( dN ) space.
Therefore the overall space complexity is O ( dN ) .
D. S TATE P ASSING C ORRECTNESS
We prove Proposition 2.
We assume that the B LOCK FFTC ONV algorithm is correct, i.e., the output y
=
B LOCK FFTC ONV ( f,u ) is equal to the output of an SSM with convolution kernel f and input u .
Proof.
Proof by induction on C .
Base case: C =1 .
WTS y =[
[Equation]
1) ] , M
[Equation]
Inthiscase, notethat N = N (cid:48) .
Then y (1) =
M xy x (0) N (cid:48)
+ B LOCK FFTC ONV ( f,u 1 )=
B LOCK FFTC ONV ( f,u 1 ) .
But u = u 1 , so y
= y (1) =[
y (1) ] .
Additionally, by the recursive deﬁnition of a state space,
[Equation]
Inductive step: C > 1 .
Assume that
[ y (1) ,...,y ( C − 1) ]=
y [: N (cid:48)
( C − 1)] , and x ( C − 1) N (cid:48)
=
x ( C − 1) N (cid:48) .
WTS that y ( C ) =
y [ N (cid:48)
( C − 1): N (cid:48) C ] , and M xx x
( C − 1) N (cid:48)
+ M ux u ( C ) =
x N .
Let t denote N (cid:48)
( C − 1) .
For i> ( C − 1) N (cid:48) , we have:
[Equation]
Thus, y ( C ) =
y [ N (cid:48)
( C − 1): N (cid:48) C ] .
Similarly,
[Equation]
Our synthetic tasks, inspired by (Olsson et al., 2022), are designed to mimic the in-context learning capability of large language models—the ability to learn from examples in the input 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 (cid:96) 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 0.0.
E. M ODEL A RCHITECTURE
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.
For the hybrid models, we use 16 attention heads for the attention layers.
The 125M hybrid model has an attention layer at layers 1 and 7, and the 355M hybrid model has attention layers at layers 1 and 13.
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-precision training, with bf16 for the MLP’s and attention.
When training language models, we use fp32 for the FFTConv.
E. O PEN W EB T EXT T RAINING
For the 125M models trained on OpenWebText, we follow the training recipe of the Megatron-LM repo.
We use an effective batch size of 512, and use gradient accumulation to ﬁt into available GPU memory.
We use the AdamW optimizer, with learning rate 6e-4 for GPT-2 small and 1.5e-4 for GPT-2 medium, and weight decay of 0.1.
All models are trained with the same hyperparameters for 100K steps.
We run all implementations with mixed-precision training (PyTorch AMP).
We train models with sequence length 1024.
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 of validation set is done once, and all models are evaluated on the same validation set.
E. T HE PILE T RAINING
For the 125M and 355M models trained on the PILE, we follow the training recipe of GPT-3.
We use batch size 256, with sequence length 2048.
We train our models for 600K steps.
We use residual dropout 0.0 and embedding dropout 0.1.
We use the AdamW optimizer, with learning rate 6e-4 for the 125M model and 3e-4 for the 355M model, and a weight decay of 0.1.
We use a cosine schedule withsteps for linear warmup, and decay the learning rate to 10% by the end of training.
We suspect that there exist better hyperparameters for H3 language models, but we did not have the resources to tune them.
For the PILE dataset, we again use 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 of validation set is done once, and all models are evaluated on the same validation set.
E. S UPER GLUE
We follow the prompts used in the GPT-3 paper (Brown et al., 2020).
For rank classiﬁcation on the binary classiﬁcation tasks, we use yes/no for WSC, WIC, MultiRC, and BoolQ, and we use true/false for RTE.
For CB, we use true/false/neither as the three choices.
For COPA and ReCoRD, we use the continuations provided by the task.
E. S EIZURE C LASSIFICATION 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 a range of psychiatric and psycho-social disorders and impact the lives of roughly one percent of the global population (Kerr, 2012).
The ﬁrst step to treating epilepsy is manual analysis of scalp EEG by board-certiﬁed neurologists.
However, the vast amount of EEG data produced by each patient (which can be up to days of data) makes manual EEG analysis a costly and time-consuming process.
To mitigate the costs associated with EEG monitoring, recent deep learning techniques have began to show promise in ﬂagging abnormal EEG segments for potential seizure 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 classiﬁcation performance Saab et al.
(2020), with the increased difﬁculty of training deep learning models on long sequences (e.g., an EEG signal sampled at 200 Hz produces 12 , 000 time steps per minute).
As a result, many techniques involve domain-specialized models and pre-processing steps, such as FFT transforms and graphical representations Tang et al.
We use the largest publicly available EEG corpus, TUSZ v1.5.2 (Shah et al., 2018), which includes 5 ,
EEG signals from 636 patients, with 3 , 050 annotated seizures.
Signals are segmented into 60-second clips, and split into train/val/test by patient.
The train set contains 39765 clips, the val set containsclips, and the test set contains 10001 clips.
F A DDITIONAL E XPERIMENTS F. LRA A CCURACY
We evaluate the accuracy of H3 on LRA.
We compare accuracy to S4D (Gu et al., 2022b), since H3 uses an S4D kernel as a 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] Table 11 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. W IKI T EXT
We train 125M-sized models on WikiText103 (Merity et al., 2016) and compare their test PPL to transformers, as well as other variatns of efﬁcient or long-range attention.
We use the same hyperparameters and setup as training on OpenWebText.
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. [Table] Table 12 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 signiﬁcantly outperforms transformers with performer, reformer
[Table]
and linear attention
[Table]
F. PG-
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 OpenWebText.
[Table]
Table 13 shows that Hybrid H3 outperforms transformers and linear attention.
F. L ENGTH E XTRAPOLATION
One property of SSMs is that they can naturally extrapolate to sequence lengths longer than those seen during training.
We use the synthetic associative recall task to demonstrate that H3 maintains this capability.
To do so, we train a two-layer H3 model on sequences of length 20 drawn from the associative recall synthetic language.
Then, we evaluate accuracy of the last token prediction on sequences of length 20 and 40.
[Table]
Table 14 shows that H3 maintains accuracy on sequences of length 40, which is twice the length of the training sequences.
F. S CALING IN N UMBER OF T OKENS
We evaluate how well a Hybrid H3 model scales with the number of tokens seen during training, compared to a Transformer.
For these experiments, we train a 125M Hybrid H3 model and a 125M Transformer on the PILE for 5B, 10B, 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 to 6e-5 by the end of training.
[Table]
Table 15 shows the results.
Both the Hybrid H3 model and Transformer model improve as the number of training tokens increases.
F. H3 L ANGUAGE M ODEL
We report the results of a pure H3 language model on NLP evaluations.
We train a 125M model on the PILE for 300B tokens.
Tables 16 and 17 show zero-shot and few-shot performance on SuperGLUE, respectively.
F. G ENERATION P ERFORMANCE
We report results on SuperGLUE for generation.
Instead of taking rank classiﬁcation, we instead let the model generate a response, and we search for the gold label (i.e., “yes” or “no” for the yes/no questions)
[Table]
[Table] [Table]
[Table]
in the output.
Tables 18 and 19 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 a parsable format.
F. F UNCTIONAL M AGNETIC R ESONANCE I MAGING D ATA
Functional Magnetic Resonance Imaging (fMRI) data are typically represented in four dimensions, indicating the measured blood-oxygen-level-dependent (BOLD) signal in temporal sequences
S =
{ V 1 ,...,V t } of 3-dimensional volumes V ∈ R x × y × z , each indicating the BOLD signal for all spatial locations of the brain (as deﬁned by three spatial dimensions x , 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 overﬁtting.
In spite of the low sample size of individual datasets, neuroimaging research can be considered as recently entering a big data era because researchers more frequently share their collected datasets publicly (Markiewicz et 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.
, enabling models to utilize the knowledge that can be learned from public functional neuroimaging data for the analysis of individual datasets.
Speciﬁcally, Thomas et al.
(2022) evaluate the performance of several self-supervised learning frameworks for functional neuroimaging data by ﬁrst pre-training models on a 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 a prose story or math problem) from measured brain activity.
The authors ﬁnd that a GPT-based model, pre-trained in a 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 a 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 a set Θ ∈ θ 1 ,...,θ n of n =1 , 024 functionally-independent brain networks θ (as deﬁned by the dictionaries of functional modes (DiFuMo) brain atlas (Dadi et al., 2020)), each describing the BOLD signal for some subset of voxels v x,y,z ∈ V , such that the resulting sequences X ∈ R t × n describe the activity pattern of each brain network n for time points t .
In line with Thomas et al.
(2022), we ﬁrst pre-train models f ( · ) to predict the distribution of brain activity for the next time point j of an input sequence X , using a mean absolute error ( L rec ) training objective given the model’s output ˆ X ∈ R t × n : L rec = 1 n (cid:80) ni
=1 | X j,i − ˆ X j,i | ; ˆ X t,n = b n + (cid:80) n f ( E X ) t,e w e,n ;
E Xt,e = E TR + E pos + b e + (cid:80) n X t,n w n,e .
Here, E TR ∈ R e and E pos ∈ R e represent learnable embeddings for each possible time point and position of an input sequence (for details, see Thomas et al.
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 lower-dimensional embedding representation
E X ∈ R t × e (with e =768 dimensions).
We evaluate two model architectures for f ( · ) , namely, the GPT variant used in Thomas et al.
(2022), with 4 hidden layers and 12 attention heads, and a corresponding H3 variant with 4 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 et 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 a 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 β 1 = 0 .
9 , β 2 = 0 .
999 , and (cid:15) =
1 e − 8 ) for 5 , 000 steps at a mini-batch size of 512 , and a learning rate of 5 e − 4 .
We apply a linear learning rate decay schedule (with a warm-up phase of 1% of the total number of training steps), gradient norm clipping at 1 . 0 , and
L 2 -regularisation (weighted by 0 . 1 ).
We also apply dropout at a rate of 0 .
1 for the GPT-based model (based on Thomas et al.
and evaluate three dropout rates for H3: 0 . 1 , 0 . 2 , and 0 . 3 .
We ﬁnd that the H3 variant trained with 0 .
2 dropout performs on par with the GPT model, in terms of mean absolute error (Fig. 3), and therefore continue all further analyses with this model variant.
We also ﬁnd that both models exhibit almost identify L rec 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).
To adapt the pre-trained models to mental state decoding, we add a learnable classiﬁcation embedding
E cls ∈ R n to the end of input sequences X and forward the model’s prediction f
( E X ) to a decoding head p ( · ) , composed of a dense hidden layer with e model units (one for each embedding dimension, with tanh activation) as well as a softmax output layer with one model unit i for each considered mental state in the data.
Accordingly, we adapt models by optimizing a standard cross entropy loss objective:
L cls = − (cid:80) i
y i log p ( f ( E X ))
i , where y i indicates a binary variable that is 1 if i is the correct mental state and 0 otherwise.
During downstream adaptation, we 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) or 10 (MDTB) distinct individuals.
We adapt models for 750 steps at a mini-batch size of 256 and a learning rate of 5 e −(otherwise using the same learning parameters as for upstream training).
Importantly, we 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-deterministic training factors (such as random initialization and data shufﬂing).
As for the upstream data
[Table]
the H3 and GPT-based models generally perform on par in their mental state decoding performances in the two left-out test datasets (Table 20)
[Table] [Table]