Hyena Hierarchy: Towards Larger Convolutional Language Models β€” part 1/5

Cleaned-text code points [0, 15472). Bibliography omitted here; source.md includes it.

\LettrineTextFont

# Hyena Hierarchy: Towards Larger Convolutional Language Models

###### Abstract

Recent advances in deep learning have relied heavily on the use of large Transformers due to their ability to learn at scale. However, the core building block of Transformers, the attention operator, exhibits quadratic cost in sequence length, limiting the amount of context accessible. Existing subquadratic methods based on low-rank and sparse approximations need to be combined with dense attention layers to match Transformers, indicating a gap in capability. In this work, we propose Hyena , a subquadratic drop-in replacement for attention constructed by interleaving implicitly parametrized long convolutions and data-controlled gating . In recall and reasoning tasks on sequences of thousands to hundreds of thousands of tokens, Hyena improves accuracy by more than 50 50 points over operators relying on state-spaces and other implicit and explicit methods, matching attention-based models. We set a new state-of-the-art for dense-attention-free architectures on language modeling in standard datasets ( WikiText103 and The Pile ), reaching Transformer quality with a 20 % 20\% reduction in training compute required at sequence length 2 2 K. Hyena operators are twice as fast as highly optimized attention at sequence length 8 8 K, and 100 Γ— 100\times faster at sequence length 64 64 K.

## 1 Introduction

Large Transformers have enabled a number of breakthrough advances in modeling language, vision, audio, biology and numerous other domains ( Vaswani et al., 2017 ) , ( Dosovitskiy et al., 2020 ) , ( Radford et al., 2022 ) , ( Cramer, 2021 ) . Much of the success of Transformers, powered by the attention operator ( Vaswani et al., 2017 ) , relies on their scaling properties ( Hoffmann et al., 2022 ) and the emergence of in-context learning ( Garg et al., 2022 ) , which allows them to generalize to unseen data and tasks given context as input. The Transformer block is a powerful tool for sequence modeling, but it is not without its limitations. One of the most notable is the computational cost, which grows rapidly as the length of the input sequence increases. Specifically, the cost scales quadratically with the length L L of the sequence, which places a strict limit on the amount of context that can be considered by the model. Breaking the quadratic barrier is a key step towards new possibilities for deep learning, such as using entire textbooks as context, generating long-form music or processing gigapixel scale images.

Efforts to reduce the computational cost of attention in models primarily involve the use of linearized, low-rank, and sparse approximations ( Child et al., 2019 ; Wang et al., 2020 ; Kitaev et al., 2020 ; Zhai et al., 2021 ; Roy et al., 2021 ; Schlag et al., 2021 ; Tu et al., 2022 ) . These approaches introduce a trade-off between expressivity and speed, requiring hybridization with standard attention layers to reach Transformer quality ( Mehta et al., 2022 ; Dao et al., 2022c ) .

A growing amount of evidence suggests that attention mechanisms only utilize a small portion of their quadratic capabilities for language processing ( Olsson et al., 2022 ; Dao et al., 2022c ) , leading us to question its role as the gold-standard operator for deep learning at scale. Specifically, we ask:

Are there subquadratic operators that can match the quality of attention at scale?

We obtain a positive answer based on a composition of efficient subquadratic primitives, such as element-wise multiplication (gating) and long convolutions i.e., convolutions with filter sizes as long as the input. We rely on a set of targeted reasoning tasks, grounded in recent work on mechanistic interpretability ( Elhage et al., 2021 ; Power et al., 2022 ; Olsson et al., 2022 ; Zhang et al., 2022 ) such as recall and induction, to distill three properties of attention correlated with its performance and the quality gap with existing subquadratic approaches: a . a. Data control: Attention implements an expressive data-controlled ( Massaroli et al., 2020 ) linear operator 1 1 1 Self-attention can be expressed as y = 𝖠 ⁑ ( k , q ) ​ v y=\mathsf{A}(k,q)v where 𝖠 \mathsf{A} is the attention matrix conditioned by linear projections k , q k,q of the input and multiplied by v v , another projection. , encoding an entire family of linear functions in a single block.

b . b. Sublinear parameter scaling: Parameter counts of attention layers are decoupled from sequence length, allowing Transformers to allocate more parameters elsewhere e.g., the feed-forward neural networks ( π–₯π–₯𝖭 \sf FFN s) between attention layers.

