Original: M EMORIZING T RANSFORMERS
Revised: M EMORIZING T RANSFORMERS

Original: A BSTRACT
Revised: A BSTRACT

Original: Language models typically need to be trained or ﬁnetuned in order to acquire new knowledge, which involves updating their weights.
Revised: Language models typically need to be trained or ﬁnetuned in order to acquire new knowledge, which involves updating their weights.

Original: We instead envision language models that can simply read and memorize new data at inference time, thus acquiring new knowledge immediately.
Revised: We instead envision language models that can simply read and memorize new data at inference time, thus acquiring new knowledge immediately.

Original: In this work, we extend language models with the ability to memorize the internal representations of past inputs.
Revised: In this work, we extend language models with the ability to memorize the internal representations of past inputs.

Original: 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).
Revised: 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).

Original: We show that the performance steadily improves when we increase the size of memory up to 262K tokens.
Revised: We show that the performance steadily improves when we increase the size of memory up to 262K tokens.

Original: On benchmarks including code and mathematics, we ﬁnd that the model is capable of making use of newly deﬁned functions and theorems during test time.
Revised: On benchmarks including code and mathematics, we ﬁnd that the model is capable of making use of newly deﬁned functions and theorems during test time.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Transformers (Vaswani et al., 2017) have led to remarkable progress in natural language processing (Devlin et al., 2019;
Revised: Transformers (Vaswani et al., 2017) have led to remarkable progress in natural language processing (Devlin et al., 2019;

Original: Brown et al., 2020), mathematical reasoning (Polu & Sutskever, 2020; Wang et al., 2020a; Rabe et al., 2021; Li et al., 2021; Hahn et al., 2021; Cobbe et al., 2021), and program synthesis (Austin et al., 2021; Chen et al., 2021; Li et al., 2022).
Revised: Brown et al., 2020), mathematical reasoning (Polu & Sutskever, 2020; Wang et al., 2020a; Rabe et al., 2021; Li et al., 2021; Hahn et al., 2021; Cobbe et al., 2021), and program synthesis (Austin et al., 2021; Chen et al., 2021; Li et al., 2022).

Original: However, transformer performance on many of these tasks is limited by the context length of attention, which is typically short.
Revised: However, transformer performance on many of these tasks is limited by the context length of attention, which is typically short.

Original: The ability to attend to far-away tokens is important in many situations.
Revised: The ability to attend to far-away tokens is important in many situations.

Original: In novels, characters and events are referenced across multiple chapters.
Revised: In novels, characters and events are referenced across multiple chapters.

Original: In source code, references to classes and functions may occur quite far from the places in which they are deﬁned.
Revised: In source code, references to classes and functions may occur quite far from the places in which they are deﬁned.

Original: In theorem proving, proofs make use of previously deﬁned lemmas.
Revised: In theorem proving, proofs make use of previously deﬁned lemmas.

Original: Attention over long sequences is also useful as a form of rapid learning.
Revised: Attention over long sequences is also useful as a form of rapid learning.

Original: Facts and information which are stored in the form of weight matrices must be slowly trained over hundreds of thousands of training steps.
Revised: Facts and information which are stored in the form of weight matrices must be slowly trained over hundreds of thousands of training steps.

Original: By using attention, however, a model can simply memorize facts (e.g. function deﬁnitions) by storing them as (key, value) pairs in long-term memory, and then retrieve those facts later by creating a query that attends to them.
Revised: By using attention, however, a model can simply memorize facts (e.g. function deﬁnitions) by storing them as (key, value) pairs in long-term memory, and then retrieve those facts later by creating a query that attends to them.

Original: In this case, attention acts as a form of information retrieval, allowing the model to look up facts that it has seen previously.
Revised: In this case, attention acts as a form of information retrieval, allowing the model to look up facts that it has seen previously.

Original: We demonstrate that a simple and effective way to increase the size of the attention context is to use approximate k -nearest-neighbor ( k NN) lookup, which is widely used in information retrieval.
Revised: We demonstrate that a simple and effective way to increase the size of the attention context is to use approximate k -nearest-neighbor ( k NN) lookup, which is widely used in information retrieval.

Original: A number of extremely scalable implementations of k NN lookup are available, such as ScaNN (Guo et al., 2020) and Faiss (Johnson et al., 2021).
Revised: A number of extremely scalable implementations of k NN lookup are available, such as ScaNN (Guo et al., 2020) and Faiss (Johnson et al., 2021).

Original: There are two things which distinguish our approach from previous work on long-range attention (c.f.
Revised: There are two things which distinguish our approach from previous work on long-range attention (c.f.

Original: Section 2).
Revised: Section 2).

Original: First, unlike some other approaches, k NN lookup does not do averaging or summarization of tokens at long distances, but retrieves exact values even from the distant context.
Revised: First, unlike some other approaches, k NN lookup does not do averaging or summarization of tokens at long distances, but retrieves exact values even from the distant context.

Original: Second, gradients are not backpropagated into the external memory, which is critical to the scalability of our technique.
Revised: Second, gradients are not backpropagated into the external memory, which is critical to the scalability of our technique.

Original: The keys and values are a function of model parameters, so attempting to backpropagate gradients into external memory would necessarily involve computing all of the keys and values with the current model parameters on every training step.
Revised: The keys and values are a function of model parameters, so attempting to backpropagate gradients into external memory would necessarily involve computing all of the keys and values with the current model parameters on every training step.

Original: 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.
Revised: 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.

Original: With
Revised: With

Original: [Figure]
Revised: [Figure]

Original: P e r p l e x i t y ( a r X i v
Revised: P e r p l e x i t y ( a r X i v

Original: M a t h )
Revised: M a t h )

Original: Transformer
Revised: Transformer

Original: Memorizing Transformer our technique, we are easily able to scale external memory up to sequence lengths of 131k or 262k tokens on a single TPU device, while maintaining a reasonable step time.
Revised: Memorizing Transformer our technique, we are easily able to scale external memory up to sequence lengths of 131k or 262k tokens on a single TPU device, while maintaining a reasonable step time.

Original: We show that model perplexity steadily improves with the size of external memory on a variety of language modelling tasks, including C4 (long documents only), Github code repositories, PGbooks, formal proofs in Isabelle, and arXiv math papers.
Revised: We show that model perplexity steadily improves with the size of external memory on a variety of language modelling tasks, including C4 (long documents only), Github code repositories, PGbooks, formal proofs in Isabelle, and arXiv math papers.

Original: We further show that models can generalize to larger memory sizes than they were trained on: models trained with a small memory show gains from using a much larger memory at inference time.
Revised: We further show that models can generalize to larger memory sizes than they were trained on: models trained with a small memory show gains from using a much larger memory at inference time.

Original: Finally, we show that our models are actually using memory in the way that we had hoped, e.g. by looking up the deﬁnitions of lemmas in a theorem proving corpus.
Revised: Finally, we show that our models are actually using memory in the way that we had hoped, e.g. by looking up the deﬁnitions of lemmas in a theorem proving corpus.

Original: 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.
Revised: 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.

Original: 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 1.
Revised: 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 1.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: A great deal of work has been done on efﬁcient long-range attention mechanisms; see Tay et al.
Revised: A great deal of work has been done on efﬁcient long-range attention mechanisms; see Tay et al.

Original: (2020b;a) recent surveys.
Revised: (2020; 2021) recent surveys.

Original: Sliding windows (Beltagy et al., 2020) use a long sequence, but attend within a smaller window, thus reducing complexity to the window size, rather than total sequence length.
Revised: Sliding windows (Beltagy et al., 2020) use a long sequence, but attend within a smaller window, thus reducing complexity to the window size, rather than total sequence length.

