Mechanistic Unlearning: Robust Knowledge Unlearning and Editing via Mechanistic Localization

TL;DR

This work tackles the challenge of editing and unlearning undesirable factual knowledge in large language models without sacrificing general performance. It contrasts two mechanistic localization strategies: output-tracing (OT) and fact-lookup (FLU), demonstrating that FLU-focused edits yield superior robustness to prompt variations and adversarial relearning across multiple datasets and models. Through latent-probing and weight-masking analyses, the authors show that FLU targets the latent sources of factual knowledge and achieves edits with greater parameter efficiency, while OT-based localizations tend to disrupt extraction mechanisms and are more vulnerable to relearning. The findings suggest a practical, interpretable framework for robust knowledge editing in LLMs and highlight the importance of targeting the factual lookup stage to disrupt latent knowledge rather than merely altering outputs.

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 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.

Figures (47)

• Figure 1: High level depiction of mechanistic unlearning. We localize components responsible for fact extraction/enrichment and modify their weights to change the associations, in order to target internal latent representations rather than targeting the output. Graph inspired by nanda2023factfinding.
• Figure 2: Spider plots illustrating the advantages of FLU for editing Sports across adversarial prompting and relearning evaluations, averaged over all three model types. (Left) The Sports-Athlete-Editing plot shows that FLU localization leads to editing that is the most robust against MCQ prompting and relearning. (Right) The plot shows that most localizations perform approximately equivalently in the Full-Sports-Editing task, with FLU localization slightly better for MCQ.
• Figure 3: Bar charts showing results of MCQ evaluations, reporting both the forget error and edit accuracy when prompted with MCQ, averaged over three model types. For both (a) Sports-Athlete-Editing and (b) Full-Sports-Editing, FLU localization answers with the original answer the least (MCQ Forget Error) and answers with the edited answer most accurately (MCQ Edit Accuracy).
• Figure 4: Spider plots illustrating the advantages of FLU for editing CounterFact across prompting evaluations, averaged over all three model types. (Left) The CounterFact-Editing plot shows that FLU localization leads to editing that is the most robust against MCQ prompting and Paraphrasing. (Right) The Sequential-CounterFact-Editing plot shows that FLU localization is the most robust against MCQ prompting.
• Figure 5: Bar charts showing results of MCQ, Paraphrase, and Neighborhood prompt evaluations, averaged over all three model types. For both (a) CounterFact-Editing and (b) Sequential-CounterFact-Editing, FLU localization has the most robust edit accuracy measured by MCQ and Paraphrase. FLU localization editing also does not incorrectly generalize to Neighborhood prompts. Sequential editing is slightly more robust than nonsequential editing in MCQ when comparing between CounterFact-Editing and Sequential-CounterFact-Editing bars.