Zoology: Measuring and Improving Recall in Efficient Language Models — part 4/14

Cleaned-text code points [46340, 61755). Bibliography omitted here; source.md includes it.

• Input-Dependence via Gating . Since pure convolution architectures with input-independent filters struggle to perform associative recall, an important task in in-context learning [ Olah, 2022 ] , subsequent work proposed to introduce input-dependence by adding a gating or element-wise multiplication operation to the architecture, where both the inputs to the operator are defined in terms of the input [ Fu et al., 2023a , Poli et al., 2023a ] . The multiplication is generally y = σ ⁡ ( W ​ x ) ⊙ x y=\sigma(Wx)\odot x . The filters remain input-independent. These architectures demonstrate promising results on the associative recall synthetics proposed in prior work Fu et al. [2023a] , Poli et al. [2023a] and provide large downstream improvements over S4. Prior work suggests the models match attention in quality [ Poli et al., 2023a , Fu et al., 2023b , Peng et al., 2023 ] and this has led to increasing use of these architectures [ Nguyen et al., 2023 , toe, 2023 ] , however our results suggest there is still a sizable gap to attention. We show the gap is largely due to the models’ recall abilities and capture this behavior in a novel synthetic task. We theoretically show that dimension scales in sequence length when solving MQAR with gated convolutions. Selective S4 Wang et al. [2023] learns which positions to mask prior to passing the input to the next long convolution (S4) layer. This resembles gating (which also controls which information flows forwards to the next layer).

• Input-Dependent Convolution Architectures Prior work introduces convolution architectures with input-dependent filters [ Yang et al., 2020 , Kosma et al., 2023 ] . For instance, Yang et al. [2020] parametrizes the convolution filter as a linear combination of n n projections of the input. Liquid S4 Hasani et al. [2022] is also motivated by introducing input-dependence to S4. To accomplish this, the work proposes including correlation terms between the input tokens during state mixing. We evaluate Liquid S4 and find the architecture lags attention on the Pile by 4.4 perplexity points in Table 3 . Overall Slices % of gap due to Model Param (M) AR Hits Other Tokens Attention 125 12.37 (2.52) 2.32 (0.84) 14.04 (2.64) — Liquid S4 145 16.80 (2.82) 22.75 (3.12) 16.42 (2.80) 52.6% Table 3: Validation perplexity on the Pile after pretraining on 5B tokens. We report log perplexity with negative log-likelihood in parentheses. AR vs. non-AR token slices are defined in Section 3 .

• Recurrent Architectures The pure long-convolution architectures can also be computed as recurrences. Letting s i s_{i} represent the hidden state at time i i , u i u_{i} be the input at time i i , and A , B , C A,B,C be projection matrices, the recurrence computes: s i + 1 = f ⁡ ( A ​ s i + B ​ u i ) s_{i+1}=f(As_{i}+Bu_{i}) y i = C ​ s i y_{i}=Cs_{i} In a linear RNN like S4, the A A , B B , and C C matrices are input-independent and f f is an identify function. However, recent work proposes input-dependent RNNs (e.g. RetNet Sun et al. [2023] ), where a subset of these matrices is input-dependent (e.g. A = x ​ W a A=xW_{a} ). Note that linear attention mechanisms can also be expressed as input-dependent RNNs [ Katharopoulos et al., 2020b ] . While these architectures improve over gated convolutions in Mqar synthetic experiments and downstream quality on the Pile ( Appendix F ), we observe and prove that the required RNN hidden state dimensionality grows with the number of key-value pairs that the model needs to recall in the sequence ( Section H.6 ). In contrast, attention complexity does not scale with the number of key-value pairs to recall.

• Hybrid Architectures Finally, we can combine architectural components that provide different capabilities. Prior work proposes architectures that hybridize sub-quadratic layers and attention [ Fu et al., 2023a , Ma et al., 2022 , Ren et al., 2023 ] , but does motivate this choice from a mechanistic design perspective. Further, H3 without attention Fu et al. [2023a] and MEGA [ Ma et al., 2022 ] , which uses blocked attention, underperform attention on language modeling. SeqBoat [ Ren et al., 2023 ] shares closest resemblance to our learned selection module. SeqBoat introduces a different module that learns where to use attention however, their model can end up using full attention at a layer, providing quadratic scaling in the worst case. Different from SeqBoat, we use an auxiliary loss to encourage sparsity in the selected attention positions and select the top- k k positions to remain sub-quadratic.