Original: Approximate mechanisms such as Linformer (Wang et al., 2020b), and Performer (Choromanski et al., 2021) refactor the attention matrix by using a different kernel than softmax to obtain O ( N ) complexity.
Revised: Approximate mechanisms such as Linformer (Wang et al., 2020b), and Performer (Choromanski et al., 2021) refactor the attention matrix by using a different kernel than softmax to obtain O ( N ) complexity.

Original: Pooling strategies such as Hierarchical 1D attention (Zhu & Soricut, 2021), and Combiner (Ren et al., 2021) apply pooling or averaging over tokens at longer distances.
Revised: Pooling strategies such as Hierarchical 1D attention (Zhu & Soricut, 2021), and Combiner (Ren et al., 2021) apply pooling or averaging over tokens at longer distances.

Original: Sparse strategies such as
Revised: Sparse strategies such as

Original: Big Bird (Zaheer et al., 2020) select only a subset of tokens to attend to; Routing Transformers (Roy et al., 2021) use clustering to select the subset, while Reformer (Kitaev et al., 2020) relies on hashing.
Revised: Big Bird (Zaheer et al., 2020) select only a subset of tokens to attend to; Routing Transformers (Roy et al., 2021) use clustering to select the subset, while Reformer (Kitaev et al., 2020) relies on hashing.

Original: Hierarchical mechanisms (Ainslie et al., 2020) combine multiple tokens into phrases or sentences to reduce sequence length.
Revised: Hierarchical mechanisms (Ainslie et al., 2020) combine multiple tokens into phrases or sentences to reduce sequence length.

Original: Expire-span (Sukhbaatar et al., 2021) prunes far-away tokens that it learns are “unimportant”.
Revised: Expire-span (Sukhbaatar et al., 2021) prunes far-away tokens that it learns are “unimportant”.

Original: (Zemlyanskiy et al., 2021) process long sequences in two passes with different encoders.
Revised: (Zemlyanskiy et al., 2021) process long sequences in two passes with different encoders.

Original: The second pass is given a lot of context by accessing summaries of the ﬁrst pass.
Revised: The second pass is given a lot of context by accessing summaries of the ﬁrst pass.

Original: Feedback transformers (Fan et al., 2020) use a recurrent architecture in which each token attends to the output of the ﬁnal layer instead of the previous layer.
Revised: Feedback transformers (Fan et al., 2020) use a recurrent architecture in which each token attends to the output of the ﬁnal layer instead of the previous layer.

Original: Recurrence does not increase the size of the attention context itself, but it expands the receptive ﬁeld at the cost of parallelism and training speed.
Revised: Recurrence does not increase the size of the attention context itself, but it expands the receptive ﬁeld at the cost of parallelism and training speed.

Original: Truncated backpropagation through time (Williams & Peng, 1990) was originally introduced as a way of training recurrent neural networks (RNN) over very long sequences, when the entire sequence does not ﬁt in memory.
Revised: Truncated backpropagation through time (Williams & Peng, 1990) was originally introduced as a way of training recurrent neural networks (RNN) over very long sequences, when the entire sequence does not ﬁt in memory.

Original: The sequence is chopped into segments, and after each training step, the ﬁnal
Revised: The sequence is chopped into segments, and after each training step, the ﬁnal

Original: RNN state for the segment is saved in a non-differentiable cache, and used as the initial state on the next training step.
Revised: RNN state for the segment is saved in a non-differentiable cache, and used as the initial state on the next training step.

Original: Neural caches (Grave et al., 2017) extend the cache to contain a record of many prior hidden states, and attend over them.
Revised: Neural caches (Grave et al., 2017) extend the cache to contain a record of many prior hidden states, and attend over them.

Original: Transformer-XL (Dai et al., 2019) applies this technique to transformers; it caches the (key,value) pairs computed from the previous training step, and uses them as a preﬁx for the tokens on the next training step, which yields signiﬁcant gains on long documents.
Revised: Transformer-XL (Dai et al., 2019) applies this technique to transformers; it caches the (key,value) pairs computed from the previous training step, and uses them as a preﬁx for the tokens on the next training step, which yields signiﬁcant gains on long documents.

Original: Rae et al.
Revised: Rae et al.

Original: (2020) improve over Transformer-XL by compressing the tokens before adding them to the cat in the hat cat in the hat the said local attention + FFN k nearest neighbor lookup.
Revised: (2020) improve over Transformer-XL by compressing the tokens before adding them to the cat in the hat cat in the hat the said local attention + FFN k nearest neighbor lookup.

Original: dense attention + FFN … more layers … input tokens external memory: cached (key, value) pairs
Revised: dense attention + FFN … more layers … input tokens external memory: cached (key, value) pairs

Original: Will be added to external memory after the current training step.
Revised: Will be added to external memory after the current training step.

Original: kNN & local attention + FFN dense attention + FFN local attention + FFN cat in the hat cat in the hat the to output predictions local context kNN attention embedding layer softmax cache.
Revised: kNN & local attention + FFN dense attention + FFN local attention + FFN cat in the hat cat in the hat the to output predictions local context kNN attention embedding layer softmax cache.

Original: In contrast, we use a very large cache without compression, combined
Revised: In contrast, we use a very large cache without compression, combined

Original: Sukhbaatar et al.
Revised: Sukhbaatar et al.

Original: (2019) make the observation that the feed-forward portion of a transformer layer functions very much like attention if one replaces the ReLU activation with softmax.
Revised: (2019) make the observation that the feed-forward portion of a transformer layer functions very much like attention if one replaces the ReLU activation with softmax.

Original: They implement a combined attention over both tokens from the input sequence and a learned (and differentiable) “memory”.
Revised: They implement a combined attention over both tokens from the input sequence and a learned (and differentiable) “memory”.

Original: Lample et al.
Revised: Lample et al.

Original: (2019) exploit this observation to replace the feed-forward layers (FFNs) with a fast k NN lookup over a much larger “memory”, and achieve large gains in model accuracy without signiﬁcant computation overhead.
Revised: (2019) exploit this observation to replace the feed-forward layers (FFNs) with a fast k NN lookup over a much larger “memory”, and achieve large gains in model accuracy without signiﬁcant computation overhead.

Original: (We use k NN lookup to approximate attention to previous tokens, not to replace the FFN.)
Revised: (We use k NN lookup to approximate attention to previous tokens, not to replace the FFN.)

Original: Non-differentiable external memory has been used in different ways by Khandelwal et al.
Revised: Non-differentiable external memory has been used in different ways by Khandelwal et al.

Original: (2020), who run a pre-trained model over an entire corpus, and construct a large table of (key, token) pairs.
Revised: (2020), who run a pre-trained model over an entire corpus, and construct a large table of (key, token) pairs.

Original: They then use that table to replace the ﬁnal softmax layer for token selection in the model, which results in signiﬁcant improvements in language modeling.
Revised: They then use that table to replace the ﬁnal softmax layer for token selection in the model, which results in signiﬁcant improvements in language modeling.

Original: Yogatama et al.
Revised: Yogatama et al.

Original: (2021) extend this approach by a gating mechanism and a process to compress the context into keys for retrieval.
Revised: (2021) extend this approach by a gating mechanism and a process to compress the context into keys for retrieval.

Original: There are several works that combine retrieval with transformers.
Revised: There are several works that combine retrieval with transformers.

Original: REALM (Guu et al., 2020),
Revised: REALM (Guu et al., 2020),