c . c. Unrestricted context: For a given input, attention has an unrestricted context i.e., it can approximate dependencies between any two inputs, without arbitrary restrictions such as locality (except in cases using masking such as autoregressive models).

#### The 𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} hierarchy

Guided by these findings, we introduce the 𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} hierarchy, an operator defined by a recurrence of two efficient subquadratic primitives: a long convolution and element-wise multiplicative gating (see Figure 1.1 ). A specified depth (i.e., number of steps) of the recurrence controls the size of the operator. For short recurrences, existing models are recovered as special cases ( Mehta et al., 2022 ; Dao et al., 2022c ) . By mapping each step in the 𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} recurrence to its corresponding matrix form, we reveal 𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} operators to be equivalently defined as a decomposition of a data-controlled matrix i.e., a matrix whose entries are functions of the input. Furthermore, we show how 𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} operators can be evaluated efficiently without materializing the full matrix, by leveraging fast convolution algorithms ( Selesnick and Burrus, 2017 ) . Empirically, 𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} operators are able to significantly shrink the quality gap with attention at scale, reaching similar perplexity and downstream performance with a smaller computational budget (Section 4.2 ) and without hybridization of attention.

#### Narrowing the capabilities gap

The design of Hyena is motivated by a quality gap between standard dense attention and alternative subquadratic operators, which we identify by focusing on reasoning tasks correlated with language modeling performance at scale. We extend the suite of basic mechanistic interpretability benchmarks ( induction and recall ) with additional tasks that probe how quickly model performance degrades when task complexity increases (e.g. vocabulary size grows). In addition, we investigate the optimal parameterization of long convolutions in 𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} . In the most challenging settings with hundreds of thousands of tokens, our implicit parameterization scheme improves over other operators leveraging state spaces ( Gu et al., 2021 ) , frequency-domain parametrizations ( Li et al., 2020 ) , or standard convolutions by over 50 % 50\% accuracy.

#### Scaling in language and vision

Next, we aim to verify whether rankings in our reasoning benchmark suite are predictive of quality at scale. We test 𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} on autoregressive language modeling at the sub-billion parameter scale, setting a new state-of-the-art for dense-attention-free architectures in standard datasets ( WikiText103 and The Pile ) and matching Transformer quality. On the The Pile at the 335 335 M parameter scale, we match Transformer perplexity with a 20 % 20\% reduction in the total count of floating point operations (FLOPs). As an extension, we investigate the generality of 𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} operators by testing on large-scale image recognition, replacing attention in the Vision Transformer (ViT) ( Dosovitskiy et al., 2020 ) . In image classification, 𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} is able to match attention in accuracy when training on ImageNet-1k from scratch.

#### Toward much longer context

Finally, we benchmark the efficiency of 𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} on long sequences. We measure 5 5 x speedups over dense self-attention at length 8192 8192 – 2 2 x over highly optimized FlashAttention 2 2 2 FlashAttention is already 2-4x faster than a standard attention implementation in PyTorch. ( Dao et al., 2022b ) – and 100 100 x speedup over FlashAttention at sequence lengths of 64 64 k, where standard attention implementation in PyTorch runs out of memory.

## 2 Preliminaries and Related Work

A discrete convolution is a function of two arguments: an input u u signal of length L L and a learnable filter h h . The linear (aperiodic) convolution of a (possibly infinitely long) measurable 3 3 3 In the L 1 ​ ( β„€ ) L^{1}(\mathbb{Z}) sense: βˆ‘ t = βˆ’ ∞ ∞ | h t | < ∞ \sum_{t=-\infty}^{\infty}|h_{t}|<\infty filter h h with a length- L L input signal u u is defined as y t = ( h βˆ— u ) t = βˆ‘ n = 0 L βˆ’ 1 h t βˆ’ n ​ u n . \displaystyle y_{t}=(h*u)_{t}=\sum_{n=0}^{L-1}h_{t-n}u_{n}. (1)

Generally, u t ∈ ℝ D u_{t}\in\mathbb{R}^{D} where D D is the width of the signal, or in deep learning parlance, the number of channels . Without loss of generality, we specialize our analysis to single input single output (SISO) layers, i.e. with D = 1 D=1 . The multiple input multiple output (MIMO) case, canonical in standard convolutional layers, follows directly.

