Propagating Knowledge Updates to LMs Through Distillation

TL;DR

This work tackles the challenge of updating knowledge stored in language model parameters so that injected facts propagate to inferences. It introduces a context distillation framework with transfer-set generation followed by KL-based distillation, enabling knowledge to influence model behavior beyond the explicit definition. Across two propagation benchmarks and three base models, the approach consistently outperforms fine-tuning and several gradient-based edits, with scalable updates to hundreds of entities and robust performance on newer information. Limitations include remaining gaps relative to in-context knowledge and computational constraints on very large models, highlighting avenues for future scaling and integration with retrieval-based or instruction-tuning paradigms.

Abstract

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 fine-tuning 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.

Figures (5)

• Figure 1: Overview of our distillation approach. Our goal is to inject the entity definition ($\mathbf{d}_e$) into the student model ($M_s$) and propagate it to make inferences based on the injected knowledge. This example uses ChatGPT as a new entity. We first generate a set of continuations of the entity's definition using a generator model (Step 1), then use these to distill the information from definition into the student model via a KL loss between the conditioned and unconditioned models (Step 2); see Section \ref{['sec:method']} for formulation.
• Figure 2: Results on GPT-Neo with varying numbers of model updates for fine-tuning and distillation approach. Left: target perplexity; right: perplexity on the definition sentence. Only distillation continues to improve in both target and definition perplexity as the number of updates increase.
• Figure 3: Per-token NLL of tokens in continuations before conditioning on definitions and after (fractional reduction). Tokens not in the definition (blue dots) are changed less but do see lower NLL when they are inferable from the definition (examples).
• Figure 4: Editing for multiple entities at once. We report the average from three runs with different random seeds for shuffling training data.
• Figure 5: Perplexity for using $n$ distinct transfer sentences during distillation, with number of updates standardized to 10. We see the benefit of having a diverse transfer set compared to repeating the same transfer sentence 10 times.