On the Role of Attention Heads in Large Language Model Safety

📄 arXiv: 2410.13708v2 📥 PDF

作者: Zhenhong Zhou, Haiyang Yu, Xinghua Zhang, Rongwu Xu, Fei Huang, Kun Wang, Yang Liu, Junfeng Fang, Yongbin Li

分类: cs.CL, cs.AI, cs.CR, cs.LG

发布日期: 2024-10-17 (更新: 2025-02-24)

备注: 28 pages, 18 figures, 7 tables. This paper has been accepted as ICLR 2025 (oral)

💡 一句话要点

提出Safety Head ImPortant Score (Ships)和Sahara算法，用于评估和归因LLM中的安全注意力头。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型安全 注意力机制 可解释性 安全头重要性评分 安全注意力头归因

📋 核心要点

  1. 现有研究忽略了多头注意力机制在LLM安全中的作用，而多头注意力在模型功能中至关重要，需要深入研究其与安全能力的关联。
  2. 论文提出Safety Head ImPortant Score (Ships)评估注意力头对安全的贡献，并设计Sahara算法来归因模型内部的关键安全注意力头。
  3. 实验表明，特定注意力头对安全有显著影响，消融单个安全头可使模型响应更多有害查询，且微调模型表现出重叠的安全头。

📝 摘要（中文）

大型语言模型(LLMs)在多项语言任务上取得了最先进的性能，但其安全防护措施可能会被规避，导致有害内容的生成。针对此问题，近期的安全机制研究表明，当安全表示或组件被抑制时，LLMs的安全能力会受到损害。然而，现有研究往往忽视了多头注意力机制对安全的影响，尽管它在各种模型功能中起着至关重要的作用。因此，本文旨在探索标准注意力机制与安全能力之间的联系，以填补安全相关的可解释性方面的空白。我们提出了一种专门为多头注意力定制的新指标，即安全头重要性评分(Ships)，以评估各个头对模型安全的贡献。在此基础上，我们将Ships推广到数据集层面，并进一步引入安全注意力头归因算法(Sahara)来归因模型内部的关键安全注意力头。我们的研究结果表明，特殊的注意力头对安全有显著影响。消融单个安全头会使对齐模型(例如，Llama-2-7b-chat)能够响应多达16倍的有害查询，而仅修改了0.006%的参数，相比之下，之前的研究需要修改约5%的参数。更重要的是，我们通过全面的实验证明，注意力头主要充当安全特征提取器，并且从同一基础模型微调的模型表现出重叠的安全头。总之，我们的归因方法和发现为揭示大型模型中安全机制的黑盒提供了一个新的视角。

🔬 方法详解

问题定义：现有的大型语言模型虽然在各种任务上表现出色，但其安全机制容易被绕过，导致生成有害内容。之前的研究主要集中在抑制安全表示或组件，而忽略了多头注意力机制在安全中的作用。因此，需要研究注意力头与模型安全能力之间的关系，从而提升LLM的安全性。

核心思路：论文的核心思路是通过量化每个注意力头对模型安全性的贡献，从而识别出关键的安全注意力头。通过分析这些安全注意力头，可以更好地理解LLM内部的安全机制，并为改进LLM的安全性提供指导。论文假设特定的注意力头在提取与安全相关的特征方面起着重要作用。

技术框架：论文提出了Safety Head ImPortant Score (Ships)和Safety Attention Head AttRibution Algorithm (Sahara)。Ships用于评估每个注意力头对模型安全性的贡献，而Sahara算法则用于识别模型内部的关键安全注意力头。整体流程包括：1) 使用Ships评估每个注意力头的重要性；2) 使用Sahara算法对注意力头进行归因，找出关键的安全注意力头；3) 通过消融实验验证这些安全注意力头的作用。

关键创新：论文的关键创新在于提出了Ships指标和Sahara算法，用于量化和归因LLM中的安全注意力头。与以往的研究不同，该方法关注的是注意力头在安全中的作用，而不是仅仅关注安全表示或组件。此外，论文还发现，从同一基础模型微调的模型表现出重叠的安全头，这表明注意力头在提取安全特征方面具有一定的通用性。

关键设计：Ships指标的设计考虑了注意力头对模型安全性的影响，它基于模型对有害查询的响应概率来计算。Sahara算法则利用Ships指标来识别关键的安全注意力头，它通过迭代地消融注意力头并观察模型性能的变化来实现。实验中使用了Llama-2-7b-chat等模型，并针对有害查询进行了评估。消融实验通过移除特定的注意力头，并观察模型对有害查询的响应变化来评估其重要性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，消融单个安全注意力头可以使Llama-2-7b-chat模型响应多达16倍的有害查询，而仅修改了0.006%的参数，相比之下，之前的研究需要修改约5%的参数。此外，实验还发现，从同一基础模型微调的模型表现出重叠的安全头，这表明注意力头在提取安全特征方面具有一定的通用性。

🎯 应用场景

该研究成果可应用于提升大型语言模型的安全性，例如，可以通过识别和加强关键的安全注意力头来提高模型对有害查询的抵抗能力。此外，该方法还可以用于分析和比较不同模型的安全机制，从而为开发更安全的LLM提供指导。未来，该研究可以扩展到其他类型的模型和任务，并与其他安全技术相结合，以构建更强大的安全防护体系。

📄 摘要（原文）

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 a novel metric which tailored for multi-head attention, the Safety Head ImPortant Score (Ships), to assess the individual heads' contributions to model safety. Based 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 the 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 more harmful queries, while only modifying 0.006% of the parameters, in contrast to the ~ 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 within large models.