Mechanistic Unlearning: Robust Knowledge Unlearning and Editing via Mechanistic Localization

Authors: Phillip Huang Guo, Aaquib Syed, Abhay Sheshadri, Aidan Ewart, Gintare Karolina Dziugaite
Primary area: interpretability and explainable AI
Keywords: Model Editing, Unlearning, Mechanistic Interpretability, Localization, Adversarial Robustness

Abstract:
Methods for knowledge editing and unlearning in large language models seek to edit or remove undesirable knowledge or capabilities without compromising general language modeling performance. This work investigates how mechanistic interpretability---which, in part, aims to identify model components (circuits) associated to specific interpretable mechanisms that make up a model capability---can improve the precision and effectiveness of editing and unlearning. 
We find a stark difference in unlearning and edit robustness when training components localized by different methods. We highlight an important distinction between methods that localize components based primarily on preserving outputs, and those finding high level mechanisms with predictable intermediate states.
In particular, localizing edits/unlearning to components associated with the \textit{lookup-table mechanism} for factual recall 1) leads to more robust edits/unlearning across different input/output formats, and 2) resists attempts to relearn the unwanted information, while also reducing unintended side effects compared to baselines, on both a sports facts dataset and the CounterFact dataset across multiple models.
We also find that certain localized edits disrupt the latent knowledge in the model more than any other baselines, making unlearning more robust to various attacks.

Decision: reject
Reviewer scores: 5, 5, 3, 8

Reviews and discussion:
<thread>
creation date:
1730583114889

thread_type:
official_review

<thread_object>
creator: reviewer_vunn

<content>
summary:
The authors study whether insights from mechanistic interpretability can improve our ability to perform unlearning or make targeted model edits. More specifically, using a collection of factual recall tasks, the authors study different ways of selecting a subset of the model's parameters to finetune for unlearning or editing certain facts. Building on prior work identifying a subset of model parameters involved in factual recall, the authors show that finetuning only these parameters results in more effective unlearning/model editing than finetuning parameters selected via other techniques. These claims are supported by a variety of analyses, such as robustness to different ways of eliciting the model's factual knowledge and robustness to retraining the model to relearn the facts.

strengths:
Overall, the results in this paper are very strong, and the suite of evaluations is impressively thorough.
1. The authors do a good job of establishing that finetuning on the "manual interp" subset of parameters results in qualitatively different unlearning dynamics and quantitatively stronger results. The field of interpretability has struggled recently to demonstrate that their insights are useful for downstream tasks in general—and for unlearning/model editing in particular as demonstrated by Hase et al. (2023)—so I expect these results will be a breath of fresh air for the interpretability community.
2. The authors perform a very thorough suite of evaluations showing that mechanistic unlearning more effectively changes knowledge stored in the model weights (see (3) and (4) for more detail). This is another place where the authors set themselves apart from the field: the unlearning literature has often struggled with thoroughly evaluating the efficacy of their methods.
3. I was especially impressed by the relearning evaluations, showing that—when training the model to relearn the unlearned facts—the mechanistically unlearned model relearns the facts much more slowly (figure 2).
4. Also impressive were the results that sweep over the number of masked parameters, revealing qualitative differences in the various unlearning techniques. Figure 5, right, which shows that mechanistic unlearning generalizes to MCQ rephrasing substantially better than any other unlearning technique, was especially striking.

