Paper: Zoology: Measuring and Improving  Recall in Efficient Language Models (ICLR-2024)

Reviewer comment (Experiments):
The proposed attention hybrid method seems to perform well in the experiment. However, it is not clear how it would perform on a larger scale.

Author response:
**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.

Reviewer comment (Experiments):
As to the training stability and the sensitivity to the training hyper-parameters, I'm wondering how the proposed look-up method and the hybrid method compare with the attention network.

Author response:
**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!

Reviewer comment (Writing):
The description of definition 3.1 could be improved; a simple example would be quite helpful.

Author response:
**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.

Reviewer comment (Writing):
In the appendix, `N` sometimes means the entire sequence length or the number of triplets in MQAR; this confusion is exacerbated by the fact that the meaning changes in the same theorem/proof.

Author response:
_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.

Reviewer comment (Theory):
Page 6 states "[e]ach BaseConv layer uses O(Nd) parameters and ... O(Nd) operations". I believe it uses `O(Nd + d^2)` parameters and `O(d N log N + Nd^2)` FLOPs. This makes me believe that the rest of the theoretical results may have to be carefully revisited by the authors.

Author response:
_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.

Reviewer comment (Theory):
Page 26, Proof of C.19 — It's unclear how `Q[i, :]` can be set to a zero vector when `i \notin Q` (also boldface 0 suffices to express a vector; no need for a superscript d; also, d is defined as `log(C)` in the same page but `C` is a set — should be `d = log(c)` small c since `c := |C|`). Similarly for `K` and `V`, because QKV is actually a linear projection of the input, and linear projection cannot implement this non-linear operation of masking some of the activations out.

Author response:
_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.

Reviewer comment (Writing):
There's a sentence that reads `... where Q, K, V ... are positional embeddings ...`. They are different projections of u.

Author response:
_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.

Reviewer comment (Theory):
The calculation of the percentage of gap due to AR hits could be better motivated and justified in the appendix.

Author response:
**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.

Reviewer comment (Novelty):
The novelty of the paper is limited. The lack of AR ability of State Space Models (which is a special kind of long convolution model with embedded recurrency) has been analyzed in the H3 paper [1] through synthetic data. The proposed architectural modification of hybridization is a simple replication of the Hybrid-H3.

Author response:
**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.

Reviewer comment (Experiments):
The scale of the experiments is limited. The authors only empirically examine their hypothesis for models with the size up to 360M number of parameters. It is not clear whether their claims still holds empirically given the shrinking trend of the performance gap that can already be observed under the current setting.

Author response:
**Experiments with increased model sizes.**

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

Reviewer comment (Reproducibility):
The paper does not provide important technical details for reproducibility. The implementation details of the proposed modification of selectively look-up is missing.

Author response:
**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.

Reviewer comment (Novelty):
The research problem that the authors are trying to solve has been alleviated with existing [1,2] or emerging solutions [3,4]. The authors do not examine the empirical AR ability of data-dependant convolutions by claiming technical difficulties, but there does exist data-dependant SSMs, such as Liquid-S4 [2], that support causal language modeling. Not to mention the latest GCM, Monarch Mixer [3], that also supports causality. On hybridization, a previous work [4] on dynamic input subset selection for attention modules has also been proposed for efficiently combining SSMs with attention. The authors should consider comparing the proposed architectural modifications with these works to avoid being outdated upon publication.

Author response:
**Comparisons to additional architectures shows the AR gap vs. attention remains large.**

To address your questions, we evaluate H3 \[Dao et al.], Liquid state space models \[Hasani et al.], M2 \[Fu et al.], and SeqBoat \[Ren et al.] on the Pile language modeling corpus. We also include a new related works section to contextualize our contributions. 

_Hungry-hungry hippos (H3)._ We evaluate H3 across parameter scales and find that the trends match those of the Hyena architecture – MQAR explains between 63.2% and 88.4% of the gap. We evaluate purely using the H3 layer, not the hybrid architecture, to focus on the behavior of the gated-convolution components. We further evaluate H3 on the MQAR synthetic task that is used throughout our submission and observe empirically that the trend closely follows Hyena and RWKV. We have updated Table 1 and Figure 2 to reflect the new experiments.

