On the Role of Attention Heads in Large Language Model Safety

Authors: Zhenhong Zhou, Haiyang Yu, Xinghua Zhang, Rongwu Xu, Fei Huang, Kun Wang, Yang Liu, Junfeng Fang, Yongbin Li
Primary area: alignment, fairness, safety, privacy, and societal considerations
Keywords: interpretability, large language model, multi-head attention, safety, harmful content

Abstract:
Large language models (LLMs) achieve state-of-the-art performance on multiple language tasks, yet their safety guardrails can be circumvented, leading to harmful generations. In light of this, recent research on safety mechanisms has emerged, revealing that when safety representations or component are suppressed, the safety capability of LLMs are compromised. However, existing research tends to overlook the safety impact of multi-head attention mechanisms, despite their crucial role in various model functionalities. Hence, in this paper, we aim to explore the connection between standard attention mechanisms and safety capability to fill this gap in the safety-related mechanistic interpretability. We propose an novel metric which tailored for multi-head attention, the Safety Head ImPortant Score (Ships), to assess the individual heads' contributions to model safety. Base on this, we generalize Ships to the dataset level and further introduce the Safety Attention Head AttRibution Algorithm (Sahara) to attribute the critical safety attention heads inside the model. Our findings show that special attention head has a significant impact on safety. Ablating a single safety head allows aligned model (e.g., Llama-2-7b-chat) to respond to **16$\times\uparrow$** more harmful queries, while only modifying **0.006\%** $\downarrow$ of the parameters, in contrast to the $\sim$ **5\%** modification required in previous studies. More importantly, we demonstrate that attention heads primarily function as feature extractors for safety and models fine-tuned from the same base model exhibit overlapping safety heads through comprehensive experiments. Together, our attribution approach and findings provide a novel perspective for unpacking the black box of safety mechanisms in large models.

Decision: accept (oral)
Reviewer scores: 6, 8, 6, 8

Reviews and discussion:
<thread>
creation date:
1730446212185

thread_type:
official_review

<thread_object>
creator: reviewer_ugyt

<content>
summary:
This work makes an essential contribution to the field of LLM interpretability and safety, especially given the increasing deployment of LLMs in sensitive areas where maintaining safety is paramount. By focusing on the attention heads, the authors target a critical yet underutilized aspect of LLM architecture. The Ships metric and Sahara algorithm also address interpretability efficiently, which is crucial for scaling such safety measures in larger models.

The results are promising; however, further research could examine whether these safety heads are consistent across more diverse LLM architectures beyond the ones tested (Llama-2-7b-chat and Vicuna-7b-v1.5). Additionally, the analysis of the trade-offs between safety and model helpfulness could be expanded to understand better how balancing these factors could affect real-world applications.

Overall, this paper offers valuable insights into safety interpretability and contributes a more resource-efficient framework for analyzing and enhancing LLM safety capabilities.

strengths:
-Great presentation with easy to understand figures and well written concise mathematical concepts/methodology are key strenghts.
-The need of only 0.006% of parameter modification to achieve a SOTA ASR is also impressive.
-I found also interesting that the ablation tests showed clearly the contribution of safety from each head, with only a few ones being critical.
-The Sahara algorithm could be relevant in verifying the safety of LLMs.

weaknesses:
Further research could examine whether these safety heads are consistent across more diverse LLM architectures beyond the ones tested (Llama-2-7b-chat and Vicuna-7b-v1.5). Additionally, the analysis of the trade-offs between safety and model helpfulness could be expanded to understand better how balancing these factors could affect real-world applications.

questions:
How could one deploy the Ships and Sahara during the development process?
Which other use cases would be possible?
Robustness, related to attacks and otherwise, are often related to OOD samples. Could you posit the results of the paper in relation to OOD? And in this case, are the number of models and datasets used in the current study sufficient to evaluate it?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730671989884

thread_type:
official_review

<thread_object>
creator: reviewer_wkxw

