Zoology: Measuring and Improving Recall in Efficient Language Models — part 6/14
Cleaned-text code points [77753, 93491). Bibliography omitted here; source.md includes it.
To construct the synthetic dataset following Procedure 1 , we set α = 0.1 \alpha=0.1 and | C | = 8,192 |C|=8,192 for all experiments. We fix N = 256 N=256 and vary D ∈ { N 32 , N 16 , N 8 , N 4 } D\in\{\frac{N}{32},\frac{N}{16},\frac{N}{8},\frac{N}{4}\} to increase the number of key-value pairs occurring per example and the range of token-interaction distances per example. For each number of key-value pairs per example, we additionally vary the model dimension ∈ { 64,128,256 } \in\{64,128,256\} .
The results are shown in Fig. 6 , validating that sliding window attention closes the gap to attention when the number of KV pairs is within the window size. I.e., 16 16 KVs means that there are 16 16 keys and 16 16 values, for 32 32 total tokens. Sliding attention with window size 32 32 performs well up until 16 16 KVs and degrades beyond this point relative to attention models of the same dimensionality. Sliding window attention with window size 8 8 does not perform well on any setting, as expected. Blocked attention also not perform well – intuitively tokens at the beginning of a block lack sufficient prior context to attend to.
### F.2 Demonstrating Mqar as a Tool for Architecture Development
In the main paper, we provide experiments showing that there is gap between the gated convolution and attention models on synthetic Mqar data and correspondingly in the downstream Pile results. In Table 4 , we show that the trends of RetNet and sparse attention on Mqar also carry to the downstream results. We hope Mqar synthetic data may be useful in future architectural development.
## Appendix G Mqar Perplexity Gap and Model Size
In this section, we investigate how the associative recall gap changes as we increase the model size to a billion parameters and beyond. Overall, the results suggest that the larger models improve in memorizing bigrams from the training data, but do not rapidly improve in in-context learning (perplexity on rare bigrams).
Below, we pretrain at the 1.4Bn parameter scale and observe that the 70M parameter attention model is a full perplexity point better on AR than this 1.4Bn Hyena model (Table 1 ). We discuss these results further in ( Section 3 ).
##### Open-source 7B parameter models
We next evaluate the RWKV (gated convolution) Peng et al. [2023] and Llama 2 (attention) pretrained models at the 7B parameter scale Touvron et al. [2023] . 9 9 9 We specifically evaluate the RWKV-Raven model downloaded from https://huggingface.co/docs/transformers/model_doc/rwkv and the Llama 2 model downloaded from https://github.com/facebookresearch/llama . These are both popular models that took a significant amount of effort to train, towards maximizing quality. We find that there is a gap between RWKV and attention at the 7B scale and that it increases as the model needs to conduct more recalls per sequence ( P P below).
We summarize the experimental protocol below. 1. Justifying the experimental protocol. Since frequent bigrams may just be memorized by the model and not require in-context recall, our work measures AR quality on infrequent bigrams in validation sequences ( Figure 1 , Section C.1 ). We do not have access to the custom training data mixtures used in training RWKV or Llama 2 to measure bigram frequencies, so we use a synthetic test to fairly measure the AR capabilities.
2. Evaluation dataset. We construct the synthetic dataset following Algorithm 1 , where each sequence contains P P token pairs, or “bigrams”. Each bigram, containing a “key” token and a “value” token, appears twice in the sequence. On the second occurrence of a “key” token, the model should look back to the prior occurrence to output the corresponding “value” token. We measure the AR perplexity of the model based on its ability to predict the correct “value” token as the next word for these repeated keys. The sequences are constructed using the models’ vocabulary tokens.
3. Evaluation details. We evaluate (inference only, no training) on sequence lengths of 1024 1024 tokens. We measure the AR perplexity when the sequences contain P ∈ { 16 , 32 , 64 , 128 , 256 } P\in\{16,32,64,128,256\} key-value pairs, using 1000 1000 samples per P P value. The tokens that do not contain a key or value are simply filled with a fixed token id (so this token is repeated frequently in the sequence). We plot perplexity for AR and non-AR tokens (fixed token) vs. P P . We find RWKV quality degrades with P P on the AR slice (blue line), while all other lines remain flat. MQAR remains problematic at scale.
Overall our results suggest that simply scaling the model size may not close the MQAR quality gap between the gated convolutions and attention.
##### Measuring AR Gap with Scale
In the main paper, we measure the AR perplexity on downstream data based on bigrams that are seen < 1,250 × <1,250\times during pretraining. However, we hypothesize that larger models may memorize a larger number of bigram pairs seen in the training dataset, but do not rapidly gain the capability to perform associative recall as well as attention in-context. In-context learning is defined as learning from examples provided in context [ Xie et al., 2021 ] .
Concretely, the gap between the Hyena / RWKV and attention models at the 350m scale is 1.85 / 1.84 perplexity points when we focus on bigrams seen < 1,250 × <1,250\times during pretraining ( Table 1 ). If we instead focus on bigrams seen 1 × 1\times during pretraining, the gap to attention quality is 12.0 / 13.2 perplexity points respectively. The gated convolutions appear to struggle in the regime of rare bigrams that require the model to use the context.
## Appendix H Details on Theoretical Analysis
This section provides proofs and extensions to the theoretical results in the main paper.
### H.1 Preliminaries and Notation
#### H.1.1 Notation
We denote the all 1 1 row vector of size k k , given by [ 1 1 … 1 1 ] \begin{bmatrix}1&1&\ldots&1&1\end{bmatrix} , and the all 0 0 row vector of size k k , given by [ 0 0 … 0 0 ] \begin{bmatrix}0&0&\ldots&0&0\end{bmatrix} , as 𝟏 k \bm{1}^{k} and 𝟎 k \bm{0}^{k} , respectively. We also construe the standard basis vector 𝐞 i \mathbf{e}_{i} as a column vector in these notes, and adhere to the following matrix indexing convention: 𝐌 [ i , j ] {\bf M}[i,j] is the entry in the i i th row and the j j th column, 𝐌 [ i , : ] ∈ 𝔽 1 × n {\bf M}[i,:]\in\mathbbm{F}^{1\times n} denotes the i i th row, and 𝐌 [ : , j ] ∈ 𝔽 m × 1 {\bf M}[:,j]\in\mathbbm{F}^{m\times 1} denotes the j j th column of 𝐌 ∈ 𝔽 m × n {\bf M}\in\mathbbm{F}^{m\times n} . We then use 𝟏 m × n , 𝟎 m × n ∈ 𝔽 m × 1 \bm{1}^{m\times n},\mathbf{0}^{m\times n}\in\mathbbm{F}^{m\times 1} to denote the matrix of all 1 1 s and 0 0 s, respectively.
Next, we denote the Hadamard product of vectors 𝐮 , 𝐯 ∈ 𝔽 n {\bf u},{\bf v}\in\mathbbm{F}^{n} as 𝐮 ⊙ 𝐯 {\bf u}\odot{\bf v} ; the operation can be extended to matrices by applying the Hadamard product column-wise across the matrices. This is commonly referred to as (element-wise) gating . For vectors 𝐮 , 𝐯 ∈ 𝔽 n {\bf u},{\bf v}\in\mathbbm{F}^{n} , we also denote their linear (or acyclic) convolution as 𝐮 ∗ 𝐯 {\bf u}\ast{\bf v} and cyclic convolution as 𝐮 ⊛ 𝐯 {\bf u}\circledast{\bf v} .
##### Polynomial Notation.
Because convolution is intimately tied to operations on polynomials, it is convenient to use them to discuss the inputs and outputs of gated convolution models. Let us define maps poly , poly ∗ : 𝔽 n → 𝔽 [ X ] / ( X n ) \operatorname{poly},\operatorname{poly}^{*}:\mathbbm{F}^{n}\to\mathbbm{F}[X]/(X^{n}) such that poly ( 𝒖 ) \displaystyle\operatorname{poly}(\bm{u}) = ∑ i = 0 n − 1 𝒖 [ i ] X i , and poly ∗ ( 𝒖 ) = ∑ i = 0 n − 1 𝒖 [ i ] X n − 1 − i . \displaystyle=\sum_{i=0}^{n-1}\bm{u}[i]X^{i},\text{ and }\operatorname{poly}^{*}(\bm{u})=\sum_{i=0}^{n-1}\bm{u}[i]X^{n-1-i}. This allows us to map between vectors and polynomial. Accordingly, we also define coeff : 𝔽 [ X ] / ( X n + 1 ) → 𝔽 n \mathrm{coeff}:\mathbbm{F}[X]/(X^{n+1})\to\mathbbm{F}^{n} as the map converting polynomials back to vectors: coeff ( 𝒖 ( X ) ) = 𝒖 \mathrm{coeff}(\bm{u}(X))=\bm{u} with 𝒖 [ i ] \bm{u}[i] defined as the coefficient in 𝒖 ( X ) \bm{u}(X) at degree i i .
These operations allow us to interpret the convolution of vectors in terms of polynomial multiplication [ Heideman and Burrus, 1988 ] . More specifically, we have 𝒖 ∗ 𝒗 \displaystyle\bm{u}\ast\bm{v} = coeff ( 𝒖 ( X ) ⋅ 𝒗 ( X ) mod X n ) , and \displaystyle=\mathrm{coeff}\left(\bm{u}(X)\cdot\bm{v}(X)\mod{X^{n}}\right),\text{ and} 𝒖 ⊛ 𝒗 \displaystyle\bm{u}\circledast\bm{v} = coeff ( 𝒖 ( X ) ⋅ 𝒗 ( X ) mod X n − 1 ) . \displaystyle=\mathrm{coeff}\left(\bm{u}(X)\cdot\bm{v}(X)\mod{X^{n}-1}\right). We can similarly interpret the Hadamard product of vectors 𝒖 ⊙ 𝒗 \bm{u}\odot\bm{v} as the Hadamard product of polynomials 𝒖 ( X ) ⊙ 𝒗 ( X ) \bm{u}(X)\odot\bm{v}(X) : 𝒖 ⊙ 𝒗 = coeff ( 𝒖 ( X ) ⊙ 𝒗 ( X ) ) = coeff ( ∑ i = 0 n − 1 ( 𝒖 [ i ] ⋅ 𝒗 [ i ] ) ⋅ X i ) . \bm{u}\odot\bm{v}=\mathrm{coeff}\left(\bm{u}(X)\odot\bm{v}(X)\right)=\mathrm{coeff}\left(\sum_{i=0}^{n-1}(\bm{u}[i]\cdot\bm{v}[i])\cdot X^{i}\right).
##### Arithmetic Circuit Notation.
We briefly introduce the notation of arithmetic circuits [ Volkovich, 2016 ] , the focus of Section H.5 . An arithmetic circuit 𝒞 \mathcal{C} with variables X ≜ { x 1 , x 2 , … , x n } X\triangleq\{x_{1},x_{2},\ldots,x_{n}\} over a field 𝔽 \mathbbm{F} is interpreted as a directed acyclic graph, where the input nodes are labelled by either the variables from X X or constants from 𝔽 \mathbbm{F} and the internal nodes are labelled by + + or × \times with the output being the polynomial computed at the output node.
We shall also refer to the size of the circuit as the number of nodes, the depth of the circuit as the length of the longest path between an input node and the output node, and the width of the circuit as the number of parallel operations in the circuit, or ‘wires’ which will be intersected by a horizontal ‘cut’ through the circuit. Moreover, the degree of a circuit is defined as the degree of the polynomial computed by the circuit. We summarize this with the following definition:
###### Definition H.1 .
An arithmetic circuit 𝒞 \mathcal{C} is an ( n , s , Δ , w ) (n,s,\Delta,w) -circuit if 𝒞 \mathcal{C} is an n n -variate arithmetic circuit of size s s and of depth at most Δ \Delta , and width w w .
##### Model Notation.
Now we introduce the notation we will be using for defining layers. In what follows, we denote 𝒖 ∈ ℝ N × d \bm{u}\in\mathbbm{R}^{N\times d} as the model input; N N as the sequence length; L L as the number of stacked layers, indexed by ℓ \ell ; and d d as the input (embedding) dimension.
#### H.1.2 Summary of the Results
The outline of our results are as follows: In Section H.2 we introduce gated convolution models and define H3, Hyena, RWKV, RetNet, and BaseConv . In Section H.3 we introduce a set of primitive operations that BaseConv can implement. We use these as tools in the subsequent proofs. Then, in Section H.5 , we show that a general arithmetic circuit of size s s and degree at most Δ \Delta can be simulated by BaseConv . We use the above results to show all models built from gating and convolutions can be simulated by BaseConv . This helps us analyze the class of gated convolutions, beyond a specific proposal (e.g. Hyena).
Next, we study the representational power that gated convolutions and attention use to solve Mqar . In Section H.7 , we derive a BaseConv model inspired from dyadic intervals to show the dimensionality for gated convolutions (with input-independent filters) solve Mqar . Next, we analyze the dimensionality for a BaseConv model with data-dependent kernels to solve Mqar in Section H.8 .
### H.2 Gated Attention-Free Models
We now present formal definitions of gated convolution and recurrence models with respect to the 5-tuple of parameters ( N , L , d , N ′ , d ′ ) \left(N,L,d,N^{\prime},d^{\prime}\right) .
###### Definition H.2 .
An ( N , L , d , N ′ , d ′ ) − Gated Convolution Model \left(N,L,d,N^{\prime},d^{\prime}\right)-\text{\text{Gated Convolution Model}} is a stacked sequence to sequence model with L L layers such that:
1. Input and output are N × d N\times d matrices,
2. Each layer’s operations consist of element-wise gating, convolution, linear projection, and
3. All the individual gated convolution layers take in N ′ × d ′ N^{\prime}\times d^{\prime} matrices and output N ′ × d ′ N^{\prime}\times d^{\prime} matrices. We refer to the tuple ( N ′ , d ′ ) (N^{\prime},d^{\prime}) as the inner dimension of the model.
We define the Hyena, RWKV, RetNet, and BaseConv layers to make step 2 more concrete. We also assume that the input 𝒖 ∈ ℝ N × d \bm{u}\in\mathbbm{R}^{N\times d} is embedded into 𝒖 ′ ∈ ℝ N ′ × d ′ \bm{u}^{\prime}\in\mathbbm{R}^{N^{\prime}\times d^{\prime}} such that
𝒖 ′ [ n , t ] = { 𝒖 [ n , t ] if n < N , t < d 0 otherwise. \bm{u}^{\prime}[n,t]=\begin{cases}\bm{u}[n,t]\ \ \text{ if }n<N,\ t<d\ \\ 0\ \ \text{ otherwise. }\end{cases}
The output from the last layer 𝒛 ∈ ℝ N ′ × d ′ \bm{z}\in\mathbbm{R}^{N^{\prime}\times d^{\prime}} is transformed into output 𝒚 ∈ R N × d \bm{y}\in R^{N\times d} by extracting the top left N × d N\times d entries in 𝒛 \bm{z} .
Next, we define the class of weight matrices that we will use in the linear projections in the models:
###### Definition H.3 .
The linear projection 𝙻𝚒𝚗𝚎𝚊𝚛 m , m {\tt Linear}_{m,m} has its matrix representation as the weight matrix 10 10 10 I.e. 𝙻𝚒𝚗𝚎𝚊𝚛 m , m ( 𝒛 ) = 𝑾 ⋅ 𝒛 {\tt Linear}_{m,m}(\bm{z})=\bm{W}\cdot\bm{z} . 𝐖 ∈ ℝ m × m \mathbf{W}\in\mathbb{R}^{m\times m} taken to be a K-matrix for 𝐖 ∈ ( ℬ ℬ ∗ ) poly- log m poly- log m \mathbf{W}\in(\mathcal{B}\mathcal{B}^{*})^{\text{poly-}\log{m}}_{\text{poly-}\log{m}} ( Definition H.17 ). Consequently, each matrix 𝐖 \mathbf{W} has 𝒪 ~ ( m ) \tilde{\mathcal{O}}(m) parameters and runtime for matrix vector multiplication, and allows us to represent general linear transformations with low-depth linear arithmetic circuits [ Dao et al., 2020 ] . We will also need linear maps 𝙻𝚒𝚗𝚎𝚊𝚛 m , n {\tt Linear}_{m,n} for m < n m<n , where each take the corresponding square matrices from 𝙻𝚒𝚗𝚎𝚊𝚛 m , m {\tt Linear}_{m,m} and note that such matrices has 𝒪 ~ ( n ) \tilde{\mathcal{O}}(n) parameters and runtime for matrix vector multiplication.
###### Remark H.4 .
We note that the weight matrix 𝐖 ∈ ℝ d × d \mathbf{W}\in\mathbb{R}^{d\times d} above is taken to be a dense matrix in the experiments. However, for our theoretical results, we restrict the linear projection in our models as per Definition H.3 . In the rest of the paper, unless mentioned otherwise all linear projections will follow Definition H.3 .
Next, we define three popular gated convolution models that we study in our work: Hyena Poli et al. [2023a] , RWKV Peng et al. [2023] , and RetNet Sun et al. [2023] .
#### H.2.1 The Hyena Layer
We will now outline the Hyena layer [ Poli et al., 2023b ] . Hyena takes a sequence 𝒖 ∈ ℝ N × d \bm{u}\in\mathbbm{R}^{N\times d} as input and produces L + 1 L+1 projections p 1 , … , p L , v p^{1},\ldots,p^{L},v by passing 𝒚 \bm{y} though a linear layer and applying a short convolution afterwards. The algorithm then recursively performs a point-wise multiplication of the projection with the convolution of the filter h l h^{l} with the previous output. We summarize this process in Algorithm 3 .
###### Remark H.5 .
In Algorithm 3 , L L is used to denute the nuber of recursive applications of the Hadamard product and convolutions, not the number of layers Note that asymptotically, the recursive step does not make a difference.