SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Thomas Hartvigsen
University of Virginia, MIT hartvigsen@virginia.edu
Swami Sankaranarayanan
Sony AI swami.sankaranarayanan@sony.com
Hamid Palangi
Microsoft Research hpalangi@microsoft.com
Yoon Kim
MIT yoonkim@mit.edu
Marzyeh Ghassemi
MIT  
mghassem@mit.edu
Deployed language models decay over time due to shifting inputs, changing user needs, or emergent world-knowledge gaps. When such problems are identified, we want to make targeted edits while avoiding expensive retraining. However, current model editors, which modify such behaviors of pre-trained models, degrade model performance quickly across multiple, sequential edits. We propose GRACE, a lifelong model editing method, which implements spot-fixes on streaming errors of a deployed model, ensuring minimal impact on unrelated inputs. GRACE writes new mappings into a pre-trained model's latent space, creating a discrete, local codebook of edits without altering model weights. This is the first method enabling thousands of sequential edits using only streaming errors. Our experiments on T5, BERT, and GPT models show GRACE's state-of-the-art performance in making and retaining edits, while generalizing to unseen inputs. Our code is available at github.com/thartvigsen/grace.
Large scale pre-trained neural networks are the state-of-the-art for many hard machine learning problems, especially in natural language processing [5, 33] and computer vision [9, 34]. But when deployed, they still make unpredictable errors [2, 43]. For example, Large Language Models (LLMs) notoriously hallucinate [17], perpetuate bias [11], and factually decay [8]. Many such errors will arise sequentially in deployment, some of which must be addressed quickly without waiting until new training data is collected [16, 21]. For example, when an LLM generates hate speech or crucial knowledge about the world changes, its behavior must be modified immediately to protect users.
While retraining or finetuning can edit a model's predictions, doing this frequently is often too computationally expensive. LLaMA [40], for instance, was trained for 21 days on 2,048 A100 GPUs, costing over $2.4M and emitting over 1,000 tons of CO $_2$ . Even repeatedly retraining or finetuning smaller models [3] quickly costs too much for most practitioners. To enable cheap, targeted updates to big, pre-trained models, we study lifelong model editing. Here, we continually make targeted edits sequentially throughout a model's deployment. Success means correcting a model's predictions on a stream of edits without decaying its performance on unrelated and previously-edited inputs.
Two possible approaches to lifelong model editing are continual learning and conventional model editing. While continual learning methods can update models sequentially, when used for sequential editing, they suffer from overfitting [25] and quickly forget previous edits and pre-training data [16, 22]. Alternatively, existing static model editors could also be run sequentially. But the editors that
Figure 1: Overview of lifelong model editing with GRACE. a) Models make important errors that must be corrected. b) GRACE makes edits by learning, caching, and selectively retrieving new transformations between layers. c) Edits appear sporadically and require quick fixes, so GRACE codebooks are curated over long sequences of edits.
rely on regularized finetuning [8, 28, 29, 37] succumb to overfitting [4], while hypernetwork methods require pre-training on hard-to-access data [30, 31]. Further, these existing approaches require large sets of representative training edits, access to the model's pre-training data, or semantically-equivalent inputs. Accessing these data is often unrealistic or impractical, and existing editor performance severely decays the pre-trained model after only a few sequential edits [12, 16].
We study lifelong editing using only singular inputs to edit models, where models are edited immediately upon arrival. This setup extends beyond recent works, which use many inputs for editing, including collected sets of semantically-equivalent inputs, training edits, and the model's pre-training data. Editing in our realistic, more cost-effective setup is hard because each edit must fix the flagged error, while generalizing to similar future inputs without impacting unrelated model behavior.
To address the challenging lifelong model editing setup, we propose General Retrieval Adaptors for Continual Editing, or GRACE. While we study transformers for natural language processing, GRACE is broadly applicable. As illustrated in Figure 1, GRACE edits a model by adding an Adaptor to a chosen layer, while never changing its weights. This Adaptor then modifies layer-to-layer transformations for select inputs. By caching embeddings for input errors and learning values that decode into desired model outputs, GRACE serves as a codebook in which edits are stored, enabling longer sequences of edits than prior works. To encourage generalization, we leverage the models' semantic similarity in its latent space by introducing  $\epsilon$ -balls around cached edits. GRACE then only applies for inputs near existing keys, differing from more-traditional Adaptors [15]. By managing the  $\epsilon$ -balls' size over time, GRACE makes immediate edits, remembers previous edits, and leaves correct model behaviors intact, making it parameter efficient. Further, since GRACE codebooks leave model weights unaltered and are fully model-agnostic, they also point towards plug-and-play, cost-effective model editing, especially to make crucial spot-fixes between bigger retraining efforts.
Our contributions are as follows:
1. We establish key metrics and comparisons for lifelong model editing, an important but understudied and challenging problem setting. We introduce two new public benchmarks for lifelong model editing: mitigating LLM hallucination [27] and addressing label shifts [6].
2. We develop GRACE, a method for lifelong model editing that avoids expensive retraining or finetuning. Unlike other works, GRACE requires no inputs beyond singular edits. GRACE works with autoregressive and non-autoregressive models alike and achieves parameter efficiency without touching the model's weights.
3. Our experiments show that GRACE outperforms seven alternatives when sequentially editing T5, BERT, and GPT models for question answering, document classification, and language generation. We show that GRACE edits can generalize to new inputs without memorizing while incurring only a small, one-time cost to inference time over long sequences of edits.
2 Methods: Lifelong Model Editing with GRACE
2.1 Problem Formulation
The Lifelong Model Editing task is to edit the same model hundreds to thousands of times in a row without forgetting upstream performance or fixes for previous edits. Let  $f_{0}$  denote a model with
frozen parameters that was pre-trained on dataset  $\mathcal{D}_{\mathrm{train}}$ . For clarity, we drop the subscript where possible. Assume that  $f$  consists of  $L$  layers, where  $f^l (\cdot)$  computes the hidden state at layer  $l$ . In this work, we assume  $f$  is a transformer architecture for natural language, though our principles are general. We then deploy  $f$  on a stream of inputs  $[x_0,x_1,\ldots ]$ , where  $x_{t}\in \mathcal{D}_{\mathrm{edit}}$ , which contains samples observed during deployment. We then monitor the model's predictions  $\hat{y}_t = f(x_t)$  over the stream. Note that the prediction tasks during training and deployment are the same. Over time, the model makes errors such that  $\hat{y}_t\neq y_t$ , where  $y_{t}$  is the true label. To continue safely deploying  $f$ , we aim to edit  $f$  such that  $f(x_{t}) = y_{t}$ . After each edit, the updated  $f$  should 1) be successfully edited such that  $f(x_{t}) = y_{t}$ , 2) retain accuracy on prior edits  $x_{< t}$ , and 3) retain its behavior on its training data:  $f(x_{i}) = f_{0}(x_{i})\forall x_{i}\in \mathcal{D}_{\mathrm{train}}$ . In contrast to prior works [8, 28, 30, 31, 37], we assume access to only input edits  $x_{t}$  and corrected labels  $y_{t}$ , since  $\mathcal{D}_{\mathrm{train}}$  is often proprietary or prohibitively large, and collecting training edits or semantically-equivalent inputs is often expensive in practice.
2.2 GRACE: General Retrieval Adaptors for Continual Editing
We propose GRACE, a method for sequentially editing a pre-trained model's behavior without altering its weights, as illustrated in Figure 1. GRACE works by wrapping a chosen layer of any pre-trained model architecture with an Adaptor. A GRACE Adaptor at model  $f$ 's layer  $l$  contains two components: (1) a codebook  $\mathcal{C}$  and (2) a deferral mechanism to decide whether to use  $\mathcal{C}$  for a given input.
GRACE codebook. A GRACE Adaptor at layer  $l$  maintains a discrete codebook, adding and updating elements over time to edit a model's predictions. The codebook contains three components:
- Keys (K): Set of keys, where each key is a cached activation  $h^{l - 1}$  predicted by layer  $l - 1$ .
- Values (V): Set of values that are randomly initialized and are updated using the model's finetuning loss for edits. Each key maps to a single, corresponding value.
- Deferral radii  $(\mathcal{E})$ : Each key has a deferral radius  $\epsilon$ , which serves as a threshold for similarity matching. The deferral mechanism uses this radius as shown in Algorithm 1. GRACE is activated at layer  $l$  only if the deferral constraint is satisfied. New entries have a default value  $\epsilon_{\mathrm{init}}$ , which is a hyperparameter.
Deferral mechanism. Before editing, GRACE layers are empty. As editing progresses, GRACE adds keys and adapts values and  $\epsilon$  entries. Conceptually, inference at layer  $l$  with GRACE entails a deferral decision, computing  $h^l$  using a similarity search over GRACE's keys:
$$
h ^ {l} = \left\{ \begin{array}{l l} \operatorname {G R A C E} \left(h ^ {l - 1}\right) & \text {i f} \min  _ {i} \left(d \left(h ^ {l - 1}, \mathbb {K} _ {i}\right)\right) <   \epsilon_ {i *}, \text {w h e r e} i _ {*} = \operatorname {a r g m i n} _ {i} \left(d \left(h ^ {l - 1}\right), \mathbb {K} _ {i}\right), \\ f ^ {l} \left(h ^ {l - 1}\right) & \text {o t h e r w i s e}, \end{array} \right. \tag {1}
$$
where  $f^{l}(h^{l - 1})$  denotes the unedited model's activation of the  $l$ -th layer.  $h^{l - 1}$  can be seen as a query to the codebook, and  $\mathrm{GRACE}(h^{l - 1})$  retrieves the value associated with its closest key.  $\epsilon_{i}^{l}$  and  $\mathbb{K}_i^l$  are the influence radius and key  $i$  in layer  $l$ , respectively, and  $d(\cdot)$  is a distance function. We follow related work [41] and use Euclidean distance for  $d(\cdot)$  in our experiments—changing this is trivial. Through explicit similarity search, we use the fact that large models encode semantic similarity with respect to their tasks in their latent spaces. This lets GRACE edits to generalize to similar inputs in the future. If a new input is unlike any cached keys, GRACE simply defers to  $f$ 's pretrained weights. This way, GRACE layers limit interference with  $\mathcal{D}_{\mathrm{train}}$  by leaving the model weightsunaltered, which helps when input distributions shift [45].
Codebook maintenance. To make an edit, a GRACE layer can perform one of two operations. Each step is described in Algorithm 1. First, if the codebook is empty or the input embedding  $h^{l-1}$  falls outside the deferral radius of all existing keys according to distance function  $d(\cdot)$ , then a new codebook entry is created and added:  $\{(h^{l-1}, v, \epsilon_{init}, y)\}$ . Thus if  $x_t$  were passed into  $f$  again,  $h^{l-1}$  would activate the codebook and value  $v$  would be passed to layer  $l+1$ . Training  $v$  is detailed below.
Sometimes, a query  $h^{l-1}$  will be close enough to an existing key that adding a new entry would cause their  $\epsilon$ -balls to overlap. To avoid this, we compare the edit label  $y$  to the model's prediction for the nearest key and distinguish two cases: 1) If the overlapping key's label is the same as the edit's label, Expand that key's  $\epsilon$  to encompass the query. 2) If the overlapping key's label is different from the edit's label, Split these keys by first decreasing the influence radius of the overlapping key, then adding a new codebook entry where the new key is simply the query  $h^{l-1}$ . We set both keys'  $\epsilon$  values to be half their distance apart.
As edits stream over long deployments, by continuously adding and updating the GRACE codebook, layer  $l$ 's latent space is partitioned according to which inputs needed edits. When not performing edits, these codebook maintenance operations are bypassed, and keys are entirely frozen. GRACE thus introduces a new model editing paradigm in which edits can be made sequentially, similar edits are encouraged to be edited similarly, and the ultimate influence of new edits can be controlled and monitored explicitly.  $\epsilon_{\mathrm{init}}$  is the sole parameter in GRACE, which sets the initial  $\epsilon$  value for new codebook entries. Intuitively, using a larger  $\epsilon_{\mathrm{init}}$  will create edits with more influence, making edits more general, but increasing the interference with unrelated inputs. In practice,  $\epsilon_{\mathrm{init}}$  could be tuned using either exogenous data or GRACE codebooks can periodically be refreshed.
Training GRACE Values. When making an edit with GRACE, either a new key-value pair is learned or an existing key-value pair is updated. To ensure that newly-learned values correct the model's behavior, we train them directly using backpropagation through the finetuning loss on the model's prediction given the edit. The learned value  $v$  then replaces  $h^{\bar{r}}$  for the rest of the forward pass. In our experiments, we train values using 100 gradient descent steps to train the values and ensure the model's behavior is updated.
Algorithm 1: Update Codebook at layer  $l$ .
Input:  $\mathcal{C} = \{(\mathbb{K}_i,\mathbb{V}_i,\epsilon_i)\}_{i = 0}^{C - 1}$  , codebook   
Input:  $f(\cdot)$  , model   
Input:  $y_{t}$  , desired label   
Input:  $x_{t}$  , edit input for which  $f(x_{t})\neq y_{t}$    
Input:  $\epsilon_{\mathrm{init}}$  , initial  $\epsilon$    
Input:  $\mathrm{d}(\cdot)$  , distance function   
Output:  $\mathcal{C}$  , updated codebook   
 $C = ||\mathcal{C}||$ $\hat{y},h^{l - 1} = f^{L}(x_{t}),f^{l - 1}(x_{t})$ $\mathrm{d}_{\mathrm{min}},i = \mathrm{min}_i(d(h^{l - 1},\mathbb{K}_i))$    