weaknesses:
Overall, this work's presentation was very poor.
1. Many important details are missing from the main body of the text. (a) The "fact loookup localization" (which is also called "manual interpetability"—why not use one term?) method is entirely explained in an appendix. While it's reasonable to put most of the FLU details in the appendix (since this is a replication of prior work) understanding at a high level what FLU is and how it assigns importance scores to various model components is essential for understanding the rest of the work. (b) The definition of the unlearning loss is described only as "the log(1 - p) measure from Mazeika et al."—this is an important part of the method and should be explained.
2. It is very difficult to follow discussions of the tasks. There are two tasks related to sports facts, one related to unlearning facts about basketball athletes (how many athletes? I think this is mentioned later, but it should be included in section 2.1) and one related to editing 16 (random?) athletes to change their sport to "golf." The authors later refer to these tasks with vague phrases like "For editing the subset of athletes..." The authors should instead give distinct names to these two tasks (e.g. "Sports-unlearning" and "Sports-editing") which they use to refer to the tasks throughout the rest of the text.
3. Although the results are strong, the authors do not make it easy to tell this from reading the work. For example, the tables of numbers in the first results section are not a reasonable way to present these results. Tables like these are good for when we want to inspect small numerical differences; in contrast here we don't care about small differences (e.g. between forget scores of 0.002 and 0.000) but about large differences (e.g. between MCQ scores of 0.110 and scores >0.5). The authors should choose a different way of presenting these results, perhaps as a bar chart. 
4. Similarly, the authors present results for all three tasks for each of their evaluations, resulting in a large number of figures which are left to the reader to synthesize. This work would be much stronger if the authors found ways of presenting their work that summarized and emphasized the key takeaways.
5. The main takeaway from the counterfact retraining experiment should have been that this experiment isn't informative, since relearning on some facts doesn't generalize to other facts for *any* of the unlearning techniques. This experiment should therefore be moved into an appendix.