Original: MARGE (Lewis et al., 2020a), RAG (Lewis et al., 2020b), and composite memory for dialog (Fan et al., 2021) retrieve documents from a knowledge base to improve question answering or dialogue.
Revised: MARGE (Lewis et al., 2020a), RAG (Lewis et al., 2020b), and composite memory for dialog (Fan et al., 2021) retrieve documents from a knowledge base to improve question answering or dialogue.

Original: The knowledge base consists of text snippets and is static and typically separate from the inputs and outputs of the models.
Revised: The knowledge base consists of text snippets and is static and typically separate from the inputs and outputs of the models.

Original: Instead, we focus on language modeling using a decoder-only model, and propose a simple model that uniﬁes attention and retrieval .
Revised: Instead, we focus on language modeling using a decoder-only model, and propose a simple model that uniﬁes attention and retrieval .

Original: k -nearest-neighbor lookup is a general-purpose technique that is used for a wide variety of machine learning and retrieval tasks, and high-performance implementations are available for various architectures (Johnson et al., 2021; Guo et al., 2020).
Revised: k -nearest-neighbor lookup is a general-purpose technique that is used for a wide variety of machine learning and retrieval tasks, and high-performance implementations are available for various architectures (Johnson et al., 2021; Guo et al., 2020).

Original: Memory-efﬁcient Transformers (Gupta et al., 2021) replace dense attention with a k NN lookup to increase speed and reduce memory usage.
Revised: Memory-efﬁcient Transformers (Gupta et al., 2021) replace dense attention with a k NN lookup to increase speed and reduce memory usage.

Original:  M ETHOD
Revised:  M ETHOD

Original: The architecture of our k NN-augmented transformer is shown in Figure 2.
Revised: The architecture of our k NN-augmented transformer is shown in Figure 2.

Original: The bulk of the model is a vanilla, decoder-only transformer (Vaswani et al., 2017).
Revised: The bulk of the model is a vanilla, decoder-only transformer (Vaswani et al., 2017).

Original: The input text is tokenized, and the tokens are embedded into vector space.
Revised: The input text is tokenized, and the tokens are embedded into vector space.

Original: The embedding vectors are passed through a series of transformer layers, each of which does dense self-attention, followed by a feed-forward network (FFN).
Revised: The embedding vectors are passed through a series of transformer layers, each of which does dense self-attention, followed by a feed-forward network (FFN).

Original: Since this is a decoder-only language model, we use a causal attention mask and the token embeddings of the last layer are used to predict the next token.
Revised: Since this is a decoder-only language model, we use a causal attention mask and the token embeddings of the last layer are used to predict the next token.

Original: Long documents are split into subsequences of 512 tokens, and each subsequence is used as the input for one training step.
Revised: Long documents are split into subsequences of 512 tokens, and each subsequence is used as the input for one training step.

Original: In contrast to standard practice, we do not shufﬂe the subsequences; instead, each long document is fed into the transformer sequentially, from beginning to end, as is done with
Revised: In contrast to standard practice, we do not shufﬂe the subsequences; instead, each long document is fed into the transformer sequentially, from beginning to end, as is done with

Original: Transformer-XL (Dai et al., 2019).
Revised: Transformer-XL (Dai et al., 2019).

Original: [Figure]
Revised: [Figure]

Original:  E
Revised: Document E

Original: Document F ...
Revised: Document F ...

Original: [Figure]

Original:  G ...
Revised: Document G ...

Original: [Figure]
Revised: [Figure]

Original: We also use  keys the
Revised: We also use a Transformer- cache, which holds keys the

Original: When doing self-attention, the cached keys and  prepended to the  a sliding-window causal mask (Beltagy et al has a local context  includes the previous
Revised: When doing self-attention, the cached keys and values are prepended to the current and values, and we use a sliding-window causal mask (Beltagy et al., 2020) so that each token has a local context that includes the previous

Original: 3. k NN- AUGMENTED A TTENTION L AYER
Revised: 3. k NN- AUGMENTED A TTENTION L AYER

Original: One of the transformer layers near the top of the stack is a k NN-augmented attention layer , which combines two forms of attention.
Revised: One of the transformer layers near the top of the stack is a k NN-augmented attention layer , which combines two forms of attention.

Original: 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.
Revised: 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.

Original: Unlike the other layers, however, it also does an approximate k -nearest-neighbor search into the external memory .
Revised: Unlike the other layers, however, it also does an approximate k -nearest-neighbor search into the external memory .

Original: The same queries are used for both the local context, and for the external memory.
Revised: The same queries are used for both the local context, and for the external memory.

Original: 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.
Revised: 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.

Original: If the document is very long, old (key, value) pairs will be dropped from the memory to make room for new ones.
Revised: If the document is very long, old (key, value) pairs will be dropped from the memory to make room for new ones.

Original: Thus, for each head, the external memory keeps a cache of the prior M (key, value) pairs, where M is the memory size.
Revised: Thus, for each head, the external memory keeps a cache of the prior M (key, value) pairs, where M is the memory size.

Original: The k NN lookup will return a set of retrieved memories , which consist of the top- k (key, value) pairs that k NN search returns for each query (i.e. each token) in the input subsequence.
Revised: The k NN lookup will return a set of retrieved memories , which consist of the top- k (key, value) pairs that k NN search returns for each query (i.e. each token) in the input subsequence.

Original: As with standard dense attention, we ﬁrst construct an attention matrix by computing the dot product of each query against the retrieved keys, then apply softmax, and ﬁnally return a weighted sum of the retrieved values.
Revised: As with standard dense attention, we ﬁrst construct an attention matrix by computing the dot product of each query against the retrieved keys, then apply softmax, and ﬁnally return a weighted sum of the retrieved values.

Original: Unlike standard dense attention, the retrieved memories contain a different set of (key, value) pairs for each query.
Revised: Unlike standard dense attention, the retrieved memories contain a different set of (key, value) pairs for each query.

Original: Attention over the local context is performed in the usual way.
Revised: Attention over the local context is performed in the usual way.

Original: The results of k NN-attention and local attention are then combined using a learned gate:
Revised: The results of k NN-attention and local attention are then combined using a learned gate:

Original: [Equation]
Revised: [Equation]

Original: where σ is the sigmoid function, and (cid:12) is element-wise multiplication.
Revised: where σ is the sigmoid function, and (cid:12) is element-wise multiplication.

Original: V a is the combined result of attention, V m is the result of attending to external memory, and V c is the result of attending to the local context.
Revised: V a is the combined result of attention, V m is the result of attending to external memory, and V c is the result of attending to the local context.

Original: The bias b g is a learned per-head scalar parameter, which allows each head to choose between local and long-range attention.
Revised: The bias b g is a learned per-head scalar parameter, which allows each head to choose between local and long-range attention.

Original: 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.
Revised: 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.

Original: We did observe that over time, most heads learned to attend almost exclusively to external memory.
Revised: We did observe that over time, most heads learned to attend almost exclusively to external memory.

Original: Position bias.
Revised: Position bias.

Original: For dense attention within the local context, we use the T5 relative position bias (Raffel et al., 2020).
Revised: For dense attention within the local context, we use the T5 relative position bias (Raffel et al., 2020).

Original: As noted by Dai et al.
Revised: As noted by Dai et al.

Original: (2019), adding a global position encoding to each token does not work well when processing long documents.
Revised: (2019), adding a global position encoding to each token does not work well when processing long documents.

Original: We don’t use a position bias for the retrieved memories.
Revised: We don’t use a position bias for the retrieved memories.

Original: Experiments on the PG19 dataset (Sun et al., 2021) 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.
Revised: Experiments on the PG19 dataset (Sun et al., 2021) 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.

Original: Batching.
Revised: Batching.

Original: Figure 3 illustrates how multiple long documents of different lengths are packed into a batch, and split into subsequences.
Revised: Figure 3 illustrates how multiple long documents of different lengths are packed into a batch, and split into subsequences.

Original: 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.
Revised: 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.

Original: 3. D ISTRIBUTIONAL S HIFT
Revised: 3. D ISTRIBUTIONAL S HIFT

Original: 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.
Revised: 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.

Original: The model parameters that produce the queries change over time, and will thus have shifted since the keys and values were stored.
Revised: The model parameters that produce the queries change over time, and will thus have shifted since the keys and values were stored.

Original: For very large memories, older records may become “stale.”
Revised: For very large memories, older records may become “stale.”

Original: Similar observations have been made for
Revised: Similar observations have been made for

Original: CrossBatch memory (Wang et al., 2020c) in the vision domain.
Revised: CrossBatch memory (Wang et al., 2020c) in the vision domain.

Original: To reduce the effects of staleness, we normalize keys and queries (Henry et al., 2020).
Revised: To reduce the effects of staleness, we normalize keys and queries (Henry et al., 2020).

Original: Normalization does not eliminate staleness, but it at least ensures that older keys and newer keys do not differ in magnitude.
Revised: Normalization does not eliminate staleness, but it at least ensures that older keys and newer keys do not differ in magnitude.

Original: We also found that normalization helps stabilize training with the Transformer-XL cache.
Revised: We also found that normalization helps stabilize training with the Transformer-XL cache.

Original: In some of our experiments, we observed that training models from scratch with a large memory sometimes resulted in worse performance than pretraining the model with a small memory of size 8192, and then ﬁnetuning it on a larger memory.
Revised: In some of our experiments, we observed that training models from scratch with a large memory sometimes resulted in worse performance than pretraining the model with a small memory of size 8192, and then ﬁnetuning it on a larger memory.

Original: This training instability could be due to staleness.
Revised: This training instability could be due to staleness.

Original: However, models seem to be able to cope with a limited degree of staleness (with the small memory) by adjusting their queries accordingly.
Revised: However, models seem to be able to cope with a limited degree of staleness (with the small memory) by adjusting their queries accordingly.

Original: 3. A PPROXIMATE k NN
Revised: 3. A PPROXIMATE k NN

Original: We employ approximate k NN search rather than exact k NN search because it signiﬁcantly improves the computational speed of our model.
Revised: We employ approximate k NN search rather than exact k NN search because it signiﬁcantly improves the computational speed of our model.

Original: We use a simple approximation of k NN for TPUs, which has a recall of about 90%, i.e. 90% of the true top k are returned in the approximate top k .
Revised: We use a simple approximation of k NN for TPUs, which has a recall of about 90%, i.e. 90% of the true top k are returned in the approximate top k .

Original: There are various other efﬁcient approximate k NN algorithms available for CPU and GPU/TPU, for example through Faiss (Johnson et al., 2021) or ScaNN (Guo et al., 2020), which can scale into the billions.
Revised: There are various other efﬁcient approximate k NN algorithms available for CPU and GPU/TPU, for example through Faiss (Johnson et al., 2021) or ScaNN (Guo et al., 2020), which can scale into the billions.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We evaluate the effect of adding external memory on ﬁve 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).
Revised: We evaluate the effect of adding external memory on ﬁve 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).