Overall, our work finds that previously proposed sub-quadratic models do not efficiently solve the Mqar task.

## Appendix B Code Listing

In this section, we include code listings for the BaseConv operator. We begin with the a standard BaseConv implementation that uses an explicitly-parameterized long convolution. Below, we also provide the implementation for BaseConv with an implicitly-parameterized long convolution.

In some of our experiments, we interleave BaseConv layers that use explicit short convolution filters (like those in torch.nn.Conv1d ) and with BaseConv layers that use implicit long convolution filters (like those described in Poli et al. [2023a] ).

Below we include the code for BaseConv with an implicit convolution.

## Appendix C Downstream Experimental Details

We use A100 80GB Nvidia GPUs to run all experiments. We use the reference training infrastructure from https://github.com/EleutherAI/gpt-neox for all pretraining runs. The Pile data is tokenized using the GPT2BPETokenizer and all models see the data in the same order.

Below we provide details on the hyperparameters and settings for training each architecture studied in the paper on the real-world Pile data. In Section C.1 we summarize and justify our method for measuring the quality gap due to associative recall between the convolution architectures and attention. In Section C.2 , we provide details on the pure gated convolution architectures, attention baseline, and hybrid architectures studied in the paper.

### C.1 Measuring the Mqar Gap on Real Language Data

Here we summarize 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 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 m , and the attention model M M , we measure the % of the quality gap between m m and M M ascribed to associative recall capability as follows. Let the average log-probability for all AR Hits across validation sequences be l H m l_{H}^{m} and l H M l_{H}^{M} for m m and M M respectively. Let the average log-probabilities of all tokens in the validation sequences be l m l^{m} and l M l^{M} respectively. Let p H p_{H} be the proportion of AR Hit tokens in the validation data. As the final gap ascribed to AR, we report: min ⁡ ( ( l H m − l H M ) ​ p H l m − l M , 1.0 ) \min(\frac{(l_{H}^{m}-l_{H}^{M})p_{H}}{l^{m}-l^{M}},1.0) Shown above, if m m is better than attention ( M M ) overall and M M is better than m 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 produced by a contextual language model. For instance, language may contain bigrams in which one word is swapped by a synonym. As another example, a model may see a sentence such as “The iPhone is outside my budget so I instead purchased an Android phone. … It was much _” and predict “cheaper” for the blank, recalling that the sequence is discussing cost. Our work does not measure such fuzzy (more abstract) recall instances.

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 × 1,250\times during training.

### C.2 Gated Convolution Downstream Architectures

We evaluate over 4 4 previously proposed architectures as well as BaseConv , the theoretically “canonical” representation for gated convolutions, introduced in Section 4 , for a total of 14 training runs. Here we provide details on the hyperaparamters and configurations used for training each architecture. We also provide details on the FLOPs computation.

• Attention [ Vaswani et al., 2017 , Touvron et al., 2023 ] We train using the the specifications in Table 6 . The parameters are sourced from the Transformer implementation in https://github.com/EleutherAI/gpt-neox .

• Hyena [ Poli et al., 2023a ] We train using the specifications in Table 8 . The parameters are sourced from the Appendix of Poli et al. [2023a] and the implementation is sourced from the provided reference at https://github.com/HazyResearch/safari .

• H3 [ Fu et al., 2023a ] We train using the specifications in Table 10 . The hyperparameters and implementation use the reference at https://github.com/HazyResearch/H3 .

• RWKV [ Peng et al., 2023 ] We train using the specifications in Table 11 . The parameters are sourced from the Appendix of Peng et al. [2023] and the details provided in the reference implementation at https://github.com/BlinkDL/RWKV-LM . We specifically evaluate RWKV-V4.

• Pure long convolution We train using the specifications in Table 13 . We evaluate a simple long convolution based model with no gating as a reference point. While this is a generic architecture, we use the reference implementation and initialziations/regularizations from recent work Fu et al. [2023c] . The implementation is provided at https://github.com/HazyResearch/safari .

• BaseConv We train using the specifications in Table 15 . The implementation, amounting to 19 lines of PyTorch, is shown in Appendix B .

