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

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

where 𝖬 q , 𝖬 k , 𝖬 v ∈ ℝ D Γ— D \mathsf{M}_{q},\mathsf{M}_{k},\mathsf{M}_{v}\in\mathbb{R}^{D\times D} are learnable linear projections and π–²π—ˆπ–Ώπ—π–¬π–Ίπ—‘ {\sf SoftMax} is intended to be applied row-wise. Attention parametrizes a family of dense linear operators and for an input u u , indexes through it via projections of u u i.e., 𝖠 ⁑ ( u ) \mathsf{A}(u) . We refer to operators of this type as data-controlled , as they encode a linear transformation u ↦ y u\mapsto y , that is, however, nonlinearly defined by u u . This approach yields expressive nonlinear operators in u u , and we hypothesize contributes, together with other mechanisms ( Olsson et al., 2022 ) , to the ability of certain operators to learn in-context i.e., to adapt to unseen tasks by leveraging context. In deep learning, the projections take on specific names: query q = u ​ 𝖬 q q=u\mathsf{M}_{q} , key k = u ​ 𝖬 k k=u\mathsf{M}_{k} and value v = u ​ 𝖬 v v=u\mathsf{M}_{v} . We often rewrite the attention operator as y = 𝖠 ⁑ ( q , k ) ​ v y=\mathsf{A}(q,k)v .

###### Remark 2.1 .

Similarly to implicit convolutions, π–²π–Ύπ—…π–Ώπ– π—π—π–Ύπ—‡π—π—‚π—ˆπ—‡ \sf SelfAttention does not entangle its ability to access distant information with the number of parameters: it looks at the whole sequence at the price of π’ͺ ⁑ ( L 2 ) \mathcal{O}(L^{2}) operations.

#### Subquadratic Operators

Existing approaches to subquadratic alternatives to attention can be summarized by altering the way the data control is implemented i.e., how the operator is nonlinearly defined by u u , and then applied to v v . For example, a layer of Attention-Free Transformers (AFTs) ( Zhai et al., 2021 ) constructs the operator through a combination of gating and π–²π—ˆπ–Ώπ—π–¬π–Ίπ—‘ \sf SoftMax (AFT full) or gating and a single explicit convolution (AFT conv). Gated State Spaces (GSS) instead compose the operator via gating and a long convolution parametrized via SSMs. Taking this idea further, Hungry Hungry Hippo (H3) ( Dao et al., 2022c ) , motivated by gaps of GSS on associative recall, extend the mechanism to include an additional gate and a short convolution obtained via a shift SSM. 𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} generalizes this body of work by introducing a recurrence of gates and implicit long convolutions, evaluated efficiently.

## 3 Hyena: Definition and Properties

In this section, we define Hyena, a class of data-controlled operators consisting of a recurrence of multiplicative gating interactions and long convolutions. Instead of seeking an approximation to attention, we guide our design by intentionally incorporating key computational properties of attention, including the decoupling of sequence length and parameter counts.

### 3.1 𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} Recurrences

At a high level, 𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} consists of the following steps (setting D = 1 D=1 for clarity): i . i. Compute a set of N + 1 N+1 linear projections of the input, similarly to attention. The number of projections ( v t , x t 1 , … , x t N ) (v_{t},x^{1}_{t},\dots,x^{N}_{t}) need not be three. One projection takes the role of value, such that a linear input-output function can be defined as y = 𝖧 ⁑ ( u ) ​ v y=\mathsf{H}(u)v for some 𝖧 ⁑ ( u ) \mathsf{H}(u) .

i ​ i . ii. The matrix 𝖧 ⁑ ( u ) \mathsf{H}(u) is defined by interleaving implicit long convolutions and element-wise multiplication with one projection x i x^{i} at a time, until all projections are exhausted. Evaluation of 𝖧 ⁑ ( u ) ​ v \mathsf{H}(u)v is done efficiently without materializing 𝖧 ⁑ ( u ) \mathsf{H}(u) . By doing so, we implicitly define a data-controlled operator as a factorization of a matrix. The long convolutions forming 𝖧 ⁑ ( u ) \mathsf{H}(u) are parametrized implicitly to retain sublinear parameter scaling in sequence length.

Next, we formally define 𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} , starting with its computational model. We leave the analysis of its data-controlled matrix form for the latter part of the section.

###### Remark 3.1 .

