Memorizing Transformers

By Yuhuai Wu, Markus Norman Rabe, DeLesley S. Hutchins, Christian Szegedy. 2022

249 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Model Architectures

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Model Architectures

Citations in the corpus

Emergent Abilities of Large Language Models (2022, neutral): “Improving model architecture and training procedures may facilitate high-quality models with emergent abilities while mitigating computational cost.”
Landmark Attention: Random-Access Infinite Context Length for Transformers (2023, methodology): “Memory Transformers [39] introduce special memory tokens that are prepended to the input, and their representation at the final layer of the model is used for the next input.”
General-purpose, long-context autoregressive modeling with Perceiver AR (2022, background): “…size can also be achieved by reducing the compute requirement on tokens that are far in the past using stop gradients (Dai et al., 2019), recurrence (Mehri et al., 2017), memory (Weston et al., 2015; Rae et al., 2019), or other forms of compression (Wu et al., 2022) to process long-term structure.”
Extending Context Window of Large Language Models via Positional Interpolation (2023, background): “Several works add memory capabilities to Transformers through recurrence, which increase the models’ capability of handling very long sequences (Bulatov et al., 2022; Wu et al., 2020; Dai et al., 2019; Wu et al., 2022; Martins et al., 2021; Mu et al., 2023).”
Training Language Models with Memory Augmentation (2022, background): “Long-term memory may have different interpretations in other contexts and we use long-term memory to refer to longrange context in modeling long sequences, following previous work (Martins et al., 2022; Wu et al., 2022).”
Learning to Compress Prompts with Gist Tokens (2023, background): “See Appendix A for a sample PyTorch implementation which can be used as a drop-in replacement for attention masking in deep learning libraries such as Hugging Face Transformers [41].”
Repository-Level Prompt Generation for Large Language Models of Code (2022, background): “Using pretrained models that allow larger context length or models that augment the context [45] offer avenues for future work.”
Long Range Language Modeling via Gated State Spaces (2022, methodology): “ArXiv Math dataset was recently collected by [Wu et al., 2022] which contains LaTeX source for articles focusing on Mathematics.”
Retrieval as Attention: End-to-end Learning of Retrieval and Reading within a Single Transformer (2022, background): “Retrieval from external datastore to improve language modeling perplexity has been explored by many works, where additional tokens are retrieved during generation based on contextual representations (Khandelwal et al., 2020; Borgeaud et al., 2021; Wu et al., 2022; Zhong et al., 2022).”
LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding (2023, background): “…2020; Beltagy et al., 2020; Zaheer et al., 2020; Wang et al., 2020; Fedus et al., 2022; Ding et al., 2023), recurrent and memory modules (Dai et al., 2019; Rae et al., 2020; Wu et al., 2022; Martins et al., 2022; Bulatov et al., 2022; Orvieto et al., 2023; Liang et al., 2023; Zhou et al., 2023).”
NaturalProver: Grounded Mathematical Proof Generation with Language Models (2022, background): “We expect improvements with increased scale, but also through improved grounding mechanisms [Wu et al., 2022], or auxiliary models, e.”
FiDO: Fusion-in-Decoder optimized for stronger performance and faster inference (2022, background): “Another approach is to pre-compute and store encoder representations in a memory and directly retrieve representations from memory, rather than re-encoding retrieved text (de Jong et al., 2022; Wu et al., 2022; Li et al., 2022).”
Natural Language to Code Generation in Interactive Data Science Notebooks (2022, background): “This finding suggests the importance of developing LLMs that could handle longer context (Wu et al., 2022a) in order to include more DataFrame information in prompts.”
Retrieval-Augmented Generation for AI-Generated Content: A Survey (2024, unknown): “For example, RAG can reduce the size of large generative models [30], provide support for long contexts [31], and eliminate certain generation steps [32].”
ProofNet: Autoformalizing and Formally Proving Undergraduate-Level Mathematics (2023, methodology): “Language-model-based theorem provers and autoformalization systems have typically been evaluated on benchmarks consisting of competition and olympiad-style problems [Zheng et al., 2022, Wu et al., 2022a].”