Revealing and Mitigating Over-Attention in Knowledge Editing

Thank you for your detailed review and valuable comments. Here are our responses to your concerns:

• W1: Why is KL divergence used to constrain the two attention weights?

Attention weights are also commonly referred to as attention distributions, representing the model’s allocation of attention across different token positions, and they sum to 1. KL divergence is a widely used metric for measuring differences between two probability distributions and has been utilized in prior works [1,2] to quantify differences in attention weights. Therefore, using KL divergence to constrain two attention weights is a natural choice and aligns with the findings in Section 3.3 of our experiments, which show that attention drift, as measured by KL divergence, is positively correlated with specificity failure.

[1] Interpreting Self-Attention Weights [2] Interpreting Attention Models with Human Visual Attention in Machine Reading Comprehension

• W2: The decline in generalization.

First, PM measures the probability assigned to the target word of the output. In the unedited model, even when the model correctly knows the knowledge (e.g., assigns the highest probability to the true object), the average probability of outputting that word is only around 10%. At the same time, the model may predict other linguistically plausible tokens, such as articles or alternative phrasings, as part of the natural language modeling process. While after editing, the probability assigned to the edited object often becomes disproportionately high, which might indicate potential overfitting. Under some setups, such as GPT-J-ROME, while PM decreases by as much as 20 points, the PS metric only drops by approximately 3 points. This demonstrates that under paraphrase tasks, the loss of knowledge in our methods remains minimal.

Regarding the decline in generalization, it is important to note that prior evaluation systems often overlooked specificity failure, achieving nearly 100% generalization at the cost of severe overfitting. What we aim to accomplish is to ensure that the knowledge retrieval process after editing remains as safe and as close to the original Transformer’s knowledge extraction mechanism as possible. We also believe that a stable knowledge editing method is more critical than achieving near 100% generalization accuracy.

• W3: Ablation on head selection

In Figure 6, we report the Edit Success and Specificity performance across 1,683 editing instances under five different $\gamma$ parameters $\gamma=50, 100, 200, 400, 800$. The p-value analysis further demonstrates that the attention head selection method achieves statistically significant improvements, as shown below:

We also evaluated the performance of randomly selecting the same number of attention heads as used in SADR. Below are the Edit Success and Specificity scores:

Edit success:

Specificity:

These results show that while generalization performance is comparable between random head selection and SADR, SADR demonstrates a consistent advantage in specificity across all $\gamma$ values (p-value less than 0.05).

Thank you for the additional experiments and explanations. I've raised my score.

• W4: Lack of efficiency analysis

In terms of memory usage, the additional variables to store in our method are the attention weights across all layers. These weights can be represented as $L \times H \times S^2$, where $L$ is the number of layers in the model, $H$ is the number of attention heads, and $S$ is the sequence length. The additional storage required is minimal compared to the overall model parameters. During our experiments, we did not observe any noticeable increase in GPU memory usage.

Regarding runtime, our method primarily involves computing a mask through comparison of attention weights and calculating the KL divergence. However, due to the use of Python loops in our current implementation, a slight runtime overhead is observed. For instance, when applying the ROME editing method to GPT-J-6B on an A100-PCIE-40GB GPU, the runtime per edit increased from 7.80 seconds (without SADR) to 9.65 seconds (with SADR).

In the revised version, we have included the efficiency analysis in Appendix F.

• Q3: Restricting attention weights of only one or a few layers?

In our early experiments, we have already evaluated the impact of restricting attention weights in different layers. In fact, the high_attn_range argument in the argparse configuration of our submitted code is designed to control which layers' attention weights are constrained. Our findings indicate that restricting only a subset of layers does not yield better results. This is primarily because over-attention occurs across different layers. Therefore, our approach identifies the attention heads exhibiting over-attention across all layers and applies constraints specifically to those heads.

Thanks for your constructive feedback on our paper. Our response to your questions is as follows:

• W1: Can SADR effectively solve specificity failure on multi-hop reasoning tasks?

We appreciate the suggestion to conduct additional experiments on multi-hop reasoning tasks. We conduct experiments on MQuake, and the results (Multi-hop reasoning score, denoted as MS) are as follows:

Our method indeed alleviates specificity failure in multi-hop reasoning to some extent. However, after analyzing the MQuake dataset, we find that most samples exhibit the following pattern: New fact: <baseball, created in, Japan>. Multi-hop question: Which political leader governs the country of origin of Mike Krukow's sport? In such cases, the edited subject is not directly mentioned in the question, limiting the impact of over-attention on failure cases. Instead, most failures arise from the inability to fully incorporate the new facts into the model’s knowledge, which is a problem of generalization, as highlighted in Section 4.2 in MQuake.

Additionally, we find a recent research[1] (released after the ICLR submission deadline) provides a multi-hop dataset where the subject is explicitly mentioned. This work attributes the failure of multi-hop questions to the edited model's over-attention on the subject. We test our method on this dataset, and the results also show a significant improvement:

These results further validate the efficacy of our approach in mitigating specificity failure under multi-hop reasoning scenarios where over-attention on the subject plays a critical role.

