Retrieval Head Mechanistically Explains Long-Context Factuality

📄 arXiv: 2404.15574v1 📥 PDF

作者: Wenhao Wu, Yizhong Wang, Guangxuan Xiao, Hao Peng, Yao Fu

分类: cs.CL

发布日期: 2024-04-24

备注: Preprint

💡 一句话要点

提出检索头机制以解释长上下文中的事实性问题

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

关键词: 长上下文 信息检索 检索头 推理能力 模型机制 自然语言处理 幻觉现象

📋 核心要点

  1. 现有的长上下文语言模型在信息检索能力上仍存在不明确的机制，导致模型在处理长文本时的表现不稳定。
  2. 论文提出了检索头的概念，指出这些特殊的注意力头在长上下文信息检索中发挥了核心作用，并具备普遍性和动态激活特性。
  3. 研究表明，检索头对链式推理有显著影响，且在去除检索头后模型的检索能力会显著下降，验证了其在模型内部信息检索中的重要性。

📝 摘要（中文）

尽管长上下文语言模型取得了近期进展，但如何从长上下文中检索相关信息仍然不明晰。本文通过系统性研究发现，特定类型的注意力头（称为检索头）在信息检索中起着关键作用。研究表明，所有具有长上下文能力的模型均存在检索头，且其数量稀疏，仅占注意力头的5%以下。这些检索头在短上下文预训练模型中已存在，并在延长上下文长度后仍能有效工作。此外，检索头的动态激活和因果性特征也被揭示，影响模型的推理能力和信息检索效果。我们的发现为减少幻觉、改善推理和压缩KV缓存提供了新的研究方向。

🔬 方法详解

问题定义：本文旨在揭示长上下文语言模型中信息检索的机制，现有方法未能清晰解释模型如何从长文本中有效提取相关信息，导致模型在实际应用中的表现不稳定。

核心思路：论文提出了“检索头”的概念，认为这些特殊的注意力头是模型进行信息检索的关键。通过系统性分析，发现所有具有长上下文能力的模型均存在这类检索头，并且它们在不同上下文中动态激活。

技术框架：研究通过对多种长上下文模型的分析，识别出检索头的特性，包括普遍性、稀疏性、内在性和动态激活。模型的整体架构包括预训练阶段和上下文扩展阶段，检索头在这两个阶段均发挥重要作用。

关键创新：最重要的创新在于提出了检索头的概念，并通过实验证明其在信息检索和推理中的重要性。这一发现与现有方法的本质区别在于，强调了特定注意力头在长上下文处理中的核心作用。

关键设计：研究中对检索头的数量和激活特性进行了详细分析，发现仅有少量（少于5%）的注意力头为检索头。此外，检索头在短上下文预训练模型中已存在，且在上下文扩展后仍能有效工作。

📊 实验亮点

实验结果显示，去除检索头会导致模型在信息检索任务中的表现显著下降，验证了检索头在长上下文处理中的关键作用。此外，研究还发现，检索头对链式推理的影响显著，而对直接生成答案的任务影响较小，这表明其在不同任务中的作用差异。

🎯 应用场景

该研究的发现对长上下文语言模型的应用具有重要意义，尤其是在信息检索、问答系统和对话生成等领域。通过优化检索头的使用，可以提高模型的推理能力和信息准确性，减少生成内容中的幻觉现象，推动自然语言处理技术的进一步发展。

📄 摘要（原文）

Despite the recent progress in long-context language models, it remains elusive how transformer-based models exhibit the capability to retrieve relevant information from arbitrary locations within the long context. This paper aims to address this question. Our systematic investigation across a wide spectrum of models reveals that a special type of attention heads are largely responsible for retrieving information, which we dub retrieval heads. We identify intriguing properties of retrieval heads:(1) universal: all the explored models with long-context capability have a set of retrieval heads; (2) sparse: only a small portion (less than 5\%) of the attention heads are retrieval. (3) intrinsic: retrieval heads already exist in models pretrained with short context. When extending the context length by continual pretraining, it is still the same set of heads that perform information retrieval. (4) dynamically activated: take Llama-2 7B for example, 12 retrieval heads always attend to the required information no matter how the context is changed. The rest of the retrieval heads are activated in different contexts. (5) causal: completely pruning retrieval heads leads to failure in retrieving relevant information and results in hallucination, while pruning random non-retrieval heads does not affect the model's retrieval ability. We further show that retrieval heads strongly influence chain-of-thought (CoT) reasoning, where the model needs to frequently refer back the question and previously-generated context. Conversely, tasks where the model directly generates the answer using its intrinsic knowledge are less impacted by masking out retrieval heads. These observations collectively explain which internal part of the model seeks information from the input tokens. We believe our insights will foster future research on reducing hallucination, improving reasoning, and compressing the KV cache.