Mechanistic Unlearning: Robust Knowledge Unlearning and Editing via Mechanistic Localization

We thank the reviewer for the detailed comments. We have updated the paper based on the feedback, and address weaknesses (W) and questions (Q) below. Due to the character limit, this first comment addresses the weaknesses, a second comment below addresses the remaining comments.

W1 Lack of Edit Accuracy Improvement Reporting: We now report the increase in accuracy relative to the new edited answers (“Edit Accuracy”), along with the original decline of the accuracy relative to the ground truth answers (“Forget Accuracy”). Our localized editing retains strong performance with this metric.

W2 Probing as a Novel Contribution: This critique misunderstands our contribution - while previous work has utilized probing to evaluate editing, our main contribution lies in utilizing these localization tools, including probing, to improve model editing performance. The localization techniques themselves may vary across tasks (and indeed we use a different methodology for the FLU localization on the CounterFact tasks than in Sports facts). Our focus is on the results of restricting editing to specific localizations rather than the localization methods themselves.

W3 (and Q4) Lack of Theoretical Guarantees: Given that theoretical guarantees remain elusive even for simpler editing approaches, this limitation should not prevent publication of empirically validated advances in the field. It is important to acknowledge that achieving guarantees for knowledge removal/editing in complex models like LLMs is very challenging, and typically only possible in very specific scenarios/under very strong assumptions (e.g., convexity, which may be achieved in infinite width limit regime, or under differentially private training, etc.). Indeed, most, if not all, approximate unlearning methods for LLMs do not provide theoretical guarantees. Our work focuses on empirically demonstrating the effectiveness of leveraging mechanistic interpretability for editing, through many adversarial evaluations that successfully attack our baselines and standard editing approaches from the literature.

While we don't offer theoretical guarantees, our experiments showcase several advantages of our approach, such as robustness to different prompting and relearning, improved trade-offs between editing quality and performance. Editing also has applications where empirical validation without guarantees may suffice, and our approach provides a valuable tool for these applications.

W4 (and Q5) On Adaptive Attacks: We now include a new experiment against model-specific optimized soft prompts. In Appendix A.8 we added results for soft prompt evaluations, where we optimize the continuous embeddings at the end of the prompt, to recover the correct answer on half of our forget set. We then evaluate the model’s performance on the other half, with this “soft prompt” in place [1]. Soft prompt evaluations can be considered to be a more narrow form of few-shot finetuning, that is closer to searching for prompts that recover the model’s knowledge. We show positive results for Sports-Athlete-Editing in Gemma-2-9b that FLU localization, Nonlocalized, and All MLPs all result in <40% recovered forget accuracy while some other localizations reach >60% recovered forget accuracy under soft prompts. In other tasks we largely fail to recover significant forget accuracy across localizations/models.

W5 (and Q6) On Sequential Unlearning:
We conducted experiments where we iteratively edited sets of 16 CounterFact facts, ultimately modifying a total of 64 facts. Our results demonstrate that: (1) FLU localizations remain robust: even with sequential editing, FLU consistently outperforms other localization methods in resisting alternative prompting schemes. (2) Potential for improved scaling: We observe a slight increase in robustness when editing sequentially compared to editing all 64 facts at once. This suggests that an adaptive, sequential approach might scale more effectively to larger numbers of facts, though further investigation with optimized hyperparameters is needed. We focused on the CounterFact dataset for sequential editing because, unlike the Sports dataset where relevant information is concentrated in a few early layers, CounterFact knowledge appears to be distributed more broadly across the model (see Appendix A.2.2). This characteristic makes it a more suitable testbed for evaluating the potential benefits of adapting localization strategies throughout the editing process.

We appreciate your feedback, and we hope we’ve improved the paper according to your suggestions. We hope you can review our changes and reevaluate your score in light of our significant improvements to the paper, as Reviewer vUnN did.

If you have any remaining concerns, we’d be glad to address them either through comments or by planning an update to the camera-ready version.

We thank the reviewer for their detailed feedback. We have made substantial revisions in response. Our revisions strengthen both the empirical foundations and accessibility of our work: we now demonstrate consistent results across three different models (Gemma-7B, Gemma-2-9B, and Llama-3-8B), introduce comprehensive evaluation scenarios including sequential editing of up to 64 facts, and add new controlled baselines that validate our mechanistic approach. We've also completely revamped our visualization strategy with intuitive spider and bar plots to better communicate these findings.

Through these substantial revisions - expanded model coverage, comprehensive editing scenarios, controlled baselines, and clearer presentation - we've demonstrated that our approach delivers consistent improvements across a broad range of practical settings. Given the pressing need for reliable knowledge editing techniques in deployed language models, we believe our work represents a significant contribution that warrants publication at this venue.