Original: The results show signiﬁcant improvements in the perplexity of the model with the addition of external memory.
Revised: The results show signiﬁcant improvements in the perplexity of the model with the addition of external memory.

Original: We experimented with various sizes of external memory, from 1536 to as high as 262K.
Revised: We experimented with various sizes of external memory, from 1536 to as high as 262K.

Original: 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.
Revised: 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.

Original: 4. E XPERIMENTAL M ETHOD

Original: 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.

Original: For all of our experiments, we used k = 32 .

Original: Unless speciﬁed otherwise, we use the 9 th layer as the k NN augmented attention layer.

Original: We used a sentence-piece (Kudo & Richardson, 2018) tokenizer with a vocabulary size of 32K.

Original: We used the Adafactor optimizer (Shazeer & Stern, 2018).

Original: In preliminary experiments, we conducted a hyperparameter search to determine the optimal learning rate among three choices ( { 3 . 0 , 1 . 0 , 3 · 10 − 1 } ), and found that 1 . 0 works best.

Original: We used a linear warmup schedule for the ﬁrststeps, followed by square root decay.

Original: We trained the models from scratch for 500K steps on all the datasets, except for the Isabelle dataset.

Original: Isabelle is small, so we stopped training after 100K steps when the model began to overﬁt.

Original: We ran all of our experiments on 32 TPU cores.

Original: Our models were implemented in JAX (Bradbury et al., 2018) and Flax (Heek et al., 2020).

Original: When training models with different context lengths, we adjusted the batch size so that the total number of tokens in a batch remains constant, at 2 17 tokens per batch.

Original: We experimented with multiple implementations of approximate k NN lookup with different tradeoffs between quality and computational cost.

Original: We did not observe a signiﬁcant degradation of the model quality when switching to lower quality approximations of k NN, so the model appears to be quite robust with respect to the quality of k NN retrieval.

Original: 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).

Original: 4. D ATASETS
Revised: 4. D ATASETS

Original: arXiv Math For the arXiv dataset, we collected a corpus of papers by downloading them via the arXiv Bulk Data Access 1 .
Revised: arXiv Math For the arXiv dataset, we collected a corpus of papers by downloading them via the arXiv Bulk Data Access 1 .

Original: We ﬁltered papers to include only articles labeled as “Mathematics” /help/ bulk_data [Table]
Revised: We ﬁltered papers to include only articles labeled as “Mathematics” and whose L A TEX source was available.

Original: None

Original: None

Original: [Table]

Original: None

Original: None

Original: [Table]

Original: None

Original: [Table]

Original: None

Original: [Table]

Original: None [Table] and whose L A TEX 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 TEX source has many special characters and the tokenizer tends to output small subwords.
Revised: 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.

Original: Github We used BigQuery 2 to obtain a large corpus of Github repositories that are published with open-source licenses.
Revised: Github We used BigQuery 2 to obtain a large corpus of Github repositories that are published with open-source licenses.

Original: We used ﬁle endings to ﬁlter for ﬁles in the languages C, C++, Java, Python (including Jupyter notebooks), Go, and TypeScript.
Revised: We used ﬁle endings to ﬁlter for ﬁles in the languages C, C++, Java, Python (including Jupyter notebooks), Go, and TypeScript.

Original: Individual source code ﬁles are often fairly short, and there are many dependencies and cross-references between ﬁles in the repository.
Revised: Individual source code ﬁles are often fairly short, and there are many dependencies and cross-references between ﬁles in the repository.

Original: To capture these dependencies, we created one long document for each Github repository by traversing the directory tree, and concatenating all of the ﬁles within it.
Revised: To capture these dependencies, we created one long document for each Github repository by traversing the directory tree, and concatenating all of the ﬁles within it.

Original: The order in which ﬁles are traversed within the repository is random, but each subdirectory is processed as a unit, so that all the ﬁles within the subdirectory are close to each other in the resulting document.
Revised: The order in which ﬁles are traversed within the repository is random, but each subdirectory is processed as a unit, so that all the ﬁles within the subdirectory are close to each other in the resulting document.

