Training Language Models with Memory Augmentation
Open accessAbstract
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.
The claims and statistics of this paper haven’t been checked yet.
Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.
Integrity signals
Registry
Cites
Statistics
Claims the paper makes
The claims and statistics of this paper haven’t been checked yet.
Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.
Claims other papers make about this one
Each is checked against this paper’s own text: does it say what the citing paper says it does?
Joint training of retrieval and generation yields better memory utilization than frozen-retriever baselines.
“Joint training of retrieval and generation [Zhong et al., 2022] yields better memory utilization than frozen-retriever baselines.”
From Memory for Autonomous LLM Agents:Mechanisms, Evaluation, and Emerging Frontiers (2026) · cites this paper as Zhong et al., 2022
Self-check
Percentages, sums and reported differences re-computed from the paper’s own numbers.
The self-check needs the paper’s full text.
Cites
Rethinking Attention with Performers
Choromanski, K., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T., Hawkins, P., Davis, J., Mohiuddin, A., Kaiser, L., et al. (2020) · arXiv · doi:10.48550/arxiv.2009.14794 · Cited by 695 · cited 1× · see citing works
“, 2020) or subquadratic-time attention (Wang et al., 2020; Choromanski et al., 2020; Peng et al., 2021; Katharopoulos et al., 2020).”
Flagged first, then by citations
Research works printed but not linked: 53
- Rami Al-Rfou, Dokook Choe, Noah Constant, Mandy Guo, and Llion Jones. 2019. Character-level lan- guage modeling with deeper self-attention. In Con- ference on Artificial Intelligence (AAAI), volume 33, pages 3159-3166.
- Alexei Baevski and Michael Auli. 2019. Adaptive in- put representations for neural language modeling. In International Conference on Learning Representa- tions (ICLR).
- Iz Beltagy, Matthew E Peters, and Arman Cohan. 2020. Longformer: The long-document Trans- former. arXiv preprint arXiv:2004.05150. · arXiv
- Sebastian Borgeaud, Arthur Mensch, Jordan Hoff- mann, Trevor Cai, Eliza Rutherford, Katie Millican, George van den Driessche, Jean-Baptiste Lespiau, Bogdan Damoc, Aidan Clark, et al. 2021. Improv- ing language models by retrieving from trillions of tokens. arXiv preprint arXiv:2112.04426. · arXiv
- Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. 2020. Language models are few-shot learners. Advances in Neural Information Process- ing Systems (NeurIPS), 33:1877-1901.
- Wenhu Chen, Pat Verga, Michiel de Jong, John Wi- eting, and William Cohen. 2022. Augmenting pre-trained language models with qa-memory for open-domain question answering. arXiv preprint arXiv:2204.04581. · arXiv
- Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. 2019. Generating long se- quences with sparse Transformers. arXiv preprint arXiv:1904.10509. · arXiv
- Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. 2022. PaLM: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311. · arXiv
- Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Car- bonell, Quoc V Le, and Ruslan Salakhutdinov. 2019. Transformer-XL: Attentive language models beyond a fixed-length context. In Association for Computa- tional Linguistics (ACL).
- Michiel de Jong, Yury Zemlyanskiy, Nicholas FitzGer- ald, Fei Sha, and William Cohen. 2022. Mention memory: incorporating textual knowledge into trans- formers through entity mention attention. In Inter- national Conference on Learning Representations (ICLR).
- Cyprien de Masson D'Autume, Sebastian Ruder, Ling- peng Kong, and Dani Yogatama. 2019. Episodic memory in lifelong language learning. Advances in Neural Information Processing Systems (NeurIPS), 32.
- Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. BERT: Pre-training of deep bidirectional Transformers for language under- standing. In North American Chapter of the Associ- ation for Computational Linguistics (NAACL).
- Angela Fan, Claire Gardent, Chloé Braud, and An- toine Bordes. 2021. Augmenting Transformers with knn-based composite memory for dialog. Transac- tions of the Association of Computational Linguis- tics (TACL), 9:82-99.
- Edouard Grave, Moustapha M Cisse, and Armand Joulin. 2017a. Unbounded cache model for online language modeling with open vocabulary. Advances in Neural Information Processing Systems (NIPS), 30.
- Edouard Grave, Armand Joulin, and Nicolas Usunier. 2017b. Improving neural language models with a continuous cache. In International Conference on Learning Representations (ICLR).
- Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pa- supat, and Ming-Wei Chang. 2020. REALM: Retrieval-augmented language model pre-training. In International Conference on Machine Learning (ICML).
- Raia Hadsell, Sumit Chopra, and Yann LeCun. 2006. Dimensionality reduction by learning an invariant mapping. In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), volume 2, pages 1735-1742.
- Junxian He, Graham Neubig, and Taylor Berg- Kirkpatrick. 2021. Efficient nearest neighbor lan- guage models. In Empirical Methods in Natural Language Processing (EMNLP).
- Gautier Izacard and Edouard Grave. 2021. Leveraging passage retrieval with generative models for open domain question answering. In European Chap- ter of the Association for Computational Linguistics (EACL).
- Haozhe Ji, Rongsheng Zhang, Zhenyu Yang, Zhipeng Hu, and Minlie Huang. 2022. LaMemo: Language modeling with look-ahead memory. In North Amer- ican Chapter of the Association for Computational Linguistics (NAACL).
- Jeff Johnson, Matthijs Douze, and Hervé Jégou. 2019. Billion-scale similarity search with GPUs. IEEE Transactions on Big Data, 7(3):535-547.
- Vladimir Karpukhin, Barlas Oguz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. 2020. Dense passage retrieval for open-domain question answering. In Empirical Methods in Natural Language Processing (EMNLP).
- Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pap- pas, and François Fleuret. 2020. Transformers are RNNs: Fast autoregressive Transformers with linear attention. In International Conference on Machine Learning (ICML), pages 5156-5165.
- Urvashi Khandelwal, Angela Fan, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. 2021. Nearest neigh- bor machine translation. In International Confer- ence on Learning Representations (ICLR).
- Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. 2020. Generalization through memorization: Nearest neighbor language models. In International Conference on Learning Representations (ICLR).
- Ankit Kumar, Ozan Irsoy, Peter Ondruska, Mohit Iyyer, James Bradbury, Ishaan Gulrajani, Victor Zhong, Romain Paulus, and Richard Socher. 2016. Ask me anything: Dynamic memory networks for natural language processing. In International Conference on Machine Learning (ICML), pages 1378-1387.
- Tao Lei. 2021. When attention meets fast recurrence: Training language models with reduced compute. In Empirical Methods in Natural Language Processing (EMNLP).
- Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Hein- rich Küttler, Mike Lewis, Wen-tau Yih, Tim Rock- täschel, et al. 2020. Retrieval-augmented generation for knowledge-intensive NLP tasks. Advances in Neural Information Processing Systems (NeurIPS), 33:9459-9474.
- Peter J Liu, Mohammad Saleh, Etienne Pot, Ben Goodrich, Ryan Sepassi, Lukasz Kaiser, and Noam Shazeer. 2018. Generating Wikipedia by summariz- ing long sequences. In International Conference on Learning Representations (ICLR).
- Pedro Henrique Martins, Zita Marinho, and André FT Martins. 2022. ∞-former: Infinite memory Trans- former. In Association for Computational Linguis- tics (ACL).
- Yuxian Meng, Shi Zong, Xiaoya Li, Xiaofei Sun, Tian- wei Zhang, Fei Wu, and Jiwei Li. 2022. Gnn-lm: Language modeling based on global contexts via gnn. In International Conference on Learning Rep- resentations (ICLR).
- Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. 2017. Pointer sentinel mixture mod- els. In International Conference on Learning Repre- sentations (ICLR).
- Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. 2019. fairseq: A fast, extensible toolkit for sequence modeling. In Proceedings of NAACL-HLT 2019: Demonstrations.
- Kishore Papineni, Salim Roukos, Todd Ward, and Wei- Jing Zhu. 2002. BLEU: a method for automatic eval- uation of machine translation. In Association for Computational Linguistics (ACL), pages 311-318.
- Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. 2021. Random feature attention. In International Confer- ence on Learning Representations (ICLR).
- Fabio Petroni, Aleksandra Piktus, Angela Fan, Patrick Lewis, Majid Yazdani, Nicola De Cao, James Thorne, Yacine Jernite, Vladimir Karpukhin, Jean Maillard, et al. 2021. KILT: a benchmark for knowl- edge intensive language tasks. In North American Chapter of the Association for Computational Lin- guistics (NAACL), pages 2523-2544.
- Matt Post. 2018. A call for clarity in reporting BLEU scores. In Proceedings of the Third Conference on Machine Translation: Research Papers, pages 186- 191, Belgium, Brussels. Association for Computa- tional Linguistics.
- Jack W Rae, Anna Potapenko, Siddhant M Jayaku- mar, and Timothy P Lillicrap. 2020. Compressive Transformers for long-range sequence modelling. In International Conference on Learning Representa- tions (ICLR).
- Abigail See, Peter J Liu, and Christopher D Manning. 2017. Get to the point: Summarization with pointer- generator networks. In Association for Computa- tional Linguistics (ACL), pages 1073-1083.
- Devendra Singh, Siva Reddy, Will Hamilton, Chris Dyer, and Dani Yogatama. 2021. End-to-end train- ing of multi-document reader and retriever for open- domain question answering. Advances in Neural In- formation Processing Systems (NIPS).
- Sainbayar Sukhbaatar, Édouard Grave, Piotr Bo- janowski, and Armand Joulin. 2019. Adaptive atten- tion span in Transformers. In Association for Com- putational Linguistics (ACL), pages 331-335.
- Sainbayar Sukhbaatar, Da Ju, Spencer Poff, Stephen Roller, Arthur Szlam, Jason Weston, and Angela Fan. 2021. Not all memories are created equal: Learning to forget by expiring. In International Con- ference on Machine Learning (ICML), pages 9902- 9912.
- PMLR.
- Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. 2020. Efficient Transformers: A survey. arXiv preprint arXiv:2009.06732. · arXiv
- Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. Advances in Neural Information Process- ing Systems (NIPS), 30.
- Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. 2020. Linformer: Self- attention with linear complexity. arXiv preprint arXiv:2006.04768. · arXiv
- Yuhuai Wu, Markus N Rabe, DeLesley Hutchins, and Christian Szegedy. 2022. Memorizing Transformers. In International Conference on Learning Represen- tations (ICLR).
- Dani Yogatama, Cyprien de Masson d'Autume, and Lingpeng Kong. 2021. Adaptive semiparametric language models. Transactions of the Association of Computational Linguistics (TACL), 9:362-373.
- Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago On- tanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. 2020. Big Bird: Transformers for longer sequences. In Advances in Neural Infor- mation Processing Systems (NeurIPS), volume 33, pages 17283-17297.
- Yury Zemlyanskiy, Joshua Ainslie, Michiel de Jong, Philip Pham, Ilya Eckstein, and Fei Sha. 2021. Readtwice: Reading very large documents with 5668 memories. In North American Chapter of the As- sociation for Computational Linguistics (NAACL).
- Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher De- wan, Mona Diab, Xian Li, Xi Victoria Lin, et al. 2022. OPT: Open pre-trained Transformer language models. arXiv preprint arXiv:2205.01068. · arXiv
- Xin Zheng, Zhirui Zhang, Junliang Guo, Shujian Huang, Boxing Chen, Weihua Luo, and Jiajun Chen. 2021. Adaptive nearest neighbor machine transla- tion. In Association for Computational Linguistics (ACL), pages 368-374.
- Yukun Zhu, Ryan Kiros, Rich Zemel, Ruslan Salakhut- dinov, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. 2015. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. In International Conference on Computer Vision (ICCV), pages 19-27.
Cited by
Sorted by citations
Retrieval-Augmented Generation for Large Language Models: A Survey
Gao, Y., Xiong, Y., Gao, X., Jia, K., Pan, J., Bi, Y., Dai, Y., Sun, J., Wang, M., & Wang, H. (2023) · arXiv · doi:10.48550/arxiv.2312.10997 · Cited by 876 · see citing works
“The preparation and curation of high-quality training datasets, defining fine-tuning objectives, and providing corresponding computational resources are necessary.”
BackgroundIn-Context Retrieval-Augmented Language Models
Ram, O., Levine, Y., Dalmedigos, I., Muhlgay, D., Shashua, A., Leyton-Brown, K., & Shoham, Y. (2023) · Transactions of the Association for Computational Linguistics · doi:10.1162/tacl_a_00605 · Cited by 449 · see citing works
ChatDB: Augmenting LLMs with Databases as Their Symbolic Memory
Hu, C., Fu, J., Du, C., Luo, S., Zhao, J., & Zhao, H. (2023) · arXiv · doi:10.48550/arxiv.2306.03901 · Cited by 13 · see citing works
“Some neural memory mechanisms have been explored (Wu et al., 2022a; Khattab et al., 2022; Zhong et al., 2022) to overcome the limited token input issue of LLMs.”
BackgroundChatLogic: Integrating Logic Programming with Large Language Models for Multi-Step Reasoning
Wang, Z., Liu, J., Bao, Q., Rong, H., & Zhang, J. (2024) · 2024 International Joint Conference on Neural Networks (IJCNN) · doi:10.1109/ijcnn60899.2024.10650138 · Cited by 8 · see citing works
LongMemEval: Benchmarking Chat Assistants on Long-Term Interactive Memory
Wu, D., Wang, H., Yu, W., Zhang, Y., Chang, K.-W., & Yu, D. (2024) · arXiv · doi:10.48550/arxiv.2410.10813 · Cited by 3 · see citing works
“A second line of research integrates differentiable memory modules into language models, proposing specialized architectural designs and training strategies to enhance memory capabilities (Weston et al., 2014; Wu et al., 2022; Zhong et al., 2022; Wang et al., 2023).”
Do as We Do, Not as You Think: the Conformity of Large Language Models
Weng, Z., Chen, G., & Wang, W. (2025) · arXiv · doi:10.48550/arxiv.2501.13381 · Cited by 1 · see citing works
“Recent efforts focus on supplementing LLMs with external, up-to-date information to improve their accuracy [71–75].”
Memory for Autonomous LLM Agents:Mechanisms, Evaluation, and Emerging Frontiers
Du, P. (2026) · arXiv · Cited by 1 · see citing works
“Joint training of retrieval and generation [Zhong et al., 2022] yields better memory utilization than frozen-retriever baselines.”
Adapting Language Models to Compress Contexts
Chevalier, A., Wettig, A., Ajith, A., & Chen, D. (2023) · arXiv · Cited by 59 · see citing works
“A number of approaches have been proposed to infuse external knowledge in the input layer (Guu et al., 2020; Shi et al., 2023), intermediate layers (Borgeaud et al., 2022) or at the output layer (Khandelwal et al., 2020; Zhong et al., 2022).”
BackgroundSelf-citationThe AI Hippocampus: How Far are We From Human Memory?
Jia, Z., Li, J., Kang, Y., Wang, Y., Wu, T., Wang, Q., Wang, X., Zhang, S., Shen, J., Li, Q., et al. (2026) · arXiv · doi:10.48550/arxiv.2601.09113
- Authors
- Zhong, Z., Lei, T., & Chen, D.
- Affiliations
- Princeton University (US)
- ORCID
- Venue
- arXiv
- Publisher
- arXiv
- Published
- 2022
- Pages
- 5657-5673
- arXiv
- 2205.12674
- OpenAlex
- W4385573102
- Cited by
- 60
- References
- 1
- Citations by year
- 2026: 5 · 2025: 15 · 2024: 18 · 2023: 20 · 2022: 2
- FWCI
- 7.60
- Type
- conference paper
- Language
- en
- Open access
- hybrid
- doi.org
- Licence
- Locations
- Conference on Empirical Methods in Natural Language Processing (EMNLP) (publishedVersion, cc-by)
- arXiv (Cornell University) (submittedVersion)
- arXiv (Cornell University)
- Topics
- Physical Sciences › Computer Science › Artificial Intelligence › Topic Modeling
- Physical Sciences › Computer Science › Artificial Intelligence › Natural Language Processing Techniques
- Physical Sciences › Computer Science › Artificial Intelligence › Speech Recognition and Synthesis
- Keywords
- language model training, memory construction, long-term memory, external memory, perplexity reduction, machine translation
About these checks
Registry, retraction, Cites and statistics results are computed, not judged, with one exception: a printed reference that no index lists is read once by a model before it counts as a citation error. Claims are extracted by a model and kept only when their quote appears verbatim in the text; each cited source is then checked, with a model’s help, for whether it says what the claim relies on. Those verdicts can be wrong — read the claim’s quote and the cited source before relying on one.
Statistics and arithmetic are recomputed, not judged.
Confidence bands are shown only for fields calibrated on our annotated test set. None is calibrated yet, so no band is shown.
What counts as a claim ›How we measure confidence ›How we check this ›
Are you an author, or do you see a problem with this page? Request a review
