Source-linked AI summary

Does Localization Inform Editing? Surprising Differences in Causality-Based Localization vs. Knowledge Editing in Language Models

Peter Hase, Mohit Bansal, Been Kim, Asma Ghandeharioun

arXiv:2301.04213v2cs.LGcs.AIcs.CL

TL;DR

The paper asks whether localizing factual knowledge identifies where model weights should be edited. It compares Causal Tracing with editing across standard and modified objectives, finding that localization generally does not identify the best edit layer. Even where tracing correlates with performance, layer choice is more predictive, within the study’s evaluated settings.

  • Problem

    It is unclear whether locating factual information in model components identifies where editing should intervene to change that knowledge.

  • Method

    The paper compares Causal Tracing localization with model-editing success across standard editing and four variants involving different inputs, targets, or objectives.

  • Results

    Causal Tracing is essentially unrelated to standard edit success; in Fact Forcing, tracing correlates best, but explains only a small fraction of performance variance.

  • Takeaways & Limitations

    Localization results from Causal Tracing do not indicate which model layer is best to edit for manipulating stored facts.

  • Takeaways & Limitations

    The conclusions may not hold beyond the evaluated CounterFact and ZSRE factual prompts and the particular localization and editing methods studied.

Abstract

from arXiv · show

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. 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. 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. Our code is available at https://github.com/google/belief-localization

1 Introduction

The paper tests whether localizing factual knowledge identifies the best weights to edit. Across standard editing and several variants, it finds that Causal Tracing localization is generally a poor guide to edit-layer choice.

  • Causal Tracing identifies where information is carried, but that location does not determine where factual knowledge should be edited.The paper distinguishes localization of information in representations from intervention sites that effectively change model behavior.
  • Causal Tracing results are statistically uncorrelated with success when injecting new facts, despite ROME and MEMIT successfully editing knowledge.The finding holds for ROME, MEMIT, and Adam-based finetuning on CounterFact with GPT-J.
  • The paper introduces four editing variants to test whether changing the editing objective or inputs can reconnect localization with edit success.The variants include Tracing Reversal, Fact Erasure, Fact Amplification, and Fact Forcing.
  • Fact Forcing shows the strongest relation between tracing effects and edit success, but edit-layer choice explains substantially more performance variation.The authors therefore conclude that tracing insights are not useful for choosing which layer to edit.
  • Many of 652 GPT-J facts appear outside layers 4-9, the range edited by methods such as ROME and MEMIT.Examples include layers 1-3 and 16-20.

2 Related Work

Related work localizes model behaviors to layers, directions, neurons, and subnetworks, and often validates localization by editing the suggested components. The paper argues that such validation alone cannot establish whether a location is necessary, proportionally predictive, or optimal for editing.

  • Localization research assigns behaviors or concepts to model components such as layers, weight matrices, latent directions, neurons, and subnetworks.
  • This paper uses layer-wise Causal Tracing, which estimates information content through denoising, focusing on MLP layers implicated in factual associations.
  • Prior studies often validate localization by editing activations, weights, or latent directions and checking whether behavior changes appropriately.
  • Editing a localized component does not by itself show that success scales with localization strength, requires that component, or is maximized there.

3 Notation and Background

The paper defines factual tuples and CounterFact variables, then describes Causal Tracing, ROME editing, and evaluation metrics for changing predictions while preserving appropriate generalization.

  • Facts are represented as tuples (s, r, o), where a subject and binary relation map to an object, with prompts completed by the object.CounterFact also supplies paraphrases, neighboring subjects, noised subjects, and false targets for editing.
  • Causal Tracing estimates information about a true fact in hidden representations by copying clean representations into a forward pass with a noised subject.The adjusted pass measures the copied representation’s effect on the probability of the true target.
  • A tracing window copies representations from multiple adjacent layers; the default window size is 5.The window-size parameter determines how many neighboring representations contribute to the layer estimate.
  • ROME applies a rank-one edit to an MLP down-projection matrix to make a false target the model’s prediction, using layer 6 by default in GPT-J.The default layer follows averaged Causal Tracing results.
  • Editing is evaluated by rewrite, paraphrase, and neighborhood scores measuring target change, paraphrase generalization, and avoidance of unrelated changes.The metrics are normalized to range from 0 to 1.