We provide the equations used to compute the FLOPs for each model, letting D D be the model width, H H the head dimension, L L the depth, N N the sequence length, V V the vocabulary size, and B B the batch size. FLOPs equations for different architecture types are provided in Table 7 (attention), Table 9 (Hyena and adapted for H3), Table 14 (pure long convolution), and Table 16 ( BaseConv ). We compute FLOPs for RWKV as in the Appendix of [ Peng et al., 2023 ] , based on the number of linear layer parameters, plus input and language modeling head FLOPs.

##### Hybrid Architectures

In the main paper and appendix, we evaluate a series of architectures with a hybrid of gated convolution and non gated convolution layers. For these architectures, we use the same number of layers as the pure gated convolution architecture (as reported in Section C.2 ). The hybridized (non gated convolution) layers are inserted as replacements of the gated convolution layer. For each architecture, we simply evenly intersperse the two types of layers. We evaluate with the following replacement layers in this work: • Full attention following the specification of attention in Section C.2 .

• Random selection We evaluate sparse attention randomly applied to tokens as introduced in Section 5 .

• Programmatic selection We apply sparse attention on AR hit tokens as introduced in Section 5 . When processing the input sequence, we can causally determine if a bigram of raw token ids is repeated in the sequence to construct the attention pattern.

• Learned selection We learn the positions on which to use attention by introducing a simple linear layer with sigmoid activation to the architecture that is trained to output high scores if attention should be applied to the token. We can take the top- k k positions to control the computational complexity of the layer. This approach is introduced in Section 5 .

We use these protocols to validate that input-dependence suffices to address the associative recall gap and highlight that there are varied approaches to incorporating input-dependence in an architecture.

## Appendix D Extended Downstream Analysis of Mqar

In this section, we provide additional analysis of how Mqar manifests in real language data. We extend our discussion of associative recall on the Pile Gao et al. [2020] . We also perform analysis on sources in RedPajama including ArXiv and StackOverflow Together [2023] .

### D.1 Additional Discussion of Mqar in the Pile

The Pile is a widely popular language modeling corpus Gao et al. [2020] . 1. First, in Section D.1.1 , we provide several real examples of how associative recall occurs in the Pile to demonstrate it’s role in language modeling. We color code the tokens to highlight differences in the next token predictions from Attention, RWKV and Hyena models.

2. Next, in Section D.1.2 , we explain where associative recall hits tend to occur in sequences. This is useful to guide the design of sequence mixers — if associative recall tends to occur in specific places, the input-dependent sequence mixer may not need to compute all N 2 N^{2} token-to-token interactions for every sequence.

#### D.1.1 Pile Examples

For the examples below, the legend is: tokens are colored as both correct , both incorrect , Attention correct and Hyena incorrect , Attention incorrect and Hyena correct . For tokens where the models disagree, the predictions of each model are provided in parentheses.

while lunch -ing at the Ma -ison Ber -gey b -ist -ro near his apartment : he had been mus -ing about (rwkv= about, attn= on) the … (723 tokens) … the young waitress -’s sigh at the Ma -ison Ber (rwkv= Bl, attn= Ber) -gey (rwkv=-nd, attn=-gey)

The second (rwkv= first, attn= second) section is all about Pixar Fest , and the (rwkv= third, attn= the) final section is all (rwkv= about, attn= all) about Pixar Pier . … (480 tokens) … -If (rwkv=-Disney, attn=-If) there wasn -âĢ -t enough Pixar at Disneyland , Pixar Fest (rwkv= would, attn= Fest) is (rwkv= at, attn= is) coming to the Disneyland Resort on April 13 , 2018 .

David Mus -a P -id -cock is an indigenous English revert to (rwkv=-ant, attn= to) Islam , (rwkv=,, attn= who) who formed the Islamic Party of Britain in 1989 , at London Central Mosque (rwkv= Mosque, attn= University) . (rwkv=., attn= in) …( 711 tokens )… -I have just found out that David Mus (rwkv= Cameron, attn= Mus) -a P (rwkv= (, attn= P) -id (rwkv=-asha, attn=-id) -cock (rwkv=-ham, attn=-cock) is (rwkv= (, attn= is) speaking at a Yorkshire Forum debate today in Bradford with