SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Peter Hase $^{1,2}$  Mohit Bansal $^{2}$  Been Kim $^{1}$  Asma Ghandeharioun $^{1}$
$^{1}$ Google Research  $^{2}$ UNC Chapel Hill
{peter,mbansal}@cs.unc.edu
{beenkim, aghandeharioun}@google.com
Language models learn a great quantity of factual information during pretraining, and recent work localizes this information to specific model weights like mid-layer MLP weights [21]. In this paper, we find that we can change how a fact is stored in a model by editing weights that are in a different location than where existing methods suggest that the fact is stored. This is surprising because we would expect that localizing facts to specific model parameters would tell us where to manipulate knowledge in models, and this assumption has motivated past work on model editing methods. Specifically, we show that localization conclusions from representation denoising (also known as Causal Tracing) do not provide any insight into which model MLP layer would be best to edit in order to override an existing stored fact with a new one. This finding raises questions about how past work relies on Causal Tracing to select which model layers to edit [21, 22]. Next, we consider several variants of the editing problem, including erasing and amplifying facts. For one of our editing problems, editing performance does relate to localization results from representation denoising, but we find that which layer we edit is a far better predictor of performance. Our results suggest, counterintuitively, that better mechanistic understanding of how pretrained language models work may not always translate to insights about how to best change their behavior.<sup>1</sup>
Language models learn a variety of facts about the world during pretraining that can be elicited via natural language prompts [28]. Recent work explores how these facts are stored in model weights and expressed in response to particular prompts, suggesting that MLP weights act as key-value memories that support factual association [12, 21, 13]. Besides improving our scientific understanding of pretrained language models, this kind of investigative work may enable the design of better model editing methods for injecting new facts into model weights, and indeed it has been used to motivate the ROME and MEMIT model-editing methods [21, 22]. These recent methods set a new state of the art for weight edits that successfully rewrite stored facts in language models. Model editing methods could be broadly useful for correcting factual errors in pretrained models, avoiding morally undesirable outputs, and updating models with changing knowledge over time.
The connection between localization (identifying components of a model responsible for a certain behavior) and editing (changing model components in order to change model behavior) is predicated on the reasonable assumption that one should go about editing a model by first localizing a behavior to a specific component and then choosing to edit that particular component. In the case of ROME and
$^{1}$ Code for all experiments is available at https://github.com/google/belief-localization
MEMIT, localization is done via Causal Tracing, which measures the information content of hidden representations, and editing is done by treating MLP weights as linear associative memories and injecting new key-value memories into the weights. Meng et al. [21, 22] choose to edit early MLP layer(s) based on results from Causal Tracing showing the largest causal effects on average in early layers.
Surprisingly, the assumption that one should change the knowledge in a model by editing the weights where it is stored turns out to be false. In fact, localization results from Causal Tracing are statistically uncorrelated with the success of an edit injecting a new fact into MLP weights. Using the CounterFact dataset from Meng et al. [21] with a GPT-J model [35], we show that (1) not only is a substantial fraction of factual knowledge stored outside of the range of layers edited by ROME/MEMIT (see Fig. 1), (2) the correlation between Causal Tracing results and edit success is near zero (for several editing methods including ROME, MEMIT, and Adam-based finetuning). We note that this is surprising largely because ROME and MEMIT do work well for editing facts, in spite of Causal Tracing often suggesting knowledge is stored elsewhere than early-to-mid-layer MLP weights.
In the face of this result, we attempt to recover the connection between tracing-based localization and editing by introducing four variants of the default model editing problem. Each variant differs in terms of the input, target, or objective used in the editing problem. One variant we introduce, called Fact Forcing, is designed to match Causal Tracing along these three factors. Specifically, Fact Forcing uses a noised input and involves maximizing the probability of the correct target output, just like Causal Tracing. We find that tracing results are related to edit success for Fact Forcing. However, even for this variant, it is still better to ignore the tracing results and always choose an early-to-mid-layer MLP weight for editing. We conclude that, although Causal Tracing is a reasonable localization method that has yielded insight into how models store factual information, this insight does not actually indicate which model layers we should edit in order to manipulate what facts are stored in language models.
Figure 1: We visualize where 652 facts known by GPT-J are stored within the model, as localized by Causal Tracing. Model editing methods like ROME and MEMIT can successfully change knowledge in LMs by editing layers 4-9. But many facts appear to be stored outside of this range, e.g. at layers 1-3 and 16-20. What about these facts?
To summarize, our conclusions are as follows:
1. We find that model edit success is essentially unrelated to where factual information is stored in models, as measured by Causal Tracing. Robustness experiments generalize this result across causal localization methods, editing methods, editing metrics, models, and datasets.
2. To reconnect localization with editing performance, we introduce four variants of a standard model editing problem, including Tracing Reversal, Fact Erasure, Fact Amplification, and Fact Forcing.
3. Edit success and tracing effects correlate best in the Fact Forcing setting. However, tracing effects explain only a small fraction of the variance in editing performance, while the choice of edit layer is a much more important factor. This suggests that, surprisingly, localization insights from Causal Tracing are not useful for choosing which model layer to edit.
Localization. A long line of work aims to interpret what certain hidden representations represent, or, in the reverse direction, to understand how a given concept is represented in a model. Both of these efforts aim to localize behaviors to specific model components. We group these methods based on the kinds of model components they consider (e.g. layers, neurons, etc.).
Many works focus on individual layers or weight matrices [37, 31, 9, 33, 11]. In this paper, we adopt the layer-wise localization method from Meng et al. [21] known as Causal Tracing, which estimates the information content of a set of representations via a denoising operation. We specifically focus on MLP layers given evidence of their role in factual association [12, 21, 13].
Related to analysis at the layer level, other work aims to localize concepts to directions in a latent space, dating back to work interpreting directions in word vector space [23, 17, 41, 14, 38, 5]. One might also place "key-value memory" theories of weight matrices in this category since a key vector represents a direction in the latent space [1, 32, 12, 21].
Neurons, meanwhile, are the most common focus of localization analysis. Past work explores the functions of groups of neurons and subnetworks [27, 6, 10, 4] or simply individual neurons [29, 40, 18, 2, 34, 26, 8, 19, 3, 16, 7, 36].
Relating Localization to Editing. Many works on localization validate the quality of their conclusions by editing neuron activations or layer weights corresponding to a particular concept, then checking that the network behavior changes appropriately. For example, Dai et al. [8] check that their "knowledge neurons" have localized a specific fact by amplifying or suppressing the expression of that fact via adjusting the corresponding neuron activations. Altogether, we find many localization analyses are validated by editing models in suggested locations [29, 18, 2, 26, 34, 8, 19, 7, 36, 4] or directions in the latent space [23, 1, 32, 21].
Changing model behavior by editing components suggested by localization seems like a reasonable validation step. However, in isolation, it paints an incomplete picture that has led to misleading interpretations about the connections between localization and editing. Such experiments alone do not show whether editing that specific component is (1) successful in proportion to the strength of the localization, (2) necessary to achieve the desired behavior, or (3) the best option for editing. In particular, these experiments do not show whether the same change in behavior can be achieved elsewhere in the network. Meng et al. [21] consider this question by measuring editing success across layers, averaged across data, then comparing the results with Causal Tracing conclusions also averaged across data. However, as we show, more fine-grained analysis at the datapoint level reveals the unexpected result that tracing results are unrelated to edit success. We are not aware of any work that primarily investigates the connection between localization and editing or that demonstrates better model editing at locations elsewhere in the network than those suggested by localization analysis.
3 Notation and Background
3.1 Data Notation
Following Meng et al. [21], we consider facts of the form  $(s,r,o)$ , where  $s$  represents a subject entity (e.g. Paris),  $r$  a binary relation (e.g. is located in), and  $o$  an object (e.g. France) for which the tuple  $(s,r,o)$  represents a factual assertion about the world. In the CounterFact dataset [21], each datapoint is a prompt  $P$  for some fact  $(s,r,o)$ . So,  $P$  might be "Paris is located in" or "Paris is situated in," to be completed by the object  $o$  to form a true statement. In an abuse of notation, we will often use  $s$  and  $r$  to refer to textual representations of a subject and relation, for instance by writing a model's conditional probability as  $p_{\theta}(\cdot | s,r)$  instead of  $p_{\theta}(\cdot | P)$ . We do so in order to more easily indicate when an input is provided where the subject or relation has been manipulated (described next).
We make use of a few variations of the data for the fact  $(s,r,o)$ . The additional variables include:
1.  $s^*$  is a "neighboring" entity to the subject  $s$  (similar to  $s$ ) for which  $(s^*, r, o)$  is a true fact like  $(s, r, o)$ . In CounterFact, "Marseille" is a neighboring entity to "Paris."
2.  $r^*$  is a paraphrase of the relation  $r$ , such as "is situated in" for "is located in."
3.  $s_{noise}$  is a noised representation of the subject  $s$ . We add Gaussian noise to the token embeddings of  $s$ , following Meng et al. [21].
4.  $o_{false}$  is an object that incorrectly completes the tuple  $(s, r, \cdot)$ . CounterFact contains an  $o_{false}$  for each datapoint, intended to be the new model output when evaluating model editing methods.
5.  $o_{true}$ , for clarity, is the object that correctly completes the fact  $(s, r, \cdot)$ , from CounterFact.
3.2 Causal Tracing
We give a brief description of Causal Tracing here and refer readers to Meng et al. [21] for more information (see Fig. 2 for an example visualization). Causal Tracing is a method for localizing information in the forward pass of an autoregressive Transformer to specific hidden representations. For a model with  $L$  layers, the input is a prompt containing  $T$  tokens (including a subject  $s$  and
relation  $r$ ). Given this input, the forward pass produces  $T \times L$  layer outputs (one representation per  $T$  tokens and  $L$  layers). The algorithm aims to estimate the amount of information about the fact  $(s, r, o_{\text{true}})$  that is contained in each of these representations. We denote the representation at token  $t$  and layer  $\ell$  as  $v_{(t,\ell)}$ .
The amount of factual information in  $v_{(t,\ell)}$  is estimated by copying this representation into a different forward pass obtained from using a noised subject in the input:
$$
\text {T r a c i n g E f f e c t} = p _ {\theta} \left(o _ {\text {t r u e}} \mid s _ {\text {n o i s e}}, r, v _ {(t, \ell)}\right) - p _ {\theta} \left(o _ {\text {t r u e}} \mid s _ {\text {n o i s e}}, r\right)
$$
where  $s_{noise}$  indicates that we add Gaussian noise with  $\sigma = 0.094$  to the token embeddings of  $s$  following Meng et al. [21], and  $v_{(t,\ell)}$  is the representation at token  $t$  and layer  $\ell$  in the forward pass on the original prompt  $P = (s,r)$ . The probability  $p_{\theta}(o_{true}|s_{noise},r,v_{(t,\ell)})$  is computed by (1) running the model forward pass on the noised prompt  $P^{*} = (s_{noise},r)$  until layer  $\ell$ , (2) overwriting the existing representation at token  $t$  and layer  $\ell$  with the representation  $v_{(t,\ell)}$ , then (3) computing the remaining  $L - \ell$  layers as normal using this adjusted set of  $T$  representations as input (adjusted at token index  $t$ ). Thus, Causal Tracing estimates the information content of a representation in terms of its effect on the probability of the true target. The results from Causal Tracing show where the representations containing information about the true target are in the model forward pass.
In practice, a set of representations from multiple adjacent layers is copied from the clean forward pass rather than a single layer's rep
presentation (for instance, ten layers in Fig. 2). The size of this set is referred to as the tracing window size. A window size of, e.g., three implies that the tracing effect at layer  $\ell$  estimates the amount of information contained in the three representations  $v_{(t,\ell -1)}$ ,  $v_{(t,\ell)}$ , and  $v_{(t,\ell +1)}$ . See Appendix Figs. 10 and 11 for analysis of the parameter's effect. In this paper, we use a tracing window size of 5 by default, and we apply Causal Tracing exclusively to MLP layers, given evidence of their role in factual association [12, 21].
MLP Tracing Effect (Window Size: 10)
Figure 2: Visualizing Causal Tracing results over MLP layers with window size 10. Tokens with an asterisk are the noised subject tokens. Here,  $p_{\theta}(o_{true}|s,r) = .923$  and  $p_{\theta}(o_{true}|s_{noise},r) = .001$ .
CounterFact Example
Figure 3: An example CounterFact datapoint.
3.3 Model Editing with ROME
We describe the ROME editing method here since we use it in our analysis in Sec. 4, and later in Sec. 5 we outline additional editing methods we consider. For mathematical detail, see Meng et al. [21].
The input to ROME includes a prompt  $P = (s, r)$  and a new desired output, which is always a false target  $o_{\text{false}}$  in the CounterFact dataset. To change the model prediction to  $o_{\text{false}}$ , ROME applies a rank one edit to the down-projection matrix in a prespecified MLP layer in the model. The default layer in GPT-J
is layer 6, following from averaged Causal Tracing results. ROME also makes use of covariance statistics of different subject representations obtained from a larger corpus as it edits individual facts. Overall, the method is designed to optimize the quantity  $p_{\theta}(o_{false}|s,r)$  while aiming to satisfy some other constraints reflecting what a desirable model edit is (described in Sec. 3.4 next).
Input Prompt: Autonomous University of Madrid, which is located in
Requested Edit: Spain  $\longrightarrow$  Sweden
Paraphrase: and Sallie Beavers Riley. Autonomous University of Madrid is located in
Neighbor: Ripollès, located in
3.4 Editing Metrics
Editing methods are typically evaluated according to their ability to (1) change the model prediction on the input  $P$  provided at runtime, (2) generalize appropriately to paraphrases of the prompt  $P$ , and (3) avoid over-generalizing to unrelated data [42, 9, 24, 15, 25]. We adopt metrics for each desideratum that we compute with available CounterFact data. Instead of the exact "magnitude" metrics from Meng et al. [21], we use normalized versions of each metric that we design to scale from 0 to 1 depending on whether the edit was maximally (un)successful, for purposes of making scores more comparable across data points. We denote the new edited weights of the LM as  $\theta^*$  and its pre-edit weights as  $\theta$ . See Fig. 3 for an example of the kinds of data these metrics are computed on.
1. Rewrite Score. The rewrite score measures how much an edit improves the target probability  $p(o_{false}|s,r)$  as a fraction of the maximum possible improvement:
$$
\frac {p _ {\theta^ {*}} \left(o _ {\text {f a l s e}} | s , r\right) - p _ {\theta} \left(o _ {\text {f a l s e}} | s , r\right)}{1 - p _ {\theta} \left(o _ {\text {f a l s e}} | s , r\right)}
$$
2. Paraphrase Score. The paraphrase score measures the target probability using syntactical paraphrases as inputs, always preserving the exact subject wording:
$$
\frac {p _ {\theta^ {*}} \left(o _ {\text {f a l s e}} | s , r ^ {*}\right) - p _ {\theta} \left(o _ {\text {f a l s e}} | s , r ^ {*}\right)}{1 - p _ {\theta} \left(o _ {\text {f a l s e}} | s , r ^ {*}\right)}
$$
which is averaged over multiple available paraphrases per input  $P$ . The score measures whether edits properly generalize across semantically equivalent prompts.
3. Neighborhood Score. The neighborhood score measures whether edits change predictions for prompts with a similar subject  $s^*$ , the same relation  $r$ , and the same (true) objects. We scale the difference in probabilities so that 1 means the probability did not change (good), and 0 means it changed to the maximum extent possible (bad):
$$
1 - \frac {\left| p _ {\theta^ {*}} \left(o _ {f a l s e} \mid s ^ {*} , r\right) - p _ {\theta} \left(o _ {f a l s e} \mid s ^ {*} , r\right) \right|}{. 5 + \left| p _ {\theta} \left(o _ {f a l s e} \mid s ^ {*} , r\right) - . 5 \right|}
$$
The score measures whether edits avoid over-generalizing from the prompt  $P$  to different subjects.
4 Does Edit Success Follow From Localization?
Ostensibly, localization results should inform editing methods because it should help to know where information is stored in a model if you are going to manipulate the model's expression of that information. More specifically, if you wanted to inject a false belief  $(s,r,o_{\text{false}})$  into a model (as defined in the ROME editing problem), it seems helpful to know which weights store the true fact  $(s,r,o_{\text{true}})$ , so that you could replace some stored representation of  $o_{\text{true}}$  with that of  $o_{\text{false}}$ . This underlying assumption about editing models appears in much past work on localization, where editing is used to verify localization analysis (see Sec. 2). In this section, we investigate the validity of this assumption as it applies to autoregressive Transformers.
4.1 Experiment Design
The goal of our experiments is to determine, for a given datapoint, whether edit success at a specific layer aligns with the results from Causal Tracing at that layer (see Causal Tracing description in Sec. 3.2). We operationalize this outcome and explanatory variable as follows:
1. Edit Success. We primarily consider Rewrite Score as our measure of edit success, given that this is the main optimization objective of ROME. Note ROME achieves an average rewrite score of  $99\%$  at layer 6 of GPT-J and above  $96\%$  at layers besides the last layer of the model.
2. Tracing Effect at layer  $\ell$ . Since the output of Causal Tracing is a  $T \times L$  grid of estimates, we obtain a single tracing effect per layer by taking the max across the  $T$  token effects at each layer (i.e., we collapse the grid in Fig. 2 down to a single curve across layers). Like our other metrics,
we use a fractional tracing effect where 0 means the intervention had no effect and 1 means it fully restored the original probability  $p_{\theta}(o_{true}|s,r)$ :
$$
\frac {p _ {\theta} \left(o _ {\text {t r u e}} \mid s _ {\text {n o i s e}} , r , v _ {(t , \ell)}\right) - p _ {\theta} \left(o _ {\text {t r u e}} \mid s _ {\text {n o i s e}} , r\right)}{p _ {\theta} \left(o _ {\text {t r u e}} \mid s , r\right) - p _ {\theta} \left(o _ {\text {t r u e}} \mid s _ {\text {n o i s e}} , r\right)}
$$
Lastly, note we use a tracing window size of 5 (smaller than the value of 10 used in Fig. 2).
4.2 Model and Data
We conduct our analysis with GPT-J [35] using the CounterFact dataset, similar to Meng et al. [21]. GPT-J is a 6 billion parameter autoregressive language model. We record editing performance at layers in  $\{1, 5, 9, 13, 17, 21, 25, 28\}$  as well as layer 6 (the default for ROME). Note ROME achieves an average rewrite score of  $99\%$  at layer 6 and above  $96\%$  at layers besides layer 28.
The CounterFact dataset includes datapoints consisting of a prompt, paraphrases, and neighboring points. For each point, a new (false) target is supplied for editing purposes. We show an example datapoint in Fig. 3. Note paraphrases intentionally include unrelated text preceding a syntactical paraphrase of the input, with the idea that this text should not affect the output. We select data for experiments from  $10\%$  of CounterFact, additionally filtering to a
subset of facts that are correctly completed by GPT-J, in order to ensure that there is knowledge to localize in the model for each point (details in Appendix A). Our final sample size is  $n = 652$ .
Figure 4: The correlation between ROME edit success and the tracing effect at layer 6 in GPT-J is not positive but in fact slightly negative ( $\rho = -0.13$ ;  $p < 1\mathrm{e} - 3$ ). The dashed red line shows a hypothetical perfect relationship.
4.3 Experiment Results
We present results in two ways. First, in Fig. 4, we show Rewrite Score as a function of the (fractional) tracing effect. The red dotted line shows a hypothetical perfect relationship between tracing and edit success. Surprisingly, there is not a positive relationship but a negative relationship between the rewrite score and the tracing effect (linear correlation of  $\rho = -0.13$ ;  $p < 1\mathrm{e} - 3$ ). This seems to fully invalidate the assumption that editing should be most effective when it occurs at a layer where information is stored about the edited fact. We wish to emphasize, however, that in most layers we simply see a near-zero rather than negative correlation, as shown in Appendix Fig. 15.
Our second mode of analysis is though linear regression models predicting rewrite score based on (1) the tracing effect, (2) the choice of edit layer treated as a categorical variable, or (3) both terms interacted, again treating edit layer as a categorical variable. The purpose of the models is to show how much of the variance in rewrite score is explained by one variable versus the other. We show the resulting  $R^2$  values in Table 1. We see that the choice of layer explains almost all of the variance in
Table 1:  $R^2$  values for predicting ROME edit success. Tracing effects explain essentially none of the variance in rewrite score, while the choice of edit layer is very important.
Method	R2Values
Layer	Tracing Effect	Both
ROME	0.947	0.016	0.948
rewrite score  $(94.7\%)$ , while adding the tracing effect to the model raises the  $R^2$  only to  $94.8\%$ . This means that the tracing effect is able to explain only  $0.1\%$  of the variance in edit success when accounting for the choice of edit layer. These results suggest that the tracing effect is essentially unrelated to the success of model editing.
This is a surprising conclusion, and it naturally raises the question of why applying ROME at layer 6 works well in the first place (see average rewrite, paraphrase, and neighborhood scores across layers in Appendix Fig. 7). We suggest a possible answer to this question in Sec. 6.