Elamur Research
☰
home / papers / 2023 / Large Language Models with Controllable Working Memory

Large Language Models with Controllable Working Memory

Written for the Research library (claude-opus-5-5) from the paper text. Original — aclanthology.org/2023.findings-acl.112.pdf

Summary

LLMs often receive a supporting context in the prompt, but it is unclear how that text interacts with facts stored in their weights. The authors argue for two behaviours. A model should follow a relevant context even when it contradicts memorized knowledge, which they call controllability. It should also disregard unrelated context and answer from memory, which they call robustness. Together, these would let factual updates be supplied as text instead of retraining or editing weights.

Focusing on question answering, the authors build benchmarks with counterfactual contexts and with irrelevant ones. They test pretrained and QA-finetuned T5 and PaLM models of several sizes. Both properties prove weak, and larger models do not fix the problem. Noise in finetuning data partly explains why models learn to ignore context.

Their remedy is knowledge aware finetuning (KAFT). It extends standard supervised QA data in two ways:

  • passages where the answer entity is swapped for a plausible alternative;
  • unrelated passages whose target is the model's own closed-book answer.

KAFT produces models that are both controllable and robust across architectures and sizes. This matters for retrieval-augmented systems, where retrieval is imperfect.

Key points

  • The target priority order is: relevant context first, then the model's parametric knowledge, then irrelevant context, which should have no effect.
  • Evaluated models include encoder-decoder T5 (e.g., T5 XXL) and decoder-only PaLM. Baselines are noisy finetuning, UQA V2 11B and the relevant-only counterfactual finetuning of Longpre et al. (2021), which lacks robustness to irrelevant context.
  • Irrelevant contexts on the same general topic as the question interfere with baseline models most strongly.
  • KAFT needs a separate dataset per model, because the irrelevant-context targets are that model's closed-book answers. The authors propose a dynamic two-forward-pass loss, akin to classifier free guidance, as future work.
  • All KAFT data is English and only QA tasks are covered. The largest run, KAFT-PaLM 540B, used 12190 TPU v4 hours.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

Fields: Natural language processing

Cited by

109 papers of our library cite it (of the 140673 whose bibliographies we read); 9 of them have a page here, 100 not yet.

Cite

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2023.findings-acl.112

