Large Language Models with Controllable Working Memory
This addresses the issue of grounding LLM predictions in context for users in NLP applications, but it is incremental as it builds on existing fine-tuning techniques.
The paper tackles the problem of how large language models (LLMs) interact with contextual information, showing that state-of-the-art models like T5 and PaLM exhibit poor controllability and robustness, which do not improve with model size. They propose Knowledge Aware FineTuning (KAFT), a method that enhances both properties by incorporating counterfactual and irrelevant contexts into training, achieving improved performance across architectures and sizes.
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.