If  $\mathrm{d}_{\mathrm{min}} > \epsilon_{i} + \epsilon_{\mathrm{init}}$  or  $C = 0$  .. #  $h^{l - 1}$  far from existing entries or empty C  $v_{\mathrm{new}} =$  finetune on  $P_{f}(y|v_{\mathrm{init}})$ $\mathcal{C}_C = (h^{l - 1},v_{\mathrm{new}},\epsilon_{\mathrm{init}})\# Add$  entry   
Else: #  $h^{l - 1}$  near existing entries If  $f^{L}(k_{i}) = y$  # Same label  $\rightarrow$  Expand  $\mathcal{C}_i\coloneqq (k_i,v_i,\epsilon_i + \epsilon_{\mathrm{init}})$  Else: # Different label  $\rightarrow$  Split  $\mathcal{C}_i = (k_i,v_i,\mathrm{d}_{\mathrm{min}} / 2)$  # Update entry i  $v_{\mathrm{new}} =$  finetune on  $P_{f}(y|v_{\mathrm{init}})$ $\mathcal{C}_C = (h^{l - 1},v_{\mathrm{new}},\mathrm{d}_{\mathrm{min}} / 2)\# Add$  entry   
return:  $\mathcal{C}$
GRACE layers with sequential inputs. For models with different representations per input token, like transformers, we must choose 1) which token should be GRACE's input query, and 2) which token to replace with a retrieved value in the subsequent layer. In practice, we find that broadcasting the value to each token is reasonable for tasks like classification, since values gain strong control over the model's behavior and makes them easy to learn. For autoregressive models however, picking the right tokens for the query and value is more nuanced. We opt for replacing only the final token of an input prompt, which we verify experimentally, since compressing future generated text into tokens is an interesting and burgeoning direction in itself [32]. Upon choosing these tokens, GRACE naturally applies to all popular transformer models.
2.3 Illustrative Example
To garner intuition about GRACE, we provide an illustrative experiment on synthetic data. As shown in Figure 2, we sample 100 instances from two 2D distributions corresponding to classes. We then train a three-layer binary classifier with two 100-dimensional hidden layers and ReLU activations. Next, we introduce edits with flipped labels, simulating local label shift at test time. Using a single
(b) Upstream Training Data + Edits
(a) Upstream Training Data
Figure 2: Illustrative example of GRACE. We train a model on separable data in (a), then introduce locally-flipped labels at test time in (b). In (c), the original model unsurprisingly misclassifies these label-flipped instances. In (d), GRACE fixes these labels without impacting other inputs.
(c) Predictions before editing
(d) Predictions after editing
SETTING	MODEL	TEST RETENTION DATA	EDITING DATA
Dataset	N	Pre-edit	Dataset	N	Pre-edit
QA	T5 (60m)	NQ [20]	1000	.72 F1	zsRE [24]	1000	.31 F1
Clf.	BERT (120m)	SCOTUS1982-1991 [6]	914	.99 Acc	SCOTUS1992-2009 [6]	931	.55 Acc
Halluc.	GPT-2 (1.5B)	OpenWebText [10]	1000	15.98 PPL	SelfCheckGPT	1392	132.7 PPL
Table 1: Dataset statistics for main results. Test retention is the testing set of each model's training data.  $N$  is the number of samples. Pre-edit is the unedited model's performance on each dataset.
key in layer two, GRACE can successfully correct these errors while barely influencing other inputs. Meanwhile, finetuning on these errors alone will clearly break the model.
3.1 Experimental Setup
We evaluate GRACE's capacity to edit models hundreds to thousands of times sequentially. This is in contrast to prior works, which largely evaluate using single edits. Further, much of the model editing literature relies on synthetic edits, often generated by flipping random labels. Beyond unrealistically assessing performance, it is well-established that learning random versus natural labels is vastly different [49]. Instead, we correct authentic mistakes made by models, each time an error is made.
Baselines. We compare against continual learning and model editing methods. First, we continually finetune (FT) [25] on streaming errors. To reduce overfitting, we also compare with Elastic Weight Consolidation (EWC) [19] and a simple form of experience replay [36] where we periodically retrain the model (Retrain) on all previous edits. Second, we compare against model editors MEND [30] and Defer, inspired by SERAC [31] but altered for our setup. We also compare against ROME [28] in our language modeling experiment. Finally, we ablate GRACE by replacing our discrete search with a Memory network containing memory module that is indexed by a soft attention mechanism. Details for all comparisons are in Appendix C.
Datasets and Pre-trained Models. We evaluate GRACE on three sequential editing tasks with corresponding pre-trained models, as shown in Table 1. 1) We edit a 60-million parameter T5 model [35] trained for context-free question-answering, as is used in [30]. We extract potential edits from the validation set of zsRE [24], following the editing literature [16, 30]. This model achieves F1 of .72 on NQ and .31 on zsRE prior to editing. Further details are in Appendix B.1. 2) We edit a 110-million BERT classifier trained for a new editing task with label shift using the SCOTUS dataset from Fairlex [6]. The prediction task is to categorize U.S. Supreme Court documents over multiple decades into 11 topics. Over time, categorization rules change, so label distributions shift. We train a BERT classifier on the 7.4k training cases from 1946-1982, then make edits on 931 cases from 1991-2009. We also introduce additional, realistic label shifts, as discussed in Appendix B.2. This model achieves Accuracy of 0.99 on the training set and .55 on the edit set. 3) We introduce a new editing task by correcting a GPT language models' Hallucination. In [27], authors prompt GPT-3 to generate 238 wikipedia-style biographies using subjects from WikiBio. They then annotate the factual accuracy of each sentence, recording which are hallucinations. We propose editing inaccurate sentences by replacing them with corresponding sentences in the true wikipedia entries. We include edits for all 238 biographies, creating 1392 sequential edits and 592 already-accurate outputs. We then edit GPT2-XL, which has 1.5B parameters. However, GPT2-XL was not trained on this task, so has high perplexity (PPL) on all sentences. Therefore, we finetune GPT2-XL on these data mixed with sentences from OpenWebText [10], a public version of GPT2's training data. Our final GPT2-XL model has PPL of 15.98 on OpenWebText (comparable to the original), 8.7 on already-accurate outputs, and 132.7 on intended edits, indicating a need for editing and room for improvement. Further details are available in Appendix B.3.
Metrics. To compare each method, we measure three main metrics that align with prior work [16, 28].
1. Edit Success (ES): We check whether an edit has been successful using ES:  $m(y, \hat{y})$ , where  $m(\cdot)$  is a task-specific measure of accuracy: standard F1 for question answering, Accuracy for classification and Perplexity (PPL) for generation.
Method	zsRE (T5; F1 ↑)	SCOTUS (BERT; Acc ↑)	Hallucination (GPT2-XL; PPL ↓)
TRR	ERR	Avg.	#E	TRR	ERR	Avg.	#E	TRR	ERR	ARR	#E	time (s)
FT [25]	.56	.82	.69	1000	.52	.52	.52	415	1449.3	28.14	107.76	1392	.26 (.07)
FT+EWC [19]	.51	.82	.66	1000	.67	.50	.58	408	1485.7	29.24	109.59	1392	.29 (.06)
FT+Retrain [36]	.27	.99	.63	1000	.67	.83	.75	403	2394.3	35.34	195.82	1392	23.4 (13.2)
MEND [30]	.25	.27	.26	1000	.19	.27	.23	672	1369.8	1754.9	2902.5	1392	.63 (.10)
Defer [31]	.72	.31	.52	1000	.33	.41	.37	506	8183.7	133.3	10.04	1392	.07 (.02)
ROME [28]	—	—	—	—	—	—	—	—	30.28	103.82	14.02	1392	.64 (.28)
Memory	.25	.27	.26	1000	.21	.20	.21	780	25.47	79.30	10.07	1392	.11 (.02)
GRACE	.69	.96	.82	1000	.81	.82	.82	381	15.84	7.14	10.00	1392	.13 (.02)
137 keys (7.30 edits/key)	252 keys (1.51 edits/key)	1341 keys (1.04 edits/key)
Table 2: Comparison of GRACE to existing methods. Metrics shown are computed after all sequential edits. For Hallucination, we also compute perplexity retention on already-accurate sentences (ARR) and the average time per edit. Parentheses in time denote standard deviation. We also count the number of keys GRACE used. GRACE achieves the best TRR/ERR balance while making edits faster than most comparisons and using small codebooks for zsRE and SCOTUS. A large codebook is required for Hallucination, since each edit label is unique.
2. Test Retention Rate (TRR): We check how well an edited model retains its performance on its original testing data using TRR:  $\frac{1}{n}\sum_{i=1}^{N}m(f(x_i),y_i)$ , where  $(x_i,y_i) \in \mathcal{D}_{\mathrm{test}}$ ,  $f$ ’s original test set. For T5,  $\mathcal{D}_{\mathrm{test}}$  is 1k random samples from NQ, for our BERT model,  $\mathcal{D}_{\mathrm{test}}$  is court documents from 1982-1991, and for GPT2-XL,  $\mathcal{D}_{\mathrm{test}}$  is WebText, for which we use the first 1k sentences of OpenWebText [10].
3. Edit Retention Rate (ERR): We check how well an edited model retains previous edits through ERR:  $\frac{1}{n}\sum_{i = 1}^{N}m(f(x_i),y_i)$  where  $(x_{i},y_{i})\in \mathcal{D}_{\mathrm{eds}}$
We also track the number of edits (#E), which may vary by editor as different editors will lead to different mistakes in practice. Further implementation details are available in Appendix A.
3.2 Results
3.2.1 Comparisons to existing methods
We first find that GRACE outperforms existing methods after long sequences of edits, as shown in Table 2. We focus here on TRR and ERR, which have a trade-off as each can be achieved in isolation, though all metrics are reported in Appendix G. On zsRE and SCOTUS, we focus on TRR, ERR, and their average. ROME is incomparable on these dataset as it is only proposed for GPT models. On Hallucination, we include Accurate Retention Rate (ARR), which is the edited model's perplexity on sentences on which it was already accurate.