@inproceedings{li2023large,
  title = {{Large Language Models with Controllable Working Memory}},
  author = {Li, Daliang and Rawat, Ankit Singh and Zaheer, Manzil and Wang, Xin and Lukasik, Michal and Veit, Andreas and Yu, Felix and Kumar, Sanjiv},
  booktitle = {Findings of the Association for Computational Linguistics: ACL 2023},
  year = {2023},
  doi = {10.18653/v1/2023.findings-acl.112},
  url = {https://doi.org/10.18653/v1/2023.findings-acl.112}
}

References

42 references, 9 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. F. Gregory Ashby, Shawn W. Ell, Vivian V. Valentin, and Michael B. Casale. 2005. FROST: A Distributed Neurocomputational Model of Working Memory Maintenance. Journal of Cognitive Neuroscience, 17(11):1728–1743. find
  2. Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chri… find
  3. Nicola De Cao, Wilker Aziz, and Ivan Titov. 2021. Editing factual knowledge in language models. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 6491–6506, Online and Punta Cana, Dominican Republic. Association for Computational Linguistics. find
  4. Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of NAACL, pages 4171– 4186. find
  5. Bhuwan Dhingra, Jeremy R Cole, Julian Martin Eisenschlos, Daniel Gillick, Jacob Eisenstein, and William W Cohen. 2022. Time-aware language models as temporal knowledge bases. Transactions of the Association for Computational Linguistics, 10:257–273. find
  6. Hady Elsahar, Pavlos Vougiouklis, Arslen Remaci, Christophe Gravier, Jonathon Hare, Frederique Laforest, and Elena Simperl. 2018. T-REx: A large scale alignment of natural language with knowledge base triples. In Proceedings of the Eleventh International Conference on Language Resources and Evaluati… find
  7. Keisuke Fukuda and Geoffrey F. Woodman. 2017. Visual working memory buffers information retrieved from visual long-term memory. Proceedings of the National Academy of Sciences, 114/20. find
  8. J M Fuster. 1973. Unit activity in prefrontal cortex during delayed-response performance: neuronal correlates of transient memory. Journal of Neurophysiology, 36(1):61–78. PMID: 4196203. find
  9. Karl H. Pribram George A. Miller, Eugene Galanter. 1960. Plans and the structure of behavior. Holt, New York. find
  10. Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. 2021. Transformer feed-forward layers are key-value memories. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 5484–5495, Online and Punta Cana, Dominican Republic. Association for Computational… find
  11. Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Ming-Wei Chang. 2020. Realm: Retrieval-augmented language model pre-training. find
  12. Jonathan Ho and Tim Salimans. 2022. Classifier-free diffusion guidance. find our card
  13. Zhiting Hu and Li Erran Li. 2022. A causal lens for controllable text generation. CoRR, abs/2201.09119. — [arXiv:2201.09119](https://arxiv.org/abs/2201.09119) find
  14. Joel Jang, Seonghyeon Ye, Sohee Yang, Joongbo Shin, Janghoon Han, Gyeonghun KIM, Stanley Jungkyu Choi, and Minjoon Seo. 2022. Towards continual knowledge learning of language models. In International Conference on Learning Representations. find
  15. Mandar Joshi, Eunsol Choi, Daniel Weld, and Luke Zettlemoyer. 2017. TriviaQA: A large scale distantly supervised challenge dataset for reading comprehension. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 1601–1611, Vancouver… find our card
  16. Mandar Joshi, Kenton Lee, Yi Luan, and Kristina Toutanova. 2020. Contextualized representations using textual encyclopedic knowledge. CoRR, abs/2004.12006. — [arXiv:2004.12006](https://arxiv.org/abs/2004.12006) find
  17. Nora Kassner and Hinrich Schütze. 2020. Negated and misprimed probes for pretrained language models: Birds can talk, but cannot fly. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 7811–7818, Online. Association for Computational Linguistics. find our card
  18. Daniel Khashabi, Yeganeh Kordi, and Hannaneh Hajishirzi. 2022. Unifiedqa-v2: Stronger generalization via broader cross-format training. find
  19. Daniel Khashabi, Sewon Min, Tushar Khot, Ashish Sabharwal, Oyvind Tafjord, Peter Clark, and Hannaneh Hajishirzi. 2020. UNIFIEDQA: Crossing format boundaries with a single QA system. In Findings of the Association for Computational Linguistics: EMNLP 2020, pages 1896–1907, Online. Association for Com… find our card
  20. Tushar Khot, Peter Clark, Michal Guerquin, Peter Alexander Jansen, and Ashish Sabharwal. 2020. Qasc: A dataset for question answering via sentence composition. ArXiv, abs/1910.11473. find
  21. Angeliki Lazaridou, Adhi Kuncoro, Elena Gribovskaya, Devang Agrawal, Adam Liska, Tayfun Terzi, Mai Gimenez, Cyprien de Masson d’Autume, Tomas Kocisky, Sebastian Ruder, Dani Yogatama, Kris Cao, Susannah Young, and Phil Blunsom. 2021. Mind the gap: Assessing temporal generalization in neural language… find
  22. Percy Liang, Rishi Bommasani, Tony Lee, Dimitris Tsipras, Dilara Soylu, Michihiro Yasunaga, Yian Zhang, Deepak Narayanan, Yuhuai Wu, Ananya Kumar, Benjamin Newman, Binhang Yuan, Bobby Yan, Ce Zhang, Christian Cosgrove, Christopher D. Manning, Christopher Ré, Diana Acosta-Navas, Drew A. Hudson, Eric… find
  23. Shayne Longpre, Kartik Perisetla, Anthony Chen, Nikhil Ramesh, Chris DuBois, and Sameer Singh. 2021. Entity-based knowledge conflicts in question answering. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 7052–7063, Online and Punta Cana, Dominican Re… find our card
  24. Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. 2022a. Locating and editing factual knowledge in gpt. arXiv preprint arXiv:2202.05262. find
  25. Kevin Meng, Arnab Sen Sharma, Alex Andonian, Yonatan Belinkov, and David Bau. 2022b. Mass-editing memory in a transformer. find
  26. Eric Mitchell, Charles Lin, Antoine Bosselut, Chelsea Finn, and Christopher D Manning. 2022. Fast model editing at scale. In International Conference on Learning Representations. find
  27. Ella Neeman, Roee Aharoni, Or Honovich, Leshem Choshen, Idan Szpektor, and Omri Abend. 2022. to appear. find
  28. Katsuhiko Ogata. 1996. Modern Control Engineering (3rd Ed.). Prentice-Hall, Inc., USA. find
  29. Marwan Omar, Soohyeon Choi, DaeHun Nyang, and David Mohaisen. 2022. Robust natural language processing: Recent advances, challenges, and future directions. CoRR, abs/2201.00768. — [arXiv:2201.00768](https://arxiv.org/abs/2201.00768) find
  30. Yasumasa Onoe, Michael JQ Zhang, Eunsol Choi, and Greg Durrett. 2022. Entity cloze by date: What lms know about unseen entities. arXiv preprint arXiv:2205.02832. find
  31. Liangming Pan, Wenhu Chen, Min-Yen Kan, and William Yang Wang. 2021. Contraqa: Question answering under contradicting contexts. CoRR, abs/2110.07803. — [arXiv:2110.07803](https://arxiv.org/abs/2110.07803) find
  32. Fabio Petroni, Patrick S. H. Lewis, Aleksandra Piktus, Tim Rocktäschel, Yuxiang Wu, Alexander H. Miller, and Sebastian Riedel. 2020. How context affects language models’ factual predictions. CoRR, abs/2005.04611. find our card
  33. Fabio Petroni, Tim Rocktäschel, Sebastian Riedel, Patrick Lewis, Anton Bakhtin, Yuxiang Wu, and Alexander Miller. 2019. Language models as knowledge bases? In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural… find our card
  34. Jack W. Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, Eliza Rutherford, Tom Hennigan, Jacob Menick, Albin Cassirer, Richard Powell, George van den Driessche, Lisa Anne Hendricks, Maribeth Rauh, Po-Sen… find
  35. Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. 2020a. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research, 21(140):1–67. find
  36. Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. 2020b. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research, 21(140):1–67. find
  37. Pranav Rajpurkar, Robin Jia, and Percy Liang. 2018. Know what you don’t know: Unanswerable questions for squad. find our card
  38. Adam Roberts, Colin Raffel, and Noam Shazeer. 2020. How much knowledge can you pack into the parameters of a language model? arXiv preprint arXiv:2002.08910. find
  39. Shaden Smith, Mostofa Patwary, Brandon Norick, Patrick LeGresley, Samyam Rajbhandari, Jared Casper, Zhun Liu, Shrimai Prabhumoye, George Zerveas, Vijay Korthikanti, Elton Zhang, Rewon Child, Reza Yazdani Aminabadi, Julie Bernauer, Xia Song, Mohammad Shoeybi, Yuxiong He, Michael Houston, Saurabh Tiwa… find our card
  40. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in NeurIPS. find
  41. Hanqing Zhang, Haolin Song, Shaoyu Li, Ming Zhou, and Dawei Song. 2022. A survey of controllable text generation using transformer-based pre-trained language models. CoRR, abs/2201.05337. — [arXiv:2201.05337](https://arxiv.org/abs/2201.05337) find
  42. Chen Zhu, Ankit Singh Rawat, Manzil Zaheer, Srinadh Bhojanapalli, Daliang Li, Felix Yu, and Sanjiv Kumar. 2020. Modifying memories in transformer models. arXiv preprint arXiv:2012.00363. find