Mechanistic Unlearning: Robust Knowledge Unlearning and Editing via Mechanistic Localization

Thank you for reading our work and providing feedback. We appreciate the opportunity to clarify our contributions and address your concerns. We believe that the assessment and low score was based on certain misunderstandings of our work: we focus on model editing of factual relations; our editing method is in fact novel and clears up misconceptions about mechanistic interpretability for model editing and unlearning, making it an important contribution to the literature; we provide an extensive set of evaluations of our method, and even stronger attacks than suggested by the reviewer.

We hope our detailed responses below will convince the reviewer to re-evaluate our work and raise the score.

On the Lack of Unlearning benchmarks

Our primary contribution lies in investigating how mechanistic interpretability can enhance the precision and robustness of knowledge editing of factual associations: i.e., replacement of certain facts with new facts. This goal underscores our choices of baselines and datasets. We use the SportsFacts dataset following work from Nanda et. al, who use it to mechanistically understand factual recall, to create a localization technique for robustly modifying factual associations. We then translate these findings to the CounterFact dataset, a benchmark widely used in the editing literature.

Most of our work focuses on editing rather than unlearning, and the baselines of RMU, GradDiff, and NPO couldn’t be used for the editing objectives without a significant reformulation of the tasks. Our inclusion of an unlearning result in A.1 of the paper was primarily to show the potential for our method to generalize to unlearning, but not a claim that our method led to state of the art unlearning performance on every benchmark. However, the current literature on unlearning is relatively unanimous in that no current method is robust against the partial relearning attack (Deeb 2024), including RMU (Li 2024) and TAR (Tamirisa 2024) which is defeated by parameter-efficient fine-tuning: we believe our method can yield progress.

In this work, we hoped to lay the groundwork for applying interpretability for unlearning. We are excited about future work that applies more sophisticated interpretability on complex datasets like WMDP: a positive outcome of the results of this paper would be to inspire further research into interpretability to achieve more robust editing and unlearning.

We will revise the paper to more accurately reflect that the majority of our current empirical results focus on fact editing. An easy fix!

On Novelty

One main novel contribution is the identification and utilization of Fact Lookup (FLU) mechanisms for robust knowledge editing. We contrast this approach with Output Tracing (OT) methods. We show that editing localized to these FLU mechanisms leads to more robust edits as measured using a number of evaluations (listed below).

We demonstrate that the relationship between localization and fact editing/unlearning is more nuanced than suggested in Hase et al. (2023), and that not all localization techniques are equally effective. This is an important contribution to the literature, showing a promising use of mechanistic interpretability.

On evaluations

Regarding the evaluation against different attacks, we would like to highlight that we do evaluate the robustness of our method against:

• Rephrasing prompts (Paraphrase Evaluation);
• Multiple-choice question extraction (MCQ Evaluation);
• Adversarial relearning attacks;
• Soft prompt attacks, which are a more challenging form of adversarial prompting since they operate in continuous space.
• Latent knowledge analysis, which trains a probe to extract the correct answer from the latent representation of the model, instead of logit lens which is a biased representation of the model’s best guess at a layer

We are confident that these evaluations provide a comprehensive assessment of the robustness of our proposed editing method.

Note that the current version of the paper is even stronger as we included additional experiments suggested by other reviewers: proposed and tested an automated version of our localization method which outperforms strong baselines; included additional tests on which components are important for localization; and demonstrated that our method can successfully edit a large number of facts (up to 1000).

References: Hase, P. et al. Does localization inform editing? surprising differences in causality-based localization vs. knowledge editing in language models, 2023.

Meng, K. et al. Locating and editing factual associations in gpt, 2023.

Nanda, N. Attribution patching: Activation patching at industrial scale, 2023.

Li, N. et al. The WMDP Benchmark: Measuring and Reducing Malicious Use With Unlearning, 2024.​

Deeb, A. et al. Do Unlearning Methods Remove Information from Language Model Weights?, 2024.​

Tamirisa, R. et al. Tamper-Resistant Safeguards for Open-Weight LLMs, 2024.​

Thank you for your detailed review and valuable proposals. In response, we proposed a way to automate our method, which allowed us to test it at scale in terms of the number of facts to be edited. We also added experiments demonstrating that adding attention heads does not improve performance.

Addressing points 1, 2:

