Revealing and Mitigating Over-Attention in Knowledge Editing

TL;DR

This work identifies a core failure mode in knowledge editing for large language models: Specificity Failure, where edited facts hijack attention to the edited entities and distort related knowledge when the edited subject appears in context. It pinpoints Attention Drift as the main trigger and demonstrates that patching or constraining drift—via the Selective Attention Drift Restriction (SADR) regularization—substantially mitigates this failure across multiple models and editing methods, with only a modest cost to edit efficacy. The authors provide a formal framework, extensive experiments on CounterFactrome and WikiData_counterfact, and ablations showing that targeting specific attention heads yields better trade-offs than blanket constraints. The approach advances reliable, safe knowledge editing for real-world deployment by preserving generalization while stabilizing context-sensitive behavior. Collectively, the work offers a practical, architecture-compatible method to improve the fidelity of edited knowledge in transformer-based models, with broad implications for model maintenance and safety.

Abstract

Large Language Models have demonstrated superior performance across a wide range of tasks, but they still exhibit undesirable errors due to incorrect knowledge learned from the training data. To avoid this, knowledge editing methods emerged to precisely edit the specific model knowledge via efficiently modifying a very small percentage of parameters. % However, those methods can lead to the problem of Specificity Failure: when the content related to the edited knowledge occurs in the context, it can inadvertently corrupt other pre-existing knowledge. However, those methods can lead to the problem of Specificity Failure, where the existing knowledge and capabilities are severely degraded due to editing. Our preliminary indicates that Specificity Failure primarily stems from the model's attention heads assigning excessive attention scores to entities related to the edited knowledge, thereby unduly focusing on specific snippets within the context, which we denote as the Attention Drift phenomenon. To mitigate such Attention Drift issue, we introduce a simple yet effective method Selective Attention Drift Restriction}(SADR), which introduces an additional regularization term during the knowledge editing process to restrict changes in the attention weight distribution, thereby preventing undue focus on the edited entity. Experiments on five frequently used strong LLMs demonstrate the effectiveness of our method, where SADR can significantly mitigate Specificity Failure in the predominant knowledge editing tasks.

Figures (13)

• Figure 1: An illustration of counterfactual knowledge editing, where the new factual association (Eiffel Tower, is located in, New York) is edited in GPT-J-6b using the ROME method (Meng et al., 2022). (a) The hidden states of the subject are enriched by the MLP with relevant information and are successfully retrieved by the attention modules. (b) The editing method modifies the MLP parameters to alter the factual association. (c) The edited MLP generates hidden states that are prone to being mistakenly focused on by the attention modules, leading to specificity failure.
• Figure 2: Example of evaluation tasks for Knowledge Editing.
• Figure 3: Visualizing "Contaminating Substitution" when replacing the output of specific modules on test prompts of the Distract Neighborhood task with window size 6.
• Figure 4: The correlation between the attention weight drift and $P(o_{edit})$ is positive on Distract Neighborhood ($\rho = 0.49; p<$1.0e-5) and Relation ($\rho = 0.62; p<$1.0e-5) tasks. The range of $l$ corresponds to the layers where attention activations have a significant impact on the final results.
• Figure 5: Impact of patching attention weights within a window size of 10 on the specificity tasks in the edited model. The performance of the original edited model is represented on the 0-th layer.