The time complexity of a 𝖧𝗒𝖾𝗇𝖺 \sf Hyena recurrence is π’ͺ ⁑ ( N ​ L ​ log 2 ⁑ L ) \mathcal{O}(NL\log_{2}L) . The input-output map can be rewritten as y = x N β‹… ( h N βˆ— ( x N βˆ’ 1 β‹… ( h N βˆ’ 1 βˆ— ( β‹― ) ) ) ) y=x^{N}\cdot(h^{N}*(x^{N-1}\cdot(h^{N-1}*(\cdots)))) where each convolution is performed through the Fourier domain in π’ͺ ⁑ ( L ​ log 2 ​ L ) \mathcal{O}(L\log_{2}L) .

Interestingly, the element-wise product in time domain corresponds to convolution in frequency domain, i.e. x t ​ u t = ( x ^ βˆ— u ^ ) t , x_{t}u_{t}=(\hat{x}*\hat{u})_{t}, where x ^ , u ^ \hat{x},\hat{u} denote the DFT of x x and u u , respectively. Thus, 𝖧𝗒𝖾𝗇𝖺 \sf Hyena is alternatively applying convolutions in the time and then the frequency domain (or alternatively applying element-wise products in the time and frequency domain). One potential explanation for the effectiveness of this procedure is that the convolution in the time domain (element-wise multiplication in the frequency domain) increases the memory length, allowing for a broader context to be taken into account. On the other hand, the element-wise multiplication in the time domain (convolution in the frequency domain) allows for more fine-grained selection of specific frequency components of the signal.

### 3.2 𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} Matrices

𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} operators build on the 𝖧πŸ₯ {\sf H3} mechanism developed by ( Dao et al., 2022c ) . For clarity of exposition, we once again consider the SISO case ( D = 1 D=1 ). Let 𝖣 q \mathsf{D}_{q} and 𝖣 k \mathsf{D}_{k} be the L L -by- L L diagonal matrices whose respective main diagonal entries are the respective entries of q q and k k . 𝖧πŸ₯ {\sf H3} realizes a surrogate attention matrix with a data-controlled, parametrized decomposition in four terms: 𝖠 ⁑ ( q , k ) \displaystyle\mathsf{A}(q,k) = 𝖣 q ​ 𝖲 ψ ​ 𝖣 k ​ 𝖲 Ο† \displaystyle=\mathsf{D}_{q}\mathsf{S}_{\psi}\mathsf{D}_{k}\mathsf{S}_{\varphi} (5) 𝖧πŸ₯ ⁑ ( q , k , v ) \displaystyle{\sf H3}(q,k,v) = 𝖠 ⁑ ( q , k ) ​ v \displaystyle=\mathsf{A}(q,k)v where 𝖲 Ο† , 𝖲 ψ \mathsf{S}_{\varphi},\mathsf{S}_{\psi} are the Toeplitz matrices of learnable causal filters Ο† , ψ \varphi,\psi parametrized via SSMs 5 5 5 For consistency with our discussion, we have swapped k k and v v compared to the notation in ( Dao et al., 2022c ) . . Alongside the q ​ k ​ v qkv -projections the filters constitute our degrees of freedom in the layer design. This decomposition allows evaluation of ( 8 ) in just π’ͺ ⁑ ( L ​ log 2 ​ L ) \mathcal{O}(L\log_{2}L) time (two FFT convolutions and two element-wise products), i.e. z t \displaystyle z_{t} = k t ​ ( Ο† βˆ— v ) t \displaystyle=k_{t}(\varphi*v)_{t} (6) y t \displaystyle y_{t} = q t ​ ( ψ βˆ— z ) t \displaystyle=q_{t}(\psi*z)_{t} 𝖧𝗒𝖾𝗇𝖺 \small\sf Hyena represents a generalization of ( 8 ) for an arbitrary number of projections – not limited to three – and with implicit free-form long filters for the convolutions. The resulting recurrence ( 4 ) can be also represented in matrix form y = 𝖧 ⁑ ( u ) ​ v y=\mathsf{H}(u)v . Let 𝖣 x n = diag ⁑ ( x n ) ∈ ℝ L Γ— L \mathsf{D}_{x}^{n}=\diag(x^{n})\in\mathbb{R}^{L\times L} and let 𝖲 h n \mathsf{S}_{h}^{n} be the Toeplitz matrix corresponding to filter h n h^{n} . The resulting 𝖧𝗒𝖾𝗇𝖺 \sf Hyena recurrence is linear in v v and can be rewritten in matrix form: y = 𝖧 ( u ) v = 𝖣 x N 𝖲 h N β‹― 𝖣 x 2 𝖲 h 2 𝖣 x 1 𝖲 h 1 v y=\mathsf{H}(u)v=\mathsf{D}_{x}^{N}\mathsf{S}_{h}^{N}\cdots\mathsf{D}_{x}^{2}\mathsf{S}_{h}^{2}\mathsf{D}_{x}^{1}\mathsf{S}_{h}^{1}v Figure 2.1 visualizes an example decomposition.

