SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Shankar Padmanabhan, Yasumasa Onoe, Michael J.Q. Zhang, Greg Durrett, Eunsol Choi  
Department of Computer Science  
The University of Texas at Austin  
shankarpadmanabhan@utexas.edu
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.
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
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 3 for formulation.
on the transfer set when it conditions on the definition (the "teacher" for distillation) and when it does not (the "student", or the language model itself). Figure 1 shows this approach.
We evaluate our approach on two knowledge propagation benchmarks: ENTITY INFERENCES [32] and Entity Cloze by Date (ECBD) [31]. We evaluate on three language models and find that our distillation approach outperforms fine-tuning and prior editing methods (MEND [26] and MEMIT [23]) across all models. To investigate the robustness of our approach, we present an ablation study focusing on the design choices during transfer set construction. Encouragingly, we find that distilling on transfer sets constructed from the base language model itself is competitive with those generated by a much larger model (GPT-3.5). This demonstrates that context distillation does not rely on distilling from a larger model, and that our approach can work across a range of model sizes. Finally, we show that our approach can be scaled to inject larger amounts of information at once: we can inject over 100 new entities into a language model with minimal performance degradation, suggesting that the distillation process performs relatively targeted editing even without additional objectives to ensure specificity as in past methods [22, 23].
To summarize, we present a new approach for propagating injected knowledge. We show that a knowledge distillation technique can effectively impart and propagate knowledge from entity definitions into the parameters of a pre-trained language model, compared to existing knowledge editing methods. Yet, we observe robust gap between providing information in-context and parameter updating methods, leaving ample room for future work. Our code and data are available at https://github.com/shankarp8/knowledge_distillation.
2 Background and Task Setup
2.1 Motivating Example
Figure 1 shows a motivating example. An LM trained on text collected prior to November 2022 will not have specific knowledge about what ChatGPT is, as ChatGPT was introduced after that time. Past retrieval-augmented generation methods [20] have shown that conditioning on information about this entity can lead to lower perplexities when evaluating on sentences like ChatGPT can respond to natural language questions [35, 32]. For example, the model assigns a higher likelihood to tokens like respond given the knowledge that ChatGPT is a chatbot.
Our approach relies on teaching a "student model" (the LM itself) to match the next-token distributions given by the model conditioned on the definition sentence even when the definition sentence is not shown. We do this via a distillation process on a set of continuations, or sampled sentences following
the definition. We impose a KL penalty between the student and teacher distributions of a set of target tokens, namely all those occurring after ChatGPT in the continuation. Because the distillation process does not make updates on tokens where the teacher and student have the same distribution (zero KL), only tokens that are in some way predictable from the definition drive parameter updates (see Section 7 for discussion).
2.2 Task Setup
We refer to language models  $M$  as  $M(\mathbf{x}) \to \mathcal{D}(\mathcal{V})$ , mapping an input context  $\mathbf{x} = (x_{1}, \ldots, x_{n})$  to a next-word distribution  $\mathcal{D}(\mathcal{V}) = p(\cdot \mid x_{1}, \ldots, x_{n})$  over a vocabulary  $\mathcal{V}$ . We will also use  $M(\mathbf{x}) \to \mathcal{D}(\mathcal{V})_{1,\dots,n}$  to represent the collection of distributions after each prefix of  $\mathbf{x}$ , which is a standard operation used in language model training. To update knowledge in the base language model  $M_{\mathrm{Base}}$ , definitional information  $\mathbf{d}_e = (d_1, \ldots, d_m)$  for an entity  $e$  is provided. We will use  $e$  both as an indicator and also a reference to the entity name string (e.g., ChatGPT in Figure 1). Our goal is to update  $M_{\mathrm{Base}}$  to  $M_s$  so that it "knows"  $\mathbf{d}_e$ , by matching  $M_s(\mathbf{x})$  with  $M_t(\mathbf{x} \mid \mathbf{d}_e)$  (the teacher model) as closely as possible with our distillation scheme, when  $\mathbf{x}$  is relevant to entity  $e$ . We set the teacher model  $M_t$  to be a copy of  $M_s$ .
We evaluate on two factors. First, propagation success measures how well the updated language model  $M_{s}$  acquired information about  $\mathbf{d}_e$  to make correct inferences in probe sentences. Crucially, our evaluation here is not just a narrow notion of whether a specific fact is injected [44, 8, 26, 22, inter alia], but captures the model's ability to make inferences on it [31, 32]. Second, specificity evaluates whether the predictions of the LM on other contexts are altered as in prior work [8, 26, 22, 23]. Ideally, edits should not impact inferences on examples unrelated to the edit.
Knowledge distillation We are not aware of prior work that uses distillation for knowledge editing. Our use of context distillation is most similar to Askell et al.'s alignment work [1]; however, they use it in a phase roughly analogous to RLHF and use a generic transfer set sampled from the language model training corpus. Our work is also related to prompt injection [5], which examines tasks like distilling a persona-conditioned language model. Unlike their work, we do not train a task-specific model to generate examples for distillation. Instead, we simply prompt existing LMs. Furthermore, while they aim to have a model memorize a particular prompt, we focus on general knowledge updates and inferences based on those. Other work has used distillation techniques for "gisting" to make shorter prompts [28] or to distill reasoning processes [37]. Similar approaches as our continuation sampling have been used for example extrapolation [19] to generate training datasets for fine-tuning.
Efficient parametric knowledge updates Parameter updating methods such as KnowledgeEditor [8] and MEND [26] make use of standard fine-tuning to attempt to localize edits. Another line of work [7, 22, 23] attempts to locate where factual information is stored in transformers and designs edit methods based on these findings. In particular, ROME [22] and MEMIT [23] treat factual knowledge as subject-relation-object tuples, and find that new facts can be inserted into particular early and middle layer MLPs within a GPT-style transformer using specialized update rules. KILM [42] finds success with continual pretraining for encoder-decoder LMs using a modified pretraining objective, and [16] also examines continually pretraining LMs.
Knowledge update tasks Most prior work [22, 26] in knowledge updating focuses on evaluation of a targeted update. Because our goal is to test propagation of knowledge, we mainly focus on two benchmarks from Onoe et al. [32]. Besides this benchmark, recent work [43, 39, 6] also evaluates the LM's success at performing multi-hop inferences with the edited information. Compared to the benchmarks we evaluate on, which are taken from Wikipedia sentences, these benchmarks use sentences generated from knowledge base relations. Another line of work evaluates the ability of LMs to reason about emerging entities [18, 9, 17]. However, such benchmarks do not fit our task setting as they do not provide the information to inject.
Algorithm 1 Knowledge Propagation Through Distillation
Input: An entity  $e$  and its definition sentence  $\mathbf{d}_e$ , a base LM  $M_{\mathrm{Base}}$ , an LM  $M_g$  to generate a transfer set, and a prompt  $\mathbf{p}$  for the transfer set generation.
Output: An updated model  $M_{s}$
1:  $\mathbf{T}_e = (\mathbf{c}_1,\dots ,\mathbf{c}_N)$  where  $\mathbf{c}_i\sim M_g([\mathbf{p};\mathbf{d}_e])$  ▷ Sample  $N$  continuations to form a transfer set  $\mathbf{T}$
2:  $M_{s}\gets M_{\mathrm{Base}}$ $\triangleright$  Create a student LM
3:  $M_{t}\gets M_{\mathrm{s}}$ $\triangleright$  Create a teacher LM
4: for  $\mathbf{c}_i\in \mathbf{T}_e$  do
5:  $\mathcal{D}_t = M_{\mathrm{t}}([d_e;c_i])$  Compute the teacher distribution for each token
6:  $\ell_i\gets \mathbf{Find}(\mathbf{c}_i,e)$  Find the end token index of the entity mention in  $\mathbf{c}_i$
7: for  $k\in \{1,\dots ,K\}$  do  $\triangleright$  Update  $M_{s}$  for  $K$  epochs
8:  $\mathcal{D}_s = M_s([c_i])$  Compute the student distribution for each token
9:  $\mathcal{L} = \frac{1}{|\mathbf{c}_i| - \ell_i}\sum_{j = \ell_i + 1}^{|\mathbf{c}_i|}D_{KL}(\mathcal{D}_{t,|\mathbf{d}_e| + j}||\mathcal{D}_{s,j})$
10:  $M_{s}\gets \nabla \mathcal{L}$  Update  $M_{s}$  w.r.t. avg. KL over the tokens after the entity mention
Our method is illustrated in Figure 1 and described formally in Algorithm 1. It consists of two steps: transfer set generation and distillation on the generated transfer set.
Transfer set generation First, we generate a transfer set corresponding to  $\mathbf{d}_e$ , written as  $\mathbf{T}_e = \{\mathbf{c}_1,\mathbf{c}_2,\dots ,\mathbf{c}_N\}$ . We do this by sampling  $N$  distinct continuations from our generator model  $M_{g}$  with a prompt  $\mathbf{p}$  followed by the entity definition  $\mathbf{d}_{\mathbf{e}}$ ; we will either use GPT-3.5 or the base LM  $M_{\mathrm{Base}} = M_s$  as the generator model  $M_g$ .
Each continuation must contain an identifiable reference to the entity string  $e$ . We describe how we ensure this in Section 5. We use  $\ell_i$  to refer to the fencepost index where this entity string ends in the continuation sentence  $\mathbf{c}_i$ ; for example, in Figure 1,  $\ell_i = 2$  with 1-based indexing to indicate the mention string ChatGPT ends before the second token. Crucially, we only want to distill losses when predicting tokens located at position  $\ell_i$  or later. Tokens before do not condition on the entity name in the student and risk making broad updates to the model, which can impact specificity negatively.
Distillation We initialize an LM  $M_{s}$  from its original pretrained checkpoint, as well as a copy of the LM,  $M_{t}$ , to serve as the teacher model during the distillation process. Then, for each continuation  $\mathbf{c}_i$  in the transfer set, we compute the student model's distributions  $M_{s}(\mathbf{c}_{i})$  (a sequence of  $|\mathbf{c}_i|$  distributions) as well as the teacher model's distributions conditioned on the definition,  $M_{t}(\mathbf{c}_{i}|\mathbf{d}_{e})$ . We compute the KL divergence summed over the tokens after  $\ell$  (line 8). Finally, we perform a gradient update on  $M_{s}$  based on this loss. This is done for  $K$  epochs.
Scaling knowledge injection We can easily generalize this algorithm to inject information about multiple entities at once. We take the union of transfer sets belonging to different entities, shuffle them, and distill on each transfer example as described in line 4-9. We evaluate this setting in Section 7.2.
4 Evaluating Knowledge Propagation
To evaluate our approach on entity knowledge propagation (EKP), we closely follow the setup laid out in Onoe et al. [32]. Here, we describe two datasets and metrics for completeness. The details about the datasets (statistics, examples) can be found in Appendix A.
Data We evaluate on two datasets. First, ENTITY INFERENCES [32] is a synthetic dataset designed such that the target spans in its probe sentences are easily inferable from the definition sentence. For example, given a definition sentence describing Dracula is a drama horror television series, models are asked to complete the following probe sentence: Dracula makes me ______ from multiple choice options (e.g., scared,atheletic,etc).
Second, Entity Cloze By Date (ECBD) [31] consists of cloze-style sentences from Wikipedia that probe for knowledge of specific entities. Examples in ECBD are separated by each entity's
origination date (e.g., when an event occurred). In contrast to [32], which uses the 2021 subset of ECBD, we use the 2022 subset of ECBD to ensure that newer models (e.g. GPT-3.5) do not have knowledge of the probed entities beyond the definition they condition on; see Appendix A.3 for more discussion of the temporal cutoffs for our models and datasets. Each example consists of a cloze-style probe sentence prefix  $\mathbf{x}$  about an entity  $e$  followed by a target span  $\mathbf{y}$ . The definition  $\mathbf{d}_e$  is taken from the first sentence of the entity's Wikipedia page.
Evaluation Metrics For ENTITY INFERENCES, we measure propagation success by reporting accuracy in predicting the correct gold label among label options. We measure specificity by evaluating the model's accuracy at predicting gold spans on similar probe sentences across all other entities.
We evaluate on ECBD by computing per-token perplexity of the continuation given the probe prefix,  $PPL(\mathbf{y} \mid \mathbf{x})$ . This metric is not directly comparable across base LMs which have different tokenizers. To evaluate propagation success, we report the decrease in perplexity from the edit,  $PPL(\mathbf{y} \mid \mathbf{x}; M_{\mathrm{Base}})$  vs.  $PPL(\mathbf{y} \mid \mathbf{x}; M_s)$ . To evaluate an edit's specificity, we randomly sample 40 examples from the "popular" subset of ECBD, ensuring that all 40 probes are about unique entities. We then report the change in perplexity on these sampled examples before and after the edit, using the same metric as above for evaluating on the target sentence.
5 Experimental Setting
Base Models We consider three autoregressive language models: GPT-Neo-1.3B [3], GPT2-XL [33] (1.5B), and LLaMA-2-7B [38]. The former two models have minimal knowledge of the entities in Entity Inferences and ECBD from their pretraining corpora as the entities in these datasets emerged after their pre-training.
Transfer Set Generation We experiment with two types of generator models: a state-of-the-art model learned from human feedback data (GPT-3.5, text-davinci-003), which can generate highly fluent transfer sentences from the definition sentence, and the base model itself, which presents a more realistic scenario in which we do not assume a better LM than the base LM that we are updating. For both models, we use a simple prompt to elicit a continuation of the definition sentence and sample five transfer sentences for each entity. For generation, we use nucleus sampling [15] with  $p = 0.9$ , a temperature of 1.0, and a max length of 40 tokens.
Table 1 summarizes the statistics of transfer sets. Upon manual inspection, we find that GPT-3.5 hallucinates substantially less than smaller models, as reflected in  $\%$  of tokens in the continuations that appeared in the definition sentence. For continuations that do not contain the entity name, we simply pretend the entity name onto the continuation. We also report the number of tokens after  $l$ , i.e., the number of tokens which we compute the distillation loss on. The exact prompt and example continuations can be found in Appendix C.
Mg	# Tokens	% Token in Ed	# Tokens after l
GPT-3.5	40.0	56.4	33.8
GPT2-XL	35.5	34.8	30.1
GPT-Neo	37.2	35.1	31.6
LLaMA-2	32.5	37.4	26.4
Table 1: Statistics for transfer set sentences generated by each generator model.
5.1 Comparison Systems
We compare against two paradigms for knowledge injection: prepending new knowledge in-context at inference time and updating the parameters of LMs. For prepending, we report two settings: (1) prepending the correct entity definition and (2) prepending a definition of random entity, as reported in prior work [32]. Next, we describe knowledge updating methods below.
Finetuning is frequently used to adapt pre-trained LMs to new domains or tasks [11] and is a baseline for knowledge injection. We train  $M_{\mathrm{Base}}$  on  $\mathbf{d}_e$  with standard negative log likelihood loss on the sequence (teacher forcing). We investigate fine-tuning the full model, as well as only the last layer.
We also compare to finetuning with the transfer set. First, we fine-tune  $M_{s}$  on the definition. Then, for each sentence in our transfer set  $\mathbf{T}_e = (\mathbf{c}_1,\dots \mathbf{c}_N)$ , we fine-tune on  $M_{s}(\mathbf{c}_{i}\mid \mathbf{d}_{e})$ , conditioning
GPT-NEO-1.3B	GPT2-XL
Pre-Edit Accuracy (↑)	34.1 Target (Δ)	34.1 Spec. (Δ)	32.9 Target (Δ)	32.9 Spec. (Δ)
Finetuning on de (full)	57.7 (+23.6)	18.3 (-15.9)	62.9 (+30.0)	24.1 (-8.8)
Finetuning on de (last only)	48.8 (+14.7)	16.4 (-17.7)	46.5 (+13.6)	35.4 (+2.5)
Finetuning on de + Te (full)	66.5 (+32.4)	28.8 (-5.3)	59.4 (+26.5)	33.8 (+0.9)
MEND	41.8 (+7.7)	34.4 (+0.3)	-	-
Distillation (Mg = Ms)	61.8 (+27.7)	32.6 (-1.6)	58.2 (+25.3)	31.4 (-1.5)
Distillation (Mg = GPT3.5)	65.9 (+31.8)	32.5 (-1.6)	65.3 (+32.4)	28.7 (-4.2)
Prepend Def.	60.0 (+25.9)	34.1	64.1 (+31.2)	32.9
Prepend Random Def.	27.7 (-6.4)	34.1	26.5 (-6.4)	32.9
Table 2: Results (accuracy) on ENTITY INFERENCES. Non-bolded lines are taken from prior work [32]. Before the edit, accuracy was 34.1 for GPT-Neo and 32.9 for GPT2-XL.
on  $\mathbf{d}_e$  and only updating the model on the tokens after the entity occurrence  $\ell$  in  $\mathbf{c}_i$  to make updates more comparable to our distillation setting. Here, we use the transfer set generated by GPT-3.5.
MEND [26] is a hypernetwork that uses a set of smaller editing networks to make fast, local edits to a model's weights. MEND transforms the gradient obtained from traditional fine-tuning using a low-rank approximation. We train MEND editors for GPT-Neo using the WikiText-103 dataset, which utilizes generated text as altered output following the configuration used in the original paper.
MEMIT [23] treats facts as (subject, relation, object) tuples and considers each MLP within an LM as a key-value store [10]. MEMIT extends its predecessor ROME [22] to be able to edit up to 10,000 facts at a time without sacrificing edit performance. Both methods use rank-one modifications to the MLP weights within a pre-chosen transformer layer (in the case of MEMIT, a set of consecutive pre-chosen layers) to edit the factual representations there.
We format the data for MEMIT as follows: For a given definition sentence  $\mathbf{d}_e$ , the subject is the name of the entity  $e$ , the relation is the part of the sentence before the masked span, and the object is the part of the sentence after the masked span, including the gold span. For details about how masked spans are defined, refer to Onoe et al. [32].
Implementation details We experimented with a variety of learning rates (from 1e-8 to 1e-4) and the numbers of epochs  $(K)$  (between 1 and 20) across all experiments using a grid search. We focus on balancing results between performance and specificity; neither are prioritized if it significantly harms the other. The specific values used can be found in Appendix B.1.
6.1 Entity Inferences
We first conduct a smaller scale study on the easier benchmark, ENTITY INFERENCES, where learning about the definition should allow us to guess the target tokens by design. Table 2 reports the results. Our distillation approach shows promising performance in two base models we test. We find that transfer sets generated from GPT-3.5 show substantially better results than transfer sets generated from the base model itself in both datasets. This sometimes even outperforms definition prepending, which might be due to GPT3.5 introducing information about the entity beyond what can be inferred from the definition sentence. Fine-tuning on the definition and transfer set using GPT-Neo does outperform distillation, at the cost of specificity. For GPT2-XL, distillation only outperforms fine-tuning on the definition sentence when using GPT3.5 as a generator model, but still shows a substantial accuracy gain using its own generated sentences (24.3%). The drop in specificity (1.6-4.2%) is substantially less severe than fine-tuning on the definition sentence. These results indicate that context distillation teaches models to make simple inferences based on injected knowledge without significantly harming the model's distribution on unrelated concepts.
<sup>1</sup>MEND is extended by a method called SERAC [27]. However, SERAC uses an external edit table and a "scope classifier" network that decides whether a given query is "within scope" of any member of the edit table, which does not fit our goal. We deliberately aim to test the queries out of scope of the definitions.