4 Does Edit Success Follow From Localization?

The paper tests whether Causal Tracing identifies the best layer for model editing and finds that tracing effects do not predict success when injecting new facts.

  • Research question: The study asks whether knowing where information is stored should guide edits that change model behavior.It evaluates whether edit success at a layer aligns with Causal Tracing at that layer.
  • Operationalization: ROME edit success is measured primarily with Rewrite Score, while tracing effects are reduced to one fractional effect per layer.The tracing effect is obtained by taking the maximum across token effects, using a tracing window size of 5.
  • Experimental setting: The experiments use GPT-J with CounterFact and evaluate edits across selected layers, including ROME’s default layer 6.ROME achieves an average rewrite score of 99% at layer 6 and above 96% at other tested layers besides the last layer.
  • Results: ρ = −0.13; p <1e−3: GPT-J ROME rewrite score is slightly negatively correlated with the tracing effect at layer 6.Most layers instead show near-zero correlations, rather than negative ones.
  • Interpretation: The results challenge the assumption that editing is most effective where Causal Tracing indicates stored information.The finding also raises the question of why ROME performs well at layer 6.
  • Results: Layer choice explains 94.7% of rewrite-score variance, while adding tracing effects raises R2 to only 94.8%.Thus, tracing effects explain only 0.1% of variance after accounting for edit layer.

5 Reconciling Localization and Editing

The paper introduces editing variants designed to resemble Causal Tracing more closely and finds that tracing effects remain weak predictors except in Fact Forcing.

  • Editing problem variants: The authors introduce four editing variants that progressively resemble Causal Tracing in their inputs, targets, or objectives.The variants are Tracing Reversal, Fact Erasure, Fact Amplification, and Fact Forcing.
  • Evaluation: The variants retain Rewrite, Paraphrase, and Neighborhood Score evaluation, changing only the target output for rewrite and paraphrase metrics.Neighborhood evaluation remains identical across variants.
  • Methods: The study compares ROME, MEMIT, constrained finetuning, and other editing settings using the same broader experimental procedure.MEMIT spreads a single-fact update across several layers, unlike ROME’s single-layer update.
  • Main results: At most 3.2% of edit-success variance is additionally explained by tracing effects, versus 58.5% on average by layer-only regressions.The choice of editing one or five layers to match restored representations makes little difference to this conclusion.
  • Fact Forcing: Fact Forcing with finetuning shows the strongest relationship, with tracing effects explaining an additional 3% of edit-success variance.This is statistically significant at p < 1e−4, but remains a weak relationship.
  • Interpretation: The Fact Forcing result is attributed to using a noised subject input, while target choice and maximizing versus minimizing probability are ruled out as explanations.The authors conclude that deeper reasons may explain why localization and editing success are generally unrelated.

6 Discussion

The discussion separates where factual information is represented from where edits work best, while proposing that information accumulated across layers may be overridden elsewhere. It also cautions that some supporting analyses remain limited.

  • ROME performs better on average when optimizing the last subject token representation than another token representation.The paper notes that this finding concerns token choice, not a direct prediction of the best edit layer.
  • Information may be accumulated across layers, allowing an edit at layer k to override information associated with another layer ℓ.The authors suggest this as a possible explanation for ROME working across a broad layer range.
  • The proposed cross-layer override hypothesis remains unresolved because arbitrarily swapping Transformer layers can greatly damage model performance.The authors leave further investigation of this hypothesis to future work.
  • Causal Tracing localizes information in representations but does not indicate which layer is best for editing factual behavior.The paper distinguishes tracing effects from the intervention choice that changes model behavior.
  • Editing experiments cannot provide further evidence for localization conclusions if localization and editing answer different questions.This challenges validation strategies that intervene on components recommended by localization analyses.
  • Datapoint-level regression would provide stronger evidence that tracing effects predict which token representation is best to optimize with ROME.The authors identify edit layer and other confounders as factors that should be ruled out.

7 Conclusion

The paper finds that Causal Tracing is essentially unrelated to factual edit success in the default setting. Across editing variants, tracing correlates best in Fact Forcing, but edit layer remains the stronger predictor.

  • Model edit success is essentially unrelated to where factual information is stored, as measured by Causal Tracing.
  • Tracing effects correlate best with edit success in the Fact Forcing setting, but explain only a small fraction of performance variance.The choice of edit layer is a much more important factor in editing performance.
  • The results suggest that better mechanistic understanding of pretrained language models may not always translate into better ways to change their behavior.