Weaknesses:

W1 (Missing details): We now elaborate on an overview of the techniques used for the FLU localization in section 2.2, and comment on the specific log(1-p) loss function used for unlearning in A.1 (where we moved unlearning results). We also now consistently use “Fact Lookup localization” throughout the text and graphs instead of “manual interpretability”.

W2 (Following results discussion): We have reformatted Section 2.1 and the results section to be far more readable and introduce consistent naming of the tasks (Sports-Unlearning, Sports-Athlete-Editing, etc).

W3 and W4 (Results readability): Our new visualizations with spider (Figure 2 and Figure 4) and bar plots (Figure 3, Figure 5, Figure 6) attempts to make the key results immediately apparent to readers.

W5 (Counterfact retraining experiment): We moved the relearning result into Figure 2 to avoid clutter with uninformative experiments, and to make clear where FLU outperforms other localization techniques. This reorganization places our strongest result in a more prominent position while reducing visual clutter.

Object level concerns:

1. (Task definitions): We added new comprehensive evaluations demonstrating robustness across both targeted edits and systematic modifications of entire knowledge categories for all models. For editing a constant set of athletes, we now vary the forget set size to be 16 and 64 athletes (Sports-Athlete-Editing), aggregating results over both sizes. Additionally, we edit all athletes who play a certain sport (basketball, baseball, or football) to golf (Full-Sports-Editing), aggregating over all sports. Similarly for counterfact, we vary the forget set size to be 16 and 64 facts (CounterFact-Editing). Additionally, we perform a sequential editing setup where we edit 16 counterfact facts at a time, resulting in 64 total facts edited (CounterFact-Sequential-Editing). We hope this constitutes a reasonable-seeming group of editing setups.
2. (Model and task pairings): We have now tested all tasks with all three models (Gemma-7B, Gemma-2-9B, and Llama-3-8B). We find robustness results translate well across models and have included them in the main text by averaging evaluation results over all model types. We hope this addresses any concerns about cherry-picking, and demonstrates that our results generalize beyond any specific architecture.
3. (Additional baseline): We have added the proposed Random MLPs baseline, as well as two baselines using only MLPs selected by Causal Tracing and MLPs selected by Attribution Patching (with the same number of layers) throughout Appendix A.7. The Random MLPs baseline is competitive with the other strong Nonlocalized and All MLPs baselines, while the Causal Tracing MLPs and Attribution Patching MLPs edits are less robust. The new controlled baselines provide further evidence that our performance gains come from better understanding and localization rather than simply from the type and quantity of components selected.

We thank the reviewer for their positive comments and very concrete proposals for improving the paper, which we’ve followed up on. We address central reviewer concerns about our paper's experimental methodology, particularly focusing on new models and ablation studies for core model components. We respond to individual questions below, and have updated the paper with a number of new results and analyses:

1. Theoretical Analysis of FLU: While a theoretical framework would be valuable, we demonstrate through extensive empirical evaluations that our approach is effective, and note that theoretical guarantees for unlearning in large language models remain an open challenge in the field. Most, if not all, approximate unlearning/editing research focuses on empirical evaluations due to the complexity of analyzing approximate unlearning methods in deep learning. Throughout Section 3, we provide clear empirical proof through many standard and adversarial evaluations across tasks and models that targeting FLU components leads to improved editing. In Section 3.3, we also provide direct evidence of our hypothesis that FLU edits disrupt the latent knowledge more at the early layers prior work has discovered this knowledge to be located in, and in Appendix A.3 and A.5 we provide evidence that baselines and OT edits target non-lookup (extraction) components rather than the FLU components. Theoretical guarantees for unlearning/editing have been primarily achieved in simplified settings (e.g., convex optimization), or by modifying the training procedure with techniques like differential privacy, which often lead to significant performance trade-offs, or by introducing architectural changes (such as some form of modularity and data compartmentalization). Neither of these techniques that yield theoretical guarantees are suitable for mechanistic unlearning analysis and also would not directly apply to this setup where we don’t have particular training data points in mind; finding the right theoretical framework would present a major breakthrough in deep learning theory.
2. Limited Experiments: We added additional experiments on Llama 3 8b, as well as additional tasks (sequential editing). Furthermore, we have evaluated all three models on both tasks, instead of one model per task. These new experiments, included in the updated paper, further validate our findings and demonstrate the generalizability of our approach.
3. Ablation Study: We have now conducted a sweep of the forget loss weights used in Section 2.3 as well as the learning rate, optimizing both parameters for the three localizations of Causal Tracing, Fact Lookup, and No Localization. The results of this sweep are included in Appendix A.6 Hyperparameters, first demonstrating a complete sweep over a task in A.6.1 which provides insight into the contribution of each loss component (learning rate is extremely important but forget/inject loss weight is not), and second providing optimized hyperparameters across models, tasks, and localizations in A.6.2.

