Zoology: Measuring and Improving Recall in Efficient Language Models

Thank you for your thoughtful review! We appreciated that you found the work novel and insightful, and the problem impactful as gated convolution sequence models are receiving increasing adoption in the ML community. Here we address your questions. 

Presentation

We take note the presentation score of 3 and have taken several steps to improve the overall delivery of our work. We hope these changes are helpful! The changes include: 

1. Section 3 and Appendix C.1: We update the sections to describe how associative recall quality is measured in the real language modeling data. We also clarify the definitions and add simple examples of prior vs. our proposed recall problem. 
2. Theory: We thoroughly checked for consistency with respect to notation and further clarified the assumptions underlying our theoretical results.
3. Section 5.1: We provide clarification of the input-dependent architectures we evaluate.
4. Appendix A: We include a new detailed related works section that discusses the history of associative recall in machine learning and an extensive list of previously proposed architectures to better contextualize our work.
5. Appendix B: We provide a simple PyTorch implementation of the BaseConv architecture to complement the theoretical statements. 
6. Appendix E: We provide a new section and algorithm box to document exactly how the MQAR synthetic data is constructed. We also provide more details on our setup for the synthetic experiments.

Please let us know if you have any additional suggestions on how we can improve the presentation and thank you for your help! 

Scaling architectures. 

Thank you for your question on scaling performance! In response, we have conducted three new experiments to better understand the scaling trends. We kindly direct you to the discussion of these experiments that is included in the main response to all reviewers as well as Appendix G of the revised submission. 

In these experiments, we show that simply scaling the gated convolution architectures does not solve the MQAR problem. We scale up to 7B parameter models in our new experiments. We hope these experiments build confidence that hybrid solutions continue to help at scale since attention layers can solve MQAR, complementing the gated-convolution backbone which struggles to solve MQAR efficiently.

Sensitivity to hyperparameters. 

You raised the important question of how the training stability and sensitivity of the selective attention methods compare to the baseline attention network. We do not apply any hyperparameter tuning (e.g., LR, Optimizer, Batch Sizes, etc.) when producing these results. We simply use the backbone gated convolution models with hyperparameters provided by the prior work and replace a few layers with the hybrid layer. Overall, training is stable and the hybrid models work even without tuning! 

Further, to back up these results we are including our training code in the revised submission. We will also release the pretrained checkpoints upon publication. Please let us know if you have further questions about this!

Thank you for your detailed review! The feedback led to significant improvements in our work. We appreciate that you found the downstream analysis insightful and here we address your questions.

Theoretical corrections, clarifications, and updates

Thank you for raising questions about the theoretical pieces. We apologize for the confusion from our original submission and have thoroughly examined the revised submission. 

Page 26. Proof of C.19. QKV matrices and dimensions. – Now page 47, Appendix H.7.

We have added clarifications on how we assumed the input u to be encoded for the proof (we have made this explicit as well as added some justification on why the encoding assumption is reasonable). The proof is updated where Q, K and V are linear projections. While fixing the argument we realized that we needed d=3c so we have made that correction as well.

Parameters and operations for BaseConv layers.

We had proven that we only need $\tilde{O}(d)$ parameters for the projections to achieve the tight general arithmetic circuit result in Appendix C.4 of our original submission (now Appendix H.5), where the projections use Kaleidoscope matrices for all architectures [Dao et al., 2020]. In all experiments, the linear projections are dense, but for all the theory results in App H, the linear projections use the restricted poly-log Kaleidoscope matrices. Doing so is what gave the claimed O(Nd) parameter count and runtime on Page 6. 

We did not specify this clearly in the main paper and update this in our revision (page 6). We also included an explicit proposition in the appendix (Proposition H.6) that argues the claim on a single BaseConv layer. We also have added explicit statements in the proofs to show that using the near linear parameters Kaleidoscope matrices are sufficient for our proof.

Clarification on the meaning of N. We had used N to represent both the sequence length and the number of (key, value, query) tuples. We have updated the number of tuples to read N/3 in the revised copy. We apologize for the confusing notation and thank you for pointing this out.

Shape of inputs in Proposition 4.3. As noted by the review, we have corrected proposition 4.3 to read $u \in \\{ 0, 1 \\}^{N \times 3c}$. 