We also note that the H3 architecture is closely related to the Hyena architecture. They differ only in the parameterization of the short and long convolutions. Each Hyena/H3 layer contains a short convolution followed by a long convolution that is sandwiched between gating. Given these similarities, we originally included just one of these two architectures, but agree with the suggestion that including the explicit comparison is important for readers. 

Below, recall that we measure the % Explained using the log-perplexities (justified in Appendix C.1). We report perplexities and log-perplexities in parentheses.

|                |              |             |              |             |
| -------------- | ------------ | ----------- | ------------ | ----------- |
| Model          | Overall PPL  | AR PPL      | Non-AR PPL   | % Explained |
| Attention 72M  | 12.99 (2.56) | 2.41 (0.88) | 14.76 (2.69) | -           |
| H3 72M         | 15.8 (2.76)  | 13.9 (2.63) | 15.9 (2.77)  | 63.2%       |
| Attention 125M | 11.01 (2.40) | 2.16 (0.77) | 12.45 (2.52) | -           |
| H3 168M        | 12.1 (2.49)  | 6.75 (1.91) | 12.60 (2.53) | 88.4%       |
| Attention 360M | 9.44 (2.25)  | 1.98 (0.69) | 10.62 (2.36) | -           |
| H3 357M        | 10.38 (2.34) | 4.81 (1.57) | 11.00 (2.40) | 65.8%       |

_Monarch-Mixer (M2, Fu et al.)._ Next, we note that M2 is an implementation of the Hyena architecture that is optimized for higher hardware utilization. M2 uses the exact same sequence mixer as Hyena, but implements the convolutions with monarch matrix multiplications (instead of \`torch.fft\`). Therefore, the analysis in our main paper pertains to M2 as well. Since Hyena lacks sufficient data-dependence for MQAR, so do current M2 architectures.

We note that the M2 authors have released code for training BERT-style language models so we train BERT with M2 vs. BERT with attention from scratch at the 130M parameter scale at a 1024 sequence length. As expected, we find that there is an 4.6% overall quality gap in the masked language modeling accuracy and 58.5% of this can be explained by the gap on repeated bigram tokens. 

We use the exact same training configurations provided in the M2 codebase. However, we note that the M2 paper compares to Devlin et al. BERT numbers. Devlin et al. trains on  Wikipedia and Books pretraining data with a 15% masking probability, while the numbers reported in the M2 paper are with C4 pretraining data and 30% masking. For fair comparison, we train both BERT and M2-BERT models using consistent data and masking probabilities (C4, 30%) for this analysis. Overall, M2 does not address the MQAR problem as it uses Hyena under the hood, supported further by our experiments. 

_Liquid state space models (Hasani et al.)._ We evaluate Liquid SSM upon suggestion and include this in Table 3 of the revision. Liquid SSM introduces a new convolution filter that is composed of the original input-independent SSM matrices as well as correlation terms between the original input tokens. The correlation terms are included to make the filter data-dependent. We evaluate Liquid SSM on the Pile after training for 5B tokens and find that quality lags attention by 4.4 ppl. 

Below, recall that we measure the % Explained using the log-perplexities (justified in Appendix C.1). We report perplexities and log-perplexities in parentheses.

|                       |              |              |              |             |
| --------------------- | ------------ | ------------ | ------------ | ----------- |
| Model after 5B Tokens | Overall PPL  | AR PPL       | Non-AR PPL   | % Explained |
| Attention (125M)      | 12.37 (2.52) | 2.32 (0.84)  | 14.04 (2.64) | -           |
| Liquid SSM  (145M)    | 16.80 (2.82) | 22.75 (3.12) | 16.42 (2.80) | 52.6%       |  

_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. 

|     |                     |                                  |
| --- | ------------------- | -------------------------------- |
|     | With Auxiliary Loss | Without Auxiliary Loss (SeqBoat) |
| PPL | 11.134              | 11.634                           |


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.