Memorizing Transformers — part 2/3

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

We evaluate the effect of adding external memory on five language modeling tasks, all of which involve long-form text: English language books (PG-19), long web articles (C4), technical math papers (arXiv Math), source code (Github), and formal theorems (Isabelle). The results show significant improvements in the perplexity of the model with the addition of external memory. We experimented with various sizes of external memory, from 1536 to as high as 262K. On most of the datasets, there was an initial sharp gain from adding a small external memory, followed by smaller but steadily increasing gains as the size of the memory was increased.

### 4.1 Datasets

##### arXiv Math

For the arXiv dataset, we collected a corpus of papers by downloading them via the arXiv Bulk Data Access 1 1 1 https://arxiv.com/help/bulk_data . We filtered papers to include only articles labeled as “Mathematics” and whose L a T e X source was available. The number of tokens per paper in this dataset is roughly comparable to the number of tokens per book in PG19, because L a T e X source has many special characters and the tokenizer tends to output small subwords.

##### Github

We used BigQuery 2 2 2 https://console.cloud.google.com/marketplace/product/github/github-repos to obtain a large corpus of Github repositories that are published with open-source licenses. We used file endings to filter for files in the languages C, C++, Java, Python (including Jupyter notebooks), Go, and TypeScript. Individual source code files are often fairly short, and there are many dependencies and cross-references between files in the repository. To capture these dependencies, we created one long document for each Github repository by traversing the directory tree, and concatenating all of the files within it. The order in which files are traversed within the repository is random, but each subdirectory is processed as a unit, so that all the files within the subdirectory are close to each other in the resulting document. Source code is usually structured so that related files are all grouped together in the same subdirectory; this traversal preserves that structure, while still shuffling files and subdirectories in random order.

##### Formal Math – Isabelle

The Isabelle corpus consists of formal mathematical proofs of theories. We collected all 627 theories available on The Archive of Formal Proofs 3 3 3 https://www.isa-afp.org/topics.html (as of October 6, 2021) and an additional 57 theories from the Isabelle standard library 4 4 4 https://isabelle.in.tum.de/ to create a corpus of 684 theories. All theories have open-source licenses. Each theory is a self-contained mathematical object, on topics such as foundational logic, advanced analysis, algebra, or cryptography, and consists of multiple files containing proofs. As with the Github corpus, all files that make up a theory are concatenated together into one long document. Unlike the Github corpus, we order the files according to their import dependencies, so that later files use sub-theorems that are proved in earlier files.

##### C4(4K + + )

C4, the colossal cleaned common crawl, is a very large collection of documents that have been scraped from the internet ( Raffel et al., 2020 ) . We filtered out all documents that have less than 4096 tokens to focus on documents where memory can have an impact.

##### PG-19

PG-19 is a large dataset of English-language books, published prior to 1919, which were retrieved from the Project Gutenberg archive ( Rae et al., 2020 ; Sun et al., 2021 ) . PG-19 is one of the few public datasets that only contains full-length books, and has become a benchmark for long-range natural language text modeling.

### 4.2 Experimental Method

We used a 12-layer decoder-only transformer (with and without Transformer-XL cache) with an embedding size of 1024, 8 attention heads of dimension 128, and an FFN hidden layer of size 4096. For all of our experiments, we used k = 32 k=32 . Unless specified otherwise, we use the 9 9 th layer as the k k NN augmented attention layer. We used a sentence-piece ( Kudo & Richardson, 2018 ) tokenizer with a vocabulary size of 32K.

We used the Adafactor optimizer ( Shazeer & Stern, 2018 ) . In preliminary experiments, we conducted a hyperparameter search to determine the optimal learning rate among three choices ( { \{ 3.0 3.0 , 1.0 1.0 , 3 ⋅ 10 − 1 3\cdot 10^{-1} } \} ), and found that 1.0 1.0 works best. We used a linear warmup schedule for the first 1000 steps, followed by square root decay. We trained the models from scratch for 500K steps on all the datasets, except for the Isabelle dataset. Isabelle is small, so we stopped training after 100K steps when the model began to overfit. We ran all of our experiments on 32 TPU cores. Our models were implemented in JAX ( Bradbury et al., 2018 ) and Flax ( Heek et al., 2020 ) .

When comparing models with different context lengths, we adjusted the batch size (the number of documents in a batch) so that there are always 2 17 2^{17} tokens in a batch. E.g., a model with a context length of 512 has a batch size of 256, while the 2048 model has a batch size of 64.

We experimented with multiple implementations of approximate k k NN lookup with different tradeoffs between quality and computational cost. We did not observe a significant degradation of the model quality when switching to lower quality approximations of k k NN, so the model appears to be quite robust with respect to the quality of k k NN retrieval. For a model with around 200M trainable parameters the step time increased from 0.2s to 0.25s when we added a memory of size 8K, and to 0.6s when we added a memory of size 65K (measured on TPUv3).

### 4.3 Effect of external memory

##### Adding external memory results in substantial gains across datasets and architectures,

as shown in Table 4 . Across all five datasets, adding external memory to either the vanilla Transformer or the Transformer-XL architecture improves perplexity by a substantial amount. For example, on C4(4K+) dataset, adding memory of size 8192 improves the perplexity of the vanilla Transformer (with context size 512) from 17.20 to 14.42, and improves Transformer-XL from 15.38 to 14.04.

##### Increasing the size of the memory increases the benefit of the memory.

The best perplexities for all datasets and architectures were obtained with a memory size of 65K.

Note that Transformer-XL with context size 2048 already has a theoretical receptive field that is quite large. Each token in a higher layer can attend up to 2048 tokens away in the layer below, so the total receptive field is 2048 ⋅ 12 2048\cdot 12 (layers) ∼ 25 \sim 25 K. Nevertheless, we still saw a substantial gain when adding an external memory of size 8192 to this model. k k NN attention into memory would appear to be a more effective way to retrieve information from the distant past than the Transformer-XL cache.

On the other hand, we also saw improvements by adding XL cache to the large-memory (65K) models. In a vanilla (non-XL) Transformer, the first few tokens in a sequence have very little context, and thus have higher perplexity. The XL cache provides additional local short-range context at the start of a sequence, which complements the long-range context provided by external memory.

Interestingly, in a vanilla Transformer, using even a small external memory of size 1536 provides a gain in perplexity which is almost as good as using a local context of size 2048 but no memory (e.g. Table 4 ). This is surprising, because the external memory is not differentiable, and is added only to one layer of the Transformer, whereas increasing the context size is differentiable and affects all layers. We conclude that the lower layers of a Transformer don’t necessarily need long-range context, and having a differentiable memory is not as important as one might suspect.

### 4.4 Scaling to larger models

We scaled up the Transformer model to sizes of 1 1 and 8 8 billion parameters. For the 1 billion parameter model, we use 8 layers, 32 heads with head dimension 128, d ​ _ ​ model d\_\textrm{model} 2048, and d ​ _ ​ ff d\_\textrm{ff} 16384. For the 8 billion parameter model, we use 64 heads, 16 layers, d ​ _ ​ model d\_\textrm{model} 4096, and d ​ _ ​ ff d\_\textrm{ff} 32768. We used a context size of 2048, memory size of 8192, and no XL cache. We ran the comparisons to the vanilla Transformer on the arXiv math dataset. Scaling plots are shown in Figure 1 .

External memory provides a consistent improvement to the model as it is scaled up. Remarkably, we found that the smaller Memorizing Transformer with just 8k tokens in memory can match the perplexity of a larger vanilla Transformer which has 5X more trainable parameters.

### 4.5 Finetuning on Larger Memories

##### Finetuning on a larger memory.

In some cases, training was unstable when using large memories, possibly due to distributional shift early in the training (See Section 3.2 ). Thus, for memories of 131K or more tokens, we first pretrain the model with a memory size of 8192 or 65K for 500K steps, and then finetune it with the larger memory for an additional 20K steps. The results of finetuning on the arXiv Math data set are shown in Table 5 . Increasing the size of external memory provided consistent gains up to a size of 262K. Note that 262K tokens is longer than almost all of the documents in arXiv, and thus we would not expect to see any gain past this point (see Appendix A ).

##### Finetuning a non-memory model to use memory

Pretraining can be very costly both in time and computational resources. Thus, a natural question to ask is: can one fine-tune a pretrained Transformer to use external memory? The answer is yes!

We took a pre-trained 1B vanilla Transformer model, and fine-tuned it to use external memory (the 1B models used in Section 4.4 ). The fine-tuning result is shown in Figure 6 . Notice that the model quickly learns to use external memory. Within 20K steps ( 4 % 4\% of the pre-training time) the fine-tuned model has already closed 85 % 85\% of the gap between it and the 1B Memorizing Transformer, and after 100k steps it has closed the gap entirely.

### 4.6 Information Retrieval Patterns

We conducted a qualitative study of what the model was actually retrieving from external memory, by finding which tokens showed the biggest improvements in cross-entropy loss when the size of the memory was increased, and then examining the top- k k retrieved memories for those tokens. We found that the model gained the most when looking up rare words, such as proper names, references, citations, and function names, where the first use of a name is too far away from subsequent uses to fit in the local context. This result is in keeping with the prior analysis of long-context Transformers on PG19 ( Sun et al., 2021 ) , which found similar lookup patterns. For this experiment, we used a slightly older version of the architecture without the gating mechanism.

##### Which tokens show a benefit from memory?

Figure 7 shows a visualization of which tokens show an improvement when the size of the external memory is increased. We selected a math paper at random, and plotted the difference in cross entropy loss for each token x i x_{i} in the paper, comparing two models with the same parameters, but with memories of different sizes. Δ i = cross-entropy 8192 ​ ( x i ) − cross-entropy 32 ​ K ​ ( x i ) \Delta_{i}=\textrm{cross-entropy}_{8192}(x_{i})\linebreak-\textrm{cross-entropy}_{32\mathrm{K}}(x_{i}) . Positive values show an improvement in loss.

The x x -axis on the chart is the token number i i , while the y y -axis is Δ i \Delta_{i} . For the first 8192 tokens, the difference between the two models is zero, since the larger capacity of the 32 32 K memory isn’t being used yet. However, after token 8193 8193 , we can see that the larger memory helps, on average, over the smaller memory. The benefit is not universal, since the predictions for some tokens become worse, possibly due to the fact that a relevant retrieved memory no longer makes it into the top- k k when the size of the external memory is increased. This figure also shows that the benefit of external memory is somewhat sparse. The improvement in perplexity seems to be mainly driven by a small percentage of tokens that obtain a large improvement in cross-entropy loss when using the larger memory.

##### What information is being looked up?

Given that only a subset of tokens shows improvement from external memory, we did a further investigation into what, exactly, those tokens are using the memory for. We took those tokens which showed the largest improvement in cross-entropy loss, and for each of them tokens, we examined the top- k k retrieved memories. We studied arXiv math, Github and Isabelle corpus. For arXiv math and Github, we found the model retrieved function and variable names. See more details with examples in Appendix B .

Query index Input Target Surrounding context Retrieved index Retrieved surrounding context 29721 mark ov rule prob_space. markov_inequality 8088 M. t \ \backslash <le> X a} \ \backslash <le> expectation X / / t" 40919 _ \_ th = ( subgraph_threshold H n / p n) 27219 threshold H n = n powr (-(1 / max_density’ 49699 S w assumes " orthonormal_system S w" 28050 definition orthonormal_system :: "

Retrieving mathematical definitions. Our case study on the Isabelle corpus provides one of the clearest illustrations of how a model can make good use of external memory. When predicting the name of a mathematical object or a lemma, the model looked up the definition from earlier in the proof. Examples of this behavior are shown in Table 8 . In example 1, the model retrieves a definition within the body of a lemma, markov_inequality . In example 2, it retrieves the definition of a previously defined concept subgraph_threshold . In example 3, it retrieves the definition of orthonormal_system . We manually checked 10 examples where the model made a prediction of lemma names, and 8 out of 10 times model found the body of the lemma it needs to predict. In the other two cases, the model also looked up materials in the immediate vicinity. To the best of our knowledge, this is the first demonstration that attention is capable of looking up definitions and function bodies from a large corpus. The Isabelle case study used a model with two memory layers of size 32 32 K.

## 5 Conclusion

We present a simple extension to the Transformer architecture, called k k NN-augmented attention, which dramatically increases the length of the context that a language model can attend to by using k k -nearest-neighbor lookup into a large external memory. We demonstrate the effectiveness of external memory in a series of language modeling experiments over a variety of long-document datasets, including LaTeX documents, source code, formal proofs, and books.

The Memorizing Transformer shows large improvements in perplexity over the baseline for all of the data sets and architectures that we studied; it is comparable to a vanilla transformer that has 5 times the number of parameters. Perplexity continues to improve with increasing memory size, although there is a point of diminishing returns. Moreover, external memory continues to provide benefits even as the transformer is scaled up from 200M to 8B parameters. Perhaps most intriguingly, a Memorizing Transformer does not need to be pre-trained from scratch; it is possible obtain large gains from adding memory to an existing pre-trained model, and then fine-tuning it.

Unlike other forms of attention, k k NN retrieval can be easily scaled up to huge memory sizes, and is thus potentially able to leverage vast knowledge bases or code repositories. How to make the best use of this capability is a topic for future work.

#### Acknowledgments