Memorizing Transformers
Authors: Yuhuai Wu
Markus Norman Rabe
DeLesley Hutchins
Christian Szegedy
Venue: ICLR 2022 Spotlight (2022)
Keywords: Transformer
architecture
memorization.

Abstract: Language models typically need to be trained or finetuned in order to acquire new knowledge, which involves updating their weights.  
We instead envision language models that can simply read and memorize new data at inference time, thus acquiring new knowledge immediately. In this work, we extend language models with the ability to memorize the internal representations of past inputs. We demonstrate that an approximate $k$NN lookup into a non-differentiable memory of recent (key, value) pairs improves language modeling across various benchmarks and tasks, including generic webtext (C4), math papers (arXiv), books (PG-19), code (Github), as well as formal theorems (Isabelle). We show that the performance steadily improves when we increase the size of memory up to 262K tokens. 
On benchmarks including code and mathematics, we find that the model is capable of making use of newly defined functions and theorems during test time.

Published as a conference paper at ICLR 2022
MEMORIZING TRANSFORMERS
Yuhuai Wu, Markus N. Rabe, DeLesley Hutchins, Christian Szegedy
(yuhuai, mrabe, delesley, szegedy / @ggogle com
ABSTRACT
Language models typically need to be trained or finetuned in order to acquire new knowledge, which involves updating their weights. We instead envision language models that can simply read and memorize new data at inference time, thus acquiring new knowledge immediately.
this work, we extend language models with the ability to memorize the internal representations of past inputs. We demonstrate that an approximate kNN lookup into a non- differentiable memory recent (key, value) pairs improves language modeling across various benchmarks and tasks, including generic webtext (C4), math papers (arXiv), books (PG-19), code (Github), well as formal theorems (Isabelle). We show that the performance steadily improves when we increase the size of memory up to 262K tokens. On benchmarks including code and mathematics, we find that the model is capable of making use of newly defined functions and theorems during test time.
INTRODUCTION
Transformers Vaswani et al., 2017 have led to remarkable progress in natural language process-
2020), mathematical reasoning (Polu & Sutskever [020 Wang et 2022 Rabe et al. 2222 Li et al., [021 Hahn et al. [202]] Cobbe et 20.1). and program synthesis (Austin et al. 2021 Chen et al. 2021 Li et al. [002]. However, transformer performance on many of these tasks limited by the context length of attention, which is typically short. The ability to attend to far-away tokens is important in many situations. In novels, characters and events referenced across multiple chapters. In source code, references classes and functions may occur quite far from the places in which they are defined. In theorem proving, proofs make use of previously defined lemmas.
ing (Devlin et al. 2019 Brown et al.
Attention over long sequences is also useful as a form of rapid learning. Facts and information which are stored in the form of weight matrices must be slowly trained over hundreds of thousands of training steps. By using attention, however, a model can simply memorize facts (e.g. function definitions) storing them (key, value) pairs in long-term memory, and then retrieve those facts later by creating query that attends to them. In this case, attention acts as form of information retrieval, allowing the model look up facts that has seen previously.
We demonstrate that a simple and effective way to increase the size of the attention context is to use approximate neighbrr (kNN) lookup, which widely used information retrieval. number of extremely scalable implementations of KNN lookup are available, such as ScaNN (Guo ct al. 2020) and Faiss (Johnson et al. 202T).
There are two things which distinguish our approach from previous work on long-range attention (c.f. Section 2. First, unlike some other approaches, ŁNN lookup does not do averaging summarization of tokens at long distances, but retrieves exact values even from the distant context.
Second, gradients are not backpropagated into the external memory, which is critical to the scalability of our technique. The keys and values are function of model parameters, attempting to backprop- agate gradients into external memory would necessarily involve computing all of the keys and values with the current model parameters every training step. However, if the external memory is not differentiable, then we can instead instead reuse keys and values that were previously computed on prior training steps, which drastically reduces the amount of computation for large memories. With

---

Published as a conference paper at ICLR 2022
2.8
Perplexity (arXiv Math)
Transformer
2.6
Memorizing Transformer
2.4
2.2
2.0
1.8 200M
1B Model size
Figure 1: Adding a memory of 8K tokens improves perplexity across different model sizes.
our technique, we are easily able to scale external memory up to sequence lengths of 131k or 262k tokens on single TPU device, while maintaining a reasonable step time.
We show that model perplexity steadily improves with the size of external memory on variety of language modelling tasks, including C4 (long documents only). Github code repositories, PG-19 books, formal proofs in Isabelle, and arXiv math papers. We further show that models generalize to larger memory sizes than they were trained on: models trained with a small memory show gains from using much larger memory at inference time. Finally, we show that models actually using memory in the way that we had hoped, e.g. by looking up the definitions of lemmas in a theorem proving corpus.
The simplicity of the changes to the Transformer architecture allows us to easily integrate this approach into existing code bases, including extremely large language models. We further show that the improvements to quality are maintained across models of increasing size, and that the model improvements gained from adding memory are even larger than increasing the size of the model by 5X or more as shown in Figure
RELATED WORK
great deal of work has been done on efficient long-range attention mechanisms; see Tay et al. (2020) 202T) recent surveys. Sliding windows (Beltagy et al. 2020) use long sequence, but attend within smaller window, thus reducing complexity the window size, rather than total sequence length.
Approximate mechanisms such as Linformer (Wang et al. 2020b), and Performer (Choromanski ct al., 2021 refactor the attention matrix by using a different kernel than softmax to obtain O(N) complexity. Pooling strategies such as Hierarchical ID attention (Zhu & Soricut 2021), and Combiner (Ren et al. 2021) apply pooling or averaging over tokens at longer distances. Sparse strategies such as Big Bird (Zaheer et al. 2020) select only subset of tokens to attend to; Routing Transformers (Roy ct al. 2021) use clustering to select the subset, while Reformer (Kitaev et al.) (2200) relies on hashing. Hierarchical mechanisms (Ainslie et al. 2020) combine multiple tokens into phrases sentences to reduce sequence length. Expirc-span (Sukhbaatar et al., 2021) prunes far-away tokens that it learns are "unimportant". Zemlyanskiy et 2021) process long sequences in two passes with different encoders. The second pass is given a lot of context by accessing summaries of the first pass.
Feedback transformers (aan al., 2020 use recurrent architecture in which each token attends to the output of the final layer instead of the previous layer. Recurrence does not increase the size of the attention context itself, but it expands the receptive field at the cost of parallelism and training speed.
Truncated backpropagation through time (Williams & Peng 1990) was originally introduced as way of training recurrent neural networks (RNN) over very long sequences, when the entire sequence does not fit in memory. The sequence is chopped into segments, and after each training step, the final RNN state for the segment saved in a non-differentiable cache, and used as the initial state on the next training step. Neural caches (Grave et al.) 2017) extend the cache to contain a record of many prior hidden states, and attend over them. Transformer-X XL (Dai et al. 2019) applies this technique to transformers; it caches the (key. pairs computed from the previous training step, and them as prefix for the tokens on the next training step, which yields significant gains on long documents. Rae al. 2020) improve over Transformer-XI. by compressing the tokens before adding them the
2

---

Published as a conference paper at ICLR 2022
oulput prodictions
the
cat
in
thce
hat
to
softmax
local aitention FFN
KNN sttcation
A nearesi neighhor lookup.
KNN & local atiention FFN
memory: cached (key, value) pairs
Will external memory after the current training step.
local attention FFN
embeddiig layer
input tokens
said
Figure 2:
We extend Transformers with access to (key, value) pairs of previously seen subsequences.
cache. In contrast, we use a very large cache without compression, combined with an approximate kNN attention mechanism over it.
Sukhbaatar et al. (019) make the observation that the feed-forward portion a transformer layer functions very much like attention if one replaces the ReLU activation with softmax. They implement a combined attention over both tokens from the input sequence and a learned (and differentiable) "memory". Lample al. (2019) exploit this observation replace the ferward layers (FFNs) with fast ANN lookup over much larger "memory", and achieve large gains in model accuracy without significant computation overhead. (We use &NN lookup to approximate attention previous tokens, not to replace the FFN.)
doe-eifferentiable external memory has been used in different ways by Khandelwal al. (2020), who run a pre-trained model over an entire corpus, and construct large table of (key, token) pairs. They then use that table replace the final softmax layer for token selection in the model, which results in significant improvements in language modeling. Yogatama al. (2021) extend this approach by gating mechanism and a process to compress the context into keys for retrieval.
There are several works that combine retrieval with transformers. REALM (Guu et al. 2020). MARGE (Lewis et al. (2020),, RAG (Lewis et al. (2020b), and composite memory for dialog (Fan et al. 2021) retrieve documents from knowledge base to improve question answering or dialogue. The knowledge base consists of text snippets and is static and typically separate from the inputs and outputs of the models. Instead, we focus language modeling using decoder-only model, and propose a simple model that unifies attention and retrieval.
d-nearest-neighhoor lookup a general-purpose technique that is used for wide variety of machine learning and retrieval tasks, and high-performance implementations are available for various architec- tures (Johnson al. 2021 Guo al., 2020). Memory-efficient Transformers (Gupta et al. 2021 replace dense attention with a kNN lookup to increase speed and reduce memory usage.
METHOD
The architecture of our aug-angee transformer is shown in Figure The bulk of the model is vanilla, decoder-only transformer (Vaswani et al. (017). The input text is tokenized, and the tokens are embedded into vector space. The embedding vectors are passed through a series of transformer layers, each of which does dense self-attention, followed by feed-forward network (FFN). Since this is a decoder-only language model, we use a causal attention mask and the token embeddings of the last layer used predict the next token.
Long documents are split into subsequences of 512 tokens, and each subsequence is used as the input for training step. In contrast to standard practice, do not shuffle the subsequences; instead, each long document is fed into the transformer sequentially, from beginning to end, as is done with Transformer-Xi XI. (Dai et al. 2019.
3

---

Published as a conference paper at ICLR 2022
subsequence
batch
Document
Document C
Document ...
dimension
Document B
Document D
Document E
Document G
Figure 3:
Our data pipeline splits documents into subsequences and packs subsequences into batches.
We also use Transformer-XL style cache, which holds the keys and values from the previous training step. When doing self-attention, the cached keys and values are prepended to the current keys and values, and we use sliding-windoon causal mask (Beltagy al. 2020) so that each token has a local context that includes the previous 512 tokens.
3.1 ANN-AUGMENTED ATTENTION LAYER
One of the transformer layers near the top of the stack is a kNN-angmented attention layer, which combines two forms of attention. Like all of the other layers, it uses standard dense self-attention on the local context, which is the input subsequence for the current training step. Unlike the other layers, however, it also does approximate k-nearest-neighhor search into the external memory.
The same queries are used for both the local context, and for the external memory. The keys and values also belong to the same distribution; after each training step, the (key, value) pairs in the local context are appended to the end of the external memory. If the document is very long, old (key, value) pairs will be dropped from the memory to make room for new ones. Thus, for each head, the external memory keeps a cache of the prior M (key, value) pairs, where M is the memory size.
The kNN lookup will return set of retrieved memories, which consist of the top-k (key, value) pairs that kNN search returns for each query (i.e. each token) in the input subsequence. As with standard dense attention, we first construct an attention matrix by computing the dot product of each query against the retrieved keys, then apply softmax, and finally return weighted sum of the retrieved values. Unlike standard dense attention, the retrieved memories contain a different set of (key, value) pairs for each query.
Attention over the local context is performed in the usual way. The results of ANN-attention and local attention are then combined using learned gate:
g a(by)
V. = Vm V
(2)
where a is the sigmoid function, and @ element-wise multiplication. V. is the combined result of attention, Vn. is the result of attending to external memory, and V. is the result of attending to the local context. The bias by is a learned per-head scalar parameter, which allows cach head to choose between local and long-range attention. In our experiments, the value of the gate g does not depend on the content of the token at each position, although that would be a trivial extension to implement. We did observe that over time, most heads learned to attend almost exclusively to external memory.
Position bias. For dense attention within the local context, we use the T5 relative position bias (Raffel al., 2020). As noted by Dai et al. (2019), adding global position encoding to each token does not work well when processing long documents. We don't 't use a position bias for the retrieved memories. Experiments the PG19 dataset (Sun al., (021) have shown that relative position does not appear to matter at long range, and the T5 relative bias puts all long-range tokens in the same bucket anyway.
Batching, Figure 3 iillustrates how multiple long documents of different lengths are packed into batch, and split into subsequences. Each subsequence in the batch comes from a different document, and thus requires a separate external memory, which is cleared at the start of each new document.
3.2 DISTRIBUTIONAL SHIFT
Because each long document is processed over multiple training steps, there is a distributional shift in the keys and values that are stored in external memory. The model parameters that produce the queries change over time, and will thus have shifted since the keys and values were stored. For very large memories, older records may become "stale." " Similar observations have been made for CrossBatch memory (Wang et al. 2020c) in the vision domain.
4

---

Published as a conference paper at ICLR 2022
To reduce the effects of staleness, normalize keys and queries (Henry al. 2020). does not eliminate staleness, but it at least ensures that older keys and newer keys do not differ in magnitude. We also found that normalization helps stabilize training with the Transformer-X XL cache.
Normalization
In some of our experiments, we observed that training models from scratch with large memory sometimes resulted in worse performance than pretraining the model with small memory of size 8192, and then finetuning on a larger memory. This training instability could be due to staleness. However, models seem to be able to cope with a limited degree of staleness (with the small memory) by adjusting their queries accordingly.
3.3 APPROXIMATE &NN
We employ approximate kNN search rather than exact kNN search because it significantly improves the computational speed of our model. We use simple approximation of INN for TPUs, which has recall of about 90%, i.e. 90% of the true top k are returned in the approximate top k. There are various other efficient approximate kNN algorithms available for CPU and GPU/TPU, for example through Faiss (Johnson et al. 202T) ScaNN (Guo al. 2020), which can scale into the billions.
EXPERIMENTS
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 small external memory, followed by smaller but steadily increasing gains as the size of the memory was increased.
4.1 DATASETS
Xiv Math For the arXiv dataset, we collected a corpus of papers by downloading them via the arXiv Bulk Data Acces{] We filtered papers to include only articles labeled as "Mathematies" and whose LTTEX 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 LTTEX source has many special characters and the tokenizer tends to output small subwords.
Github We used BigQuery to obtain 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 TypeSeript. 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, 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, 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.
created one long document for each Github repository by traversing the
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 (as of October 6, 2021) and an additional 57 theories from the Isabelle standard libraryE to create a corpus of 684 theories. All theories have open-source licenses. Each theory is self-contained mathematical object, on topics such as foundational logic, advanced analysis, algebra. cryptography, and consists of multiple files containing proofs. As with the Github corpus, all files that make up a theory are concatenated
ht pttpp////arttoo.o///elppp/p/ data
repos
attos: atsos://wwwieo-af..com
Rtpa:///saane htt.:aaa....pporoiiiiocssesee
5

---

Published as a conference paper at ICLR 2022
Context
Memory
XI. cache
arXiv
PG19
C4(4K+)
GitHub
Isabelle
512
None
None
3.29
13.71
17.20
3,05
3.09
2048
None
None
2.69
12.37
14.81
2.22
2.39
512
None
512
2.67
12.34
15.38
2.26
2.46
2048
None
2048
2.42
11.88
14.03
2.10
2.16
512
1536
None
2.61
12.50
14.97
2.20
2.33
512
8192
None
2.49
12.29
14.42
2.09
2.19
512
8192
512
2.37
11.93
14.04
2.03
2.08
512
65K
512
2.31
11.62
14.04
1.87
2.06
2048
8192
2048
2.33
11.84
13.80
1.98
2.06
2048
65K
2048
2.26
11.37
13.64
1.80
1.99
Table 4: Average token-level perplexities of each model when trained for 500k steps.
together into one long document. Unlike the Github corpus, we order the files according to their import dependencies, that later files use sub-theorems that are proved in earlier files.
C4(4K+) C4, the colossal cleaned common crawl, is 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- is large dataset of laglish-language books, published prior to 1919, which were retrieved from the Project Gutenberg archive (Rae et |l.] [020 (unn et al. (002)) 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 12-layer decoder-only transformer (with and without Transformer-XL. cache) with 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 & 32. Unless specified otherwise, we use the 9th layer as the kNN augmented attention layer. We used sentence-piece (Kudo & Richardson 2018) tokenizer with vocabulary size of 32K.
We used the Adafactor optimizer (Shazeer & Stern 2018). In preliminary experiments, we conducted hyperparameter search determine the optimal learning rate among three choices (f3.0, 1.0, 3 1)- and found that 1.0 works best. We used 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,
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 batch) that there always 2 tokens in batch. E.g., model with 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 ANN 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 KNN, the model appears to be quite robust with respect to the quality of kNN retrieval. For a model with around 200M trainable parameters the step time increased from 0.2s 0.25s when we added 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 Across all five datasets, adding external memory to either the vanilla Transformer the Transformer-X architecture improves perplexity by substantial amount. For example,

---

Published as a conference paper at ICLR 2022
Context
Pretrain
Fine-tune
Perplexity
512
8192
None
2.37
512
65K
None
2.31
512
8192
65K
2.32
512
8192
131K
2.30
512
8192
262K
2.26
2048
8192
None
2.33
2048
65K
None
2.26
2048
65K
131K
2.23
2048
65K
262K
2.21
Table 5: Finetuning for 20K steps to make use of a larger memory on the arXiv data set.
C4(4K+) dataset, adding memory of size 8192 improves the perplexity the vanilla Transformer (with context size 512) from 17.20 to 14.42, and improves Transformer-XI from 15.38 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 memory size of 65K.
Note that Transformer-Xi with context size 2048 already has theoretical receptive field that is quite large. Each token in higher layer attend up 2048 tokens away in the layer below, the total receptive field is 2048 12 (layers) 25K. Nevertheless, we still saw a substantial gain when adding external memory of size 8192 to this model. ANN attention into memory would appear to be more effective way to retrieve information from the distant past than the Transformer- XI. cache.
On the other hand, also saw improvements by adding XL cache to the large-memory (65K) models. In a vanilla (non-XL) Transformer, the first few tokens 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 vanilla Transformer, using a gain in perplexity which is almost as good as using a local context of size 2048 but no memory (c.g. Table This 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 Transformer don't necessarily need long-range context, and having a differentiable memory is not important as one might suspect.
small external memory of size 1536 provides
4.4 SCALING TO LARGER MODELS
We scaled up the Transformer model to sizes and 8 billion parameters. For the billion parameter model, we layers, 32 heads with head dimension 128, d__modl 2048, and d_ff 16384. For the 8 billion parameter model, we use 64 heads, 16 layers, d d__odel 4096, and d 32768. We used a context size of 2048, memory size of 8192, and XL cache. We ran the comparisons to the vanilla Transformer on the arXiv math dataset. Scaling plots are shown in Figure
External memory provides consistent improvement the model found that the smaller Memorizing Transformer with just 8k tokens in memory can match the perplexity of larger vanilla Transformer which has 5X more trainable parameters.
is scaled up. Remarkably,
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 13IK more tokens, we first pretrain the model with memory size of 8192 or 65K for 500K steps, and then finetune it with the larger memory for an additional 20K steps. The results finetuning on the arXiv Math data set consistent gains up to a size of 262K. Note that 262K tokens is longer than almost all of the documents in arXiv, and thus would not expect any gain past this point (see Appendix A).
shown in Table 5 Increasing the size of external memory provided
7

---

Published as a conference paper at ICLR 2022
Memory Fine-tuning
Perplexity (arXiv Math)
Transfomer
Memorizing Transfomer
2.4
2.2
300K
500K 600K
Training steps
Figure 6: Finetuning a IB vanilla Transformer model to use external memory of size 65K.
Finetuning non-memory model use memory Pretraining can be very costly both in time and computational resources. Thus, natural question to ask is: can one fine-tune a pretrained Transformer to use external memory? The answer yes!
We took pre-trained IB vanilla Transformer model, and fine-tuned it to use external memory (the IB 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% of the pre-training time) the fine-tuned model has already closed 85% of the gap between and the Memorizing Transformer, and after 100k steps it has closed the gap entirely.
4.6 INFORMATION RETRIEVAL PATTERNS
We conducted 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 increased, and then examining the top-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 of name is too far away from subsequent uses to fit in the local context. This result in keeping with the prior analysis of long-context Transformers on PG19 (Sun et al.) 202T), which found similar lookup patterns. For this experiment, used slightly older version of the architecture without the gating mechanism.
Which tokens show benefit from memory? Figure 7 shows 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 æ; in the paper, comparing two models with the same parameters, but with memories of different sizes. A, cross-entropy (192(1;) (æ:) -cross-nttropyssk (z+). Positive values show an improvement in loss.
The r-axis on the chart the token number i, while the y-axis is A,. For the first 8192 tokens, the difference between the two models is zero, since the larger capacity of the 32K memory isn't 'n''t being used yet. However, after token 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 relevant retrieved memory no longer makes it into the top-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 small percentage of tokens that obtain large improvement in cross-entropy loss when using the larger memory.
What information being looked up? Given that only a subset of tokens shows improvement from external memory, we did 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 retrieved memories. We studied arXiv math, Github and Isabelle corpus. For arv math and Github, we found the model retrieved function and variable names. See more details with examples in Appendix B
8

---

Published as a conference paper at ICLR 2022
5
20k
Figure 7: Difference in loss for each token in randomly chosen paper, using the same model once with a memory size of 8K and once with 32K. Higher numbers mean the longer memory helped in comparison to the shorter memory. This paper is 22K tokens long.
Query index
Input
Target
Surrounding context
Retrieved index
Retrieved surrounding context
mark
prob_ space. markov iineuuulttyy
|xic Xa) <<le> expectation X
40919 th = subgraph, threshold a/ 27219
threshold H a powr (-(1 max density'
49699
assumes'
orthonormal system 28050
definition orthonormal. system
Table 8: Examples of memory retricval in the Isabelle dataset. The model is able to find the definition of a lemma from reference to it. The retrieved surrounding context (highlighted) is the definition body of the mathematical object highlighted in the querying context.
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 mathematical object or a lemma, the model looked up the definition from earlier in the proof. Examples of this behavior are shown in Table In example 1, the model retrieves definition within the body of lemma, markov previously defined concept markov_i ieeuualiiiaaa In example 2, it retrieves the definition of aebbgraph_threshoee
In example 3, it retrieves the definition of orthonorma syssssssssssssssssmsssm 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 predict. In the other two cases, the model also looked up materials in the immediate vicinity. To the best of knowledge, this is the first demonstration that attention capable of looking up definitions and function bodies from a large corpus. The Isabelle case study used model with two memory layers of size 32K.
CONCLUSION
We present simple extension to the Transformer architecture, called &NN-augmented attention, which dramatically increases the length of the context that a language model can attend to by using t-nearest-neighhor lookup into a large external memory. We demonstrate the effectiveness of external memory in series of language modeling experiments over 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; comparable to 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; possible obtain large gains from adding memory an existing pre- pre-trained model, and then fine-tuning it.
Unlike other forms of attention, &NN retrieval can easily scaled up to huge memory sizes, and is thus potentially able to leverage vast knowledge bases code repositories. How to make the best use of this capability is topic for future work.
ACKNOWLEDGMENTS
We want to thank Charles Staats for the many fruitful discussions and detailed comments, Henryk Michalewski for early version of of the memory implementation, Petros Maniatis for his help with our code datasets, Aitor Lewkowyez for his help with larger scale memorizing transformer experiments, Behnam Neyshabur for his comments finetuning non-memory models, Imanol Schlag for his proofread and detailed comments, and Dennis Lee and Manzil Zaheer for discussions about large-scale attention and retrieval.

---

Published as a conference paper at ICLR 2022
ETHICS
The ability to memorize large databases of facts could have potential ramifications for society, especially if those databases include sensitive personal information copyrighted works. However, one advantage of using an external memory is that the memory can be easily cleared of all such information, we do the end of each document that we train on. The same not true of differentiable model parameters, which is what most existing architectures use to store facts and information that they are trained on.
REPRODUCIBILITY
Details of our architecture and training hyperparameters are given in Section 4.2 The datasets for C4 and PG-19 are publicly available. Our additional datasets, Github, Isabelle, and ArXiv Math are derived from publicly available data buckets, which we link in the main part of the paper. Subsection [..] include details on how we constructed the datasets from those datasets. We plan to release our code open source.
REFERENCES
Joshua Ainslie, Santiago Ontañón, Chris Alberti, Vaclav Cvicek, Zachary Fisher, Philip Pham, Anirudh Ravula, Sumit Sanghai, Qifan Wang, and Li Yang. ETC: encoding long and structured inputs transformers. In EMNLP. 2020.
Jacob Austin, Augustus Odena, Maxwell Nye. Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang. Carrie J. Cai, Michael Terry, Quoc V. Le, and Charles Sutton. Program synthesis with large language models. CoRR, abss22008.07732, 2021. URL httttttss///xxxiy org/abss.rg//b/.....
2108.077 07732]
Iz. Beltagy, Matthew F. Peters, and Arman Cohan. Longformer: The long-document transformer. CoRR, 0bbs/2004.05150, 2020. URL
05150
James Bradbury, Roy Frostig. Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Miin and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
http: httt:///githb. com/google/// jax
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voon Voss. Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In NeurlPS, 2020.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harrison Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, Dave Cummings, Matthias Plappert, Fotios Chantzis, Elizabeth Barnes, Ariel Herbert- Voss. William Hebgen Guss, Alex Nichol, Alex Paino, Nikolas Tezak, Jie Tang, Igor Babuschkin, Suchir Balaji, Shantanu Jain, William Saunders, Christopher Hesse, Andrew N. Carr, Jan Leike, Joshua Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew Knight, Miles Brundage. Mira Murati, Katie Mayer, Peter Welinder, Bob McGrew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojeiech Zaremba. Evaluating large language models trained on code. CoRR, abx/2107.033 033744 2021. URL pottps://arse 2/55/bb//250 wwt.aa/////2.... 03374
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamás Sarlós, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser. David Ben- jamin Belanger, Lucy J. Colwell, and Adrian Weller. Rethinking attention with performers. In
ICL.R, 2021.
10

---

Published as a conference paper at ICLR 2022
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. CoRR, abs/2110.14168, 2021. URL http:tt:: /tttpps//arxiie org/
Zihang Dai, Zhilin Yang, Yiming Yang. Jaime Carbonell, Quoc Viet Le, and Ruslan Salakhutdinov. Transformer-Xie XL: Attentive language models beyond fixed-length context. In ACL, 2019.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In ACL, 2019.
Angela Fan, Thibaut Lavril, Edouard Grave, Armand Joulin, and Sainbayar Sukhbaatar. Addressing some limitations of transformers with feedback memory. arXiv preprint wrXiv:2002.09402 2020.
Angela Fan, Claire Gardent, Chloé Braud, and Antoine Bordes. Augmenting transformers with KNN-based composite memory for dialog. Transactions of the Association for Computational Linguistics, 9:82-99, 2021.
Edouard Grave, Armand Joulin, and Nicolas Usunier. Improving neural language models with continuous cache. In ICLR, 2017.
Ruiqi Guo, Philip Sun, Erik Lindgren, Quan Geng, David Simcha, Felix Chern, and Sanjiv Kumar. Accelerating large-scale inference with anisotropic vector quantization. In JCML, 2020.
Ankit Gupta, Guy Dar, Shaya Goodman, David Ciprut, and Jonathan Berant. Memory-efficient transformers via top-k attention. CoRR, 0bss/200600899 2021. URL https://axxiv.oo//
05s/21006068993
Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Ming- Wei Chang. Retrieval augmented language model pre-training. In ICML, 2020.
Christopher Hahn, Frederik Schmitt, Jens U. Kreber, Markus Norman Rabe, and Bernd Finkbeiner. Teaching temporal logics to neural networks. In ICLR, 2021.
Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Mare van Zee. Flax: A neural network library and ecosystem for JAX, 2020. URL http:://ithub: com wow.://oole////oooooooooooo flax
Alex Henry, Prudhvi Raj Dachapally, Shubham Shantaram Pawar, and Yuxuan Chen, Query-key normalization for transformers. In EMNLP, 2020.
Jeff Johnson, Matthijs Douze, and Hervé Jégou. Billion-scale similarity search with GPUs. IEEE Transactions on Big Data, 2021.
Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Generalization through memorization: Nearest neighbor language models. In ICLR, 2020.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In ICLR, 2020.
Taku Kudo and John Richardson. Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. In EMNLP, 2018.
Guillaume Lample, Alexandre Sablayrolles, Mare' Aurelio Ranzato, Ludovic Denoyer, and Hervé Jégou. Large memory layers with product keys. In NeurIPS, 2019.
Mike Lewis, Marjan Ghazvininejad, Gargi Ghosh, Armen Aghajanyan, Sida Wang, and Luke Zettlemoyer. Pre-training via paraphrasing. In NeurlPS, 2020a.
Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal,
Heinrich Küttllr, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, et al. Retrieval-augmenter genera- tion for inoeeevee-iteensive NLP tasks. In NewrIPS, 2020b.
Wenda Li, Lei Yu, Yuhuai Wu, and Lawrence C. Paulson. Isarstep: a benchmark for high-level mathematical reasoning. In ICLR, 2021.
11

---

Published as a conference paper at ICLR 2022
Yujia Li, David Choi, Junyoung Chung, Nate Kushman, Julian Schrittwieser, Rémi Leblond, Tom Eccles, James Keeling, Felix Gimeno, Agustin Dal Lago, Thomas Hubert, Peter Choy, Cyprien de Masson d' Autumme Igor Babuschkin, Xinyun Chen, Po-Sen Huang, Johannes Welbl, Sven Gowal, Alexey Cherepanov, James Molloy, Daniel J. Mankowitz, Esme Sutherland Robson, Pushmeet Kohli, Nando de Freitas, Koray Kavukcuoglu, and Oriol Vinyals. Competition-level code generation with alphacode. DeepMind, 2022.
Stanislas Polu and Ilya Sutskever. Generative language modeling for automated theorem proving. CoRR, 0bss/2009003393, 2020. URI.httss: www.a. wwtp:/wwwwwg.aa/////oooo 2209..009..0 03393
Markus Norman Rabe, Dennis Lee, Kshitij Bansal, and Christian Szegedy. Mathematical reasoning via self-supervised skip-tree training. In ICLR, 2021.
Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, Chloe Hillier, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling. In ICLR, 2020.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research, 2020.
Hongyu Ren, Hanjun Dai, Zihang Dai, Mengjiao Yang, Jure Leskovec, Dale Schuurmans, and Bo Dai. Combiner: Full attention transformer with sparse computation cost. CoRR, abs/2107.00766. 2021.
URL https::////mm / waaaaaivii.. org/ http://bb//21..... 05766777.67788
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 9:53-68, 2021.
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost.
In ICML, 2018.
Sainbayar Sukhbaatar, Edouard Grave, Guillaume Lample, Herve Jegou, and Armand Joulin. Aug- menting self-attention with persistent memory. arXiv preprint 1907:190070017000 2019.
Sainbayar Sukhbaatar, Da Ju, Spencer Poff, Stephen Roller, Arthur Szlam, Jason Weston, and Angela Fan. Not all memories are created equal: Learning to forget by expiring. In ICML, 2021.
Simeng Sun, Kalpesh Krishna, Andrew Mattarella-Micke, and Mohit lyyer. Do long-range language models actually use long-range context? In EMNLP, 2021.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. arXiv preprint arXiv:2009.061: 06732, 2020.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: benchmark for efficient transformers. In ICLR, 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. In NeurIPS, 2017.
Qingxiang Wang, Chad Brown, Cezary Kaliszyk, and Josef Urban. Exploration of neural machine translation in autoformalization of mathematics in mizar. In International Conference on Certified Programs and Proofs, 2020a.
Sinong Wang, Belinda 7 Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04776 2020b.
Xun Wang, Haozhi Zhang, Weilin Huang, and Matthew R. Scott. Cross-batch memory for embedding learning. In CVPR, 2020c.
Ronald J. Williams and Jing Peng. An efficient gradient-based algorithm for on-line training of recurrent network trajectories. Neural Computation, 1990.
Dani Yogatama, Cyprien de Masson d' Autume, and Lingpeng Kong. Adaptive semiparametric language models. ACI., 9:362-373, 2021.
12

---

Published as a conference paper at ICLR 2022
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big bird: Transformers for longer sequences. In NeurIPS, 2020.
Yury Zemlyanskiy, Joshua Ainslie, Michiel de Jong, Philip Pham, Iya Eckstein, and Fei Sha. Readtwice: Reading very large documents with memories. In ACL: Human Language Technologies, 2021.
Zhenhai Zhu and Radu Soricut. H-transformer- 1d: Fast one-dimensional hierarchical attention for sequences. In ACL, 2021.
13

---

Published as a conference paper at ICLR 2022
A
LENGTH OF INPUTS
04
50c
100%
150K
200k
250k
50ok
350k
400k
450k
500c
Figure Histogram of the number tokens arXiv math papers dataset. We tuncated the histogram at 500k tokens. The maximum paper had almost 1.6M tokens.
200%
400k
300k
600K
/0ok
BOlk
900k
1000k
Figure 10: Histogram of the number of tokens in Github repositories dataset. We cut off the long tail of this plot. The repository with the maximum length has just over 9M tokens.
0c
30k
100k
150k
230k
250k
300K
J50k
450k
500c
Figure 11: Histogram of the number of tokens in Isabelle proof scripts dataset.
400k
Figure 12: Histogram of the number of tokens in PG19 books dataset.
30k
100k
Figure 13: Histogram of the number of tokens in C4 documents filtered by documents that have less than 4096 tokens.
14

---

Published as a conference paper at ICLR 2022
A.1 ABLATION STUDIES
In the following section, we performed ablation studies to investigate the effects of various hy- perparameters. Unless otherwise specified, carried these experiments with memorizing transformer with context size 512, XL cache 512 with a memory size of 8192.
Multiple NNN layers. We experimented with using two kNN layers, rather than just one. However, we did not see further benefits brought by more than multiple retrieval layers.
KNN layer index We experimented with adding the external memory to layer 3, 6, 9 and 12 in a 1 12-layer transformer, with results shown in Table We found that adding memory the middle of the layer stack will obtain the best result, whereas adding memory to layers either too close to the input or the output obtained less gains.
Table 14: Different layer index.
Layer index
Perplexity
3
2.40
6
2.36
9
2.37
12
2.43
Number of neighbors We studied the effects of the number of neighbors we retrieve from memory. with results shown in T5] We found that even with 32 number of neighbors, we can already obtain a comparable results with 128 or 256 neighbors.
Table 15: Number of neighbors.
Number of neighbors
Perplexity
32
2.38
128
2.37
256
2.37
Random seeds We measured the statistical significant of the results reported. We did runs with 3 random seeds for Transformer XL. of size 512, and also a memorizing transformer with memory size 81 8192. We measured the standard deviation of perplexities after 500K steps of training, shown in Table We saw the standard deviation between different runs of the same experiment appears to be much smaller than the gap between different models.
Table 16: Random seeds.
Models
Perplexity
Transformer XL.
2.67 + 0.01
Memorizing Transformer
2.37 + 0.005
15

---

Published as a conference paper at ICLR 2022
B WHAT DOES THE MODEL RETRIEVE FROM MEMORY?
Retrieving citation names On arXiv math, several examples are shown in Table 17 which includes both the retrieved token and its surrounding context. We observe that many of the gains in cross- entropy loss took place when trying to predict the name of bibitems, citations, or references, by looking up the references and citations used previously in the paper. Such lookups usually span over the entire paper, which is much longer than 8192 tokens, providing plausible explanation for the gain beyond memory size of 8192.
Table 17:
The table shows several examples which tokens were retrieved during language modelling of arXiv math dataset. The model is retrieving names of the references from previous passages.
Query index
Input
Target
Surrounding context
Retrieved index
Retrieved surounding context
20389
Mon
thus
bibitem!
ComtedMonthus Yor
2208
Brownian motion |cite| ComtecMonthus Yor r-1.2 of (citel chakrabarti
16623
cha
kra
|cite( chakrabarti }-
4677
14747
d
legreff asdig
which
3365
begin (equieiqu))) (n |Jabel| asdfg .1)
Retrieving function names from the codebase As with the arXiv papers, we also studied which tokens the model retrieved from memory. As might be expected, the model is often looking up the names of functions, and variables, as shown in Table 18
Table 18: Examples of memory retrieval in the Github dataset. The model looks up how functions are used elsewhere in the repository.
Query index
Input
Target
Surrounding context
Retrieved index Retrieved surrounding context
23837
Fo
nic
menu menu_play-> sctarFonte
14607
menu load-> sctarFonte
23825
35
hscreen/2-50, 50, 200.
14599
20, y+40, 200,
14546
adi
panel-> adicionaCompooeente
5205
panel-> adicionaCommoneett
16

---

Published as a conference paper at ICLR 2022
B.1
MORE RETRIEVING EXAMPLES FORMAL THEOREM PROVING CORPUS
Example
Input token index: 64604
Input token:
Target token: "pair"
Surrounding context:
)) by (simp add: Fourier_sum_lingg Laiimn_lmii_ppaaaa pair [OF f, symmetric] Fourier' pair
Name needs to be predicted:
Retrieved token: "Four"
Retrieved token index: 64412
Retrieved context: 2 Fourier f trigonometric_set k
Definition of the name:
Lenna Fourier sum Limit pair:
assunes absolutely integrable on {ip'..pi)'
shows "(An. ks
Fourier_coefficien Fourier coefficient k trigononetric set
(An. ksn. Fourier coefficient _eefier.coefficient
trigonometric set
(is "?ths ?rhs")
Figure 19: Definition of Fourier_ssi sum 1imit pair.
Example 2
Input token index: 46175
Input token: "tri"
Target token: "gon"
Surrounding context: <le>n. k * trigonometrie_sete ket kx)
Name needs to be predicted:
Retrieved token: "gon"
Retrieved token index: 35457
Retrieved context: lemma orthonormal sastnoomml_ssstemmmm systsyysssssses trigonometric_sett set::, "orthonormal
Definition of the name:
Lemma orthonormal system trigonometric set:
"orthonormal system {pp trigonometric_se set"
Figure 20: Definition of orthonormal www.hooomalssssteeeeeeeeeoeeeeee system systeen.trignnmmrric..e...
17

---

Published as a conference paper at ICLR 2022
Example 3
Input token index: 49760
Input token: "sum""
Target token: "m"
Surrounding context: nusing squrier squrresssseree summable [OF assms, of"
Name needs to be predicted: Fourier Fourlerr_serees_suurre.._nn
Retrieved token: "sum"
Retrieved token index: 35457
Retrieved context: lemma Fourier_scries square summable n assumes:
Definition of the name:
series square
assunes os: oorthnnomal systen and f: "f square integrable ş*
and w: 'A1. (M 1) square integrable
shows *surnable (confine (orthanormal coeff S i) I)"
Figure 21: Definition of Fourier_series_ summable..
Example 4
Input token index: 49697
Input token:
Target token: "system"
Surrounding context: lemma Riemann Riemann_lebesgur Jebesgue_suuareeeeeeeoeeeeeeeale square, squaree integrable:
nassumes "orthonormal_ systtmm_smmmmmm w
Name needs to predicted: arthenormal_sysee
Retrieved token: "system"
Retrieved token index: 28052
Retrieved context: definition orthonormalsystemm :: "\ daaceuclidean"
Definition of the name:
definition orthonormal system :: 12product (M then else 0)"
space set = ("b "a real) + bool"
system
Figure 22: Definition of orthonormai system.
18

---

Published as a conference paper at ICLR 2022
Example 5
Input token index: 34817
Input token: "*.
Target token: "b"
Surrounding context: shows "integrable (lebesgue_oo (a..b))
Retrieved token I: "!"
Retrieved token index 1: 2416
Retrieved context 1: lebesgue__og on fa..
Retrieved token 2:
Retrieved token index 2445
Retrieved context 2: (lebesgue_on (a-c..bcc)) (
Retrieved token 3: "L"
Retreived token index 3: 6479
Retrieved context 3: (lebesgue_on (-pi..pi))
Example
Input token index: 49759
Input token:
Target token: "sum"
Surrounding context: On using Fourier Fourier_series, peries_square, summable [OF assms
Retrieved token "set"
Retrieved token index 35044
Retrieved context 1: definition trigonometric_set:: "nat \<Rightarrow>
Retrieved token 2: "ier"
Retrieved token index 2: 47272
Retrieved context 2: definition Fourier nnwer jnwheee
Retrieved token "ine"
Retrieved token index 3: 18160
Retrieved context 3: lemma Schwartz Schwartz_inequaltyy strriittyssrotg''" nassumes "f"
Retrieved token 4: "system"
Retrieved token index 4: 28052
Retrieved context 4: definition orthonormall systelssssmm__ :: "I "Vaseuucideann
Retrieved token "<"
Retrieved token index 5: 47241
Retrieved context 5: subsection's <open>Connvrrencce wrt the L'
Retrieved token 6: "n"
Retrieved token index 6: 40835
Retrieved context subsection) -cobeee <open>> bit of extra'
19


Reviews:
[{"review_id": "r2wKYBSrEPu", "reviewer": "Reviewer_eMDi", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper proposes a novel method to connect a k-NN based memory with transformers for language modeling. The idea seems interesting and somewhat novel and evaluated empirically. The previous work is nicely identified, although the work of Rae et al. (ICLR 2020) could be described there as well. \n\nStrong points:\n- The memory seems to be usefully working and having a positive effect.\n- Practically, seems to work.\n- Relevant work to this community.\n\nWeak points:\n- Section 3 could see more formalization and better description on what are the input for searching the kNN, how the model is trained (maybe, an algorithm figure).\n- Experimentation could be greatly improved: no hyperparameter search or reasoning behind the selection, no baselines with other memory architectures showing differences either performance or specific use-cases (beyond plain transformers).\n- Memory sizes are too big for the data selected. It is unclear why such big memories are still improving the model.\n- The results on PG-19 doesn't seem to replicate previous work (original paper shows 33.6 for Compressive transformer, whereas this work shows ~19 for plain transformers). \n- Results are very close for each memory size and there are no error measurements (e.g., std dev). It is hard to conclude anything about memory sizes and their effectiveness. \n- Hard to replicate: several proposed datasets are downloaded from the web, it will be hard to download the same files.\n- No details on the specific hardware and software used.\n\nMinor:\n- What do you refer with \"16 attention heads of size 64\"?\n- Describe the content of the figures in the caption (e.g., what is the architecture about?)", "summary_of_the_paper": "This work proposes to deal with long-range dependencies in sequences using a transformer language model augmented with memory. The memory comes in the form of a k-Nearest Neighbor (kNN). The memory stores the last M  (Key, Values) from the previous-to-last layer, and proposes k neighbors from the memory per element in the sequence to attend to the last layer, in addition to the current ones in the sequence. The last layer is trained to utilize the memories. The memory is non-differentiable, but it is still possible to be updated, and used an \"input\" to a layer in the network. The authors test their model with math paper from arXiv (in Latex), PG-19, C4, Github code, and Isabelle (theorem proving). \nThe experiments measure performance of the language models with perplexity. The context (input) length is typically fixed to a relatively small number (512). The paper studies effects of memory vs no-memory,  different sizes of memory, and fine-tuning by training on small memory first followed by an increase in memory size. An alternative with 2 k-NN memories for the last 2 layers was tested too. Further, a few examples show how the memory is most usefully used (mostly for predicting rare tokens). The experiments show improvements in perplexity when using memory, and some saturation points.", "summary_of_the_review": "This is an empirical work that aims at improving the processing of long-range dependencies in sequential data. \nI've updated my score  recommendation is to accept the paper given that the additional experiments and standard deviations are finalized and included in the paper. I would recommend to further include a word-level tokenization of the PG-19 dataset (and maybe others) to make them comparable to previous work [Rae et al.] and more challenging to the model. This would allow readers understand whether such benefits are indeed beyond the compressive transformer or other work. Also, as the authors are using character-level language modeling, it would be much better to use the bits-per-character metric instead of perplexity. \n\nMy recommendation is to reject this work given the poor quality of the experimental section (see above and below). Also, there will be some challenges to reproduce this work with respect to the datasets that should be addressed as well. \n\nQuestions:\n- Please, could you clarify how did you choose the hyperparameters and justify them?\n- What is being used as search in k-NN? The output of the layer connected to the memory? Are the keys used in the search or both key/values?\n- Why did you choose k=128 in k-NN? What would happen if these value changes?\n- Why do you use one layer of memory, if two improve the results?\n- Could you add std dev. values to all your experiments? \n- Please explain why does the arXiv dataset have such a low perplexity?\n- How many sequences are used for training, validation and testing?\n\n==================================\n\nUpdate:\n\nI've updated my score based on the rebuttal discussion with the authors, and the additional elements presented. The new score is based on the existing merit, however, the contributions are limited and the experimental results require further work. The additional experiments could strongly benefit the understandings of the benefits of this method. For example, finalizing the experiment with 2K vs 500 + 1.5K memory with transformer and transformer-XL, would add lots of value. The presented results are difficult to compare with previous work (other than the included transformer-XL) due to the selected tokenizer. Therefore, including other models, e.g., compressive transformer, is missing. This could have been avoided by reproducing the experiments from those works. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "HL_YDGg7lSw", "reviewer": "Reviewer_gVad", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\n- Integrating kNN module into Transformer is an interesting modeling technique, which would be potentially useful for language generation especially when the context is quite long.\n\n- The kNN layer is integrated into the attention layer, while previous work usually integrates it into FFN blocks.\n\n- The LM experiments are conducted on several corpora.\n\nWeaknesses:\n\n- The method is not compared with previous models in the experiments, such as Transformer-XL, Compressive Transformers, and other efficient Transformers for LM.\n\n- There is no ablation study in terms of modeling.\n\n- I didn't find clear comparisons of training speed and GPU memory usage between different models. It is valuable for the readers to have a sense of what kind of resource is needed to train or host such models.\n\n- Is the application restricted to sequence decoding? Can we finetune the proposed model as BERT on downstream tasks, and how about the performance?\n\n- The above question is related to generalization. Such kind of kNN method seems to suffer from too much memorization instead of generalization. If the model cannot generalize well, it would not be good at zero- or few-shot learning. I agree that language modeling is a meaningful task, but it is more like a testbed for modeling techniques that improve generalization. Otherwise, the shortcut of memorizing context would hide the real challenge. \n\n- The code seems unavailable. It is unclear how many engineering efforts are required to reproduce the results. The paper writing ignores many details that are helpful for reproducing the method. Even the code would be released, the standalone paper should be informative enough.\n\n- The kNN layer is densely triggered for each position, which seems quite heavy. How about using a gate to determine whether we need to perform kNN?\n\n- Is there any evaluation about the performance of the learned kNN search? For example, we could conduct some human evaluations to see whether the retrieved knowledge makes sense for prediction.\n\n- The model is only evaluated on the language modeling task. More tasks would be helpful to indicate whether the proposed method is general enough in terms of data distribution.\n", "summary_of_the_paper": "The paper plugs a kNN memory module into a Transformer for long-distance knowledge. Specifically, the bottom layers are still Transformer blocks, then a kNN-augmented attention layer is used to store the processed (key, value) pairs into external memory, and another Transformer block is put at the top layer for aggregation. The LM experiments are conducted on webtext (C4), math papers (arXiv), books (PG-19), code (Github), as well as formal theorems (Isabelle).", "summary_of_the_review": "I am okay to accept the submission. But it can be improved as mentioned in the above review (e.g., including ablation studies, more tasks, comparisons with previous models, and analysis).", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "86eil-gt7wU", "reviewer": "Reviewer_75xx", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\n1. The method proposed is novel, intuitive and simple to implement. This is a huge advantage in the language modeling space. We’ve seen many memory-augmented transformers in the past few years, and most of them have failed to have significant impact because they are too complicated to replicate or take too many resources to train. \n2. The results on the C4 dataset seem significant and show that  memory augmentation could really improve transformer language modeling. \n3. The analysis is super interesting, and it shows the potential of this simple approximate attention method to find very relevant contexts to retrieve from the recent past, and how these retrievals aid in improving the LM. \n4. The paper is well-written and easy to understand. \n\nWeaknesses:\n\n1. The authors do not compare their model to any previously published model. Throughout the paper the authors only compare their model to baselines that they designed and trained. Doing this for some of the datasets is fine, but not even having one dataset in which you compare to strong, previously published models makes it harder to convince the reader that you have found something interesting. You did show results for PG19, which has been used before in LMing papers, so it’s not clear why you haven’t compared to previously published models on that dataset, such as the compressive transformer? The numbers in the compressive transformer are very different from the numbers your models obtain, so I’m assuming that you preprocessed that dataset differently and so have a different vocabulary, making those results incomparable? Could you re-train your model using the same settings as the compressive transformer so that we could compare the results of your method and theirs?\n2. The baseline has just 6 layers, which is significantly smaller than the current transformer LMs that are being discussed by the community. I’m definitely not expecting to see GPT-3 sized models in every research paper (or even close to that) but I think 16 layers would make more sense for such an exploration. I’m afraid the gains shown here would not transfer to larger models. On the other hand, larger models might produce better representations resulting in even bigger gains. \n3. The authors do not show standard deviation numbers for any of the baselines on the various datasets. This makes understanding the results on each dataset, for the different models, harder. For example, on the arxiv dataset the authors talk about a drop from 2.927 to 2.877 perplexity between two different models, but is a 0.05 perplexity difference really statistically significant here? \n4. The authors don’t provide any speed or memory usage statistics. Of course, it’s totally fine if this model runs slower or uses more memory than the baseline- I expect such preliminary investigations to potentially be nonoptimal, with later work making it an actually viable model. I would not penalize this paper for having slow runtime or using too much memory. But I think it's important to discuss these things. \n\n\nMisc:\n\n1.  Although the perplexity improvements on many of the datasets are quite small (<=.1 perplexity) this could be a non-issue, and more of a problem with the perplexity metric than with the model. The perplexity metric, and most of the data that we train our LMs on, heavily incentivizes models to get short term dependencies right. Getting the model to look back into the past is both hard and barely incentivized, and so these small initial gains that are attained in this paper might be improved when we find a better loss function that incentivizes looking back more, or when we find data that has more long-term dependencies and less short-term ones. \n2. The way I see it, there are two different ways to augment the transformer with a memory module. The memory module could be ‘global’ or ‘local’. Local memory transformers have a memory module that can only access the recent past. So models like the compressive transformer or Transformer-XL are local memory transformers. On the other hand, models like the kNN-LM or REALM are ‘global’ memory models, since they can access the entire datastore for every prediction, without any kind of limitation that is dependent on the current timestep and the recent previous timesteps.  \nThe model proposed here is clearly a ‘local’ memory transformer, since it performs approximate attention over the recent past. I believe that that is part of what makes this model work, and I think the paper might benefit from a discussion of this and maybe even an empirical test. Does this model improve performance on WikiText-103, where all the articles are in a shuffled ordering? I don’t think it would, and that’s totally fine, but I think readers would benefit from a discussion of this property of the model.  \n3. I think it could be helpful if the authors had more comparisons where their model and the baseline can access the same amount of tokens for every prediction. For example, if you train one of the memory models with 2k tokens of memory and 512 tokens of local context it would be interesting to compare it to the baseline model trained on 2.5k tokens. Of course, since your model performs approximate attention we do expect to see worse perplexity values, but it would be very insightful to see precisely how much using this approximate attention method degrades performance.  \n4. For section 3.2, consider referencing [CrossBatch] which discusses this distributional shift in section 3.2 of their paper, in the vision domain. \n\nReferences:\n\n[Compressive Transformer] Jack W. Rae, Anna Potapenko, Siddhant M. Jayakumar, Timothy P. Lillicrap. “Compressive Transformers for Long-Range Sequence Modelling” \n\n[REALM] Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, Ming-Wei Chang. “REALM: Retrieval-Augmented Language Model Pre-Training”\n\n[CrossBatch] Xun Wang, Haozhi Zhang, Weilin Huang, Matthew R. Scott. “Cross-Batch Memory for Embedding Learning” \n", "summary_of_the_paper": "\nTransformer language models attend to the previous N tokens, and so do not have any access or knowledge of the tokens that appeared more than N tokens ago. Extending N would allow language models to have more knowledge which might increase performance. But there’s a tight limit as to how far we can naively extend N, since memory usage grows quadratically with N, and so in SoTA transformer LMs N is limited to a few thousand tokens. \n\nThis paper shows a novel, simple and intuitive method for extending the number of tokens a transformer can access at every timestep. The authors augment the model with a kNN datastore, so that in the top layer, the model attends to the local context using normal attention but then also attends to up to 262k previous tokens by using an approximate nearest-neighbors datastore. This is basically equivalent to approximately attending over those previous tokens. No gradients are passed to the representations in the datastore in any way, making this easy to implement and relatively efficient to train. \n\nThe authors present results on multiple language modeling datasets that have long-range dependencies. On most of the datasets the gains from using the proposed method are quite small (<=.1 perplexity), and it’s not clear that they are statistically significant, but on the C4 dataset the model reduces perplexity by 2 which I believe to be a strong result. I believe that this idea is promising and that giving it a stage at ICLR will push the community to try further improving upon it, hopefully leading to larger gains. \n\n", "summary_of_the_review": "This paper introduces a new method for extending the transformer model with a simple datastore. The insert and query operations into this datastore do not have to be trained which is a big advantage for efficiency and simplifies the implementation. Results are shown on multiple datasets, with the model showing big gains on one and smaller gains on the others. I believe the idea presented is reproducible, and extensible and I’m looking forward to seeing how this research direction develops. The analysis presented is insightful and could guide further research. I believe this paper should be accepted. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "7fv6o2kjgT2", "reviewer": "Reviewer_zRRc", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros:\n- propose a non-differentiable cache mechanism to reuse keys and values coming from prior training steps.\n- comprehensive experiments are conducted to evaluate the effectiveness of the proposed memory-based method.\nCons:\n- there are some benchmarks that require models to be able to handle long inputs, such as document-level QA datasets and document-level retrieval. Besides the LM task, it could be better if more document-level downstream tasks could be evaluated to verify the usefulness of the proposed approach.\n", "summary_of_the_paper": "This paper presents a memory-based Transformer where a memory is used to leverage the external inputs for scenarios that have long inputs like documents or codes. The key contribution is that it describes how external contexts are integrated into the representations of the current inputs. By introducing such external contexts with the memory mechanism, perplexity of LM can be improved on every dataset studied in the paper.", "summary_of_the_review": "This work is highly related to those scenarios, such as document-level QA/retrieval and code completion, where the models should consider and leverage long inputs. The key contribution, based on my understanding, is the joint KNN/local attention part, where both external context and local context are combined. In general, the paper is well written and all details are clear to me. One thing could be further improved is that I suggest to include more downstream tasks to verify the effectiveness of the proposed method.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper studies the problem of dealing with long contexts within a Transformer architecture.
The key contribution is a  kNN memory module that works in concert with  a Transformer by integrating upper layers with additional retrieved context.

The idea is simple  but the execution is good.  While the  idea is reminiscent of other recent work on this topic, and novelty is somewhat borderline, it is practically useful.
Overall, though ambivalent, my recommendation is that the paper should probably be accepted
Decision: Accept (Spotlight)