We originally ignored the attention heads as candidates for localization as a result of work done by Nanda et. al who identified attention heads as playing a “fact extraction” role in the recall mechanism. We don’t claim that we found the precisely optimal localization for editing, but rather that this somewhat coarse localization was sufficient enough to yield significant robustness improvements: future work with more sophisticated interpretability techniques could further strengthen these results. However, we agree that empirical evidence to prove this would be valuable.

To address this, along with your second point for a more automated strategy for localization and your third point for a more challenging editing scenario, we run an experiment scaling up the number of CounterFact facts edited from 64 to 1000, and try a new localization. We stick to CounterFact as the dataset contains factual questions in varying formats. Our difficulty increase thus comes from scaling the number of facts edited: our technique now has to maintain robustness to orders of magnitude greater number of facts. We do this experiment for the Gemma-7B model, and plan to replicate this setup across Gemma-2-9B and Llama-3-8B by the camera-ready deadline.

This scaling of facts also means we have to use a more automated technique for localization. We localize per-fact, employing the same technique as described in A.2.2. We pick the components by utilizing a heuristic, picking all important MLPs that affect the final logit difference by > 2 stds. This differs from our original manual localization where we analyzed a group of facts, took the average contribution for each MLP, and assigned a group of MLPs to be the localization.

We test our technique, based on your suggestion, to also include relevant attention heads. We compare this against a strong baseline of picking all the MLPs. We measure the robustness of these techniques to prompt changes (using an MCQ prompt format as in section 3.1) and to latent attacks by measuring probe accuracies, as in Section 3.3.

