Zoology: Measuring and Improving Recall in Efficient Language Models — part 2/14
Cleaned-text code points [15717, 31059). Bibliography omitted here; source.md includes it.
To better understand the differences between attention and gated convolutions, we perform a fine-grained analysis of next token predictions and observe that convolution-based models struggle to recall associations previously seen in context. For example, in Fig. 1 , the model must recall the association between “ Tim ” and the last name “ Rice ”. Following a long line of prior work, we call this skill associative recall (AR) ( Willshaw et al., 1969 ; Hopfield, 1982 ) (see Appendix A for an extended discussion of AR’s history in machine learning). In Section D.1.1 , we provide annotated Pile examples to demonstrate the phenomenon qualitatively.
Quantifying AR performance. It is challenging to derive a quantitative measure of associative recall performance on the Pile because we don’t know which next token predictions in raw text require associative recall. We use a simple heuristic to identify these tokens, which we refer to as AR Hits . An AR Hit is the last token of an n n -gram repeated in context ( e.g. the second occurence of ` ` R i c e " ``Rice" in Fig. 1 ). However, some common n n -grams ( e.g. “of the” ) could have been memorized during training, so we factor in the frequency with which an n n -grams appeared in the training data. This heuristic enables us to scale our analysis to over 10 million tokens of Pile validation data.
We stratify Pile tokens into two slices based on this heuristic and report perplexity on each in Table 1 : 1. AR Hits : (6.4% of tokens) Tokens in the final position of a bigram (a pair of consecutive tokens) which previously appeared in context, but ≤ 1250 × \leq 1250\times during training.
2. Other tokens: (93.6% of tokens) Tokens in the final position of a bigram which did not previously appear in context or it appeared > 1,250 >1,250 times during training.
In Figure 1 , we visualize these slices by plotting log-perplexity against the frequency with which bigrams appear during training. Strikingly, the gap between attention and gated convolutions is the largest on AR hits with the fewest occurrences. On the other tokens, there is no gap.
In Table 1 , we also compute the percentage of the difference in perplexity between attention and each model that is due to AR tokens: Δ log ( ϕ AR ) ⋅ | T AR | Δ log ( ϕ ) ⋅ | T | \frac{\Delta\log(\phi_{\text{AR}})\cdot|T_{\text{AR}}|}{\Delta\log(\phi)\cdot|T|} , where ϕ \phi is the perplexity and T T is the set of tokens in the test set. This quantity can also be interpreted as the fraction of the overall gap that would close if a model matched attention on the AR slice. We find that the AR slice accounts for 82% of the average quality gap between the gated convolutions and attention.
To evaluate the AR capacity of larger models, we train two 1.4 billion parameter attention and Hyena models for 50 billion tokens on the Pile and repeat this analysis (see Table 5 ). Strikingly, a 70 million parameter attention model is a full perplexity point better in the AR slice than this 1.4B Hyena model that is 20 × \times its size (2.41 vs. 3.43 ppl.). In Appendix G , we also evaluate open-source RWKV and attention models trained up to 7 billion parameters. We use a controlled semi-synthetic dataset to measure AR capacity and show that RWKV’s performance degrades sharply as we increase the number of queries in an example while attention performs consistently well Appendix G .
### 3.2 Formalizing the problem: Multi-Query Associative Recall
This gap in associative recall perplexity is very surprising because prior work shows gated-convolutions can perfectly solve a formalized version of the task ( Fu et al., 2023a ; Poli et al., 2023a ; Olsson et al., 2022 ) . In this synthetic task, the input 𝒙 \bm{x} contains a sequence of bigrams representing key-value pairs from a random dictionary followed by a single query token. For example, the correct output for the input below would be 3 3 : A 4 B 3 ⏟ Key-Value C 6 E 2 F 1 C 6 G 8 → B ? ⏟ Query \text{A 4}\underbrace{\text{B 3}}_{\mathclap{\textbf{Key-Value}}}\text{C 6 E 2 F 1 C 6 G 8 }\rightarrow\underbrace{\text{B ?}}_{\mathclap{\textbf{Query}}} Gated convolutions ( e.g. H3, Hyena, RWKV) can solve this task perfectly for most sequence lengths.
These conclusions are inconsistent with our findings on the Pile, as described above, so we ask how this formulation of AR differs from the way AR manifests in real language. We identify a major difference. In real world inputs, the language model often needs to perform multiple associative recalls in a single forward pass, at varying positions in the sequence ( e.g. “Tim Rice” and “March 2018” in Fig. 1 . We refer to this as Multi-Query AR ( Mqar ). We formally define the Mqar problem as follows: 4 4 4 In Section H.7.1 we define a formal, general form of this definition, used in the theoretical analysis.
###### Definition 3.1 (Multi-Query-AR ( Mqar ) \left(\textsc{Mqar}\right) ) .
We are given an input sequence 𝒙 = { x 0 , … , x N − 1 } \bm{x}=\{x_{0},\ldots,x_{N-1}\} where each x i ∈ C x_{i}\in C is a token drawn from a vocabulary of size c = | C | c=|C| . The task is to check, for every query 1 ≤ i < N 1\leq i<N , whether there exists a 0 ≤ j < i 0\leq j<i such that 𝒖 i ≡ 𝒖 j \bm{u}_{i}\equiv\bm{u}_{j} . If so, output 𝒖 j + 1 \bm{u}_{j+1} .
For example, the correct output for input below would be 4, 6, 1, 2, 3 : A 4 B 3 C 6 F 1 ⏟ Key-Value E 2 → A ? C ? F ? ⏟ Query E ? B ? \text{A 4 B 3 C 6}\underbrace{\text{F 1}}_{\mathclap{\textbf{Key-Value}}}\text{E 2}\rightarrow\text{A ? C ?}\underbrace{\text{F ?}}_{\mathclap{\textbf{Query}}}\text{E ? B ?} In Section 4 , we use Mqar to explain the quality gap between gated convolutions and attention.
## 4 Explaining the associative recall problem
In this section, we provide an explanation for the gap in associative recall performance by analyzing the formal Mqar task theoretically and empirically. In Section 4.1 , we define a simple gated-convolution architecture, called BaseConv , which we show can simulate a broad class of architectures built from gating and convolutions. This allows us to make general statements that apply to popular gated-convolution architectures like Hyena, RWKV, or H3. In Section 4.2 , we show that there exist theoretical solutions to Mqar that could in principle be learned by BaseConv , and we analyze their complexity in terms of model width and depth. In Section 4.3 , we use experiments on synthetic data to show that solving Mqar with BaseConv (and other gated-convolution architectures) requires model dimension to scale linearly with the sequence length. In contrast, attention solves Mqar consistently in our experiments with model dimension scaling independently of sequence length. These empirical scaling laws provide a potential explanation for the AR gap and, alongside our theoretical analysis, point to the potential solutions discussed in Section 5 .
### 4.1 BaseConv : a minimal gated convolution operator
In this section, we define our minimal gated-convolution architecture, called BaseConv . Given a function, we would like to know the most efficient model ( e.g. parameters, FLOPs) that can represent the solution. In this work, we show we can precisely reason about this question for representing polynomial functions with gated convolutions as gating and convolutions are both polynomial operations. The standard model defining computational complexity for polynomials is by the size of the smallest arithmetic circuit that can compute the polynomial. We define the BaseConv gated convolution operator which is exciting because (1) it is universal in that it that can simulate any arithmetic circuit 𝒞 \mathcal{C} (with only a poly-log blowup in the corresponding parameters) and (2) it is simple to implement efficiently (19 lines of pure PyTorch including imports, see Appendix B ).
###### Definition 4.1 ( BaseConv Operator) .
Given an input 𝒖 ∈ ℝ N × d \bm{u}\in\mathbbm{R}^{N\times d} , the BaseConv operator for layer ℓ \ell is defined as: 𝒚 \displaystyle\bm{y} : = ( 𝒖 ⋅ 𝑾 ℓ + 𝒃 1 ℓ ) ⏟ Linear Projection ⊙ ( 𝒉 ℓ ∗ 𝒖 + 𝒃 2 ℓ ) ⏟ Convolution \displaystyle:=\underbrace{\left(\bm{u}\cdot\bm{W}^{\ell}+\bm{b}_{1}^{\ell}\right)}_{\mathclap{\textbf{Linear Projection}}}\odot\underbrace{\left(\bm{h}^{\ell}\ast\bm{u}+\bm{b}_{2}^{\ell}\right)}_{\mathclap{\textbf{Convolution}}}\quad (1) where the layer is parameterized by learnable filters 𝒉 ∈ ℝ N × d \bm{h}\in\mathbbm{R}^{N\times d} , a linear projection 𝑾 ℓ ∈ ℝ d × d \bm{W}^{\ell}\in\mathbbm{R}^{d\times d} , and ‘bias’ matrices 𝒃 1 , 𝒃 2 ∈ ℝ N × d \bm{b}_{1},\bm{b}_{2}\in\mathbbm{R}^{N\times d} . The ⊙ \odot is component-wise product and convolution of two matrices is computed as convolution of the corresponding columns.
In our experiments, each BaseConv layer uses 𝒪 ~ ( N d + d 2 ) \tilde{\mathcal{O}}(Nd+d^{2}) parameters 5 5 5 We use 𝒪 ~ ( ⋅ ) \tilde{\mathcal{O}}{(\cdot)} to hide poly-log factors. and can be computed in 𝒪 ~ ( N d 2 ) \tilde{\mathcal{O}}(Nd^{2}) operations. For our theoretical results, we can assume the weight matrix 𝑾 ℓ \bm{W}^{\ell} is restricted to a class of matrices that support near-linear time matrix multiplication ( e.g. Kaleidoscope matrices, see Definition H.3 ). Under this assumption, BaseConv uses 𝒪 ~ ( N d ) \tilde{\mathcal{O}}(Nd) parameters and 𝒪 ~ ( N d ) \tilde{\mathcal{O}}(Nd) FLOPs ( Proposition H.6 ). We now state the equivalency result between arithmetic circuits and BaseConv a “canonical” representation of arithmetic circuits ( Theorem H.21 in Section H.5 ):
###### Theorem 4.2 (Equivalency to Arithmetic Circuits) .
For an arithmetic circuit 𝒞 \mathcal{C} of size s s and depth Δ \Delta that takes 𝐮 ∈ ℝ N × d \bm{u}\in\mathbbm{R}^{N\times d} as input, there exists an equivalent BaseConv operator that uses 𝒪 ~ ( s Δ ) \tilde{\mathcal{O}}(s\Delta) parameters and 𝒪 ~ ( Δ ) \tilde{\mathcal{O}}(\Delta) layers. 6 6 6 The formal statement in the Appendix has a sharper version of this result in terms of the circuit ‘width’.
In other words, any gated convolution model with small number of layers can be simulated by BaseConv with only a (poly)logarithmic blowup in parameters and layers. We note that arithmetic circuits are a very well studied computation model in computational complexity Bürgisser et al. (1996) . Many well-known efficient algorithms on matrices (e.g. the FFT or the current best known matrix-matrix multiplication algorithm) in fact give small arithmetic circuits. However, arithmetic circuits are inherently discrete objects – we cannot learn them via gradient descent. Theorem 4.2 shows that (up to poly-log loss in parameters), we can instead learn over BaseConv models. This result generalizes a similar result from Dao et al. (2020) for the special class of linear functions: we generalize the earlier result to the class of all polynomials.
For specific gated convolution layers, we can get rid of the poly-logarithmic factor blowup–we observe in the appendix that BaseConv and Hyena models can simulate each other with only a small constant blowup in parameters ( Proposition H.12 in Section H.5 ).
### 4.2 Theoretical analysis of gated convolution capacity and associative recall
In this section, we provide theoretical Mqar solutions that could in principle be learned by each architecture and analyze their complexity in terms of model width and depth. First, we note that attention solves Mqar with parameters independent of sequence length ( Proposition H.27 ).
###### Proposition 4.3 (Attention) .
Given an input 𝐮 ∈ { 0 , 1 } N × 3 c \bm{u}\in\{0,1\}^{N\times 3c} , Attention (even without using soft-max) solves Mqar for 𝐮 \bm{u} using 𝒪 ( c 2 ) {\mathcal{O}}(c^{2}) parameters, 𝒪 ( N c 2 + N 2 c ) \mathcal{O}(Nc^{2}+N^{2}c) time complexity and 𝒪 ( 1 ) {\mathcal{O}}(1) layers.
It is natural to wonder then if all pairwise comparisons among tokens are necessary to solve Mqar . Indeed, in the RAM setting, a sequential algorithm can simply utilize N N logarithmic insertion and membership queries to solve Mqar in subquadratic time. Unfortunately, any model attempting to emulate this would require Ω ( N ) \Omega(N) layers. Instead, we observe that we can parallelize this algorithm using dyadic intervals and achieve a depth of 𝒪 ~ ( 1 ) \tilde{\mathcal{O}}(1) ( Proposition H.30 ). We then convert this algorithm into an arithmetic circuit and apply Theorem 4.2 to derive an equivalent BaseConv model. This allows us to prove new upper bounds for BaseConv models applied to Mqar , which improves upon the quadratic time complexity of attention to near-linear runtime at the cost of using poly-log layers ( Theorem H.37 in Section H.7 ).
###### Theorem 4.4 (Data-Independent Filters 7 7 7 We note here that existing architectures also use data-independent convolution filters, meaning the filter is defined as a function of the model parameters, independent of the input. ) .
Given an input 𝐮 ∈ { 0 , 1 } N × 𝒪 ( log c ) \bm{u}\in\{0,1\}^{N\times\mathcal{O}(\log{c})} to Mqar (where we assume that distinct tokens are embedded into distinct vectors in { 0 , 1 } 𝒪 ( log c ) \{0,1\}^{\mathcal{O}(\log{c})} ), there exists a BaseConv operator that solves Mqar for 𝐮 \bm{u} using 𝒪 ~ ( N log c ) \tilde{\mathcal{O}}(N\log{c}) parameters as well as time complexity and 𝒪 ~ ( 1 ) \tilde{\mathcal{O}}(1) layers.
Nevertheless, the poly-logarithmic number of layers in the above result is undesirable in practice. But, we show that using input-dependent convolution filters, one can get constant many layers (for a sub-class of inputs). Towards that end, we define the interaction distance between a query 𝒒 i \bm{q}_{i} and the matching key 𝒌 j \bm{k}_{j} as i − j i-j . This then allows us to present the corresponding upper bound for data-dependent mixing ( Theorem H.38 in Section H.8 ).
###### Theorem 4.5 (Input-Dependent Filters) .
Given an input 𝐮 ∈ { 0 , 1 } N × c \bm{u}\in\{0,1\}^{N\times c} to Mqar (where we assume that the tokens are embedded as one-hot encoding in { 0 , 1 } c \{0,1\}^{c} and there exists at most t t distinct interaction distances), 8 8 8 Note that the interaction distances can be arbitrary: there is just a bounded number of distinct distances. there exists a BaseConv operator that uses input-dependent kernels to solve the above case of Mqar using 𝒪 ( t ⋅ N c ) \mathcal{O}(t\cdot Nc) parameters and 𝒪 ( 1 ) {\mathcal{O}}(1) layers.
### 4.3 Empirical analysis of gated convolution capacity and associative recall
In this section, we measure empirically how model dimension must scale in order for different sequence mixers to solve Mqar .
##### Setup
We train and evaluate models on a synthetic Mqar with vocabulary size 8,192 8,192 , varying model dimension and sequence length from 64 64 to 512 512 . Appendix E provides further details on the formulation and construction of this synthetic task. Following Olsson et al. (2022) , we train two layer models with a Transformer backbone that interleaves sequence mixing and state mixing (MLPs). For each architecture, we sweep four learning rates from log ( − 4 ) \log(-4) to log ( − 2 ) } \log(-2)\} for each architecture, and report maximum test accuracy.
Our results, which are summarized in Figure 2 , support two main claims: