Propagating Knowledge Updates to LMs Through Distillation

Modern language models have the capacity to store and use immense amounts of knowledge about real-world entities, but it remains unclear how to update such knowledge stored in model parameters. While prior methods for updating knowledge in LMs successfully inject atomic facts, updated LMs fail to make inferences based on injected facts. In this work, we demonstrate that a context distillation-based approach can both impart knowledge about entities and propagate that knowledge to enable broader inferences. Our approach consists of two stages: transfer set generation and distillation on the transfer set. We first generate a transfer set by prompting a language model to generate continuations from the entity definition. Then, we update the model parameters so that the distribution of the LM (the 'student') matches the distribution of the LM conditioned on the definition (the 'teacher') on the transfer set. Our experiments demonstrate that this approach is more effective at propagating knowledge updates than finetuning and other gradient-based knowledge-editing methods. Moreover, it does not compromise performance in other contexts, even when injecting the definitions of up to 150 entities at once. Introduction As large language models (LLMs) are used for a wider variety of applications, it is crucial to ensure that they contain up-to-date information about the world. One potential solution is retrieval augmentation, which prepends retrieved texts to the language model's context [20, 29, 35, 34] . However, this raises inference costs and becomes impractical when updating large amounts of information. An alternative approach, and our goal in this work, is to internalize the new knowledge into the language model via parameter updates [36, 44, 8, 26, 22, 12] . Recent work on injecting LLMs with information about emerging entities [32] demonstrates that updating parameters effectively enables models to acquire updated facts (Rishi Sunak is the prime minister of the UK), but struggles to teach models how to propagate this knowledge, or make inferences based on it (what might Rishi Sunak do tomorrow?). This contrasts with results from retrieval augmentation [20, 35] and chain-of-thought prompting [40] , which show that LLMs can make such inferences when information is placed in the prompt. This work aims to bridge the gap between the two approaches in knowledge injection. We use a form of knowledge distillation [13] called context distillation [1] that updates an LM to act like it is conditioned on a given context, even when that context is not shown. Our approach consists of two steps: transfer set generation and distillation on the generated transfer set. The transfer set consists of continuations of the entity definition sentence generated by prompting a language model. To distill on this transfer set, we minimize the Kullback-Leibler (KL) divergence between the model's predictions 37th Conference on Neural Information Processing Systems (NeurIPS 2023).