We see that our fact lookup localization technique (localizing MLP layers only) maintains its MCQ forget error as we scale the facts edited, outperforming the baselines (fig: https://imgur.com/a/60azzEo). Our localization also outperforms other methods on its MCQ edit accuracy (fig: https://imgur.com/a/L7TnBwu), being the only localization to generalize to an alternative prompt setting. Finally, we see that latent knowledge attacks continue to fail when using our localization, as the probe accuracies are as accurate as random chance (fig: https://imgur.com/a/Fn6xWp1). All figures are anonymized and do not contain author information.

Addressing point 3: In this work, we hoped to lay the groundwork for applying interpretability for unlearning. We chose to work with CounterFact and Sports Facts datasets because the mechanisms of factual recall in these datasets have been well studied in the literature, and because CounterFact is a benchmark widely used in the editing literature, including in the seminal paper by Meng et. al. We are excited about future work that applies more sophisticated interpretability analyses on complex datasets like MMLU: a positive outcome of the results of this paper would be to inspire further research into interpretability to achieve more robust editing and unlearning.

Addressing point 4: We will add error bars and also results for each model across the various evaluations in the camera-ready version. For latent knowledge we do present results for each model in Appendix A.7.4 since each model has a different number of layers. For the latent knowledge analysis, Fact Lookup localization is most robust in Gemma-7b and Llama-3-8b, while some other methods are competitive with Fact Lookup in Gemma-2-9b.

Thank you again for your feedback. We hope that these strong additional empirical results, automation of our method, and our clarifications addressed all of your concerns.

References: [1] Nanda, N., Rajamanoharan, S., Kram ́ar, J., and Shah, R. Fact finding: Attempting to reverse engineer factual recall on the neuron level, Dec 2023. URL https://www.alignmentforum.org/posts/iGuwZTHWb6DFY3sKB/fact-finding-attempting-to-reverse-engineer-factual-recall.

Thank you for your comments, we will make sure to improve the description of the FLU localization technique for CounterFact in Sec 2.2, and we report new positive experimental results editing significantly more facts. The methodology is briefly summarized here:

An important pre-requisite is “path patching”, described in section 3.1 of Wang et. al. This allows us to measure the importance of the direct edge between two components in a model. We first measure the direct effect of each attention head on the final output, using the logit difference between the original and edit answer as our measure. Nanda et. al show that these attention heads “extract” relevant associations from the residual stream. We call these the fact extraction heads.

We then measure the effect of each MLP on the final logit difference as mediated only through these fact extraction heads. That is, we iterate through the MLPs and path patch the MLP -> {all fact extraction head} edges, and measure the logit difference. MLPs that cause a large change in this measure do so by introducing the factual association into the residual stream via a lookup mechanism, enabling the fact extraction heads to parse the association. We call these our fact lookup MLPs, which are the localization.

Below we report new experimental results in response to your concern about the number of facts we edit.

We acknowledge your point about the relatively small size of the dataset for editing. Our reasoning for maintaining a smaller set of facts was primarily to facilitate the creation of precise manual mechanistic localizations. However, we recognize the importance of understanding how our findings might generalize to a larger set of edited facts.

To address this, we slightly modify our technique to localize per fact rather than averaging over a set of facts. Then, similar to our original technique in A.2.2, we pick the relevant components by identifying the MLPs that have a >2 std impact on the logit difference in the fact. Using this, we can precisely localize each fact at scale. This allowed us to run an additional evaluation scaling up the number of facts on the CounterFact dataset to be edited all the way to 1000 facts.

We do this evaluation only on our Gemma-7b model due to time constraints, but we can expand this to Gemma-2-9B and Llama-3-8B as well by the camera-ready deadline. We compare our localization technique to using all the MLPs as a localization (a strong baseline) and a localization that uses our technique but also includes attention heads (as suggested by reviewer MshQ). We measure the robustness of these techniques to prompt changes (using an MCQ prompt format as in section 3.1) and to latent attacks by measuring probe accuracies as in section 3.3.

We see that our fact lookup localization technique maintains its MCQ forget error as we scale the facts edited, outperforming the baselines (fig: https://imgur.com/a/60azzEo). Our localization also outperforms on its MCQ edit accuracy (fig: https://imgur.com/a/L7TnBwu), being the only localization to generalize to an alternative prompt setting. Finally, we see that latent knowledge attacks continue to fail when using our localization, as the probe accuracies are as accurate as random chance (fig: https://imgur.com/a/Fn6xWp1). All figures are anonymized and do not contain author information.

We hope our additional experiments have convinced you of the robustness of our method at scale.

References:

Nanda, N., Rajamanoharan, S., Kram ́ar, J., and Shah, R. Fact finding: Attempting to reverse engineer factual recall on the neuron level, Dec 2023. URL https://www.alignmentforum.org/posts/iGuwZTHWb6DFY3sKB/fact-finding-attempting-to-reverse-engineer-factual-recall.

Wang, Kevin, et al. "Interpretability in the wild: a circuit for indirect object identification in gpt-2 small." arXiv preprint arXiv:2211.00593 (2022).

Thank you for a thorough read of our submission. Below we answer your question regarding the generalization of relearning in the Sports-Athlete-Editing and Full-Sports-Editing tasks, and explain our rationale behind the task choice for relearning attacks.

Our relearning experiments in Section 3.2 are performed on the Sports-Athlete-Editing task. In this task, for a randomly selected set of athletes, we edit their associated sport by assigning them a new sport chosen uniformly at random from the existing set of sports.

In this Sport-Athlete-Editing task, we do not expect relearning to “generalize”: there is no inherent correlation between an athlete, their original sport, and the newly assigned sport. Therefore, relearning a subset of edited facts should not provide the model with a basis to correctly infer the previously unlearned facts for other athletes, even if they were originally associated with the same sport like basketball (unless this old information is still stored in the model after editing).

On the other hand, relearning attack is uninformative in some other setups, where the forget set accuracy could go up by the model simply learning a fixed mapping (e.g., always respond with “Golf” instead of “Basketball”). For example, this concern would apply for relearning the Full-Sports-Editing task: here we reassign one sport to another for all athletes playing that sport. For example, we might change all associations of "basketball" to "golf." If we then retrain the model on half of these edited facts, it's highly probable that the model would simply learn the global reassignment (e.g., all "golf" becomes "basketball") rather than truly relearning the specific original associations. This makes it difficult to distinguish between a newly learned association and the relearning of previously stored knowledge.

It is important to note that the relearning paper we base our evaluations off (Deeb 2024) also makes efforts to avoid retraining on a set which does not have independent information from the rest of the forget set: when the information used for retraining is independent from the rest of the forget set, we don’t expect any information recovery given perfect unlearning/editing. However, when the retrained information is not independent from the rest of the forgotten facts, it is unclear what baseline amount of recovery is expected.

We hope this clarifies our approach and the rationale behind our experimental design. We would also like to highlight that we ran additional experiments requested by other reviewers to further strengthen the paper: we showed that the localization part of our method can be automated, and that our method can be scaled to successfully edit a large number of facts.