# Training Language Models with Memory Augmentation

arXiv: https://arxiv.org/abs/2205.12674 · HTML: https://mlsanity.com/paper/2205.12674/

## Key facts

- Cited by 113 papers in the map (data through 14 Oct 2026).
- Builds on 49 of its 61 references that are in the map.
- First posted on arXiv 25 May 2022.

## Builds on (most cited first, 30 of 49)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2004.04906/) (2020)
- [OPT: Open Pre-trained Transformer Language Models](https://mlsanity.com/paper/2205.01068/) (2022)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [Billion-scale similarity search with GPUs](https://mlsanity.com/paper/1702.08734/) (2017)
- [REALM: Retrieval-Augmented Language Model Pre-Training](https://mlsanity.com/paper/2002.08909/) (2020)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://mlsanity.com/paper/1901.02860/) (2019)
- [A Call for Clarity in Reporting BLEU Scores](https://mlsanity.com/paper/1804.08771/) (2018)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [fairseq: A Fast, Extensible Toolkit for Sequence Modeling](https://mlsanity.com/paper/1904.01038/) (2019)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [Improving language models by retrieving from trillions of tokens](https://mlsanity.com/paper/2112.04426/) (2021)
- [Get To The Point: Summarization with Pointer-Generator Networks](https://mlsanity.com/paper/1704.04368/) (2017)
- [Leveraging Passage Retrieval with Generative Models for Open Domain Question Answering](https://mlsanity.com/paper/2007.01282/) (2020)
- [Aligning Books and Movies: Towards Story-like Visual Explanations by Watching Movies and Reading Books](https://mlsanity.com/paper/1506.06724/) (2015)
- [Generalization through Memorization: Nearest Neighbor Language Models](https://mlsanity.com/paper/1911.00172/) (2019)
- [Efficient Transformers: A Survey](https://mlsanity.com/paper/2009.06732/) (2020)
- [Compressive Transformers for Long-Range Sequence Modelling](https://mlsanity.com/paper/1911.05507/) (2019)
- [KILT: a Benchmark for Knowledge Intensive Language Tasks](https://mlsanity.com/paper/2009.02252/) (2020)
- [Random Feature Attention](https://mlsanity.com/paper/2103.02143/) (2021)
- [Generating Wikipedia by Summarizing Long Sequences](https://mlsanity.com/paper/1801.10198/) (2018)
- [Adaptive Input Representations for Neural Language Modeling](https://mlsanity.com/paper/1809.10853/) (2018)

## Built on by (most cited first, 30 of 113)

- [Retrieval-Augmented Generation for Large Language Models: A Survey](https://mlsanity.com/paper/2312.10997/) (2023)
- [REPLUG: Retrieval-Augmented Black-Box Language Models](https://mlsanity.com/paper/2301.12652/) (2023)
- [LongMemEval: Benchmarking Chat Assistants on Long-Term Interactive Memory](https://mlsanity.com/paper/2410.10813/) (2024)
- [In-Context Retrieval-Augmented Language Models](https://mlsanity.com/paper/2302.00083/) (2023)
- [Enabling Large Language Models to Generate Text with Citations](https://mlsanity.com/paper/2305.14627/) (2023)
- [A Survey on RAG Meeting LLMs: Towards Retrieval-Augmented Large Language Models](https://mlsanity.com/paper/2405.06211/) (2024)
- [LeanDojo: Theorem Proving with Retrieval-Augmented Language Models](https://mlsanity.com/paper/2306.15626/) (2023)
- [Augmented Language Models: a Survey](https://mlsanity.com/paper/2302.07842/) (2023)
- [Adapting Language Models to Compress Contexts](https://mlsanity.com/paper/2305.14788/) (2023)
- [Retrieval-Augmented Generation for AI-Generated Content: A Survey](https://mlsanity.com/paper/2402.19473/) (2024)
- [Check Your Facts and Try Again: Improving Large Language Models with External Knowledge and Automated Feedback](https://mlsanity.com/paper/2302.12813/) (2023)
- [Adaptive Chameleon or Stubborn Sloth: Revealing the Behavior of Large Language Models in Knowledge Conflicts](https://mlsanity.com/paper/2305.13300/) (2023)
- [Personal LLM Agents: Insights and Survey about the Capability, Efficiency and Security](https://mlsanity.com/paper/2401.05459/) (2024)
- [RA-DIT: Retrieval-Augmented Dual Instruction Tuning](https://mlsanity.com/paper/2310.01352/) (2023)
- [Augmenting Language Models with Long-Term Memory](https://mlsanity.com/paper/2306.07174/) (2023)
- [Chain-of-Note: Enhancing Robustness in Retrieval-Augmented Language Models](https://mlsanity.com/paper/2311.09210/) (2023)
- [A Comprehensive Study of Knowledge Editing for Large Language Models](https://mlsanity.com/paper/2401.01286/) (2024)
- [Dynamic Cheatsheet: Test-Time Learning with Adaptive Memory](https://mlsanity.com/paper/2504.07952/) (2025)
- [LM vs LM: Detecting Factual Errors via Cross Examination](https://mlsanity.com/paper/2305.13281/) (2023)
- [Infini-gram: Scaling Unbounded n-gram Language Models to a Trillion Tokens](https://mlsanity.com/paper/2401.17377/) (2024)
- [ChatDB: Augmenting LLMs with Databases as Their Symbolic Memory](https://mlsanity.com/paper/2306.03901/) (2023)
- [Focused Transformer: Contrastive Training for Context Scaling](https://mlsanity.com/paper/2307.03170/) (2023)
- [Unlimiformer: Long-Range Transformers with Unlimited Length Input](https://mlsanity.com/paper/2305.01625/) (2023)
- [xRAG: Extreme Context Compression for Retrieval-augmented Generation with One Token](https://mlsanity.com/paper/2405.13792/) (2024)
- [Complex Claim Verification with Evidence Retrieved in the Wild](https://mlsanity.com/paper/2305.11859/) (2023)
- [Lift Yourself Up: Retrieval-augmented Text Generation with Self Memory](https://mlsanity.com/paper/2305.02437/) (2023)
- [Training Trajectories of Language Models Across Scales](https://mlsanity.com/paper/2212.09803/) (2022)
- [In-context Pretraining: Language Modeling Beyond Document Boundaries](https://mlsanity.com/paper/2310.10638/) (2023)
- [Grokked Transformers are Implicit Reasoners: A Mechanistic Journey to the Edge of Generalization](https://mlsanity.com/paper/2405.15071/) (2024)
- [Reliable, Adaptable, and Attributable Language Models with Retrieval](https://mlsanity.com/paper/2403.03187/) (2024)

## Abstract (arXiv)

Recent work has improved language models (LMs) remarkably by equipping them with a non-parametric memory component. However, most existing approaches only introduce mem-ories at testing time or represent them using a separately trained encoder, resulting in suboptimal training of the language model. In this work, we present TRIME, a novel yet simple training approach designed for training LMs with memory augmentation. Our approach uses a training objective that directly takes in-batch examples as accessible memory. We also present new methods for memory construction and data batching, which are used for adapting to different sets of memories--local, long-term, and external memory--at testing time. We evaluate TRIME on multiple language modeling and machine translation benchmarks and show that it is able to achieve significant improvements across all the settings. Concretely, TRIME reduces the perplexity from 18.70 to 15.37 on WIKITEXT-103, by effectively leveraging a large memory set from the training corpus. Compared to standard LM training, TRIME adds negligible computational overhead and is compatible with different neural architectures, making it a versatile solution for training memory-augmented LMs.