[1] Zhang M, Ye X, Liu Q, et al. Uncovering Overfitting in Large Language Model Editing[J]. arXiv preprint arXiv:2410.07819, 2024.

• W2,3 & Missing References

Thank you for pointing out the typos and reminding us of the missing references. In the revised version, we have corrected these issues.

• Q1: The reasons that lead to attention drift.

This is a good question. In traditional editing methods (e.g., ROME discussed in Section 2.2), the optimization objective explicitly trains the model to predict the new $o_{\text{edit}}$ given $(s, r)$. This process can unintentionally shape the subject’s hidden state in a way that makes it disproportionately prioritized by the attention mechanism, creating a shortcut. As a result, whenever the edited subject appears, the model overemphasizes $o_{\text{edit}}$, achieving the optimization goal but causing specificity failure.

To further investigate whether factors such as the norm of the hidden layer vector or the distance between hidden state vectors pre- and post-editing contribute to attention drift, we conducted a correlation analysis. Specifically, we examined the relationships between hidden state norm post-editing, L2 distance between hidden states pre- and post-editing, and cosine similarity of hidden states pre- and post-editing with attention drift. The results are as follows:

The results suggest that the shift or norm of the hidden state vector is weakly correlated with attention drift. In fact, the implemention of ROME already constrains the shift in hidden state vectors during optimization by introducing the clamp_norm_factor. The primary cause of attention drift likely lies in the optimization objective, which hard-codes the knowledge into the model’s forward propagation rather than enabling a more natural and reasonable assimilation of new knowledge.

-Q2: Additional runtime overhead of SADR.

Regarding runtime, our method only involves comparing attention weights to create a mask and calculating the KL divergence. However, due to the use of Python loops in our implementation, there is a slight increase in runtime. For instance, on an A100-PCIE-40GB GPU, when applying the ROME editing method on GPT-J-6B, the time required for each edit with and without SADR was 9.65 seconds and 7.80 seconds, respectively. In the revised version, we have included the efficiency analysis in Appendix F.

Tanks for your reply! Some of my concern has been addressed. I wish I could raise the score by one point, but it is very unfortunate that there is no option for 7 in ICLR. However, this does not affect my belief that this is a good paper.

Also, here are some additional responses:

The results suggest that the shift or norm of the hidden state vector is weakly correlated with attention drift. In fact, the implemention of ROME already constrains the shift in hidden state vectors during optimization by introducing the clamp_norm_factor.

In a previous reply, I pointed out that the norm of the model might be one reason affecting the performance. This is because I have conducted experiments in the past and found that: with the increase in the number of edits, although the clamp_norm_factor is used, the norm of the model will inevitably become larger, some existing results can also corroborate this view [1, 2].

The primary cause of attention drift likely lies in the optimization objective, which hard-codes the knowledge into the model’s forward propagation rather than enabling a more natural and reasonable assimilation of new knowledge.

I hope that the author can validate this view in the final version.

In summary, I like your views.

$Ref$:

[1] Model Editing Harms General Abilities of Large Language Models: Regularization to the Rescue. (ACL 2024)

[2] Model Editing at Scale leads to Gradual and Catastrophic Forgetting. (EMNLP 2024)

Thank you for your response! My doubts have been basically resolved! Finally, perhaps the authors can further validate the results with a pearson coefficient between attention drift and editing difficulty, just as you did before.

Tanks for your reply! Finally, I decide to raise my score! I believe this paper will have a significant impact on this field.

• W3: Comparison between the attention-based and decoding-based constraint methods.

DeCK [1] primarily encourages the model to output words with greater differences from the original distribution by employing contrastive decoding, thereby increasing the model's confidence in the edited facts. However, DeCK will increase the model's confidence in its outputs, which leads to more severe specificity failure. Therefore, the formula needs to be modified to constrain the divergence between the model’s output and the original distribution. By modifying the Equation 6 in [1] as follows:

$\mathcal F(\mathbb P_{\text{edited}}(x), \mathbb P_{\text{unedited}}(x)) = \log \mathbb P_{\text{edited}}(x) - 0.5 \log \left(\frac{\mathbb P_{\text{edited}}(x)}{\mathbb P_{\text{unedited}}(x)}\right),$

we can constrain the difference between the edited and unedited distributions during decoding. We refer to this approach as constrained decoding, and the results are as follows:

Although this approach significantly improves the metrics for the Relation task, its performance on all other tasks is decreased. Using decoding methods alone cannot fundamentally address the model's specificity failure and may also significantly harm the success rate of edits.

$Ref:$

[1] Decoding by Contrasting Knowledge: Enhancing LLMs’ Confidence on Edited Facts.

• W2: Analysis on the reason for sepecificity failure.

This is indeed a very meaningful and important question for our work. First, we would like to emphasize that the specificity problem we investigate in this paper refers to cases where the model’s ability is negatively affected when content related to the edited knowledge appears in the context (as noted in lines 44–46).

