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

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

-The most common map elements (hyena= in, attn= elements) in Sub -Space are prizes , or ” -g (hyena=-points, attn=-g) -reens ” (hyena=-",, attn=") ( -for their green color -). Pri -zes allow players to upgrade their ships and gain (hyena= other, attn= gain) special weapons or abilities (hyena= abilities, attn= upgrades) . While prizes are generally plent -ifully scattered throughout the map (hyena= map, attn= game) , the upgrades or abilities (hyena= bonuses, attn= abilities) they award are randomly selected by the zone . -Energy -Rather than dealing with ammunition counts and hit points separately , Sub -Space combines both of these elements into a single unit of measure : energy . Each ship is equipped with a certain amount of energy , from which (hyena= which, attn= the) it must draw its health as well as its weapons power . … (1,526 tokens) … Speed Zone proved to be less popular than the Jack -pot -/ -Running , Chaos , or ” -flag ” zone games and support was discontinued shortly after Sub -Space (hyena=-space, attn=-Space) went to retail .

Thus far , no systematic study has been published on is -olation (hyena=-olation, attn=-ot) via sequential centrif -ug -ation , and no systematic analysis is … (491 tokens) … Fig . 1 (hyena= 1, attn=-Âł) -is -olation via sequential (hyena= centrif, attn= sequential) centrif -ug -ation .

Miss -ouri (hyena=-ouri, attn=-iss) Southern has had 14 Major League Baseball draft selections since the draft began in 1965 . … (149 tokens) … Fred G . Hughes Stadium (hyena= Stadium, attn= Field) ( -opened in 1975 ) is named (hyena= the, attn= named) after former (hyena= the, attn= former) J -op -lin Globe publisher and (hyena= and, attn= Fred) Missouri Southern (hyena= State, attn= Southern) board of reg -ents member

#### D.1.2 Distribution of Mqar Hit Positions in Sequences

In Figure 3 , we compute and plot the distances between AR hits and its prior bigram (key-value) occurrence in the sequence across the Pile training data. The distances follow a power law distribution where most AR hits are within 100 token positions from the prior bigram occurrence, and a long tail of AR hits requires long-range interactions.

##### Implications

This suggests that architectures which compute token-to-token interactions within windows, where the window size is less than the full sequence length, may suffice to handle Mqar in real data. For instance, sliding window attention may help with AR [ Beltagy et al., 2020 , Sun et al., 2023 , inter alia.] . Prior work leverages the observation that local attention captures most of the token dependencies to improve efficiency Sukhbaatar et al. [2019] .

### D.2 Analyzing Mqar across Varied Language Distributions

We use the HuggingFace sample of the popular RedPajama language modeling corpus Together [2023] to understand how the prevalence of Mqar hits and and ways in which the hits appear vary across language distributions. We analyze text sequences from the following sources: ArXiv papers, Books, C4, Common Crawl, GitHub, Stack Exchange, and Wikipedia.

First we profile the prevalence of Mqar hits by computing the number of repeated bigram occurrences per sequence. Each sequence is 2,048 2,048 tokens in length and exclude bigrams containing NLTK stop words or punctuation from the computation. The prevalence of hits and the distances between hit tokens and the prior bigram occurrence in context vary across distributions as shown in Figure 4 and 5 .

Arxiv, Github, and Stack Exchange contain relatively structured or richly formatted langauge data. We observe that the prevalence of Mqar hits is relatively high in these three sources in Figure 4 . We inspect why these documents contain many bigrams. In Arxiv, we find domain-specific terminology, groups of related citations, Latex commands, and mathematical equations are frequently repeated within the same document. In GitHub and Stack Exchange, we find function and variable names are frequently reused. Meanwhile, the C4, CC, Books, and Wikipedia bigrams are highly variable. The topics and documents in these corpora are relatively diverse and unstructured.

If the placement V V of the hyperplanes is not generic, then the induced subdivision Σ \Sigma is not a triangulation. However, the above bijection is unaffected and, in particular, the fine type of a vertex of C ​ D A CD_{A} can be read off the corresponding facet of Σ \Sigma . We define the crosscut complex c ​ c ​ u ​ t ​ Σ \ ccut{\Sigma} ⊆ 2 [ n ] × [ d ] \subseteq 2^{[n]\times[d]} to be the unique simplicial complex with the same vertices-in-facets incidences as the polyhedral complex Σ \Sigma . The crosscut complex is a standard notion in combinatorial topology and can be defined in more generality (see Björner). The following observation is immediate from the definitions. Proposition D.1 . For V V an ordered sequence of n n points in T d − 1 T^{d-1} let A = A ⁡ ( V ) A=A(V) be the corresponding tropical arrangement and let Σ \Sigma be the regular subdivision of Δ n − 1 × Δ d − 1 \Delta_{n-1}\times\Delta_{d-1} induced by V V . Then the fine cotype ideal f ​ c ​ I \ fcI is Alexander dual to the Stanley-Reisner ideal of c ​ c ​ u ​ t ​ Σ \ ccut{\Sigma} . The crosscut complex encodes the information of which collections of vertices lie in a common face. Hence, the crosscut complex is a purely combinatorial object and does not see the affine structure of the underlying polyhedral complex .