I also have some object-level concerns about various choices made by the authors. I mostly expect these to be easy to address with follow-up experiments (and I'll be happy to raise my score once I see these follow-ups).
1. The task definitions involve various arbitrary-seeming choices: only basketball athletes are targeted for unlearning, only golf is uses as the target relation for editing, only a particular set of 16 facts is used for CounterFact, etc. This makes it hard to tell if these results are due to cherry-picking. The authors should sweep over options for these choices (e.g. also targeting athletes for different sports, or rerandomizing for different unlearned facts) and present averages over the conditions.
2. While it is good that the authors study multiple models, they pair specific models with specific tasks, again making the reader wonder about cherry-picking. Unless there is a good reason not to do so, the authors should test all three tasks for *both* Gemma-7B and Gemma-2-9B.
3. In addition to the all MLPs baseline, the authors should also include a baseline for the same number of MLPs as in the "manual interp" condition, but randomly selected.

questions:
(All of my questions were asked in the "weaknesses" section.)
</content>
</thread_object>

</thread>
<thread>
creation date:
1730648532681

thread_type:
official_review

<thread_object>
creator: reviewer_reey

<content>
summary:
This article investigates the performance of various localization methods in unlearning/editing tasks, particularly focusing on their limitations in adapting to shifts in prompting/output distributions and adversarial relearning scenarios. It compares three main approaches: output tracing, attribution patching, and FLU. Through experiments conducted on two models and two datasets, the findings reveal that the component set identified by FLU localization is more closely tied to the factual query process, demonstrating greater robustness and generalization when fine-tuned. Additionally, the authors achieve more efficient parameter editing by controlling model modifications through weight masking.

strengths:
1. The motivation and approach of this article are interesting, as it breaks down factual recall into more granular steps to enhance the accuracy and generalization of editing methods. 

2. The extensive experiments demonstrate the effectiveness of the proposed approach. The findings indicate that fine-tuning the FLU-related components identified through manual localization effectively eliminates specific knowledge from the model and makes it less susceptible to re-learning.

3. Using multiple-choice questions (MCQs) can help eliminate the influence of input patterns while allowing for a more effective exploration of knowledge deletion. This approach can provide clearer insights into how specific knowledge is affected by unlearning processes.

weaknesses:
1. Regarding the editing task, only the decline in the original answers is reported, which aligns with the unlearning aspect. However, results seem to lack a demonstration of improvement in the correctness of the new answers. This is significantly related to the performance of the editing task. 

2. Additionally, for the analysis of intermediate representations using probing, this concept is derived from existing work and does not represent a novel contribution to this research.

3. [Critical] In the unlearning task, there is no theoretical proof or guarantee that the knowledge is fully forgotten. Since approximate unlearning can be easily exploited, this method is vulnerable. I believe that a theoretical guarantee is crucial for the unlearning task because the security issue is fundamentally a "to be or not to be" problem.

4. [Critical] No adaptive attack experiments were conducted. The authors performed only standard unlearning/editing experiments, without testing for membership inference or adaptive attacks, despite the fact that approximate unlearning methods are particularly susceptible to adaptive attacks.

5. [Critical] There is a lack of experiments on adaptive unlearning, which would involve sequentially unlearning specific types of knowledge—for instance, first basketball, then football, and finally table tennis. Would adaptive unlearning impact the efficiency of the unlearning methods?

questions:
1. In the section defining the method, it mentions, "In practice, we fix τ such that Cτ contains the same number of parameters in OT, FLU, and random localizations." How should this statement be understood? For example, in the counterfact dataset, what are the MLP results using OT and EAP, given that our analysis highlights layers 3-5, 7-10, and 14-17 as the critical MLPs?

2. Could you explain in more detail the process of direct path patching on the counterfact dataset? For example, after obtaining the set of components related to the fact extraction mechanism, how do we replace all edges from each MLP to all components in the set?

3. In the manual localization process, why is only the MLP considered as the localization component, while other methods like EAP and OT do not also set the form to only consider MLP? Instead, they assess both attention heads and MLP components simultaneously?

4. [Critical] Can the authors provide theoretical proof or guarantee to show that the knowledge is forgotten?

5. [Critical] Can the authors provide the experiments of adaptive attack (attackers that easily conquer approximate unlearning )?

6. [Critical] Can the authors provide the experiments of sequential unlearning?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730713251663

thread_type:
official_review

<thread_object>
creator: reviewer_dovc

<content>
summary:
This paper investigates the effectiveness of mechanistic interpretability techniques in improving the precision and robustness of knowledge editing and unlearning in LLMs. The authors mainly discuss two types of localization methods, i.e., OT techniques that focus on preserving outputs, and mechanistic localization, which identifies high-level mechanisms with predictable intermediate states. They claim in the paper that localizing edits to components associated with the FLU mechanism leads to more robust unlearning across different input/output formats and resists attempts to relearn unwanted information. They conduct experiments on the Sports Facts and CounterFact dataset using Gemma-7B and Gemma-2-9B models.

strengths:
This paper contains the following several strengths:
+ The paper addresses an important topic by attempting to improve the robustness of knowledge unlearning in LLMs through  localizing edits to components associated with the FLU.
+ The authors provide a in-depth analysis and comparison between mechanistic unlearning and previous OT methods.

weaknesses:
+ The paper would benefit from a more in-depth theoretical analysis to explain why FLU could inherently lead to more robust unlearning. While the authors claim that targeting the fact lookup components is more effective, they do not provide analysis or proof to support this.
+ The experiments are limited, especially limiting itself to Gemma-7B and Gemma-2-9B models and two datasets. The authors could provide a larger variety of models and unlearning tasks in order to better demonstrate the consistency of their findings.
+ Can the author provide more ablation study, for example on the loss weights parameter being used in 2.3, so that we could better understand the contribution of each loss in the finetuning process.

questions:
See above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730721694325

thread_type:
official_review

<thread_object>
creator: reviewer_uzky

<content>
summary:
This work explores methods for knowledge editing and unlearning in large language models, focusing on how mechanistic interpretability can enhance the precision and effectiveness of these processes. The study reveals that localizing edits to components associated with lookup-table mechanisms for factual recall leads to more robust unlearning, resisting unwanted information relearning and minimizing side effects. Additionally, certain localized edits disrupt latent knowledge more effectively than other methods, resulting in increased resilience against various attacks.

strengths:
* Studying unlearning methods from the perspective of knowledge storage and mechanistic interpretability is indeed a very important and promising direction.

* This paper further confirms that causal tracing-based localization methods are not suitable for editing and unlearning tasks.

* The paper is well presented and the literature review is thorough.

* The experimental design is generally comprehensive.

weaknesses:
1. The test dataset appears to be limited to this triplet format; is it constrained by the knowledge format, and could it be applied to more broadly and flexibly expressed knowledge sentences, such as continuous text, etc.?

2. Manually analyzing and then selecting layers for operations seems to lack convenience and flexibility in the context of large-scale data editing/unlearning.

3. The proposed new unlearning method lacks sufficient originality. There are already some works that attempt unlearning directly from the perspective of mechanistic interpretability, including [2].

4. Knowledge is not necessarily stored entirely in the MLP; there are certain cases where it exists in the attention mechanism [1], yet the method described in the paper only considers knowledge stored in the MLP.

5. There is a lack of discussion on unlearning methods in Representation Engineering [3, 4].

6. Can the proposed method achieve performance advantages on other representative series of transformers, such as LLaMA?


---
**References:**

[1] Dissecting Recall of Factual Associations in Auto-Regressive Language Models

[2] Intrinsic Evaluation of Unlearning Using Parametric Knowledge Traces

[3] The WMDP Benchmark: Measuring and Reducing Malicious Use With Unlearning

[4] Improving Alignment and Robustness with Circuit Breakers

questions:
Please see the Weaknesses section above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732738786387

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Comments

comment:
We thank all the reviewers for their constructive comments. We have addressed all concerns raised, resulting in a submission that is not only technically sound with strong empirical results showcasing the role of mechanistic interpretability for localized model editing but also clearly and effectively communicates our contributions.

Biggest Changes:

1) Significantly expanded experiments: 
* New Model: Included LLaMA-3-8B, demonstrating consistency across architectures (now includes Gemma-7B, Gemma-2-9B, and LLaMA-3-8B).
* All Tasks on All Models:  Replicated original results across all three models and tasks, addressing concerns about cherry-picking.
* Tried multiple forget set sizes (16 and 64 facts) for both Sports and CounterFact.
* Added systematic edits of entire knowledge categories in Sports (e.g., changing all basketball players to golfers), for each of the three target sports.
* Introduced a sequential editing task, editing 16 facts at a time in 4 rounds in CounterFact dataset.
* New Baselines: Added a "Random subset of MLPs" baseline and “MLPs from Causal Tracing”, “MLPs from Attribution Patching” OT localizations for more rigorous comparison.
* Moved unlearning results to appendix - the goal of unlearning these basic factual associations is unclear, since a model who has never seen information about the factual association would probably still assign >0 probability to the correct answer on priors especially in Sports. Thus, we decided to focus on editing for our new sweeps.