Given these substantial additions - new empirical results across model families, and comprehensive sweeps and ablation studies - we hope that the reviewer can confidently increase the score to a strong accept.

We appreciate your feedback, and we hope we’ve improved the paper according to your suggestions. We hope you can review our changes and reevaluate your score in light of our significant improvements to the paper, as Reviewer vUnN did.

If you have any remaining concerns, we’d be glad to address them either through comments or by planning an update to the camera-ready version.

We thank the reviewer for the suggestions.

We address six key reviewer concerns, including our triplet format dataset choice, manual analysis limitations, and originality. We strengthen our argument by running additional LLaMA experiments, while maintaining our core finding that targeting MLPs enables robust knowledge editing despite knowledge potentially being distributed elsewhere in the model.

1. Triplet format dataset: Our novel contribution lies in the mechanistic interpretability techniques and editing approach itself, which can be readily extended beyond triplet data once appropriate localization methods are developed, making this limitation primarily an implementation detail rather than a conceptual barrier. In the paper, the triplet format dataset was chosen because: (1) it aligns with prior work on mechanistic interpretability, allowing us to directly leverage their localization techniques, and (2) the structured format facilitates cleaner analysis of knowledge editing and unlearning. However, we emphasize again that our proposed editing approach is not restricted to triplet data. It can be applied to any dataset where localization of relevant model components is possible.

2. Manual Analysis: We argue that this is primarily an engineering limitation rather than a fundamental barrier, since our method can be automated by existing techniques like probing and edge patching, and the results show our approach works robustly even with this current semi-manual pipeline. Our technique for sports facts is a simple probe across MLPs and for CounterFact involves a series of patches on the model's edges (between the MLPs and the extraction heads). We set up the interpretability technique once and for all experiments in the same task, we simply run the interpretability technique automatically. The efficacy of unlearning could even be used as a method to evaluate automated interpretability methods, which are known to have a lack of ground truth.

3. Originality: We appreciate the reviewer pointing out the connection to [2]. We have now included this work in our related work section and clarified the distinctions between our approaches. While both works utilize mechanistic interpretability, we differ fundamentally in our results/findings, localization strategies, and unlearning targets. [2] focuses on identifying and manipulating "concept vectors" within MLPs, while our method directly targets the mechanistic components responsible for knowledge integration. Furthermore, their NPO+KL method, which is most similar to ours, showed limited effectiveness in their experiments. As explained in their section 4.3, their successful Needle method is an oracle skyline, since they evaluate all methods on a split specifically chosen for high Needle performance, whereas our method succeeds without such curated splits. We believe our focus on directly modifying the "fact lookup enrichment" stage via mechanistic interpretability contributes a novel perspective to the unlearning literature, and most importantly demonstrates substantial robustness improvements in a fair comparison.

4. Knowledge Storage Beyond MLPs: The reviewer rightly points out that knowledge is not solely stored in MLPs and can also reside in attention mechanisms, as highlighted in [1]. We focused on MLPs based on prior work suggesting their primary role in enriching factual knowledge [1]. Our results demonstrate that targeting MLPs is sufficient to achieve significant robustness gains. However, we agree that incorporating attention mechanisms in the FLU localization could potentially further enhance editing effectiveness, while increasing complexity of the method.

5. Lack of Discussion on Representation Engineering: We thank the reviewer for bringing this to our attention. We have now included a discussion of relevant unlearning methods in Representation Engineering, specifically addressing works like [3, 4] that focus on mitigating malicious use and improving model alignment through unlearning.

6. Additional Experiments: We ran experiments on Llama 3 8b, and have incorporated the results into the paper. The new Llama experiments validate our original claims, making this a strong addition that reinforces rather than modifies our paper's conclusions. We also pair every editing task with every model and average results over models, as opposed to our previous setup of pairing models with individual tasks.

We believe the additional results and updates to the paper address the reviewer's concerns, and will result in a strong accept. We would appreciate an opportunity to respond if any concerns remain.

We appreciate your feedback, and we hope we’ve improved the paper according to your suggestions. We hope you can review our changes and reevaluate your score in light of our significant improvements to the paper, as Reviewer vUnN did.

If you have any remaining concerns, we’d be glad to address them either through comments or by planning an update to the camera-ready version.