Hyena Hierarchy: Towards Larger Convolutional Language Models β part 4/5
Cleaned-text code points [46922, 62896). Bibliography omitted here; source.md includes it.
Training hyperparameters match those of standard GPT training pipelines, and are thus likely suboptimal for new attention-free architectures such as Hyena . We run some preliminary experiments and find that e.g., some modifications to the learning rate schedule, currently involving linear warmup and cosine decay, to improve perplexity at convergence of Hyena models (we recommend slightly lower learning rates for Hyena models compared to GPT of a similar size). Despite these findings, we use standard GPT hyperparameters for both GPT and Hyena .
#### PG-19
We also report results of additional training runs on other datasets. We train a π§ππΎππΊ {\sf Hyena} 153 153 M model on the standard PG-19 long-range corpus ( Rae et al., 2019 ) , with a context length of 16 16 k tokens, reaching a test perplexity of 14.6 14.6 (using the standard GPT2 tokenizer) in 8 8 epochs.
#### Architectures
Architectural hyperparameters for Hyena are shown in Table A.4 . We use sine as an activation function for the FFN of Hyena filters.
#### FLOP computation
The number of floating point operations (FLOPs) reported in the main text are computed using the same strategy as in ( Hoffmann et al., 2022 ) . For GPT, we do not use the approximation, opting instead for the more accurate formula based on FLOP counts of individual layers. In the case of Hyena , FLOPs are computed using the same method, except attention layers are replaced by: i. Projections: order Γ \times d_model Γ \times d_model Γ \times seq_len.
ii. Short conv on projections: order Γ \times d_model Γ \times seq_len Γ \times filter_len (usually 3 3 ).
iii. FFTConv: 5 5 Γ \times (order - 1) Γ \times d_model Γ \times log β‘ ( seq_len ) \log(\text{seq\_len}) Γ \times seq_len.
iv. Output: d_model Γ \times d_model Γ \times seq_len.
with a leading factor 2 2 to account for both additions and multiplications.
### A.3 Downstream Evaluation
#### SuperGLUE:
We evaluate models on the SuperGLUE ( Wang et al., 2019 ) with the parsing pipeline of ( Arora et al., 2022 ) . For all tasks except WIC, CB and BoolQ, we generate a response using greedy decoding, then check for the gold label. WIC, CB and BoolQ use logit scoring instead of generation.
#### Models
The models considered are the open-source checkpoint of GPTNeo 125 125 M trained for 300 300 B tokens The Pile , and the RWKV-v4 169 169 M checkpoint trained for 332 332 B tokens on The Pile . Hyena is a 153 153 M model trained for 137 137 B tokens on The Pile .
#### LAMBADA:
We evaluate π§ππΎππΊ {\sf Hyena} on the LAMBADA ( Paperno et al., 2016 ) task. We apply a stop word filter and check whether predictions for all tokens corresponding to the last word agree with the ground truth. The small π§ππΎππΊ {\sf Hyena} model trained on 137 137 B tokens reaches 44.64 % 44.64\% accuracy.
### A.4 Image Classification
#### a
ImageNet: We use ImageNet-1k which consists of 1000 classes and 1.3M images and train from scratch with no outside data on 8 Nvidia A100 GPUs. In our ViT benchmark, we swap the attention layers with the Hyena operator defined in our language experiments, and remove the class token and positional embeddings, similar to S4ND ( Nguyen et al., 2022 ) . The parameter count is kept similar at 87M ViT-B (base) vs 88M Hyena-ViT. The training procedure from T2T-ViT ( Yuan et al., 2021 ) is used, including augmentations such as RandAugment ( Cubuk et al., 2020 ) , Mixup ( Zhang et al., 2017 ) , and AugMix ( Hendrycks et al., 2019 ) . See table A.5 for hyperparameter settings used.
#### CIFAR-10:
We use CIFAR-10 in sequential and 2D experiments. For sequential, we use the Hyena operator defined in our language tasks and compare with an S4 model ( Gu et al., 2021 ) of the same size by swapping layers in the residual blocks. In 2D, we learn Hyena filters (in both x x and y y dimensions) that are equal to the size of the input shape, and forgo the gating mechanism from our language experiments. We window (i.e., apply a soft mask spatially to) the Hyena filters with a decay term. The rate of decay varies across channels, ensuring different sizes of the filters at initialization. We compare with another implicit 2D convolution, S4ND ( Nguyen et al., 2022 ) , by swapping the model layers with the 2D Hyena filters. The "isometric" model consists of 4 residual blocks of model dimension 128. We use basic image augmentations, 0.1 dropout, 0.03 weight decay and train for 100 epochs using a Nvidia T4 GPU.
## Appendix B Theoretical Results and Details
### B.1 Proofs
#### Proof of Proposition 3.1
###### Proof.
A discrete L L -by- L L operator is causal if it is lower triangular, i.e., when there is no leakage of future input information to the output. The π§ππΎππΊ {\sf Hyena} operator π§ \mathsf{H} is the product of alternating diagonal and Toeplitz matrices. Thus, if all the Toeplitz matrices π² h n \mathsf{S}_{h}^{n} are lower triangular then π§ \mathsf{H} is lower triangular. In turn, each π² h n \mathsf{S}_{h}^{n} is lower triangular if and only if the filter h h is causal, concluding the proof. β
### B.2 Analysis of Data-Controlled Mechanisms
We discuss the surrogate attention mechanism of Hyena - 2 2 : q , k , v β¦ y q,k,v\mapsto y : z t \displaystyle z_{t} = k t β ( Ο β v ) t \displaystyle=k_{t}(\varphi*v)_{t} (8) y t \displaystyle y_{t} = q t β ( Ο β z ) t \displaystyle=q_{t}(\psi*z)_{t} If Ο \varphi and Ο \psi are convolutions parametrized via state-space models (SSMs), the above resembles the H3 mechanism ( Dao et al., 2022c ) . We investigate the effect of the convolutional kernels Ο \varphi and Ο \psi on the attention layer. We start by introducing a matrix representation of the layer, and we isolate the attention matrix π Ο Ο β ( q , k ) \mathsf{A}_{\varphi}^{\psi}(q,k) such that y \displaystyle y = π Ο Ο β ( q , k ) β v . \displaystyle=\mathsf{A}_{\varphi}^{\psi}(q,k)v. (9)
#### Isolating the surrogate attention matrix
In the case of length- L L discrete sequences z t \displaystyle z_{t} = k t β β m = 0 L β 1 Ο t β m β v m \displaystyle=k_{t}\sum_{m=0}^{L-1}\varphi_{t-m}v_{m} (10) y t \displaystyle y_{t} = q t β β m = 0 L β 1 Ο t β m β z m \displaystyle=q_{t}\sum_{m=0}^{L-1}\psi_{t-m}z_{m} Therefore we can rewrite ( 8 ) as y t \displaystyle y_{t} = q t β β m = 0 L β 1 Ο t β m β k m β β n = 0 L β 1 Ο m β n β v n \displaystyle=q_{t}\sum_{m=0}^{L-1}\psi_{t-m}k_{m}\sum_{n=0}^{L-1}\varphi_{m-n}v_{n} (11) = q t β β m = 0 L β 1 β n = 0 L β 1 Ο t β m β k m β Ο m β n β v n \displaystyle=q_{t}\sum_{m=0}^{L-1}\sum_{n=0}^{L-1}\psi_{t-m}k_{m}\varphi_{m-n}v_{n} Move Ο , k inside inner sum \displaystyle\text{Move $\psi$, $k$ inside inner sum} = q t β β n = 0 L β 1 β m = 0 L β 1 Ο t β m β k m β Ο m β n β v n \displaystyle=q_{t}\sum_{n=0}^{L-1}\sum_{m=0}^{L-1}\psi_{t-m}k_{m}\varphi_{m-n}v_{n} Index shift \displaystyle\text{Index shift} = β n = 0 L β 1 q t β β m = 0 L β 1 Ο t β m β k m β Ο m β n β v n \displaystyle=\sum_{n=0}^{L-1}q_{t}\sum_{m=0}^{L-1}\psi_{t-m}k_{m}\varphi_{m-n}v_{n} And we can define the surrogate attention matrix π Ο Ο β ( q , k ) \mathsf{A}_{\varphi}^{\psi}(q,k) t , t β² \displaystyle{}_{t,t^{\prime}} = q t β β m = 0 L β 1 Ο t β m β k m β Ο m β t β² . \displaystyle=q_{t}\sum_{m=0}^{L-1}\psi_{t-m}k_{m}\varphi_{m-t^{\prime}}. (12)
#### Operator decomposition of the surrogate attention matrix
We can decompose the linear map v β¦ y ; y = π Ο Ο β ( q , k ) β v v\mapsto y;~y=\mathsf{A}_{\varphi}^{\psi}(q,k)v into a sequence of factors, each dependent on a projection of the input π Ο Ο β ( q , k ) = π Ο β ( q ) β π Ο β ( k ) \mathsf{A}_{\varphi}^{\psi}(q,k)=\mathsf{A}^{\psi}(q)\mathsf{A}_{\varphi}(k) . 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 . Then, we have that π Ο β ( q ) \displaystyle\mathsf{A}^{\psi}(q) = π£ q π² Ο , π£ q = diag ( q ) , \displaystyle=\mathsf{D}_{q}\mathsf{S}_{\psi},\qquad\mathsf{D}_{q}=\diag(q), (16) π Ο β ( k ) \displaystyle\mathsf{A}_{\varphi}(k) = π£ k π² Ο , π£ k = diag ( k ) . \displaystyle=\mathsf{D}_{k}\mathsf{S}_{\varphi},\qquad\mathsf{D}_{k}=\diag(k). The matrix has been decomposed into two terms π Ο β ( q ) \mathsf{A}^{\psi}(q) and π Ο β ( k ) \mathsf{A}_{\varphi}(k) constructed by multiplying the diagonal matrices π£ q \mathsf{D}_{q} and π£ k \mathsf{D}_{k} with the Toeplitz matrices π² Ο \mathsf{S}_{\psi} and π² Ο \mathsf{S}_{\varphi} . π² Ο \mathsf{S}_{\psi} and π² Ο \mathsf{S}_{\varphi} are the kernels of the convolution operators with filterβs impulse responses Ο \psi and Ο \varphi respectively. In the current applications of interest, Ο \psi and Ο \varphi are chosen to be causal, i.e. Ο β‘ [ t ] = 0 β for β t < 0 \psi[t]=0\text{ for }t<0 and Ο β‘ [ t ] = 0 β for β t < 0 \varphi[t]=0\text{ for }t<0 . This results in π² Ο \mathsf{S}_{\psi} and π² Ο \mathsf{S}_{\varphi} to be lower triangular matrices π² Ο = [ Ο 0 0 β― 0 Ο 1 Ο 0 β― 0 β± β± Ο L β 1 Ο L β 2 β― Ο 0 ] , π² Ο = [ Ο 0 0 β― 0 Ο 1 Ο 0 β― 0 β± β± Ο L β 1 Ο L β 2 β― Ο 0 ] . \mathsf{S}_{\psi}=\begin{bmatrix}\psi_{0}&0&\cdots&0\\ \psi_{1}&\psi_{0}&\cdots&0\\ \vdots&\ddots&\ddots&\vdots\\ \psi_{L-1}&\psi_{L-2}&\cdots&\psi_{0}\end{bmatrix},\qquad\mathsf{S}_{\varphi}=\begin{bmatrix}\varphi_{0}&0&\cdots&0\\ \varphi_{1}&\varphi_{0}&\cdots&0\\ \vdots&\ddots&\ddots&\vdots\\ \varphi_{L-1}&\varphi_{L-2}&\cdots&\varphi_{0}\end{bmatrix}. (17) The surrogate attention matrix is then given by π Ο Ο β ( q , k ) = π£ q β π² Ο β π£ k β π² Ο \mathsf{A}_{\varphi}^{\psi}(q,k)=\mathsf{D}_{q}\mathsf{S}_{\psi}\mathsf{D}_{k}\mathsf{S}_{\varphi} (18) We can expand the matrix multiplications in ( 16 ) in the case of causal filters Ο \varphi and Ο \psi as π£ q [ q 0 q 1 β± q L β 1 ] β π² Ο [ Ο 0 Ο 1 Ο 0 β± β± Ο L β 1 Ο L β 2 β― Ο 0 ] β π£ k [ k 0 k 1 β± k L β 1 ] β π² Ο [ Ο 0 Ο 1 Ο 0 β± β± Ο L β 1 Ο L β 2 β― Ο 0 ] \displaystyle\underset{\begin{bmatrix}q_{0}&&&\\ &q_{1}&&\\ &&\ddots&\\ &&&q_{L-1}\end{bmatrix}}{\displaystyle\mathsf{D}_{q}}\underset{\begin{bmatrix}\psi_{0}&&&\\ \psi_{1}&\psi_{0}&&\\ \vdots&\ddots&\ddots&\\ \psi_{L-1}&\psi_{L-2}&\cdots&\psi_{0}\end{bmatrix}}{\displaystyle\mathsf{S}_{\psi}}\underset{\begin{bmatrix}k_{0}&&&\\ &k_{1}&&\\ &&\ddots&\\ &&&k_{L-1}\end{bmatrix}}{\displaystyle\mathsf{D}_{k}}\underset{\begin{bmatrix}\varphi_{0}&&&\\ \varphi_{1}&\varphi_{0}&&\\ \vdots&\ddots&\ddots&\\ \varphi_{L-1}&\varphi_{L-2}&\cdots&\varphi_{0}\end{bmatrix}}{\displaystyle\mathsf{S}_{\varphi}} (19) = [ q 0 β Ο 0 q 1 β Ο 1 q 1 β Ο 0 β± β± q L β 1 β Ο L β 1 q L β 1 β Ο L β 2 β― q L β 1 β Ο 0 ] π Ο β ( q ) β [ k 0 β Ο 0 k 1 β Ο 1 k 1 β Ο 0 β± β± k L β 1 β Ο L β 1 k L β 1 β Ο L β 2 β― k L β 1 β Ο 0 ] π Ο β ( k ) \displaystyle\underset{\displaystyle\mathsf{A}_{\psi}(q)}{=\begin{bmatrix}q_{0}\psi_{0}&&&\\ q_{1}\psi_{1}&q_{1}\psi_{0}&&\\ \vdots&\ddots&\ddots&\\ q_{L-1}\psi_{L-1}&q_{L-1}\psi_{L-2}&\cdots&q_{L-1}\psi_{0}\end{bmatrix}}\underset{\displaystyle\mathsf{A}_{\varphi}(k)}{\begin{bmatrix}k_{0}\varphi_{0}&&&\\ k_{1}\varphi_{1}&k_{1}\varphi_{0}&&\\ \vdots&\ddots&\ddots&\\ k_{L-1}\varphi_{L-1}&k_{L-1}\varphi_{L-2}&\cdots&k_{L-1}\varphi_{0}\end{bmatrix}}
## Appendix C Discussion and Additional Results
#### Vocabulary size scaling
Table C.1 showcases interesting correlation between associative recall performance for varying vocabulary sizes and loss on the The Pile . In this case, we fix sequence length for associative recall to be 2048 2048 , the same sequence length used to train all models on the The Pile .
We observe a similar phenomenon on other slices of tasks from our mechanistic design benchmarks, indicating that it may be possible to derive predictive laws for performance at scale, based on fast experimentation on synthetic tasks with models of 1 1 or 2 2 layers. Surprisingly, performance on our language synthetics appears to be further linked to performance as attention replacement in other domains (Appendix A.4 for results on image classification).
#### Single layer recall
All experiments on our synthetic tasks default to 2 2 layer models. We choose 2 2 as it is the canonical number for mechanistic analysis of Transformers ( Elhage et al., 2021 ) based on circuits . Interestingly, a single layer of Hyena (width 64 64 ) is capable of performing associative recall, solving the task completely even in the challenging setting with vocabulary size 40 40 . Reverse engineering exactly how the single Hyena operator is able to perform recall is left for future work.
### C.1 Learning Arithmetic
We showcase an additional task in our mechanistic design benchmark: learning arithmetic. We train Hyena models of increasing depth ( 1 1 , 2 2 and 3 3 layers) on a dataset of D n D_{n} -digit addition. As an example, a 3 3 -digit addition input sample is given by the sequence π· , πΈ , πΉ , πΏ , π» , πΊ , π· , πΆ , π½ , π½ {\tt 1,2,3,9,5,4,1,0,7,7} where the first 6 6 digits contain the two 3 3 digits numbers to add, and the last 4 4 the result. Our models are optimized using standard autoregressive training i.e., predicting the next token, since they are causal. In particular, we optimize models to learn a map x β¦ y x\mapsto y where x x is the original prompt without the last element, and y y equal to x x shifted right by one position. We mask the first 2 β D n β 1 2D_{n}-1 elements of the loss for each sequence since they contain predictions for addends and not results.
We report results in Figure C.1 . A single layer of Hyena is able to learn to perform addition with up to 4 4 digits. Longer numbers require deeper models. In our experiments, alternative architectures such as AFT-conv struggle to learn arithmetic, signaling a cap in capability.
## Appendix D Samples and Visualizations
### D.1 Hyena Matrices
We provide visualizations of attention and π§ππΎππΊ {\sf Hyena} matrices activated by test strings. In D.1 , D.2 , we compare GPTNeo ( Black et al., 2021 ) attention matrices with Hyena matrices extracted by our pre-trained small Hyena model. In D.3 and D.4 , we provide additional Hyena matrices for the 355 355 M model, activated by test strings of different length.
For attention, we visualize the raw post-softmax matrix. For π§ππΎππΊ {\sf Hyena} matrices, we plot the (element-wise) absolute value of π§ β‘ ( u ) \mathsf{H}(u) : π§ β‘ ( u ) \displaystyle\mathsf{H}(u) = π£ x N π² h N β― π£ x 2 π² h 2 π£ x 1 π² h 1 \displaystyle=\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} π§ ^ β ( u ) i β j \displaystyle\hat{\mathsf{H}}(u)_{ij} = | π§ β ( u ) i β j | \displaystyle=\left|\mathsf{H}(u)_{ij}\right| Since Hyena does not normalize the entries of its matrices with e.g., softmax, there are notable differences with attention: (1) the entries of π§ β‘ ( u ) \mathsf{H}(u) can be either positive and negative, and (2) the magnitude is unconstrained. We observe the magnitude of matrices in pre-trained Hyena models to be around 10 β 3 10^{-3} .
### D.2 Hyena Filters
Figure D.5 provides a visualization of Hyena long convolution filters at initialization and after training to completion on The Pile .
We find a substantial performance difference (up to 5 % 5\% perplexity) between initialization schemes. If the filters at initialization are excessively smooth (see Appendix D.3 for a discussion of positional encoding and activation), the model finds a worse solution and takes longer to converge. Further, we observe initialization schemes that regularize filters towards typical filters learned at convergence to decrease performance. These observations are in line with performance gaps between convolution parametrization schemes discussed in main text and Appendix A.1 . In particular, the performance improvements obtained via Hyena filters could be due to easier optimization in the space of convolutional filters.
At convergence, Hyena learns a collection of lower-order filters with a similar structure, which can be exploited to further speed up inference after training.
### D.3 Positional Encoding and Filters Initialization