Importantly, these expanded experiments consistently demonstrate the effectiveness of our approach, yielding strong and positive results across all models, tasks, and evaluation metrics.

2) Included extensive hyperparameter sweeps: Conducted comprehensive sweeps for learning rates and forget loss weights across all models and editing tasks, ensuring optimized performance and fair comparisons.

3) Strengthened adversarial evaluations: 
* Relearning:  Aligned our relearning methodology with [1] for a more standard and challenging evaluation: we relearn the ground truth answer on half of the forget set, and evaluate on the other half of the forget set. FLU localizations are hardest to relearn on.
* Latent knowledge: modified probe training to reflect model’s internal representations of the answer throughout layers, which highlights significant differences in the forget set answers represented by models edited with different localizations. FLU localizations consistently represent the edited answer throughout layers, while other localizations start by representing the original answer in early layers and then switch to representing the edited answer in later layers.
* Adaptive Attacks: Introduced soft prompt evaluations to test resilience against adaptive attacks.

4) Improved readability and presentation: 
* Reorganized and clarified task descriptions in Section 2.1.
* Introduced spider plots and bar charts to clearly visualize key results in Section 3, making the paper much easier to parse.
* Moved less informative experiments to the appendix to reduce clutter and emphasize core findings.

We believe these revisions significantly strengthen the paper and address the reviewers' concerns. We are excited about these new results and strongly believe this work represents a valuable contribution to the field of model editing and unlearning, showcasing the practical benefits of mechanistic interpretability. We urge the reviewers to reconsider their scores in light of these substantial improvements.