Beyond the comments in the review, we thoroughly checked for consistency with respect to notation and further clarified the assumptions underlying our theoretical results.

Clarifying Definition 3.1 and Descriptions of MQAR

We have updated the manuscript in the following ways to address this feedback:

1. We streamlined definition 3.1 and added a simple example per the reviewer suggestions. We also include a simple example of how AR is formulated in prior work for clear comparison. 
2. We add Appendix D containing several examples of MQAR that we sourced from real language modeling data.
3. We add Appendix E (Algorithm 1) detailing exactly how MQAR synthetic data gets constructed.  

We hope these collectively clarify the MQAR contribution.

Explaining the procedure for measuring the AR gap

In response to the review, we add Appendix C.1. to justify our method for computing the amount of quality gap between the convolution and attention models that is ascribed to associative recall capability (e.g., in Table 1):

1. Given an input sequence, we identify recurring bigrams (i.e. bigrams that have already appeared in the sequence at a prior position). Since bigrams that appear frequently during training may be memorized by the model, rather than requiring the model to perform recall at inference-time, we only measure AR log-probabilities with respect to bigrams that are seen fewer than a threshold number of times during training. The threshold used in all the experiments in our submission is $1,250$ training occurrences in the 10B tokens of pretraining data.
2. We measure the log-probability assigned to the true bigram completion. This bigram completion is referred to as an AR Hit in our work. This protocol assumes that the model can produce the completion by recalling the prior occurrence of the bigram in the sequence. 
3. For the model being evaluated $m$, and the attention model $M$, we measure the % of the quality gap between $m$ and $M$ ascribed to associative recall capability as follows. 
4. Let the average log-probability for all AR Hits across validation sequences be $l\_{H}^m$ and $l\_{H}^M$ for $m$ and $M$ respectively. Let the average log-probabilities of all tokens in the validation sequences be $l^m$ and $l^M$ respectively. Let $p\_{H}$ be the proportion of AR Hit tokens in the validation data. As the final gap ascribed to AR, we report:

$\min (\frac{(l\_{H}^m - l\_{H}^M) p\_{H}}{l^m - l^M}, 1.0)$

Shown above, if $m$ is better than attention ($M$) overall and $M$ is better than $m$ at AR, we ascribe 100% of the gap to AR.

We briefly discuss two important decisions in this protocol. 

First, we only measure explicit bigrams, i.e. bigrams are identified based on token ids in the sequence. However, intuitively, models may also perform associative recall between related concepts. For instance, language may contain bigrams in which one word is swapped by a synonym. We note that our work does not measure recall in higher dimensional recall in favor of a more transparent procedure. 

Next, we measure the gap based on log-probabilities rather than perplexity. This is simply because we want to make our metrics independent of the number of tokens in each of the slices of the validation set. Approximately 6.4% of validation tokens are AR Hits with the threshold set to consider bigrams seen less than $1,250 \times$ during training.

Additional Comments

Overall we have made clarifications to the submission, which we hope makes the work more understandable.

We include Appendix C and E, which provide several additional details on our experimental protocols for synthetic and downstream training across models. We also include a zip file of our code for reproducing our results with all models. Upon publication, we will also release all the pretrained checkpoints analyzed in our paper. We hope these materials build confidence in the reproducibility and thoroughness of our work.

References

[1] Dao et al., Kaleidoscope: An Efficient, Learnable Representation For All Structured Linear Maps, 2020.

SeqBoat (Ren et al.). We thank the reviewer for this pointer! The simple learned selection function we evaluate in Section 5 is similar to SeqBoat. One major difference is that we explicitly limit the number of activated tokens to ensure sub-quadratic scaling, while SeqBoat allows all tokens to activate. This occurs in some of experiments in the SeqBoat (See Figure 3 of their paper). While SeqBoat use the sparse attention at every layer, we simply replace three of the BaseConv layers with the sparse attention mechanism. Also, unlike SeqBoat, we use an auxiliary loss that guides the model to sparsely select attention positions. Based on the reviewer’s pointer, we evaluate SeqBoat’s approach of using no auxiliary loss and find that it increases perplexity by 0.5. 

