Source-linked AI summary
Memorizing Transformers
Yuhuai Wu, Markus N. Rabe, DeLesley Hutchins, Christian Szegedy
TL;DR
Language models ordinarily acquire new knowledge through weight updates, motivating a way to memorize information immediately from newly read data. The paper adds approximate kNN access to a non-differentiable memory of past representations and finds improved language modeling across diverse tasks, including use of newly defined code functions and mathematical theorems. Large memories remain beneficial, though stale representations can make large-memory training unstable.
Problem
Language models usually require training or fine-tuning with weight updates to acquire new knowledge, rather than memorizing newly read information immediately.
Method
The paper extends decoder-only Transformers with a kNN-augmented attention layer that searches a non-differentiable external memory of previously seen key-value representations alongside local attention.
Results
The method improves perplexity across C4, Github, PG-19, Isabelle, and arXiv, with gains comparable to a vanilla Transformer using five times as many parameters.
Takeaways & Limitations
External memory provides benefits across model scales and can let models retrieve newly encountered lemma definitions and other task-relevant information during use.
Takeaways & Limitations
Changing model parameters shifts stored keys and values, making older records stale; training from scratch with large memory sometimes performs worse than small-memory pretraining followed by large-memory fine-tuning.
Abstract
from arXiv · showhide
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 kNN 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.
1 INTRODUCTION
Transformer context lengths limit access to distant information, while attention can support rapid learning by retrieving previously seen facts. The paper proposes scalable approximate kNN memory and reports improvements across models and language-modeling tasks.
- Short attention contexts limit Transformers when references span distant parts of novels, source code, and theorem-proving proofs.
- Attention can rapidly retrieve memorized facts, such as function definitions, stored as key-value pairs instead of slowly learned weight updates.
- Approximate k-nearest-neighbor lookup increases attention context size using scalable information-retrieval implementations.
- The approach retrieves exact distant values rather than averaging or summarizing them, distinguishing it from some long-range-attention methods.
- 8K-token memory improves perplexity across different model sizes, while external memory scales to 131K or 262K tokens on one TPU with reasonable step time.
- Perplexity steadily improves with memory size across C4, Github, PG-19, Isabelle, and arXiv, and models retrieve lemma definitions during theorem proving.
- Memory improvements remain across increasing model sizes and can exceed the gains from increasing model size by 5X or more.
2 RELATED WORK
Prior work extends context through windows, approximation, recurrence, caches, learned memories, and retrieval. This paper instead combines a very large uncompressed cache with approximate kNN attention over previously seen tokens.
- Long-range attention methods include sliding windows, kernel approximations, pooling, and recurrent architectures with different efficiency and parallelism trade-offs.
- Truncated backpropagation and neural caches preserve prior states in non-differentiable memory, while Transformer-XL applies related cache-based recurrence.
- The paper uses a very large cache without compression together with approximate kNN attention over it.
- Learned-memory approaches combine attention with differentiable memory, whereas this work uses kNN to retrieve previous tokens rather than replace feed-forward layers.
- Other external-memory methods retrieve token pairs for final token selection or compress context into retrieval keys.
- Retrieval-augmented systems typically use static text knowledge bases for question answering or dialogue, while this paper targets decoder-only language modeling with a different memory setup.
- kNN lookup is a general-purpose retrieval technique with high-performance implementations and has also been used to replace dense attention for efficiency.
3 METHOD
The method augments a decoder-only Transformer with local self-attention and approximate kNN retrieval over a non-differentiable memory of past key–value pairs. Documents are processed sequentially in subsequences, while learned gating, normalization, and memory management address long-range retrieval and training stability.
- Architecture: A decoder-only Transformer uses causal dense self-attention and final-layer token embeddings to predict the next token.
- Training and stability: Long documents are processed sequentially in 512-token subsequences, with separate external memories for batched documents and normalization used to reduce—but not eliminate—staleness.Older records can become stale as model-produced queries, keys, and values shift during multi-step processing; large-memory training sometimes performs worse from scratch than small-memory pretraining followed by large-memory finetuning.
- kNN-augmented attention: One upper Transformer layer combines local dense self-attention with approximate k-nearest-neighbor search over external memory.The same layer therefore accesses both the current subsequence and distant retrieved representations.
- External memory: The external memory stores the prior M key–value pairs per head, appending each step’s local pairs and dropping old pairs when necessary.Queries are shared between local and external attention, and the two key–value sources belong to the same distribution.
- kNN retrieval: Retrieved memories are scored with query–key dot products, normalized with softmax, and combined as a weighted sum of retrieved values.
- Attention combination: A learned per-head sigmoid gate combines external-memory attention with local attention, allowing heads to favor local or long-range information.The gate is content-independent in these experiments, and most heads learned to attend almost exclusively to external memory over time.
4 EXPERIMENTS
Across five long-form language-modeling tasks, external memory consistently improves perplexity, with larger memories generally yielding further gains. The experiments also examine scaling, finetuning, token-level effects, and retrieval of mathematical definitions.
- External-memory effects: External memory significantly improves perplexity across PG-19, C4, arXiv Math, Github, and Isabelle, with initial sharp gains followed by steadier improvements as memory grows.The experiments vary memory from 1536 to 262K tokens.
- External-memory effects: On C4(4K+), memory size 8192 lowers vanilla Transformer perplexity from 17.20 to 14.42 and Transformer-XL perplexity from 15.38 to 14.04.
- External-memory effects: The best perplexities across datasets and architectures occur with a memory size of 65K, and increasing memory size increases the benefit.
- Scaling models: External memory continues improving scaled models: an 8K-memory Memorizing Transformer matches the perplexity of a vanilla Transformer with 5X more trainable parameters.
- Finetuning: A pretrained 1B vanilla Transformer closes 85% of its gap to the 1B Memorizing Transformer within 20K finetuning steps and closes it entirely after 100K steps.
- Information retrieval patterns: With larger memory, improvements are generally concentrated in a small percentage of tokens, although some predictions worsen when relevant memories fall outside the top-k.
5 CONCLUSION
The paper introduces kNN-augmented attention, which expands a Transformer’s usable context through large external memory and improves language modeling across long-document datasets. Performance improves with memory and remains beneficial as models scale, while the approach can also leverage existing pretrained models.
- 5 CONCLUSION: kNN-augmented attention increases the context available to language models through k-nearest-neighbor lookup into a large external memory.The extension is designed for language modeling over LaTeX documents, source code, formal proofs, and books.
- 5 CONCLUSION: Memorizing Transformers show large perplexity improvements over baseline models across all studied datasets and architectures, comparable to a vanilla Transformer with 5 times more parameters.Perplexity continues improving as memory grows, although gains eventually diminish.
- 5 CONCLUSION: External memory continues to provide benefits as Transformer size increases from 200M to 8B parameters.The reported improvements are maintained across models of increasing size.
- 5 CONCLUSION: kNN retrieval can scale to huge memory sizes, potentially supporting vast knowledge bases or code repositories.The paper identifies making the best use of this capability as future work.
ETHICS
Large factual memories may create societal risks when they contain sensitive personal information or copyrighted works. External memory offers a way to clear such information more easily than differentiable model parameters.
- ETHICS: Memorizing large databases of facts could have societal ramifications when they include sensitive personal information or copyrighted works.The concern applies to the contents stored in the external memory.
- ETHICS: External memory can be cleared of stored information, unlike differentiable model parameters used by most existing architectures to store learned facts.The authors describe clearing the memory at the end of each training document.
REPRODUCIBILITY
The paper provides implementation and dataset-construction details and indicates that its datasets derive from publicly available sources. The authors also plan to release the code as open source.
- REPRODUCIBILITY: Architecture and training hyperparameters are detailed in Section 4.2.
- REPRODUCIBILITY: C4 and PG-19 datasets are publicly available, while Github, Isabelle, and arXiv Math are derived from publicly available data sources.The paper links the additional source buckets in the main paper.
- REPRODUCIBILITY: The authors plan to release the code as open source.
A LENGTH OF INPUTS
The evaluated corpora contain highly variable document lengths, with especially long tails in the arXiv math and Github datasets. The reported maxima reach almost 1.6M tokens for a math paper and just over 9M for a repository.
- A LENGTH OF INPUTS: ArXiv math papers have a broad token-length distribution, with the histogram truncated at 500k tokens and a maximum paper of almost 1.6M tokens.
- A LENGTH OF INPUTS: Github repositories have a long-tailed token-length distribution, with the maximum-length repository containing just over 9M tokens.The figure cuts off the long tail of the plot.
- A LENGTH OF INPUTS: The Isabelle proof-scripts dataset has a documented histogram of token counts.
- A LENGTH OF INPUTS: The PG19 books dataset has a documented histogram of token counts.
- A LENGTH OF INPUTS: C4 documents in this analysis are filtered to contain fewer than 4096 tokens.
A.1 ABLATION STUDIES
The ablations examine retrieval-layer placement, neighbor count, multiple retrieval layers, and run-to-run variability. Memory works best in the middle of the transformer, while relatively few neighbors are sufficient and additional retrieval layers provide no further benefit.
- Experimental setup: The ablations used a memorizing transformer with context size 512, XL cache 512, and memory size 8192 unless otherwise specified.These settings define the default experimental configuration for the ablation studies.
- Multiple kNN layers: Using more than one kNN layer did not provide further benefits over a single retrieval layer.The experiment compared two kNN layers with one layer.
- kNN layer index: Middle-layer memory placement produced the best result, while layers near the input or output produced smaller gains.The study tested layers 3, 6, 9, and 12 in a 12-layer transformer.
- Number of neighbors: Thirty-two retrieved neighbors already achieved results comparable to using 128 or 256 neighbors.The comparison is reported in the number-of-neighbors ablation.
- Random seeds: Run-to-run perplexity variation was much smaller than the gap between different models.The measurement used three random seeds for Transformer XL and a memorizing transformer with memory size 8192 after 500K training steps.
B WHAT DOES THE MODEL RETRIEVE FROM MEMORY?
The model retrieves recurring symbolic identifiers and their surrounding usage contexts from earlier material. In arXiv math, retrieval often targets citation and reference names, while in Github it commonly targets function and variable names.
- arXiv math: On arXiv math, many cross-entropy gains came from predicting bibitem, citation, and reference names retrieved from earlier passages.These lookups can span an entire paper, exceeding 8192 tokens.
- Github: On Github, the model often retrieves function and variable names from elsewhere in the codebase.The examples concern how functions are used elsewhere in the repository.
B.1 MORE RETRIEVING EXAMPLES IN FORMAL THEOREM PROVING CORPUS
The formal-theorem-proving examples show retrieval of theorem, definition, and syntactic context from earlier Isabelle material. Retrieved tokens include names and fragments associated with Fourier analysis, integrability, and proof statements.
- Fourier analysis: The examples include retrieval of Fourier_sum_limit_pair and surrounding Fourier_coefficient context.The associated context also references convergence-related material.
- Formal theorem names: The model retrieves the name orthonormal_system_trigonometric_set together with nearby lemma context.The example includes a retrieved occurrence of the named theorem-related token.
- Analytic contexts: The retrieval examples include fragments about integrability over intervals such as {a..b}, {a-c..b-c}, and {-pi..pi}.These contexts appear alongside retrieved punctuation and token fragments.
- Proof contexts: Retrieved proof contexts reference Fourier_series_square_summable, Schwartz_inequality_strong, and convergence-related subsections.The examples also contain fragments from formal proof structure and subsection names.
- Formal definitions: Other retrieved contexts contain definitions involving orthonormal_system and trigonometric_set.These examples connect retrieved fragments to formal definitions in the corpus.