<content>
summary:
The paper proposes the Safety Head Important Score (Ships) to evaluate the contribution of individual attention heads to the safety of language models. It presents the Safety Attention Head AttRibution Algorithm (Sahara), which identifies groups of attention heads that contribute to safety. The authors demonstrate through experiments that certain attention heads are crucial for safety, showing how their ablation can significantly increase the model's susceptibility to harmful queries.

strengths:
1. The paper introduces a novel approach to understanding the safety mechanisms within large language models (LLMs) by presenting the Safety Head Important Score (Ships) and the Safety Attention Head AttRibution Algorithm (Sahara).
2. It effectively shifts the focus from generalized model parameters to specific attention heads that have a direct impact on the model's ability to reject harmful queries.
3. This work addresses a significant gap in the literature by systematically exploring the role of multi-head attention mechanisms in ensuring model safety, an area that has been relatively underexplored.
4. The discussions in Sections 4 and 5 are comprehensive and logically coherent, incorporating several insightful analyses and discussions.

weaknesses:
1. Lines 256 and Appendix B.3 indicate that the ASR metric used in this paper employs a keyword-detection method, which is noted in [1] as having limitations that “lead to false positive and false negative cases.” Why is the GPT4-judge method, validated in [1] as a more comprehensive and accurate metric, not used? This method is commonly employed in 2024 LLM safety papers to measure ASR. The inaccuracies of ASR based on keyword detection in assessing successful attacks weaken the experimental data and analysis presented in the paper.
2. The analysis of Figure 2 is relatively weak (Lines 261-267). For instance, according to Figure 2, the improvement in ASR for Vicuna on Advbench and Jailbreakbench (direct) is quite limited, yet these two datasets are highly mainstream in the field of LLM safety. Does this not weaken the conclusions drawn in Lines 262-263?
3. There are still some writing issues in the paper, such as the figure numbering in Line 496, which should refer to Figure 6.
4. Is Figure 6b derived from Llama2 or Vicuna? I couldn't find this information in the paper. Additionally, Figure 6 shows an average decline of about 0.1 in Zero-Shot Task Scores, indicating a 15% decrease compared with the vanilla model. Given that the ASR determination in the paper only detects refusal keywords, I am concerned that the decline in helpfulness may reduce the model's understanding of harmful queries, potentially leading to responses that are affirmative but do not align with the harmful query. This could be counted towards the ASR, further undermining the discussion surrounding Figure 2.

[1] Fine-tuning aligned language models compromises safety, even when users do not intend to!

questions:
1. Could you consider adding a new ASR metric, such as the GPT4-judge referenced in [1]? Alternatively, an analysis of the accuracy of the ASR based on keyword detection (in relation to human judgments) could be included.
2. In light of the second weakness, would it be possible to incorporate additional models, such as Gemma (instruct), to enhance the generalizability of the conclusion that “ablating the attention head with the highest Ships score significantly reduces safety capability”?

[1] Fine-tuning aligned language models compromises safety, even when users do not intend to!

I will reconsider my score if all these issues are adequately addressed.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730721459910

thread_type:
official_review

<thread_object>
creator: reviewer_zypy

<content>
summary:
The authors study how heads contribute to LLM safety, and find that by modifying less than 0.006% of the parameters (a significantly smaller number of parameters than before), safety alignment is degraded significantly. The authors propose Safety Head ImPortant Score (Ships) which measures individual heads’ contributions to model safety. On the dataset level, the authors propose an algorithm Sahara that iteratively choose important heads, creating a group of heads that degrades safety performance for a dataset. The method is efficient in compute hours needed, and impact safety at greater granularity than before. The authors study the effect of safety attention heads, including experiments of concatenating heads to a pre-trained model, and observes that safety capability is close to that of aligned model.

strengths:
- The paper proposes a novel method for mechanistically locating and ablating heads that are important to safety alignment, with greater granularity and less compute than prior methods. The method of head ablation is well motivated and the experiments are detailed, considering course correction (reverting back to safety) as well.
- The paper is well written and organized