8 Limitations

The experiments are limited to particular factual datasets, two autoregressive Transformer models, and selected localization and editing methods. The authors caution that conclusions may not generalize beyond this experimental setting.

  • The study uses only CounterFact and ZSRE short English prompts with factual completions for specific subject–object relations.This represents a basic form of factual knowledge.
  • Localization and editing analysis may yield different trends for forms of knowledge beyond the basic factual setting studied.
  • The study uses two autoregressive Transformers, and conclusions may not generalize to models larger than GPT-J with prompting-related phase changes.
  • The conclusions cover selected layer-level localization and MLP editing methods rather than the full breadth of related methods.The authors recommend caution when applying the findings beyond their experimental setting.

9 Broader Impacts

The paper notes that mechanistic understanding could eventually improve editing, including potentially harmful editing applications. It therefore frames interpretability and editing as relevant to controlling language models.

  • Mechanistic understanding could eventually improve model editing, including the injection of harmful beliefs or dangerous knowledge.The paper identifies this as a possible dual-use concern.
  • The authors hope that mechanistic interpretability and model editing will improve the ability to control language models.

A Experiment Details

The experiments use CounterFact and ZSRE data with GPT-J and GPT2-XL, evaluating several editing methods, layer choices, and tuned hyperparameters.

  • CounterFact and ZSRE are publicly available datasets used in the experiments.
  • Data Filtering: GPT-J data are filtered to facts the model completes correctly, producing a final sample of n = 652.GPT-J achieves 32.6% completion accuracy under the filtering scheme.
  • Compute: Computations use a single NVIDIA A6000 GPU, with GPT-J editing and causal-tracing runs taking about eight and twelve hours, respectively.
  • Hyperparameters: Editing-method hyperparameters include norm constraints for finetuning and regularization weights for ROME and MEMIT.
  • Experiments evaluate Error Injection and Fact Forcing for GPT2-XL, while the broader setup tunes methods across additional editing variants.

B Additional Results

Additional analyses examine datasets, metrics, localization methods, concentrated tracing effects, and essence drift. They continue to find weak links between localization and editing, while edit-layer choice remains more informative.

  • ZSRE: ZSRE results with ROME and GPT-J match the CounterFact conclusion that tracing effects are weakly related to edit success.
  • Representation Zeroing: Representation zeroing, which sets clean MLP output representations to zero, also yields near-zero correlations with edit success across layers.
  • Concentrated Tracing Effects: For concentrated tracing effects, adding tracing improves ROME Error Injection prediction by 0.2%, while the maximum added R2 effect is 3.2% for constrained finetuning with Fact Forcing.
  • Essence Drift: Essence score measures changes in perplexity on sampled texts describing a subject’s known properties after editing.
  • Regression Analysis: Tracing effects have some predictive value for Fact Forcing, but the choice of edit layer explains substantially more variance in rewrite score.
  • Essence Drift: Later edit layers generally improve essence scores, flattening and shifting the combined overall-score trend toward mid-range layers.

C Robustness Experiments

Robustness experiments vary metrics, tracing windows, models, token positions, and metric scaling. Across these changes, tracing remains weakly predictive, whereas edit-layer choice is more informative.

  • Alternative Metrics: Using paraphrase, neighborhood, and overall scores, the added predictive effect of tracing remains very small across conditions.The added effect is less than 3% for paraphrase and 2% or less for neighborhood and overall scores.
  • Tracing Window Size: Increasing the tracing window size to 10 produces results nearly identical to the main analysis.
  • GPT2-XL: GPT2-XL shows similar trends, with the largest explanatory effects occurring for Fact Forcing with constrained finetuning but remaining about 2%.
  • Unscaled Metrics: Original unscaled editing and tracing metrics preserve the conclusion that edit-success correlations with tracing remain near zero.
  • Regression Error: Across regression error metrics, predicting edit success from edit-layer choice yields much lower errors than using tracing effects.
  • Token Position: The last-subject-token analysis examines the representation choice directly motivated by ROME’s causal-tracing rationale.
Loading 2301.04213v2…