Original: Source code is usually structured so that related ﬁles are all grouped together in the same subdirectory; this traversal preserves that structure, while still shufﬂing ﬁles and subdirectories in random order.
Revised: Source code is usually structured so that related ﬁles are all grouped together in the same subdirectory; this traversal preserves that structure, while still shufﬂing ﬁles and subdirectories in random order.

Original: Formal Math – Isabelle The Isabelle corpus consists of formal mathematical proofs of theories.
Revised: Formal Math – Isabelle The Isabelle corpus consists of formal mathematical proofs of theories.

Original: We collected all 627 theories available on The Archive of Formal Proofs 3 (as of October 6, 2021) and an additional 57 theories from the Isabelle standard library 4 to create a corpus of 684 theories.
Revised: We collected all 627 theories available on The Archive of Formal Proofs 3 (as of October 6, 2021) and an additional 57 theories from the Isabelle standard library 4 to create a corpus of 684 theories.

Original: All theories have open-source licenses.
Revised: All theories have open-source licenses.

Original: Each theory is a self-contained mathematical object, on topics such as foundational logic, advanced analysis, algebra, or cryptography, and consists of multiple ﬁles containing proofs.
Revised: Each theory is a self-contained mathematical object, on topics such as foundational logic, advanced analysis, algebra, or cryptography, and consists of multiple ﬁles containing proofs.

Original: As with the Github corpus, all ﬁles that make up a theory are concatenated together into one long document.
Revised: As with the Github corpus, all ﬁles that make up a theory are concatenated [Table] None None

Revised: [Table]

Revised: None

Original: Unlike the Github corpus, we order the ﬁles according to their import dependencies, so that later ﬁles use sub-theorems that are proved in earlier ﬁles.
Revised: None [Table] None [Table]

Revised: None

Revised: [Table]

Revised: [Table]

Revised: None

Revised: [Table]

Revised: together into one long document.

Revised: Unlike the Github corpus, we order the ﬁles according to their import dependencies, so that later ﬁles use sub-theorems that are proved in earlier ﬁles.

Original: 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).
Revised: 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).

Original: We ﬁltered out all documents that have less than 4096 tokens to focus on documents where memory can have an impact.
Revised: We ﬁltered out all documents that have less than 4096 tokens to focus on documents where memory can have an impact.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: 4. E FFECT OF EXTERNAL MEMORY
Revised: 4. E XPERIMENTAL M ETHOD

Revised: 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.

Revised: For all of our experiments, we used k = 32 .

Revised: Unless speciﬁed otherwise, we use the 9 th layer as the k NN augmented attention layer.

Revised: We used a sentence-piece (Kudo & Richardson, 2018) tokenizer with a vocabulary size of 32K.

Revised: We used the Adafactor optimizer (Shazeer & Stern, 2018).

Revised: In preliminary experiments, we conducted a hyperparameter search to determine the optimal learning rate among three choices ( { 3 . 0 , 1 . 0 , 3 · 10 − 1 } ), and found that 1 . 0 works best.

Revised: We used a linear warmup schedule for the ﬁrststeps, followed by square root decay.

Revised: We trained the models from scratch for 500K steps on all the datasets, except for the Isabelle dataset.

Revised: Isabelle is small, so we stopped training after 100K steps when the model began to overﬁt.

Revised: We ran all of our experiments on 32 TPU cores.

Revised: Our models were implemented in JAX (Bradbury et al., 2018) and Flax (Heek et al., 2020).

Revised: 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 tokens in a batch.

Revised: 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.

Revised: We experimented with multiple implementations of approximate k NN lookup with different tradeoffs between quality and computational cost.

Revised: We did not observe a signiﬁcant degradation of the model quality when switching to lower quality approximations of k NN, so the model appears to be quite robust with respect to the quality of k NN retrieval.

Revised: 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).

Revised: 4. E FFECT OF EXTERNAL MEMORY

Original: Adding external memory results in substantial gains across datasets and architectures, as shown in Table 4.
Revised: Adding external memory results in substantial gains across datasets and architectures, as shown in Table 4.

Original: Across all ﬁve datasets, adding external memory to either the vanilla Transformer or the Transformer-XL architecture improves perplexity by a substantial amount.
Revised: Across all ﬁve datasets, adding external memory to either the vanilla Transformer or the Transformer-XL architecture improves perplexity by a substantial amount.

Original: For example, on
Revised: For example, on

Original:  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.[Table]
Revised: [Table] 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.

Original: Increasing the size of the memory increases the beneﬁt of the memory.
Revised: Increasing the size of the memory increases the beneﬁt of the memory.

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

Original: Note that Transformer-XL with context size 2048 already has a theoretical receptive ﬁeld that is quite large.
Revised: Note that Transformer-XL with context size 2048 already has a theoretical receptive ﬁeld that is quite large.

Original: Each token in a higher layer can attend up to 2048 tokens away in the layer below, so the total receptive ﬁeld is 2048 · 12 (layers) ∼ 

Revised: Each token in a higher layer can attend up to 2048 tokens away in the layer below, so the total receptive ﬁeld is 2048 · 12 (layers) ∼

Original: K. Nevertheless, we still saw a substantial gain when adding an external memory of size 8192 to this model.
Revised: K. Nevertheless, we still saw a substantial gain when adding an external memory of size 8192 to this model.

Original: 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.
Revised: 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.

Original: On the other hand, we also saw improvements by adding XL cache to the large-memory (65K) models.
Revised: On the other hand, we also saw improvements by adding XL cache to the large-memory (65K) models.

Original: In a vanilla (non-XL) Transformer, the ﬁrst few tokens in a sequence have very little context, and thus have higher perplexity.
Revised: In a vanilla (non-XL) Transformer, the ﬁrst few tokens in a sequence have very little context, and thus have higher perplexity.

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 4. S CALING TO LARGER MODELS
Revised: 4. S CALING TO LARGER MODELS

Original: We scaled up the Transformer model to sizes of 1 and 8 billion parameters.
Revised: We scaled up the Transformer model to sizes of 1 and 8 billion parameters.

Original: For the 1 billion parameter model, we use 8 layers, 32 heads with head dimension 128, d _model 2048, and d _ff 16384.
Revised: For the 1 billion parameter model, we use 8 layers, 32 heads with head dimension 128, d _model 2048, and d _ff 16384.

Original: For the 8 billion parameter model, we use 64 heads, 16 layers, d _model 4096, and d _ff 32768.
Revised: For the 8 billion parameter model, we use 64 heads, 16 layers, d _model 4096, and d _ff 32768.

Original: We used  context size of 2048, memory size of 8192, and no XL cache.
Revised: We used a context size of 2048, memory size of 8192, and no XL cache.

Original: We ran the comparisons to the vanilla
Revised: We ran the comparisons to the vanilla

Original: Transformer on the arXiv math dataset.
Revised: Transformer on the arXiv math dataset.

Original: We show scaling plots in in Figure 1.
Revised: Scaling plots are shown in Figure 1.

Original: External memory provides a consistent improvement to the model as it is scaled up.
Revised: External memory provides a consistent improvement to the model as it is scaled up.

Original: 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.
Revised: 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.

Original: 4. F INETUNING ON L ARGER M EMORIES
Revised: 4. F INETUNING ON L ARGER M EMORIES

Original: Finetuning on a larger memory.
Revised: Finetuning on a larger memory.

Original: In some cases, training was unstable when using large memories, possibly due to distributional shift early in the training (See Section 3.2).
Revised: In some cases, training was unstable when using large memories, possibly due to distributional shift early in the training (See Section 3.2).