weaknesses:
1. The helpfulness / utility measurement is done with lm-eval, which mostly consists of single-turn question and answering utility measurement. More comprehensive utility measurement would benefit the paper. 
2. Sahara uses heuristic to choose group size, and group size is important to how safety capability is affected. Such size heuristics (more than 3) might not hold for different models with different number of parameters. 
3. The paper is overall well-written with some small typos:  "Bottom. Results of attributing attributing specific harmful queries using Ships" (line 445). Consider changing the color for axis label for Figure 13. It's currently quite hard to read.

questions:
Q1. What lm-eval helpful tasks were experimented on for Figure 6b?

Q2: Why do you think there is minimal overlap between the top-10 attention heads via UA ablation and SC ablation (line 477-478)? In combination with results from Appendix E, does this suggest SC is not a good ablation method?

Q3: What do you think contribute to safety capability improve when ablating a small head group between 1 and 3 (line 371)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730822711873

thread_type:
official_review

<thread_object>
creator: reviewer_beky

<content>
summary:
This paper introduces a method named Ships, which assesses each attention head’s contribution to LLM safety at both the query and dataset levels. It also proposes a heuristic search-based method, Sahara, to identify a group of safety-related attention heads. Ablating these heads led to a notable increase in ASR on Llama-2-7b-chat and Vicuna-7b-v1.5. The approach is claimed to be more parameter and computing efficient than previous methods.

strengths:
1. To my knowledge, no other work has attempted to interpret each attention head’s contributions to LLM safety. I'm glad to see such work.

2. The method for assessing the importance of attention heads to safety is intuitive and reasonable.

weaknesses:
1. Other previous works, such as those identifying safety-related parameters through probing [1], could also be discussed.

2. There are inconsistencies in notation usage that need to be addressed. To name a few:
- In Eq 2, 7 and 8, $d_k$ denotes the model dimension. However in line 297, $d$ is used instead.
- In Appendix A.1, $N = d/n$ should be clarified.
- Throughout the paper, $L$ and $n$ are used to denote the number of layers and the number of heads, respectively, but Algorithm 1 uses $\mathbb{L}$ and $\mathbb{N}$ for the same.
- Equation 5 defines parameter importance on safety as $\Delta p = p (\theta_O) - p (\theta_O \backslash \theta_c)$. However, $\Delta p$ is never mentioned again. Instead, Equation 9 uses KL divergence between two probabilities. I guess $\Delta p$ means to represent the difference between two probabilities computed by some function.
- Typo: Eq.6, $h_i^m$ → $h_i$

3. The impact of ablating safety-related heads on model utility could be evaluated.

4. Additional clarifications needed:
- In Table 3, are the results at the bottom obtained by, using Ships on each query to identify and ablate the most important head, then assessing whether the attack was successful for that query? If so, why does query-level Ships perform worse than dataset-level Ships?
- Line 264 notes that 0.006% of all parameters corresponds to the number of parameters in a single attention head. However, Section 4 suggests that ablating 0.006% of parameters (potentially one head) can achieve an ASR of approximately 0.72 on Llama. Yet, Figure 4(a) indicates that ablating at least two heads is necessary to surpass an ASR of 0.7. What is the group size used in Table 1?

[1] A Mechanistic Understanding of Alignment Algorithms: A Case Study on DPO and Toxicity. ICML 2024.

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

</thread>


Metareview:
metareview:
This paper introduces SHIPS, a method for interpreting the contribution of individual attention heads to LLM safety. It proposes the Safety Attention Head AttRibution Algorithm (SAHARA), utilizing Undifferentiated Attention and Scaling Contribution to achieve this.

The paper presents a novel approach to analyzing attention heads in the context of LLM safety. The writing is clear and easy to follow. 

The evaluation is limited, focusing primarily on LLaMA and Vicuna. Further evaluation on multi-turn inputs and a broader range of models would strengthen the findings. 

The novelty of the proposed method and clear presentation justify acceptance.