⬇

## Appendix E Synthetic Mqar Experimental Details

In this paper, we propose Mqar as a useful tool to help explain gaps between three popular sequence modeling layers — Transformers or Sparse Transformers, Convolutions, and Recurrences. In this section, we detail the motivation behind the design of Mqar and include extended experiments on additional architectures. We provide a procedure for generating Mqar synthetic data, which can help in the development of new architectures.

### E.1 Mqar Generation Procedure

Here we provide additional discussion on the properties of our Mqar synthetic data. The objective of the synthetic analysis is to help explain the differences in language modeling behavior between different classes of language modeling architectures, as observed on real-world data Section 3 . Synthetic recall tasks were used in the development of Hyena Poli et al. [2023a] and H3 Fu et al. [2023a] , and Olsson et al. [2022] to study Transformer in-context learning behavior. 1. Attention : Attention can perform recall trivially ( Proposition 4.3 ). The bound is independent of the sequence length N N .

2. Convolutions : The gated convolution models use input-independent filters. We show in Section 4 that the dimensionality for solving recall grows with the input sequence length. Real language modeling can require performing 𝒪 ⁡ ( N ) \mathcal{O}(N) recalls in one forward pass. Our intuition is that it is difficult to efficiently compute all token-interaction distances with such a filter.

3. Recurrences : Recurrent models include a single hidden state as an information bottleneck. We show in Section H.6 that the gated recurrence requires Ω ⁡ ( N ) \Omega(N) bits to solve Mqar for d ≤ N d\leq\sqrt{N} for model dimension d d . Our intuition is that it is difficult to store large numbers of key-value pairs seen in the sequence in low-dimensional hidden states.

Motivated by this analysis, we propose Procedure 1 for generating synthetic Mqar data. This procedure allows toggling two simple properties of the synthetic dataset, specified below, to reveal the differences between the sequence modeling architecture families. We also discuss deviations from the prior work that uses synthetic recall data. 1. Size of the Key-Value Map : The number of unique key-value pairs appearing in each example in the dataset. In Procedure 1 , we use the input parameter D D to toggle this value.

2. Number of Gaps in the Data : The number of unique token-interaction distances required to perform the recall task. Prior work assumes there is a single query token that requires recall per example, failing to test whether the architecture can support multiple token-interaction distances. In Procedure 1 , we use the parameters N N (sequence length), D D (size of the key-value map per example), and α \alpha to toggle this property. We enforce that the token-interaction distances appearing in the synthetic examples follow a power-law distribution specified by α \alpha (based on Figure 3 ). Keeping α \alpha and D D constant while varying N N changes the number of unique token-interaction distances appearing in the example.

Finally, noting that the vocabulary size appears in the theoretical bounds for all three architecture families ( Section 4 ), we increase the vocabulary size to be much larger than the model dimension as typically seen in language modeling (30k - 50k tokens). Prior work uses vocab sizes ≤ 40 \leq 40 tokens.

### E.2 Training Details

We first describe the architectures evaluated on Mqar synthetic data and then the training hyperparameters. We evaluate the following four architecture categories in the experiments: 1. Attention Standard GPT-2 style multi-headed Transformer architecture with learned positional embeddings Brown et al. [2020] . Attention, the core building block of Transformers, is defined in Section 2 . The architecture for synthetics has 1 attention head.

2. Gated convolutions (Hyena Poli et al. [2023a] , RWKV Peng et al. [2023] ). Gating, convolutions, and the class of gated convolutions are defined in Section 2 .

3. Gated recurrences (RetNet Sun et al. [2023] ). RetNet proposes computing attention over chunks of the sequence and combining this with a recurrence over the sequence. We evaluate this architecture (“Chunked RNN”) with chunk sizes of 32 32 and 8 8 in Figure 6 . In RetNet, hidden states are updated as: S n = γ ​ S n − 1 + A n T ​ V n S_{n}=\gamma S_{n-1}+A_{n}^{T}V_{n} Outputs at each timestep are defined as O n = C n ​ S n O_{n}=C_{n}S_{n} where γ ∈ ℝ 1 \gamma\in\mathbb{R}^{1} , A = x ​ W a A=xW_{a} for input x ∈ ℕ × 𝕕 x\in\mathbb{N\times d} and learned weight matrix W a ∈ ℝ d × d W_{a}\in\mathbb{R}^{d\times d} , C = x ​ W c C=xW_{c} for W c ∈ ℝ d × d W_{c}\in\mathbb{R}^{d\times d} , V = x ​ W v V=xW_{v} for W v ∈ ℝ d × d W_{v}\in\mathbb{R}^{d\times d} . RetNet similar to a state-space-model (SSM) Gu et al. [2021] , except for that the matrices A A and C C are input-dependent (i.e. they are functions of the input x x ). We note that RetNet is a special case of the recently proposed Mamba Anonymous [2023a] and GateLoop Anonymous [2023b] architectures, which replace the γ \gamma term in RetNet with yet another input-dependent matrix. The RetNet model is formally defined in Section H.2.3 .