In this case, the input signal can be represented as a vector u ∈ ℝ L u\in\mathbb{R}^{L} and the convolution as a matrix-vector product between the input and the Toeplitz kernel matrix 𝖲 h ∈ ℝ L Γ— L \mathsf{S}_{h}\in\mathbb{R}^{L\times L} induced by the filter h h : ( h βˆ— u ) = [ h 0 h βˆ’ 1 β‹― h βˆ’ L + 1 h 1 h 0 β‹― h βˆ’ L + 2 β‹± h L βˆ’ 1 h L βˆ’ 2 β‹― h 0 ] ​ [ u 0 u 1 u L βˆ’ 1 ] \displaystyle(h*u)=\begin{bmatrix}h_{0}&h_{-1}&\cdots&h_{-L+1}\\ h_{1}&h_{0}&\cdots&h_{-L+2}\\ \vdots&\vdots&\ddots&\vdots\\ h_{L-1}&h_{L-2}&\cdots&h_{0}\end{bmatrix}\begin{bmatrix}u_{0}\\ u_{1}\\ \vdots\\ u_{L-1}\end{bmatrix} (2)

### 2.1 Explicit and Implicit Convolutions

Parametrizing and optimizing convolution filters h t h_{t} is a standard procedure in deep learning and more broadly signal processing. The classical approach of convolutional neural networks (CNNs) ( Fukushima and Miyake, 1982 ; LeCun et al., 1998 ; Ronneberger et al., 2015 ; He et al., 2016 ) is to optimize directly the values h t h_{t} of the filter’s response at M M prescribed steps, a parametrization we call explicit . M M is referred to as the filter size and is typically much shorter than the input sequence length M β‰ͺ L M\ll L . Such filters are denoted in signal processing as finite impulse response (FIR).

FIR filters are local and can capture dependencies between inputs separated at most by M M steps. Their main advantage is their speed, with complexity π’ͺ ⁑ ( M ​ L ) \mathcal{O}(ML) . However, the number of parameters of FIR filters scales linearly with filter size, which can be computationally prohibitive. To disentangle the parameter count from the filter size, we can instead represent the filter h t h_{t} as a parametric function of the time step t t , i.e. h t = Ξ³ ΞΈ ​ ( t ) h_{t}=\gamma_{\theta}(t) , where ΞΈ \theta are the parameters of the function Ξ³ ΞΈ \gamma_{\theta} . This parametrization is called implicit . The class of functions Ξ³ ΞΈ \gamma_{\theta} is a design choice with a significant impact on the expressivity and computational complexity of the layer.

