Large Language Models with Controllable Working Memory

TL;DR

This paper investigates how large language models manage internal knowledge versus contextual information, revealing controllability and robustness issues, and proposes KAFT to improve these properties by incorporating counterfactual contexts.

Contribution

The paper introduces KAFT, a novel fine-tuning method that enhances LLM controllability and robustness by integrating counterfactual and irrelevant contexts during training.

Findings

State-of-the-art models show poor controllability and robustness.

KAFT improves model performance across architectures and sizes.

Controllability and robustness do not scale with model size.

Abstract

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.