Zoology: Measuring and Improving Recall in Efficient Language Models — part 1/14
Cleaned-text code points [0, 15717). Bibliography omitted here; source.md includes it.
# Zoology: Measuring and Improving Recall in Efficient Language Models
###### Abstract
Attention-free language models that combine gating and convolutions are growing in popularity due to their efficiency and increasingly competitive performance. To better understand these architectures, we pretrain a suite of 17 attention and gated-convolution language models, finding that SoTA gated-convolution architectures still underperform attention by up to 2.1 perplexity points on the Pile. In fine-grained analysis, we find 82% of the gap is explained by each model’s ability to recall information that is previously mentioned in-context, e.g. Hakuna Matata means no worries Hakuna Matata it means no → \rightarrow ?? . On this task, termed associative recall , we find that attention outperforms gated-convolutions by a large margin: a 70M parameter attention model outperforms a 1.4 billion parameter gated-convolution model on associative recall. This is surprising because prior work shows gated convolutions can perfectly solve synthetic tests for AR capability. To close the gap between synthetics and real language, we develop a new formalization of the task called multi-query associative recall ( Mqar ) that better reflects actual language. We perform an empirical and theoretical study of Mqar that elucidates differences in the parameter-efficiency of attention and gated-convolution recall. Informed by our analysis, we evaluate simple convolution-attention hybrids and show that hybrids with input-dependent sparse attention patterns can close 97.4% of the gap to attention, while maintaining sub-quadratic scaling. Our code is accessible at: https://github.com/HazyResearch/zoology .
## 1 Introduction
Two advances – gating and long convolutions – have catalyzed a wave of excitement around gated-convolution language models ( Fu et al., 2023a ; Ma et al., 2022 ; Wang et al., 2022 ; Poli et al., 2023a , inter alia.) . These architectures combine gating ( i.e. element-wise multiplication) with long convolutional filters ( i.e. the length of the sequence) to enable interactions between distant tokens ( Dauphin et al., 2017 ; Gu et al., 2021 ) . Recent work suggests that these models, which exhibit better asymptotic scaling in input sequence length than attention, can match attention in language modeling quality ( Poli et al., 2023a ; Peng et al., 2023 ; Fu et al., 2023b ) .
We pretrain and evaluate 17 language models across 4 scales (70M - 1.4Bn) and 5 architectures on the same data and infrastructure setup. Surprisingly, we find that there is still a perplexity gap of up to 2.1 points between state-of-the-art convolution-based architectures and strong Transformer baselines in language modeling on the Pile (Table 1 ). Through fine-grained analysis, we find a single, simple capability is responsible for much of the gap: recalling information seen in-context. Consider the example below, where some tokens can be predicted by recalling an earlier association:
Hakuna Matata! ⏟ Key-Value It means no worries ⏟ Key-Value for the rest of your days! Hakuna ⏟ Query Matata ⏟ AR Hit means no ⏟ Query → worries ⏟ AR Hit \underbrace{\text{Hakuna Matata!}}_{\mathclap{\textbf{Key-Value}}}\text{It means}\underbrace{\text{no worries}}_{\mathclap{\textbf{Key-Value}}}\text{for the rest of your days!}\underbrace{\text{Hakuna}}_{\mathclap{\textbf{Query}}}\underbrace{\text{Matata}}_{\mathclap{\textbf{AR Hit}}}\text{means}\underbrace{\text{no}}_{\mathclap{\textbf{Query}}}\rightarrow\underbrace{\text{worries}}_{\mathclap{\textbf{AR Hit}}}
We find that errors on “AR Hits” ( e.g. worries above) account for 82% of the perplexity gap to attention on average, despite only representing 6.4 % 6.4\% of all tokens in the Pile dataset. 1 1 1 We measure AR on real data using a simple heuristic: n n -gram tokens that are repeated in context. A 70M parameter Transformer can predict AR Hits better than a 1.4Bn parameter Hyena gated convolution model ( 20 × 20\times larger) (Table 1 , Table 5 ). The AR gap persists at the 7Bn parameter scale when comparing RWKV and Llama-2 ( Appendix G ).
This task, associative recall (AR), has a long history in machine learning ( Graves et al., 2014 ; Ba et al., 2016 , inter alia.) ( Section A.1 ). Prior work argues that a model’s ability to perform AR is predictive of in-context learning quality ( Elhage et al., 2021 ; Olsson et al., 2022 ) . As a result, AR has been adopted as a tool in designing new architectures and prior work has shown that gated convolution architectures match attention on synthetic AR tasks used as proxies for real language modeling ( Fu et al., 2023a ; Poli et al., 2023a ; Lutati et al., 2023 ) . For this reason, the downstream AR perplexity gaps are surprising.
Through measuring recall on real data, we learn the key disparity is that prior synthetic formulations assume there is one query per input, at a fixed position in the sequence, where tokens come from a small vocabulary size ( e.g. | V | < |V|< 50, less than model dimension). Yet, language modeling often requires performing multiple recalls ( e.g. for both “Hakuna Matata” and “no worries” above, in a single forward pass), at varying positions , with tokens from a large vocabulary (larger than model dimension). We thus propose the study of multi-query AR ( Mqar ). Compared to the prior AR formulations, Mqar better captures the persisting quality gaps on synthetic and real world data (Section 4 ). However, it is not clear why Mqar elucidates the gap.
We formalize the the Mqar gap across the architecture classes. Gated-convolutions process variable sequences using fixed filters defined by the model weights rather than as functions of the input data (See Figure 1 ). We find it is inefficient for gated-convolutions to perform the variable distance token-to-token interactions (e.g. at a distance of 10 10 tokens for Hakuna Matata and 9 9 for no worries ) in a parameter and FLOPs efficient way compared to attention. Attention achieves input-dependence since it computes all token-to-token interactions when determining how to mix information in the sequence. We formally describe the limitation of gated convolutions via theory and experiments in the rest of the paper.
We first introduce a simple operator, BaseConv , that can provably simulate the class of architectures built from gating and convolution primitives. This includes architectures such as H3, Hyena, RWKV, and RetNet. We show with theory and experiments that the model dimension for BaseConv (and thus the aforementioned architectures) to solve Mqar grows with the input sequence length (Theorem 4.4 ) while attention can solve Mqar with model dimension independent of sequence length (Proposition 4.3 , Section 4.3 ). 2 2 2 Note that the runtime of Attention is quadratic, but gated convolutions is near linear. In practice, gated-convolutions appear to encode approximate solutions to AR that support only a subset of token-interaction distances (which affects the ability to identify matching keys given Mqar queries). While theoretically analyzing deep learning architectures is challenging ( Hahn, 2020 ; Merrill et al., 2022 ; Keles et al., 2023 ) , the fact that gating and convolutions are polynomial operations facilitates our precise analysis.
We show that input-dependent sequence mixing is important to solve Mqar efficiently. The scaling for gated convolutions with input-independent convolutions is undesirable so we next ask which architectural choices close the gap. We show that data-dependent sequence mixing helps an architecture solve Mqar efficiently (Theorem 4.5 ). The model needs to adapt the sequence mixing weights based on the token-interaction distances required for each new example.
Several architectural modifications could satisfy the input-dependence property. Based on our analysis, we evaluate minimal modifications that only add input-dependent operations on exact-match repeated bigram tokens (our heuristic for measuring tokens that require recall). Note that language models often need to perform recall between fuzzier substrings (e.g. synonymous bigrams) or higher-dimensional concepts. However, we show that simply inserting input-dependent operator — e.g., a convolution filter that shifts the sequence based on the bigram positions or sparse attention placed only on repeated bigram positions — to the BaseConv architecture at < 10 % <10\% of layers suffices to outperform the Transformer baseline on Pile language modeling ( Section 5 ) and succeed on Mqar synthetic tasks ( Section 4.3 ). Moreover, this closes > 80 % >80\% of the Mqar perplexity gap on the Pile validation data. Finally, we prototype solutions that learn the positions at which to use input-dependent operations, validating that they also close the gap to attention on the Pile.
In this work, we analyze the increasingly-popular convolution-based architectures and identify fundamental limitations. We hope the Mqar framework and our analysis of the role of input-dependence for Mqar inform the design of future architectures. We release our code for reproducability: https://github.com/HazyResearch/zoology .
## 2 Background and Preliminaries
In this section, we describe the setting, introduce notation, and discuss important related work. See Appendix A for a broader discussion of related work.
Language modeling. We study auto-regressive language models trained on the task of next token prediction ( Gao et al., 2020 ) . Given a sequence of N N tokens 𝒙 = { x 0 , … , x N − 1 } \bm{x}=\{x_{0},...,x_{N-1}\} drawn from a vocabulary C C , the model outputs a probability distribution over C C for each of x i x_{i} given the preceding tokens P ( x i | x 0 , … , x i − 1 ) \mathrm{P}(x_{i}|x_{0},...,x_{i-1}) . The language models in this work share the same high-level architecture. First, each token x i x_{i} in the input is embedded in d d -dimensional space yielding a matrix 𝒖 ∈ ℝ N × d \bm{u}\in\mathbb{R}^{N\times d} . Next, 𝒖 \bm{u} is passed through a stack of L L layers, with layer ℓ \ell outputting 𝒖 ℓ ∈ ℝ N × d \bm{u}^{\ell}\in\mathbb{R}^{N\times d} . Finally, the embeddings 𝒖 L \bm{u}^{L} output by the last layer are mapped back to logits over C C with a linear projection. Each layer transforms 𝒖 \bm{u} with a sequence mixer ( e.g. attention) followed by a state mixer ( e.g. MLP). Unless specified, our models adhere to the implementation details of the LLaMA architecture (except the sequence mixer, which we vary throughout) ( Touvron et al., 2023 ) .
Sequence mixers. Our work evaluates how the choice of sequence mixer affects the quality and behavior of language models. Most sequence mixers aggregate the token embeddings in a sequence via a weighted sum. For example, 𝒚 [ i , : ] = ∑ j = 0 N − 1 ω ( i , j ) 𝒖 [ j , : ] ) \bm{y}[i,:]=\sum_{j=0}^{N-1}\omega(i,j)\bm{u}[j,:]) , where ω \omega is a function outputting scalar weights. We study the differences between two classes of sequence mixers, discussed next.
Attention. ( Vaswani et al., 2017 ) We review attention, the de facto language model sequence mixer. An attention layer is parameterized by three learnable projection matrices 𝐐 , 𝐊 , 𝐕 ∈ ℝ N × d \mathbf{Q},\mathbf{K},\mathbf{V}\in\mathbb{R}^{N\times d} . To compute the output 𝒚 \bm{y} given inputs 𝒖 \bm{u} , attention applies the projections to the input: 𝒒 = 𝐐𝐮 \bm{q}=\mathbf{Q}\mathbf{u} , 𝒌 = 𝐊𝐮 \bm{k}=\mathbf{K}\mathbf{u} , 𝒗 = 𝐕𝐮 \bm{v}=\mathbf{V}\mathbf{u} . The projected embeddings are aggregated according to: 𝒚 = softmax ( 1 d 𝒒 𝒌 ⊤ ) 𝒗 \bm{y}=\text{softmax}(\frac{1}{\sqrt{d}}\bm{q}\bm{k}^{\top})\bm{v} (shown in Fig. 1 ) in 𝒪 ( N 2 d ) \mathcal{O}(N^{2}d) time, which is expensive for long sequences (large N N ).
Gated-convolutions. A more efficient alternative to attention is the convolution, which is defined as 𝒚 [ i , : ] = ∑ j = 0 N − 1 𝐤 [ j , : ] ⊙ 𝒖 [ i − j , : ] \bm{y}[i,:]=\sum_{j=0}^{N-1}\mathbf{k}[j,:]\odot\bm{u}[i-j,:] where the kernel 𝐤 ∈ ℝ N × d \mathbf{k}\in\mathbb{R}^{N\times d} is a learnable weight matrix. Convolutions can be computed in time O ( N d log N ) O(Nd\log N) using the Fast Fourier Transform (FFT) and the convolution theorem: 𝒚 = 𝐮 ∗ 𝐤 = FFT − 1 ( FFT ( 𝐮 ) ⊙ FFT ( 𝐤 ) ) \bm{y}=\mathbf{u}\ast\mathbf{k}=\mathrm{FFT^{-1}(FFT(\mathbf{u})\odot FFT(\mathbf{k}))} ( Cooley and Tukey, 1965 ) . While purely convolutional architectures match or outperform attention in certain domains ( e.g. vision ( Tay et al., 2022 ; Gu et al., 2021 ) , audio ( Goel et al., 2022 ) , and time-series ( Zhang et al., 2023 ) ), they trail by a large margin on language ( Fu et al., 2023a ) .
Recent work has closed much of this gap by combining convolutions with gating ( i.e. elementwise multiplication of the input with a transformed version of itself). Gating was first proposed in the context of convolutions by Dauphin et al. (2017) , but more recently it has played a central role in state-of-the-art sub-quadratic language models, some of which claim attention-level quality ( Fu et al., 2023a ; Wang et al., 2022 ; Poli et al., 2023a ; Peng et al., 2023 ; Zhai et al., 2021 ; Fu et al., 2023b , inter alia) . Though they may appear different on the surface, these sub-quadratic architectures can all be expressed in terms of convolutions and gating. See Section H.2 for detailed descriptions of these architectures and their similarities. This work analyzes the differences between attention and the broad class of gated convolution mixers.
## 3 Identifying the associative recall problem
In this section, we measure the perplexity gap between gated convolutions and attention and show that single skill termed associative recall accounts for 82% of the gap on average. This is surprising because prior work shows gated convolutions solve a synthetic version of associative recall perfectly. Informed by our analysis, we define a new synthetic formulation that better reflects real data. This task facilitates our analysis of why the gap occurs (Section 4 ) and how to fix it (Section 5 ).
### 3.1 Fine-grained analysis of downstream quality
##### Perplexity Gap
We pretrain a suite of large language models with different sequence mixers across 3 scales (70M-360M) for 10B tokens on the standard Pile language modeling setting using the EleutherAI GPT-NeoX training infrastructure ( Gao et al., 2020 ; Andonian et al., 2023 ) . In the main paper, we compare attention to three state-of-the-art gated convolution sequence mixers: H3, Hyena, and RWKV ( Fu et al., 2023a ; Poli et al., 2023a ; Peng et al., 2023 ) . 3 3 3 RWKV is commonly referred to as an RNN. We show it can be viewed as a convolution ( Section H.2.2 ). We further include a pure long-convolution model to underscore the importance of gating. In Appendix F , we include results on additional sequence mixers ( Hasani et al., 2022 ; Sun et al., 2023 ) . We use a strong Transformer attention baseline with rotary embeddings and SwiGLU MLPs following the Llama architecture ( Touvron et al., 2023 ) . We also use this strong training recipe when training the attention-free sequence mixers. For experimental details and hyperparameters for all architectures see Appendix C .
Across scales, we find that attention outperforms the gated convolutions by at least a third of a perplexity point on average: the minimum gaps are + 2.14 +2.14 , + 0.59 +0.59 , + 0.35 +0.35 PPL at 70M, 160M, and 360M parameter scales, respectively. We report overall test perplexity in Table 1 . Though these gaps are relatively small on average, models may still perform very differently on different subsets of data ( Eyuboglu et al., 2022 ) .
##### Associative Recall Perplexity