4. Sparse attention Several works have proposed methods for efficiently computing attention Tay et al. [2022] . We evaluate two classical methods: sliding window attention Beltagy et al. [2020] , Zaheer et al. [2020] and blocked window attention Child et al. [2019b] , Qiu et al. [2019] . Consider a window size w w . Sliding window attention permits each token to attend to the prior w w tokens in the sequence. Blocked window attention first splits the sequence into blocks of size w w tokens, and then computes standard causal attention within the block.

For all synthetic runs in the main paper and the appendix, we use the following training protocol: • Optimizer and schedule : Weight decay 0.1 0.1 , warmup duration 10 % 10\% , linear warmup, AdamW optimizer. For each run, we sweep the learning rates in np . logspace ⁡ ( − 4 , − 2 , 4 ) \mathrm{np.logspace(-4,-2,4)} . We train for 64 64 epochs.

• Training duration : The global batch size is 8 8 for input sequence length or model dimension ≥ 512 \geq 512 , 16 16 for input sequence length or model dimension ≥ 256 \geq 256 , and 64 64 otherwise.

• Width and depth : For all synthetic experiments, we use exactly two layers (each with one sequence mixer and one MLP, interleaved with layer normalization). The model dimension, sequence length, and number of KV pairs are varied based on the relevant experiment (see Appendix F ).

• Position information : No position embeddings are used for the pure convolution runs (input-dependent nor input-independent filters). Position embeddings are used for the runs with any attention variant.

• Data : We train and evaluate each model on 100,000 100,000 and 3,000 3,000 data points respectively. The data and data order for all runs are constant.

## Appendix F Extended Results on Mqar Across Architectures

In this section, we provide further validation that Mqar is a useful diagnostic tool to explain the behaviors of a wide variety of architecture families. We show: 1. Gated convolutions In the main paper, we claim the gated convolution models with input-independent filters use larger dimensionality than attention as the number of token-interaction distances required for the task increases ( N N , holding α \alpha and D D constant in Procedure 1 ).

2. Gated recurrences In this section, we show gated recurrrent models use larger dimensionality than attention to solve the task as the number of unique key-value pairs to store ( D D in Procedure 1 ) increases.

3. Sparse attention In this section, we validate that sliding window attention helps close the Mqar gap when the token-interaction distances are relatively short, and degrades on longer distances. Meanwhile, blocked window attention struggles to close the gap (intuitively tokens earlier in a block are able to interact with few other tokens in the sequence). Concretely, in Appendix D , we observe that repeated bigrams in the Pile and RedPajama corpora tend to occur within neighboring tokens in the context.

We finally show the architectures that perform well on Mqar also perform well on downstream real-world AR on the Pile. We use the experimental protocols outlined in Appendix E for the synthetic experiments and Appendix C for the downstream experiments in this section.

### F.1 Synthetic Experiments

##### Experiment 1: Increasing the number of token-interaction distances per example.

To construct the synthetic dataset following Procedure 1 , we set α = 0.1 \alpha=0.1 and | C | = 8,192 |C|=8,192 for all experiments. We vary N ∈ { 64,128,256,512 } N\in\{64,128,256,512\} to increase the number of token-interaction distances that will appear in the examples and we vary the model dimension ∈ { 64,128,256,512 } \in\{64,128,256,512\} . The results of this experiment are shown in Fig. 2 , validating that the gated convolutiosn use larger dimensionality than attention as the number of token-interaction distances required for the task increases.

##### Experiment 2: Increasing the number of key-value pairs occurring per example.

To construct the synthetic dataset following Procedure 1 , we set α = 0.1 \alpha=0.1 and | C | = 8,192 |C|=8,192 for all experiments. We fix N = 256 N=256 and vary D ∈ { N 32 , N 16 , N 8 , N 4 } D\in\{\frac{N}{32},\frac{N}{16},\frac{N}{8},\frac{N}{4}\} to increase the number of key-value pairs occurring per example. For each number of key-value pairs per example, we additionally vary the model dimension ∈ { 64,128,256 } \in\{64,128,256\} . The results are shown in Fig. 6 . We validate that as the number of associative recall keys and values exceeds the window size, the model again uses increased dimensionality relative to full O ⁡ ( N 2 ) O(N^{2}) attention to solve the task. While Fig. 6 highlights RetNet, we note that it is a special case of the recently released Mamba Anonymous [2023a] and GateLoop Anonymous [2023b] architectures. Corresponding to this experiment, we theoretically show in Section H.6 that the gated recurrence uses Ω ⁡ ( N ) \Omega(N) bits to solve Mqar for d ≤ N d\leq\sqrt{N} . Intuitively, it is difficult to store a large number of key-value pairs in low dimensional hidden states.

##### Experiment 3: Increasing the range of token-interaction distancs per example.