Original: Thus, for memories of 131K or more tokens, we ﬁrst pretrain the model with a memory size of 8192 or 65K for 500K steps, and then ﬁnetune it with the larger memory for an additional 20K steps.
Revised: Thus, for memories of 131K or more tokens, we ﬁrst pretrain the model with a memory size of 8192 or 65K for 500K steps, and then ﬁnetune it with the larger memory for an additional 20K steps.

Original: The results  are shown in Table 5.
Revised: The results of ﬁnetuning on the arXiv Math data set are shown in Table 5.

Original: We observed that the  external memory provided consistent gains up to  size of 262K. Remarkably, a 200M Memorizing Transformer with memory size 262K obtained comparable results to a vanilla
Revised: 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). [Figure]

Original:  Transformerof 8 billion parameters (2.21 vs. 2.13), which is 40X larger in model size.
Revised: e r p l e x i t y ( a r X i v M a t h ) Memory Fine-tuning Transformer Memorizing Transformer

Original: Finetuning a non-memory model to use memory Pretraining can be very costly both in time and computational resources.
Revised: Finetuning a non-memory model to use memory Pretraining can be very costly both in time and computational resources.

Original: Thus, a natural question to ask is: can one ﬁne-tune a
Revised: Thus, a natural question to ask is: can one ﬁne-tune a

Original: pretrained
Revised: pretrained

Original: Transformer to use external memory?
Revised: Transformer to use external memory?

Original: The answer is yes!
Revised: The answer is yes!

Original: [Figure]

Original: P e r p l e x i t y ( a r X i v M a t h )

Original: Memory Fine-tuning Transformer Memorizing Transformer We took a pre-trained 1B vanilla Transformer model, and ﬁne-tuned it to use external memory (the 1B models used in Section 4.4).
Revised: We took a pre-trained 1B vanilla Transformer model, and ﬁne-tuned it to use external memory (the 1B models used in Section 4.4).

Original: The ﬁne-tuning result is shown in Figure 6.
Revised: The ﬁne-tuning result is shown in Figure 6.

Original: Notice that the model quickly learns to use external memory.
Revised: Notice that the model quickly learns to use external memory.

Original: Within 20K steps ( 4% of the pre-training time) the ﬁne-tuned model has already closed 85% of the gap between it and the 1B Memorizing Transformer, and after 100k steps it has closed the gap entirely.
Revised: Within 20K steps ( 4% of the pre-training time) the ﬁne-tuned model has already closed 85% of the gap between it and the 1B Memorizing Transformer, and after 100k steps it has closed the gap entirely.

Original: 4. I NFORMATION R ETRIEVAL P ATTERNS
Revised: 4. I NFORMATION R ETRIEVAL P ATTERNS

Original: We conducted a qualitative study of what the model was actually retrieving from external memory, by ﬁnding which tokens showed the biggest improvements in cross-entropy loss when the size of the memory was increased, and then examining the top- k retrieved memories for those tokens.
Revised: We conducted a qualitative study of what the model was actually retrieving from external memory, by ﬁnding which tokens showed the biggest improvements in cross-entropy loss when the size of the memory was increased, and then examining the top- k retrieved memories for those tokens.

Original: We found that the model gained the most when looking up rare words, such as proper names, references, citations, and function names, where the ﬁrst use of a name is too far away from subsequent uses to ﬁt in the local context.
Revised: We found that the model gained the most when looking up rare words, such as proper names, references, citations, and function names, where the ﬁrst use of a name is too far away from subsequent uses to ﬁt in the local context.

Original: This result is in keeping with the prior analysis of long-context Transformers on PG19 (Sun et al., 2021), which found similar lookup patterns.
Revised: This result is in keeping with the prior analysis of long-context Transformers on PG19 (Sun et al., 2021), which found similar lookup patterns.

Original: For this experiment, we used a slightly older version of the architecture without the gating mechanism.
Revised: For this experiment, we used a slightly older version of the architecture without the gating mechanism.

Original: Which tokens show a beneﬁt from memory?
Revised: Which tokens show a beneﬁt from memory?

Original: Figure 7 shows a visualization of which tokens show an improvement when the size of the external memory is increased.
Revised: Figure 7 shows a visualization of which tokens show an improvement when the size of the external memory is increased.

Original: We selected a math paper at random, and plotted the difference in cross entropy loss for each token x i in the paper, comparing two models with the same parameters, but with memories of different sizes.
Revised: We selected a math paper at random, and plotted the difference in cross entropy loss for each token x i in the paper, comparing two models with the same parameters, but with memories of different sizes.

Original: ∆ i = cross-entropy 8192 ( x i ) − cross-entropy 32K ( x i ) .
Revised: ∆ i = cross-entropy 8192 ( x i ) − cross-entropy 32K ( x i ) .

Original: Positive values show an improvement in loss.
Revised: Positive values show an improvement in loss.

Original: The x -axis on the chart is the token number i , while the y -axis is ∆ i .
Revised: The x -axis on the chart is the token number i , while the y -axis is ∆ i .

Original: For the ﬁrst 8192 tokens, the difference between the two models is zero, since the larger capacity of the 32 K memory isn’t being used yet.
Revised: For the ﬁrst 8192 tokens, the difference between the two models is zero, since the larger capacity of the 32 K memory isn’t being used yet.

Original: However, after token 8193 , we can see that the larger memory helps, on average, over the smaller memory.
Revised: However, after token 8193 , we can see that the larger memory helps, on average, over the smaller memory.

Original: The beneﬁt 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 when the size of the external memory is increased.
Revised: The beneﬁt 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 when the size of the external memory is increased.

Original: This ﬁgure also shows that the beneﬁt of external memory is somewhat sparse.
Revised: This ﬁgure also shows that the beneﬁt of external memory is somewhat sparse.

Original: 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.
Revised: 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.

Original: What information is being looked up?
Revised: What information is being looked up?

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: We studied arXiv math, Github and Isabelle corpus.
Revised: We studied arXiv math, Github and Isabelle corpus.

Original: For arXiv math and Github, we found the model retrieved function and variable names.
Revised: For arXiv math and Github, we found the model retrieved function and variable names.

Original: See more details with examples in Appendix B.
Revised: See more details with examples in Appendix B.

Revised: [Figure]

Original: Retrieving mathematical deﬁnitions.
Revised: Retrieving mathematical deﬁnitions.

Original: Our case study on the Isabelle corpus provides one of the clearest illustrations of how a model can make good use of external memory.
Revised: Our case study on the Isabelle corpus provides one of the clearest illustrations of how a model can make good use of external memory.

Original: When predicting the name of a mathematical object or a lemma, the model looked up the deﬁnition from earlier in the proof.
Revised: When predicting the name of a mathematical object or a lemma, the model looked up the deﬁnition from earlier in the proof.

Original: Examples of this behavior are shown in Table 8.
Revised: Examples of this behavior are shown in Table 8.

Original: In example 1, the model retrieves a deﬁnition within the body of a lemma, markov_inequality .
Revised: In example 1, the model retrieves a deﬁnition within the body of a lemma, markov_inequality .

Original: In example 2, it retrieves the deﬁnition of a[Figure] previously deﬁned concept subgraph_threshold .
Revised: In example 2, it retrieves the deﬁnition of a previously deﬁned concept subgraph_threshold .

Original: In example 3, it retrieves the deﬁnition of orthonormal_system .
Revised: In example 3, it retrieves the deﬁnition of orthonormal_system .

Original: 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.
Revised: 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.

Original: In the other two cases, the model also looked up materials in the immediate vicinity.
Revised: In the other two cases, the model also looked up materials in the immediate vicinity.

