# Large Language Models with Controllable Working Memory

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

## Key facts

- Cited by 169 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 51 references that are in the map.
- First posted on arXiv 9 Nov 2022.

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

- [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)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension](https://mlsanity.com/paper/1705.03551/) (2017)
- [REALM: Retrieval-Augmented Language Model Pre-Training](https://mlsanity.com/paper/2002.08909/) (2020)
- [Language Models as Knowledge Bases?](https://mlsanity.com/paper/1909.01066/) (2019)
- [Know What You Don't Know: Unanswerable Questions for SQuAD](https://mlsanity.com/paper/1806.03822/) (2018)
- [Transformer Feed-Forward Layers Are Key-Value Memories](https://mlsanity.com/paper/2012.14913/) (2020)
- [Scaling Language Models: Methods, Analysis & Insights from Training Gopher](https://mlsanity.com/paper/2112.11446/) (2021)
- [Mass-Editing Memory in a Transformer](https://mlsanity.com/paper/2210.07229/) (2022)
- [How Much Knowledge Can You Pack Into the Parameters of a Language Model?](https://mlsanity.com/paper/2002.08910/) (2020)
- [Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model](https://mlsanity.com/paper/2201.11990/) (2022)
- [Editing Factual Knowledge in Language Models](https://mlsanity.com/paper/2104.08164/) (2021)
- [Fast Model Editing at Scale](https://mlsanity.com/paper/2110.11309/) (2021)
- [UnifiedQA: Crossing Format Boundaries With a Single QA System](https://mlsanity.com/paper/2005.00700/) (2020)
- [Entity-Based Knowledge Conflicts in Question Answering](https://mlsanity.com/paper/2109.05052/) (2021)
- [Time-Aware Language Models as Temporal Knowledge Bases](https://mlsanity.com/paper/2106.15110/) (2021)
- [QASC: A Dataset for Question Answering via Sentence Composition](https://mlsanity.com/paper/1910.11473/) (2019)
- [Negated and Misprimed Probes for Pretrained Language Models: Birds Can Talk, But Cannot Fly](https://mlsanity.com/paper/1911.03343/) (2019)
- [Mind the Gap: Assessing Temporal Generalization in Neural Language Models](https://mlsanity.com/paper/2102.01951/) (2021)
- [Modifying Memories in Transformer Models](https://mlsanity.com/paper/2012.00363/) (2020)
- [How Context Affects Language Models' Factual Predictions](https://mlsanity.com/paper/2005.04611/) (2020)
- [Towards Continual Knowledge Learning of Language Models](https://mlsanity.com/paper/2110.03215/) (2021)
- [UnifiedQA-v2: Stronger Generalization via Broader Cross-Format Training](https://mlsanity.com/paper/2202.12359/) (2022)
- [Entity Cloze By Date: What LMs Know About Unseen Entities](https://mlsanity.com/paper/2205.02832/) (2022)
- [A Causal Lens for Controllable Text Generation](https://mlsanity.com/paper/2201.09119/) (2022)
- [Contextualized Representations Using Textual Encyclopedic Knowledge](https://mlsanity.com/paper/2004.12006/) (2020)

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

- [A Survey on Hallucination in Large Language Models: Principles, Taxonomy, Challenges, and Open Questions](https://mlsanity.com/paper/2311.05232/) (2023)
- [Large Language Models Can Be Easily Distracted by Irrelevant Context](https://mlsanity.com/paper/2302.00093/) (2023)
- [A Survey on RAG Meeting LLMs: Towards Retrieval-Augmented Large Language Models](https://mlsanity.com/paper/2405.06211/) (2024)
- [Making Retrieval-Augmented Language Models Robust to Irrelevant Context](https://mlsanity.com/paper/2310.01558/) (2023)
- [Hallucination is Inevitable: An Innate Limitation of Large Language Models](https://mlsanity.com/paper/2401.11817/) (2024)
- [Benchmarking Large Language Models in Retrieval-Augmented Generation](https://mlsanity.com/paper/2309.01431/) (2023)
- [Knowledge Conflicts for LLMs: A Survey](https://mlsanity.com/paper/2403.08319/) (2024)
- [Rethinking Machine Unlearning for Large Language Models](https://mlsanity.com/paper/2402.08787/) (2024)
- [Continual Learning of Large Language Models: A Comprehensive Survey](https://mlsanity.com/paper/2404.16789/) (2024)
- [Survey on Factuality in Large Language Models: Knowledge, Retrieval and Domain-Specificity](https://mlsanity.com/paper/2310.07521/) (2023)
- [Knowledge Editing for Large Language Models: A Survey](https://mlsanity.com/paper/2310.16218/) (2023)
- [RoleLLM: Benchmarking, Eliciting, and Enhancing Role-Playing Abilities of Large Language Models](https://mlsanity.com/paper/2310.00746/) (2023)
- [Domain Specialization as the Key to Make Large Language Models Disruptive: A Comprehensive Survey](https://mlsanity.com/paper/2305.18703/) (2023)
- [A Survey of Context Engineering for Large Language Models](https://mlsanity.com/paper/2507.13334/) (2025)
- [AI Transparency in the Age of LLMs: A Human-Centered Research Roadmap](https://mlsanity.com/paper/2306.01941/) (2023)
- [To Believe or Not to Believe Your LLM](https://mlsanity.com/paper/2406.02543/) (2024)
- [Empowering Biomedical Discovery with AI Agents](https://mlsanity.com/paper/2404.02831/) (2024)
- [RPTQ: Reorder-based Post-training Quantization for Large Language Models](https://mlsanity.com/paper/2304.01089/) (2023)
- [R-Tuning: Instructing Large Language Models to Say `I Don't Know'](https://mlsanity.com/paper/2311.09677/) (2023)
- [The Earth is Flat because...: Investigating LLMs' Belief towards Misinformation via Persuasive Conversation](https://mlsanity.com/paper/2312.09085/) (2023)
- [InstructRAG: Instructing Retrieval-Augmented Generation via Self-Synthesized Rationales](https://mlsanity.com/paper/2406.13629/) (2024)
- [DisentQA: Disentangling Parametric and Contextual Knowledge with Counterfactual Question Answering](https://mlsanity.com/paper/2211.05655/) (2022)
- [BadEdit: Backdooring large language models by model editing](https://mlsanity.com/paper/2403.13355/) (2024)
- [Large Language Models and Games: A Survey and Roadmap](https://mlsanity.com/paper/2402.18659/) (2024)
- [Pearl: Personalizing Large Language Model Writing Assistants with Generation-Calibrated Retrievers](https://mlsanity.com/paper/2311.09180/) (2023)
- [HalluLens: LLM Hallucination Benchmark](https://mlsanity.com/paper/2504.17550/) (2025)
- [Cutting Off the Head Ends the Conflict: A Mechanism for Interpreting and Mitigating Knowledge Conflicts in Language Models](https://mlsanity.com/paper/2402.18154/) (2024)
- [Enhancing Noise Robustness of Retrieval-Augmented Language Models with Adaptive Adversarial Training](https://mlsanity.com/paper/2405.20978/) (2024)
- [Context-faithful Prompting for Large Language Models](https://mlsanity.com/paper/2303.11315/) (2023)
- [Learning to Route LLMs with Confidence Tokens](https://mlsanity.com/paper/2410.13284/) (2024)

## Abstract (arXiv)

Large language models (LLMs) have led to a series of breakthroughs in natural language processing (NLP), owing to their excellent understanding and generation abilities. Remarkably, what further sets these models apart is the massive amounts of world knowledge they internalize during pretraining. While many downstream applications provide the model with an informational context to aid its performance on the underlying task, how the model's world knowledge interacts with the factual information presented in the context remains under explored. As a desirable behavior, an LLM should give precedence to the context whenever it contains task-relevant information that conflicts with the model's memorized knowledge. This enables model predictions to be grounded in the context, which can then be used to update or correct specific model predictions without frequent retraining. By contrast, when the context is irrelevant to the task, the model should ignore it and fall back on its internal knowledge. In this paper, we undertake a first joint study of the aforementioned two properties, namely controllability and robustness, in the context of LLMs. We demonstrate that state-of-the-art T5 and PaLM (both pretrained and finetuned) could exhibit poor controllability and robustness, which do not scale with increasing model size. As a solution, we propose a novel method - Knowledge Aware FineTuning (KAFT) - to strengthen both controllability and robustness by incorporating counterfactual and irrelevant contexts to standard supervised datasets. Our comprehensive evaluation showcases the utility of KAFT across model architectures and sizes.