Indeed, the edit vector's direction, space, and norm [1,2,3] can influence the model's specificity performance. However, the referenced works primarily focus on preserving general knowledge and capabilities, rather than addressing the specificity failure that arises when the edited subject appears in the context. To explore the relevance of these factors to the specificity failure problem studied in our work, we conducted a correlation analysis. Specifically, we compared four factors—attention drift, hidden state norm post-editing, L2 distance between hidden states pre- and post-editing, and the cosine similarity of hidden states pre- and post-editing—with the probability of $P(o_{\text{edit}})$ in specificity tasks.

The results show that, compared to the direction or norm of the edit vector, attention drift has a more direct and significant impact on specificity failure. We have also included this experiment in Appendix D.3 of the revised version.

Intuitively, the attention mechanism is likely a key factor contributing to specificity failure. Previous studies have shown that when language models recall factual associations, the attention mechanism extracts answers from the hidden states of the subject[4,5]. Editing methods primarily modify the hidden states of the edited subject, which then influence the final output through the attention mechanism. In traditional editing methods (e.g., ROME discussed in Section 2.2), the optimization objective explicitly trains the model to predict the new $o_{\text{edit}}$ given $(s, r)$. This may create a shortcut, where the subject’s hidden state is shaped in a way that makes it prone to being overly prioritized by the attention mechanism. Consequently, whenever the edited subject appears, the model disproportionately outputs $o_{\text{edit}}$, satisfying the optimization objective while inadvertently causing specificity failure.

Experimentally, as demonstrated in Section 3.3, we show that attention weights are a necessary condition for specificity failure. By replacing only the post-editing attention weights with the pre-editing attention weights—while keeping all other components unchanged (e.g., MLP outputs)—we observe a significant reduction in the probability of incorrect answers and a corresponding increase in the probability of correct ones during specificity tasks. This result strongly suggests that attention drift is a primary driver and a necessary cause of specificity failure.

$Ref:$

[1] AlphaEdit: Null-Space Constrained Knowledge Editing for Language Models

[2] Knowledge Circuits in Pretrained Transformers

[3] Perturbation-Restrained Sequential Model Editing

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

[5] Locating and Editing Factual Associations in GPT

Thank you for your valuable feedback and constructive comments. We provide the following responses to address your concerns:

• W1: Performance drop in the generalization.

First, we would like to clarify that our method is designed to constrain the model's excessive focus when the attention mechanism over-focuses on the subject compared to its normal levels, rather than to make the model directly focus less on the subject. In Equation (2), $H_l(S_j)$ explicitly defines that the constraint is applied only when the attention exceeds the maximum attention of the original model.

An ideal knowledge editing process should not drastically alter the Transformer’s mechanism for knowledge extraction (e.g., the norm of attention weights on specific words). Changing the location of the Eiffel Tower, for instance, should not lead the model to over-focus on the Eiffel Tower while neglecting other contexts (given that attention weights sum to 1). For this reason, imposing restrictions only when the attention weights surpass their normal levels is both necessary and beneficial, as this deviation can cause side effects like specificity failure (discussed in Section 3).

For the slight drop in the generalization metric caused by the SADR method (less than 3%), it is important to note that prior evaluation systems often ignored specificity failure, achieving near 100% generalization at the cost of severe overfitting. What we aim to accomplish is to ensure that the knowledge retrieval process after editing remains as safe and close to the original Transformer’s knowledge extraction mechanism as possible. We also believe a stable knowledge editing method is more critical than achieving near 100% generalization accuracy.

Thanks for the author's response, and they have dealt with some of my concerns, but I'm still not convinced of the generalization failure and I think this is the limitation. Anyway, the author's response makes sense; and it is a good work. I raised my score.

We greatly appreciate your insightful feedback and constructive suggestions, and thank you for your recognition of our work. Here are our responses to your concerns:

• Weakness: Why the experiments are limited to locate-then-edit methods? Typo line 47: Paris

Due to paper length constraints, we focuse our analysis of specificity failure and experiments with SADR on locate-then-edit methods to maintain consistency and clarity. We choose the locate-then-edit approach in the main text because it is a mainstream method in knowledge editing, offering state-of-the-art performance across many benchmarks with low computational demands. To illustrate the generalizability of the specificity issue and the proposed SADR method, we conduct extended experiments across more editing methods and datasets in the appendix.

The typo issue has been corrected in the revised version of the paper.

• Questions: Have parameter preserving or meta-learning methods also been investigated? What's the RS/RM and DNS/DNM scores for methods like GRACE or ICE? Adding the scores for MEMIT and PMET to Table 1.

In Appendix E2, we provide metrics for WISE and MEND, which represent these two categories of methods. Our analysis shows that specificity failure remains a significant issue across these approaches, as evidenced by their RS and DNS scores compared to the original model.

We have also evaluated the GRACE method on our tasks, with the results summarized below:

The results indicate that GRACE does not perform well on the paraphrase task in the Counterfact editing benchmark. For specificity tasks, GRACE shows some success in the Distracting Neighborhood scenario (likely due to its limited generalization capability) but still suffers from substantial overfitting in the Relation task. We acknowledge the importance of further exploring methods like ICE and plan to include them in future investigations.

Additionally, to provide a more comprehensive illustration of specificity failure, we have added the scores for MEMIT and PMET to Table 1.