Original: To the best of our knowledge, this is the ﬁrst demonstration that attention is capable of looking up deﬁnitions and function bodies from a large corpus.
Revised: To the best of our knowledge, this is the ﬁrst demonstration that attention is capable of looking up deﬁnitions and function bodies from a large corpus.

Original: The Isabelle case study used a model with two memory layers of size 32 K.
Revised: The Isabelle case study used a model with two memory layers of size 32 K.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We present a simple extension to the Transformer architecture, called k NN-augmented attention, which dramatically increases the length of the context that a language model can attend to by using k -nearest-neighbor lookup into a large external memory.
Revised: We present a simple extension to the Transformer architecture, called k NN-augmented attention, which dramatically increases the length of the context that a language model can attend to by using k -nearest-neighbor lookup into a large external memory.

Original: 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.
Revised: 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.

Original: The Memorizing Transformer shows large improvements in perplexity over the baseline with all of the data sets and architectures  we studied; it is comparable to a vanilla transformer that has 5 times the number of parameters.
Revised: 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.

Original: Perplexity continues to improve with increasing memory size, although there is a point of diminishing returns.
Revised: Perplexity continues to improve with increasing memory size, although there is a point of diminishing returns.

Original: Moreover, external memory continues to provide beneﬁts even as the transformer is scaled up from 200M to 8B parameters.
Revised: Moreover, external memory continues to provide beneﬁts even as the transformer is scaled up from 200M to 8B parameters.

Original: 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 ﬁne-tuning it.
Revised: 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 ﬁne-tuning it.

Original: Unlike other forms of attention, 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.
Revised: Unlike other forms of attention, 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.

Original: How to make the best use of this capability is a topic for future work.
Revised: How to make the best use of this capability is a topic for future work.

Original: A CKNOWLEDGMENTS
Revised: A CKNOWLEDGMENTS

Original: We want to thank Charles Staats for the many fruitful discussions and detailed comments, Henryk
Revised: We want to thank Charles Staats for the many fruitful discussions and detailed comments, Henryk

Original: Michalewski for early version of of the memory implementation, Petros Maniatis for his help with our code datasets, Aitor Lewkowycz for his help with larger scale memorizing transformer experiments,
Revised: Michalewski for early version of of the memory implementation, Petros Maniatis for his help with our code datasets, Aitor Lewkowycz for his help with larger scale memorizing transformer experiments,

Original: Behnam Neyshabur for his comments on ﬁnetuning 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.
Revised: Behnam Neyshabur for his comments on ﬁnetuning 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.

Original: E THICS
Revised: E THICS

Original: The ability to memorize large databases of facts could have potential ramiﬁcations for society, especially if those databases include sensitive personal information or copyrighted works.
Revised: The ability to memorize large databases of facts could have potential ramiﬁcations for society, especially if those databases include sensitive personal information or copyrighted works.

Original: However, one advantage of using an external memory is that the memory can be easily cleared of all such information, as we do at the end of each document that we train on.
Revised: However, one advantage of using an external memory is that the memory can be easily cleared of all such information, as we do at the end of each document that we train on.

Original: The same is not true of differentiable model parameters, which is where most existing architectures  store facts and information that they are trained on.
Revised: The same is not true of differentiable model parameters, which is what most existing architectures use to store facts and information that they are trained on.

Original: R EPRODUCIBILITY
Revised: R EPRODUCIBILITY

Original: Details of our architecture and training hyperparameters are given in Section 4.1.
Revised: Details of our architecture and training hyperparameters are given in Section 4.2.

Original: The datasets for C4 and PG-19 are publicly available.
Revised: The datasets for C4 and PG-19 are publicly available.

Original: Our additional datasets, Github, Isabelle, and ArXiv
Revised: Our additional datasets, Github, Isabelle, and ArXiv

Original: Math are derived from publicly available data buckets, which we link in the main part of the paper.
Revised: Math are derived from publicly available data buckets, which we link in the main part of the paper.

Original: Subsection 4.2 include details on how we constructed the datasets from those datasets.
Revised: Subsection 4.1 include details on how we constructed the datasets from those datasets.

Original: We plan to release our code as open source.
Revised: We plan to release our code as open source

Original: A L ENGTH OF INPUTS
Revised: A L ENGTH OF INPUTS

Original: A. A BLATION STUDIES
Revised: A. A BLATION STUDIES

Original: In the following section, we performed ablation studies to investigate the effects of various hyperparameters.
Revised: In the following section, we performed ablation studies to investigate the effects of various hyperparameters.

Original: Unless otherwise speciﬁed, we carried out these experiments with a memorizing transformer with context size 512, XL cache 512 with a memory size of 8192.
Revised: Unless otherwise speciﬁed, we carried out these experiments with a memorizing transformer with context size 512, XL cache 512 with a memory size of 8192.

Original: Multiple k NN layers.
Revised: Multiple k NN layers.

Original: We experimented with using two k NN layers, rather than just one.
Revised: We experimented with using two k NN layers, rather than just one.

Original: However, we did not see further beneﬁts brought by more than multiple retrieval layers.
Revised: However, we did not see further beneﬁts brought by more than multiple retrieval layers.

Original: k NN layer index We experimented with adding the external memory to layer 3 , 6 , 9 and 12 in a 12-layer transformer, with results shown in Table 14.
Revised: k NN layer index We experimented with adding the external memory to layer 3 , 6 , 9 and 12 in a 12-layer transformer, with results shown in Table 14.

Original: We found that adding memory to the middle of the layer stack will obtain the best result, whereas adding memory to layers either too close to the input or to the output obtained less gains.
Revised: We found that adding memory to the middle of the layer stack will obtain the best result, whereas adding memory to layers either too close to the input or to the output obtained less gains.

Original: [Table]
Revised: [Table]

Original: Number of neighbors We studied the effects of the number of neighbors we retrieve from memory, with results shown in Table 15.
Revised: Number of neighbors We studied the effects of the number of neighbors we retrieve from memory, with results shown in Table 15.

Original: We found that even with 32 number of neighbors, we can already obtain a comparable results with 128 or 256 neighbors.
Revised: We found that even with 32 number of neighbors, we can already obtain a comparable results with 128 or 256 neighbors.

Original: [Table]
Revised: [Table]

Original: Random seeds We measured the statistical signiﬁcant of the results reported.
Revised: Random seeds We measured the statistical signiﬁcant of the results reported.

Original: We did 3 runs with 3 random seeds for Transformer XL of size 512, and also a memorizing transformer with memory size 8192.
Revised: We did 3 runs with 3 random seeds for Transformer XL of size 512, and also a memorizing transformer with memory size 8192.

Original: We measured the standard deviation of perplexities after 500K steps of training, shown in
Revised: We measured the standard deviation of perplexities after 500K steps of training, shown in

Original: Table 16.
Revised: Table 16.

Original: We saw the standard deviation between different runs of the same experiment appears to be much smaller than the gap between different models.
Revised: We saw the standard deviation between different runs of the same experiment appears to be much smaller than the gap between different models.

Original: [Table]
Revised: [Table]

Original: B W HAT DOES THE MODEL RETRIEVE FROM MEMORY ?
Revised: B W HAT DOES THE MODEL RETRIEVE FROM MEMORY ?

Original: Retrieving citation names On arXiv math, several examples are shown in Table 17, which includes both the retrieved token and its surrounding context.
Revised: Retrieving citation names On arXiv math, several examples are shown in Table 17, which includes both the retrieved token and its surrounding context.

Original: We observe that many of the gains in crossentropy 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.
Revised: We observe that many of the gains in crossentropy 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.