###### Remark 3.2 ( 𝖧𝗒𝖾𝗇𝖺 \sf Hyena generalizes 𝖧πŸ₯ \sf H3 and 𝖦𝖲𝖲 \sf GSS .) .

The 𝖧πŸ₯ \sf H3 mechanism ( Dao et al., 2022c ) corresponds to 𝖧𝗒𝖾𝗇𝖺 2 {\sf Hyena}_{2} and 𝖦𝖲𝖲 \sf GSS ( Mehta et al., 2022 ) is 𝖧𝗒𝖾𝗇𝖺 1 {\sf Hyena}_{1} , with a particular choice of parametrization for the long convolutions (SSMs).

Analysis of the 𝖧πŸ₯ \sf H3 mechanism as a decomposition 𝖣 q ​ 𝖲 ψ ​ 𝖣 k ​ 𝖲 Ο† \mathsf{D}_{q}\mathsf{S}_{\psi}\mathsf{D}_{k}\mathsf{S}_{\varphi} of its surrogate attention matrix 6 6 6 Some of this analysis is reported in the Appendix. clarifies a connection to fast evaluation algorithms for matrix-vector multiplications. In particular, the generalization of ( 8 ) to an arbitrary order is inspired by fast evaluation algorithms for structured dense matrices based on butterfly decompositions ( Li et al., 2015 ; Dao et al., 2019 ; Dao et al., 2022a ) , with length of the decomposition closely tied to its expressivity (in the classes of matrices it can represent). The 𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} operator blends data control with a special case of butterfly decomposition.

###### Remark 3.3 .

𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} operators have unbounded context. Namely, they are not artificially restricted by e.g., locality, and can learn long-range dependencies between any of the elements of v v via long convolutions, which we discuss next.

### 3.3 𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} Filters

Here we provide details on the convolution parametrization. We represent the filters of each 𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} operator as a map from the time (or space) domain t t to values h t h_{t} , and learn it with a shallow feed-forward neural network ( π–₯π–₯𝖭 \sf FFN ): h t = π–Άπ—‚π—‡π–½π—ˆπ— ⁑ ( t ) β‹… ( π–₯π–₯𝖭 ∘ π–―π—ˆπ—Œπ—‚π—π—‚π—ˆπ—‡π–Ίπ—…π–€π—‡π–Όπ—ˆπ–½π—‚π—‡π—€ ) ​ ( t ) h_{t}={\sf Window}(t)\cdot({\sf FFN}\circ{\sf PositionalEncoding})(t) (7) This approach builds on the neural implicit representation literature ( Mildenhall et al., 2021 ; Sitzmann et al., 2020 ) , which has found application in long convolution layers ( Romero et al., 2021b ; Romero et al., 2021a ) . One advantage of ( 7 ) is given by the decoupling of filter length and parameter cost.

#### Specializing filters in Hyena

The window and positional encoding functions are used to specialize filters in 𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} operators, biasing them towards a specific type. Figure 3.1 provides an important example: we choose at least one of the convolutions in 𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} to be shaped towards exponential decay, mirroring the findings of ( Li et al., 2022 ) in other applications.

Interestingly, we find that long exponentially decaying filters display synergy with high-frequency filters, as they enable the operator to select specific inputs at specific steps 7 7 7 This observation finds mirrors in the parametrization of the convolutions in H3 ( Dao et al., 2022c ) as a shift SSM and a diagonal SSM. . Similarly to ( Romero et al., 2021b ) , we use high-frequency periodic activations (sine) in the π–₯π–₯𝖭 \sf FFN . This allows ( 7 ) to learn filters with high-frequency content, addressing the low-frequency bias of neural networks ( Basri et al., 2020 ) . Owing to the π–₯π–₯𝖭 \sf FFN , the parametrization in ( 7 ) can approximate filters obtained through other means, such as S4 ( Gu et al., 2020 ; Gu et al., 2021 ) , CKConv ( Romero et al., 2021b ) , SGConv ( Li et al., 2022 ) and Fourier Neural Operator (FNO) ( Li et al., 2020 ) .

#### Preserving causality

Causality is necessary to train autoregressive language models, in order for the output at a given position to depend only on the past. For example, Transformers mask the attention matrix to be lower triangular. In the case of 𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} , causality can be guaranteed by parametrizing causal convolutions:

###### Proposition 3.1 (Causal Hyenas) .

