Large Language Models with Controllable Working Memory (2023)

Abstract: Large language models (LLMs) have led to a series of breakthroughs in natural language processing (NLP), partly owing to the massive amounts of world knowledge they memorize during pretraining. While many downstream applications provide the model with an informational context to aid its 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 then facilitates updating specific model predictions without frequently retraining the model. 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 models (both pretrained and finetuned) could exhibit low controllability and robustness that does not improve with increasing the model size. As a solution, we propose a simple yet effective method – knowledge aware finetuning (KAFT) – to strengthen both controllability and robustness by injecting counterfactual and irrelevant contexts to standard supervised datasets. Our comprehensive evaluation showcases the utility of KAFT across model architectures and sizes.

## introduction large language models (llms) pretrained on large scale datasets have shown promising results across natural language tasks @xcite @xcite @xcite @xcite . however, as models scale ever larger, they become more expensive to train, making it unrealistic to frequently update model parameters. on the other hand, many real world applications often necessitate adjusting model behavior. this dilemma is especially sharp in the case of factual (world) knowledge that plays important role in realizing impressive performance of llms. it is well known that llms memorize large amounts of factual knowledge in their parameters @xcite @xcite , which could potentially be out-dated or incorrect. even for moderate-size models, it is prohibitively expensive to retrain every time an update happens or a mistake is uncovered. even if resources are ample, it is difficult to ensure that the modification of model parameters do not affect unrelated skills or knowledge. in human cognition, working memory @xcite provides the biological brain with the ability to hold information temporarily to perform tasks such as conversation, reasoning, and mathematics in a way that is adaptive to the ever changing environment. as shown both experimentally and theoretically @xcite , working memory is stored in sustained activations of neurons, as opposed to the long term memory which is stored in weights. working memory is also the immediate information buffer that is accessed while performing conscious tasks. in particular, it is where the fusion of perceptual inputs and long term memory happens @xcite . this suggests that a potential method to solve llms' pointwise knowledge update and correction problem is to control the working memory stored in activations, rather than editing the long term memory stored in the model weights. as demonstrated by their powerful in-context few shot learning abilities @xcite , llm could utilize different activation patterns resulting from different contexts during inference to solve a diverse set of tasks without any changes in controllability robustness question dave gilmour and roger waters were in which rock group? how has british art survived in normandy? context george roger waters (born 6 september 1943) is an english singer, . . . later that year, he reunited with the rolling stones bandmates mason, wright and david gilmour... in britain, norman art primarily survives as stonework or metalwork, such as capitals and baptismal fonts... ## methods for concreteness, consider a reading comprehension qa task where the model takes question q together with a context c as its input. the question has an answer label a. we also need a relevance label r denoting whether c entails a. starting with a pretrained lm m , we would like to build a model m such that when the context c is relevant, its answer is always grounded in c, when c is irrelevant, it sticks to the pretrained model's answer. in equations: in the kaft data, @xmath0, where a is the ground truth or counterfactual answer, respectively; @xmath1. here the label is given by the pretrained model's answer to the same question in a few-shot closed book setting, reflecting the model's parametric knowledge. to provide more interpretability, we make the model output its classification of the context's relevance along side the answer itself. see ## results in this section we measure the controllability and robustness of kaft with the metrics defined in sec. 3.7 and compare with baselines in sec. 3.8. ## conclusion in this work, we analyzed the interaction between llms' parametric knowledge (stored in its model parameters) and knowledge contained in informational contexts provided as a part of the input sequence. we find that models are prone to ignoring the context, especially when the context is in conflict with the parametric knowledge. in addition, the model's output can be swayed by irrelevant context even when there is no logical link between such context and the model's task at hand. we quantitatively characterize these behaviours as controllability and robustness of llms when one attempts to control their working memory with noisy context. we proposed a new finetuning method, kaft, that utilizes data augmentations to substantially boost the controllability and robustness of an llm without significantly affecting its performance on standard qa tasks. with kaft, we can build llms with a clear order of priority when utilizing information from difference sources, including its own parametric knowledge.