We include a discussion of our differences from this work in Appendix A.3 (Related Works). We also note that our objective in Section 5 is to analyze simple, common-sense approaches for closing the gap sub-quadratically. We do not claim that this simple method is novel or that we are the first to propose it, but rather our contribution lies in the explanation of why these mechanisms work. 

Finally, please see synthetic MQAR and downstream Pile experiments on additional architectures (RetNet, Sliding window attention like Mistral, and Blocked attention) in Appendix F.

Thank you for your thoughtful review! The experiments you suggested and questions you raised helped us greatly improve the work! We appreciated that you found our work comprehensive across theory, synthetic experiments, and downstream experiments. Below, we address your questions and highlight new experiments on an additional set of 6 baseline architectures (H3, Liquid S4, M2, Sliding window attention as in Mistral, Block attention, and RetNet) and a new analysis of associative recall as model size increases (up to the 7B parameter scale).

Novelty and differences from prior work (e.g H3)

Thank you for your thoughtful comments about H3. You are correct that the H3 work studies synthetic associative recall tasks and previously proposes the use of hybrid models. However, we make fundamentally different claims. They argue that gated-convolutions (like H3) can solve AR. In contrast, our claim is that gated-convolutions cannot solve AR without scaling the hidden dimension with N. To underscore this point, we have added evaluations of H3 in the main paper (Table 1) and show that the trends we documented for other gated-convolutions also hold for H3 (Figure 2). Further, we find that the synthetic task studied in H3 and Hyena don’t reflect recall in real language, so we develop and release a new synthetic better aligned with actual data. 

We use this synthetic to explain why adding attention helps with recall. The H3 paper does not provide this explanation. Finally, in downstream hybrid experiments, we go beyond the application of full quadratic attention layers to show that sparse-attention, localized to potential AR tokens, is sufficient to close the gap to the Transformer baseline. We hope this clarifies our differences and contributions. 

Please also see our summary of contributions in the main response to all reviewers. 

Experiments with increased model sizes.

Thank you for raising this question! Please find our new experimental results in the main response to all reviewers. 

New Related Works Section

Beyond the three architectures discussed above, we include an extended related works section in Appendix A. In response to your feedback that recent or concurrent architectures could alleviate the problem, the section discusses relevant architectures. 

Details for Reproducibility.

In response to your feedback on the reproducibility of our work, we provide several additional details on our experimental protocols for synthetic (details in Appendix E.2) and downstream training (details in Appendix C) across models. We also include a zip file of our code for reproducing our results with all models. Upon publication, we will also release all the pretrained checkpoints analyzed in our paper. We hope these materials build confidence in the reproducibility of our work.

Analyzing off-the-shelf large language models. 

We next evaluate the RWKV-Raven (gated-convolution) and Llama 2 (attention) pre-trained models at the 7B parameter scale. These are both popular models that took a significant amount of effort to train, towards maximizing quality. We find that there is a large gap between RWKV and attention at the 7B scale on AR tokens and that it increases as the model needs to conduct more recalls per sequence (P below). On the non-AR tokens, quality is comparable.

We summarize the experimental protocol for the 7B scale. 

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). We do not have access to the custom training data mixtures used in training RWKV/Llama 2 to measure bigram frequencies, so we use a synthetic test to fairly measure the AR capabilities. 
2. Evaluation dataset. We construct a synthetic dataset where each sequence contains 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 synthetic construction procedure is shown in Algorithm 1 (same procedure as Figure 2) and the sequences are constructed using the models’ vocabulary tokens. 
3. Evaluation details. We evaluate (inference only, no training) on sequence lengths of 1024 tokens. We measure the AR perplexity when the sequences contain $P \in \\{4, 8, 16, 32\\}$ key-value pairs, using 1000 samples per 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. We find RWKV quality degrades with P on the AR slice (blue line), while all other lines remain flat. MQAR remains problematic at scale. We release code to reproduce this experiment.

Analysis of the AR gap with scale

In the main paper, we measure the AR perplexity on downstream data based on bigrams that are seen $< 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\times$ during pretraining (Table 1). If we instead focus on bigrams seen $1\times$ during pretraining, the gap to attention quality is 12.0 / 13.19 perplexity points respectively. The gated convolutions appear to struggle in the regime of rare bigrams that require the model to use the context. 