If each filter h n , n = 1 , … , N h^{n},~n=1,\dots,N is causal, then the corresponding 𝖧𝗒𝖾𝗇𝖺 N {\sf Hyena}_{N} operator is causal.

In practice, we need not constrain the learning of the filter ( 7 ) to ensure its numerical causality. If we use FFT-based convolution algorithms, all we need is to evaluate the filter at t = 0 , … , L βˆ’ 1 t=0,\dots,L-1 and zero-pad the input and filter sequences to 2 ​ L βˆ’ 1 2L-1 before taking FFT.

#### Efficiency

One bottleneck of long convolution models can be their low utilization of hardware accelerators, especially when they involve iterative numerical methods to materialize the filter 8 8 8 In contrast, deep learning primitives are designed for high GPU utilization, with FFNs and attention usually reaching 50 βˆ’ 70 % 50-70\% or higher, if optimized. . Evaluation of 7 is fast, since it involves a single forward pass of an π–₯π–₯𝖭 \sf FFN , and can be performed in parallel across sequence length and all orders of an 𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} operator as displayed in Algorithm 3 , increasing hardware utilization. An additional source of low utilization is the FFT, which is also shared by other long other convolutional layers. This bottleneck can be partially addressed by blocking ( Selesnick and Burrus, 2017 ) , and optimization of the underlying routines ( Dao et al., 2022c ) . We benchmark runtime in Section 4.5 .

### 3.4 Hyena Algorithm

A forward pass of Hyena is summarized below.

###### Proposition 3.2 (Computational Complexity) .

The computational cost of processing an input u ∈ ℝ L Γ— D u\in\mathbb{R}^{L\times D} with an order- N N 𝖧𝗒𝖾𝗇𝖺 \sf Hyena operator is π’ͺ ⁑ ( N ​ D ​ L ​ ( log 2 ​ L + D ) ) \mathcal{O}(NDL(\log_{2}L+D))

## 4 Experiments

### 4.1 Shrinking the gap on in-context learning

We begin by empirically motivating the 𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} design, including the choice of long convolution parametrization. We consider the suite of tasks described in Table 4.1 .

Our evaluation is grounded in recent work on mechanistic interpretability of Transformers ( Elhage et al., 2021 ; Power et al., 2022 ; Olsson et al., 2022 ; Zhang et al., 2022 ) . Recently, associative recall, in particular, has been successfully used to guide the design of H3 ( Dao et al., 2022c ) . We extend the suite of tasks from these works and include benchmarking more challenging versions of each task . For example, solving associative recall with a vocabulary size of only 10 10 reveals whether a model is structurally capable of performing recall. Testing on much longer sequences and larger vocabularies reveals additional gaps in performance that are otherwise hidden.

#### How to parametrize long convolutions

We compare the performance of the following long convolution parametrizations for S 1 S^{1} and S 2 S^{2} in an order 2 2 Hyena: β€’ Conv1d: Explicit convolutions (regular convolution layers with fixed filter size).

β€’ FNO: Filters parametrized explicitly in the frequency-domain ( Li et al., 2020 ) .

β€’ H3: Implicit parametrization using state-space models (SSMs), in particular the standard S4 ( Gu et al., 2021 ) .

β€’ TransferFunc: Implicit parametrization via transfer functions, a classical system-theoretic generalization of SSMs 9 9 9 Transfer functions roughly correspond to a frequency-domain representation of SSMs.

β€’ CKConv: Implicit parametrization using π–₯π–₯𝖭 \sf FFN s ( Romero et al., 2021b ) .

β€’ 𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} : Combination of implicit parametrizations via π–₯π–₯𝖭 \sf FFN s (with exponential decay modulation as shown in Figure 3.1 ), and short explicit filters.

All models have the same width and 2 2 layers. Figure 4.1 shows implicit approaches based on FFNs outperform other long convolutions, with the gap widening on longer sequences and larger vocabulary sizes. We train a different model on each setting of sequence length and vocabulary size. The ranking is correlated with the ability to decouple sequence length from parameter count ( 𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} , CKConv, TransferFunc, H3) and expressivity (Hyena, CKConv). We observe similar trends on the other tasks.

#### Pushing sequence length to the limit

Next, we evaluate associative recall performance on extremely long sequences of length 131 131 k. To the best of our knowledge, these represent the first empirical display of attention-free in-context learning on sequences of this length. The gap between parametrization schemes widens as shown in Appendix A, with 𝖧𝗒𝖾𝗇𝖺 {\sf Hyena} outperforming CKConv by 80 80 points.

#### Comparing operators