[1] Aghyad Deeb and Fabien Roger. (2024). Do unlearning methods remove information from language model weights?
</content>
</thread_object>

</thread>
<thread>
creation date:
1733277127219

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Message to Area Chairs

comment:
We are writing to request a careful review of the feedback and discussion submitted for our paper. While we appreciate the reviewers' time and feedback, we believe some reviews may not accurately reflect the current state of our manuscript **after revisions addressing every initial concern raised**.

Specifically, we highlight the following:

1. *Reviewer vUnN*: This reviewer provided a thorough assessment, acknowledging the **"extremely impressive revision"** which addressed all prior concerns and improved presentation. They rated our paper 8 overall, with 4/4 on both Contribution and Soundness, and even **encouraged other reviewers to reconsider their scores**. We believe this review best reflects the quality of our current submission.
2. *Reviewer uzky*: While initially providing a score of 5, this reviewer **has not responded to our extensive revisions**, which include substantial new experiments in response to their initial concerns. Given that Reviewer vUnN increased their score from 6 to 8 after seeing these revisions, we believe Reviewer uzky's outdated rating no longer applies and should be adjusted upwards accordingly.
3. *Reviewer DovC*: This reviewer acknowledges that we have addressed all their initial concerns with additional explanations and experiments. They explicitly state that our paper "addresses an important topic" and provides "in-depth analysis." However, they increase their score to only 5 without providing justification. **Given the absence of remaining concerns and their positive assessment of our work's significance, we find the low scores on Soundness (2) and Contribution (2) puzzling**, and we believe their rating should also be adjusted upwards.
4. *Reviewer REEY*: While we addressed all initial concerns raised by this reviewer with extensive experiments, they subsequently **introduced new criteria** without ever improving their rating from 3, arguing that our evaluation of adversarial attacks is insufficient without specifically including the computationally expensive GCG attack. This seems unreasonable given our inclusion of even more powerful adaptive attack methodologies, including the similar but more efficient and powerful soft prompt method and the significantly more powerful attack of adversarial few-shot relearning.

Furthermore, Reviewer REEY's subsequent comments bring up a **completely new justification** for why our paper deserves a rating of 3: they refer to the discussion of a different unlearning paper “Intrinsic Evaluation of Unlearning Using Parametric Knowledge Traces”. They argue that inherent limitations in that paper (to do with how it is unclear how to measure unlearning and how causal unlearning is not defined) must also apply to our work, despite our distinct focus on model editing (which in comparison has a much clearer goal for our tasks) and uniquely strong empirical results (the other paper does not present a successful general method for unlearning). This comparison disregards the significant contributions of our work. Given that our work presents more positive results and does not have the limitations of unlearning discussed in that paper, we believe that Reviewer REEY's rating of 3 does not accurately reflect the quality and contributions of our research. We thus believe that Reviewer REEY’s rating should be adjusted upwards.

We have made extensive revisions to address all reviewer concerns, including a number of additional experiments, hyperparameter sweeps, adversarial evaluations, and massively improved presentation (summarized in the general comment in more detail). **Importantly, these expanded experiments consistently demonstrate the effectiveness of our approach, yielding strong and positive results across all models, tasks, and evaluation metrics.**

Our paper tackles the crucial issue of applying mechanistic interpretability to enhance robustness in real-world scenarios, with a particular focus on model editing and unlearning. We believe our work makes a valuable contribution to the field, showcasing the practical benefits of mechanistic interpretability.

We kindly request that you consider these points during the discussion phase and ensure our paper receives a fair and comprehensive evaluation.

Thank you for your time and consideration.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper explores the use of mechanistic interpretability to improve knowledge unlearning and editing in LLMs, with a focus on localizing edits. While it demonstrates robustness gains and improved resistance to relearning unwanted information, reviewers highlighted concerns regarding limited novelty, insufficient evaluation on diverse tasks and models, and the lack of theoretical guarantees as major weaknesses. The authors' response did not fully address these concerns, and I recommend rejection.