Original: Such lookups usually span over the entire paper, which is much longer than 8192 tokens, providing a plausible explanation for the gain beyond memory size of 8192.
Revised: Such lookups usually span over the entire paper, which is much longer than 8192 tokens, providing a plausible explanation for the gain beyond memory size of 8192.

Original: [Table]
Revised: [Table]

Original: Retrieving function names from the codebase As with the arXiv papers, we also studied which tokens the model retrieved from memory.
Revised: Retrieving function names from the codebase As with the arXiv papers, we also studied which tokens the model retrieved from memory.

Original: As might be expected, the model is often looking up the names of functions, and variables, as shown in Table 18.
Revised: As might be expected, the model is often looking up the names of functions, and variables, as shown in Table 18.

Original: [Table]
Revised: [Table]

Original: B. M ORE R
Revised: B. M ORE R

Original: Example
Revised: Example

Original: • Input token index:• Input token: “_” • Target token: “pair” • Surrounding context: )) by (simp add: Fourier_sum_limit_pair [OF f, symmetric] Fourier’ • Name needs to be predicted: Fourier_sum_limit_pair • Retrieved token: “
Revised: • Input token index:• Input token: “_” • Target token: “pair” • Surrounding context: )) by (simp add: Fourier_sum_limit_pair [OF f, symmetric] Fourier’ • Name needs to be predicted: Fourier_sum_limit_pair • Retrieved token: “

Original: Four” • Retrieved token index:• Retrieved context: 2 * n. Fourier_coefﬁcient
Revised: Four” • Retrieved token index:• Retrieved context: 2 * n. Fourier_coefﬁcient

Original: f k * trigonometric_set k t) • Deﬁnition of the name:
Revised: f k * trigonometric_set k t) • Deﬁnition of the name:

Original: Example
Revised: Example

Original: • Input token index:• Input token: “tri”’ • Target token: “gon” •
Revised: • Input token index:• Input token: “tri”’ • Target token: “gon” •

Original: Surrounding context: <le>n. a k * trigonometric_set k x) • Name needs to be predicted:
Revised: Surrounding context: <le>n. a k * trigonometric_set k x) • Name needs to be predicted:

Original: orthonormal_system_trigonometric_set • Retrieved token: “gon” • Retrieved token index:• Retrieved context: lemma
Revised: orthonormal_system_trigonometric_set • Retrieved token: “gon” • Retrieved token index:• Retrieved context: lemma

Original: orthonormal_system_trigonometric_set: \ n "orthonormal_system • Deﬁnition of the name:
Revised: orthonormal_system_trigonometric_set: \ n "orthonormal_system • Deﬁnition of the name:

Original: Example
Revised: Example

Original: • Input token index:• Input token: “sum”’ • Target token: “
Revised: • Input token index:• Input token: “sum”’ • Target token: “

Original: m” • Surrounding context: nusing Fourier_series_square_summable
Revised: m” • Surrounding context: nusing Fourier_series_square_summable

Original: [OF assms, of’ • Name needs to be predicted:
Revised: [OF assms, of’ • Name needs to be predicted:

Original: Fourier_series_square_summable • Retrieved token:
Revised: Fourier_series_square_summable • Retrieved token:

Original: “sum” • Retrieved token index:• Retrieved context: lemma Fourier_series_square_summable \ n assumes: • Deﬁnition of the name:
Revised: “sum” • Retrieved token index:• Retrieved context: lemma Fourier_series_square_summable \ n assumes: • Deﬁnition of the name:

Original: Example
Revised: Example

Original: • Input token index:• Input token: “_”’ • Target token: “system” • Surrounding context: lemma Riemann_lebesgue_square_integrable: nassumes "orthonormal_system S w • Name needs to be predicted: orthonormal_system • Retrieved token: “system” • Retrieved token index:• Retrieved context: deﬁnition orthonormal_system :: " \ ’a::euclidean’ • Deﬁnition of the name:
Revised: • Input token index:• Input token: “_”’ • Target token: “system” • Surrounding context: lemma Riemann_lebesgue_square_integrable: nassumes "orthonormal_system S w • Name needs to be predicted: orthonormal_system • Retrieved token: “system” • Retrieved token index:• Retrieved context: deﬁnition orthonormal_system :: " \ ’a::euclidean’ • Deﬁnition of the name:

Original: Example
Revised: Example

Original: • Input token index:• Input token: “.”’
Revised: • Input token index:• Input token: “.”’

Original: • Target token: “b” • Surrounding context: shows "integrable (lebesgue_on {a..b}) •
Revised: • Target token: “b” • Surrounding context: shows "integrable (lebesgue_on {a..b}) •

Original: Retrieved token 1: “.”
Revised: Retrieved token 1: “.”

Original: • Retrieved token index 1:• Retrieved context 1: lebesgue_on {a..b})
Revised: • Retrieved token index 1:• Retrieved context 1: lebesgue_on {a..b})

Original: f i • Retrieved token 2: “-” • Retrieved token index 2:• Retrieved context 2: (lebesgue_on {a-c..b-c}) ( • Retrieved token 3: “
Revised: f i • Retrieved token 2: “-” • Retrieved token index 2:• Retrieved context 2: (lebesgue_on {a-c..b-c}) ( • Retrieved token 3: “

Original: -” • Retreived token index 3:• Retrieved context 3: (lebesgue_on {-pi..pi}) (
Revised: -” • Retreived token index 3:• Retrieved context 3: (lebesgue_on {-pi..pi}) (

Original: Example
Revised: Example

Original: • Input token index:• Input token: “_”’ • Target token: “sum” •
Revised: • Input token index:• Input token: “_”’ • Target token: “sum” •

Original: Surrounding context: 0" \ n using Fourier_series_square_summable [OF assms • Retrieved token 1: “set” • Retrieved token index 1:• Retrieved context 1: deﬁnition trigonometric_set :: "nat \ <Rightarrow> • Retrieved token 2: “ier” • Retrieved token index 2:• Retrieved context 2:
Revised: Surrounding context: 0" \ n using Fourier_series_square_summable [OF assms • Retrieved token 1: “set” • Retrieved token index 1:• Retrieved context 1: deﬁnition trigonometric_set :: "nat \ <Rightarrow> • Retrieved token 2: “ier” • Retrieved token index 2:• Retrieved context 2:

Original: deﬁnition Fourier_coefﬁcient \ nwhere • Retrieved token 3: “ine” • Retrieved token index 3:• Retrieved context 3: lemma Schwartz_inequality_strong: \ nassumes “f’ • Retrieved token 4: “system” • Retrieved token index 4:• Retrieved context 4: deﬁnition orthonormal_system :: “ \ ’a::euclidean’ • Retrieved token 5: “<” • Retrieved token index 5:•
Revised: deﬁnition Fourier_coefﬁcient \ nwhere • Retrieved token 3: “ine” • Retrieved token index 3:• Retrieved context 3: lemma Schwartz_inequality_strong: \ nassumes “f’ • Retrieved token 4: “system” • Retrieved token index 4:• Retrieved context 4: deﬁnition orthonormal_system :: “ \ ’a::euclidean’ • Retrieved token 5: “<” • Retrieved token index 5:•

Original: Retrieved context 5: subsection \ <open>Convergence wrt the L’ • Retrieved token 6: “n” • Retrieved token index 6:• Retrieved context 6: \ n subsection \ <open>
Revised: Retrieved context 5: subsection \ <open>Convergence wrt the L’ • Retrieved token 6: “n” • Retrieved token index 6:• Retrieved context 6: \ n subsection \ <open>

Original: A bit of extra’
Revised: A bit of extra’