[1] Xie et al. An Explanation of In-context Learning as Implicit Bayesian Inference, 2021.

Discussion of new experiments

To address questions from 9Su6 and 5EW4 about the size of the MQAR gap as the attention and gated convolution model sizes increase, we include three new sets of experiments:

1. Pretraining scaled up. We pretrain and evaluate new attention and Hyena (gated convolution) models at the 1.4B parameter scale. Even at 1.4B parameters, Hyena still performs worse on AR than attention models 20x smaller (70M parameters). 
2. Analyzing off-the-shelf large language models. We then scale up to 7B pre-trained RWKV-Raven (gated convolution) model downloaded from HuggingFace and Llama 2 (attention) model downloaded from Meta. We note that RWKV is the only gated-convolution architecture that has been scaled up and widely disseminated thus far. Both RWKV 7B and Llama 2 are popular amongst users. 
3. Analysis of the AR gap with scale. Finally, in further analysis of our pretrained checkpoints, we find that for bigrams seen $<1,250\times$ during training the gap between 350M parameter attention and Hyena / RWKV on AR is 1.85 / 1.84 perplexity points respectively (reported in Table 1). However, if we consider bigrams seen $1 \times$, the gaps to attention are 12.0 / 13.2 perplexity points respectively (Appendix G). This suggests that the larger models may get better at memorizing bigrams, but improve less quickly on in-context learning. Overall we find that MQAR remains problematic even as we scale up the gated convolution and attention architectures. We detail these new experiments in Appendix G.

Pretraining scaled up.

We train Hyena and attention models on the Pile at 1.4Bn parameters and observe the following. We observe that the 20x larger 1.4Bn Hyena model still underperforms the 70M parameter attention model on the AR slice. Overall, the % Explained by AR is less than at the smaller scales, but still sizeable. We note that it is also possible for the model to perform AR between higher dimensional concepts (e.g. fuzzy bigram matches) and these are not captured by the bigram heuristic used in our work to measure AR. The heuristic may be less effective for larger models. We include the new results and discussion in Section 3 and Appendix G. Below, recall that we measure the % Explained using the log-perplexities (justified in Appendix C.1). We report perplexities and log-perplexities in parentheses.

Summary of Changes

Revision Summary. In our revision, we made extensive reviewer-suggested changes that strengthen our work and highlight our contributions. These changes include:

1. Experiments on larger scale models [9Su6, 5EW4] (Section 3, Appendix G). In response to comments on the model sizes used in our experiments,  we extend our empirical analysis to 1.4 and 7 billion parameter language models. We find that the AR gap between attention and gated convolutions persists at these scales.
2. Experiments with new architectures [9Su6, 5EW4] (Section 3, Appendix F). We include five new architectures in our analyses including H3, Liquid S4, RetNet, Sliding Window Attention, and Blocked Attention. 
3. Improved presentation and reproducibility [9Su6, 5EW4]. Based on feedback from 9Su6 and 5EW4 on missing technical details, we’ve restructured Sections 2, 3 and 5 for improved clarity. We additionally include code for reproducibility and will publicly release code and pre-trained checkpoints from this work upon publication. 
4. Clarifications on novelty [9Su6]. In response to questions on the difference between our work and H3, we clarify in Section 3.2 how our work (1) makes a fundamentally different claim than H3 and (2) uses a novel formalization of associative recall, MQAR, which remains unsolved by gated-convolutions with constant dimensionality. We also add an extended related works section that (1) details the prior use of associative recall in the literature and (2) details the efficient architectures that have been proposed. 
5. Updates to the theory [xdZW]. In response to concerns, typos, and clarity issues in our theoretical analysis raised by xdZW, we’ve made updates to Proposition 4.3 and the corresponding proof H.27, though the overarching claim and proof structure are unchanged. We also provide additional discussion in all proof sections to improve clarity for readers. Finally, we also made a thorough pass through our proofs checking for correctness and typos. We also note that we provide empirical validation of our theoretical statements, using simple synthetics, supporting the correctness.