One choice of implicit parametrization is to select h h as the response function of a linear state-space model (SSM) ( Chen, 1984 ) , described by the first-order difference equation: x t + 1 \displaystyle x_{t+1} = 𝖠 ​ x t + 𝖑 ​ u t \displaystyle=\mathsf{A}x_{t}+\mathsf{B}u_{t} state equation \displaystyle\text{state equation} y t \displaystyle y_{t} = 𝖒 ​ x t + 𝖣 ​ u t \displaystyle=\mathsf{C}x_{t}+\mathsf{D}u_{t} output equation \displaystyle\text{output equation} Here, the convenient choice of x 0 = 0 x_{0}=0 renders the input-output map to a simple convolution y t \displaystyle y_{t} = βˆ‘ n = 0 t ( 𝖒𝖠 t βˆ’ n ​ 𝖑 + 𝖣 ​ Ξ΄ t βˆ’ n ) ​ u n \displaystyle=\sum_{n=0}^{t}\left(\mathsf{C}\mathsf{A}^{t-n}\mathsf{B}+\mathsf{D}\delta_{t-n}\right)u_{n} where Ξ΄ t \delta_{t} denotes the Kronecker delta. We can then identify the filter h h as t ↦ h t = { 0 t < 0 𝖒𝖠 t ​ 𝖑 + 𝖣 ​ Ξ΄ t t β‰₯ 0 t\mapsto h_{t}=\begin{cases}0&t<0\\ \mathsf{C}\mathsf{A}^{t}\mathsf{B}+\mathsf{D}\delta_{t}&t\geq 0\end{cases} where the entries of 𝖠 , 𝖑 , 𝖒 \mathsf{A},\mathsf{B},\mathsf{C} and 𝖣 \mathsf{D} are the learned parameters of the filter. In terms of layer design, the degrees of freedom of SSMs are the dimension of the state and the structure of the matrices. SSMs are a canonical example of how long convolutions with sub-linear parameter counts can improve deep learning models for long sequences ( Gu et al., 2020 ; Gu et al., 2021 ) . Other implicit approaches include parametrizing filters as maps from (a positional encoding of) t t to the filter response i.e. Ξ³ ΞΈ : t ↦ h t = Ξ³ ΞΈ ​ ( t ) \gamma_{\theta}:t\mapsto h_{t}=\gamma_{\theta}(t) , for example with feed-forward neural networks ( Romero et al., 2021b ; Romero et al., 2021a ) .

#### Fast Methods for Convolutions

One of the first applications of the Cooley-Tukey fast Fourier transform (FFT) algorithm was to implement convolution faster than the direct evaluation of ( 1 ). At first glance ( 1 ) comes with O ⁑ ( L 2 ) O(L^{2}) an asymptotic time complexity. A common approach to achieve fast long convolutions in subquadratic time is through the FFT algorithm. The method first converts the aperiodic convolution into a circular convolution Selesnick and Burrus (2017) by appropriate zero-padding of input and filter sequences. The resulting kernel 𝖲 ^ h \hat{\mathsf{S}}_{h} is a circulant matrix and is diagonalized by the discrete Fourier basis 𝖲 ^ h = 𝖢 βˆ’ 1 ​ 𝖣 H ​ 𝖢 \hat{\mathsf{S}}_{h}=\mathsf{W}^{-1}\mathsf{D}_{H}\mathsf{W} where 𝖢 \mathsf{W} is the DFT matrix, 𝖢 t ​ t β€² = z βˆ’ t , z = e i ​ 2 ​ Ο€ ​ t β€² / L \mathsf{W}_{tt^{\prime}}=z^{-t},z=e^{i2\pi t^{\prime}/L} and H H is the DFT of the padded filter h h , H = 𝖢𝗉𝖺𝖽 ⁑ ( h ) H=\mathsf{W}{\sf pad}(h) . Thus, the calculation of such convolutions is performed as 𝗉𝖺𝖽 ⁑ ( y ) \displaystyle{\sf pad}(y) = 𝖲 ^ h ​ 𝗉𝖺𝖽 ​ ( u ) \displaystyle=\hat{\mathsf{S}}_{h}{\sf pad}(u) = 𝖢 βˆ’ 1 ​ 𝖣 H ​ 𝖢 ​ 𝗉𝖺𝖽 ​ ( u ) \displaystyle=\mathsf{W}^{-1}\mathsf{D}_{H}\mathsf{W}~{\sf pad}(u) = 𝗂π–₯π–₯𝖳 ⁑ ( 𝖣 H ​ π–₯π–₯𝖳 ​ ( 𝗉𝖺𝖽 ⁑ ( u ) ) ) \displaystyle={\sf iFFT}(\mathsf{D}_{H}{\sf FFT}({\sf pad}(u))) where 𝖣 H \mathsf{D}_{H} is the matrix with 𝖢 ​ h \mathsf{W}h on its diagonal. The above is known as the convolution theorem of DFT ( Oppenheim et al., 1997 ) . In this π–₯π–₯π–³π–’π—ˆπ—‡π— {\sf FFTConv} form the convolution can be performed without materializing the operator 𝖲 h \mathsf{S}_{h} with the same asymptotic cost O ⁑ ( L ​ log 2 ​ L ) O(L\log_{2}L) of FFT.

### 2.2 The Self-Attention Operator

At the heart of Transformers is the multi-head attention (MHA) mechanism. Given a length- L L sequence u ∈ ℝ L Γ— D u\in\mathbb{R}^{L\times D} , each head of scaled self-attention ( Vaswani et al., 2017 ) is a map from ℝ L Γ— D \mathbb{R}^{L\times D} to ℝ L Γ— D \mathbb{R}^{L\times D} which performs the following operations 𝖠 ⁑ ( u ) \displaystyle\mathsf{A}(u) = π–²π—ˆπ–Ώπ—π–¬π–Ίπ—‘ ⁑ ( 1 D ​ u ​ 𝖬 q ​ 𝖬 k ⊀ ​ u ⊀ ) \displaystyle={\sf SoftMax}\left(\tfrac{1}{\sqrt{D}}u\mathsf{M}_{q}\mathsf{M}^{\top}_{k}u^{\top}\right) (3) y \displaystyle y = π–²π–Ύπ—…π–Ώπ– π—π—π–Ύπ—‡π—π—‚π—ˆπ—‡ ⁑ ( u ) \displaystyle={\sf SelfAttention}(u) = 𝖠 ⁑ ( u ) ​ u ​ 𝖬 v , \displaystyle=\mathsf{A}(u)u\mathsf{M}_{v},