Retrieval Head Mechanistically Explains Long-Context Factuality

尽管长上下文语言模型的研究取得了进展，但是transformer-based模型如何从长上下文中的任意位置检索相关信息的能力仍然难以捉摸。本文旨在回答这个问题。我们对各种模型进行了系统研究，发现一种特殊类型的注意力头负责检索信息，我们称之为检索头。我们确定了检索头的有趣特性：(1)通用性：所有具有长上下文能力的探索模型都有一组检索头；(2)稀疏性：只有少部分(不到5%)的注意力头是检索头；(3)内在性：检索头已经存在于预先训练的短上下文模型中。通过持续的预训练扩展上下文长度时，执行信息检索的仍然是同一组头。(4)动态激活：以Llama-2 7B为例，12个检索头始终关注所需信息，无论上下文如何改变，其余的检索头在不同的上下文中被激活。(5)因果关系：完全剪枝检索头会导致无法检索相关信息，导致产生幻觉，而剪枝随机非检索头不会影响模型的检索能力。我们进一步表明，检索头强烈影响了链式思维(Chain-of-Thought, CoT)推理，其中模型需要频繁地参考问题和先前生成的上下文。相反，直接使用其内在知识生成答案的任务受到屏蔽检索头的影响较小。这些观察结果共同解释了模型的哪个内部部分从输入标记中寻找信息。我们相信，我们的洞见将促进未来关于减少幻觉、改进推理和压缩KV缓存的研究。

论文旨在探究transformer-based模型如何从长文本中检索相关信息。

论文发现一种特殊类型的attention heads，即retrieval heads，是检索信息的关键。这些heads普遍存在于具有长文本能力的模型中，只占所有attention heads的很小一部分。这些heads在预训练模型中已经存在，随着继续预训练，仍然是同一组heads执行信息检索。这些发现有助于理解模型内部哪些部分从输入token中寻找信息。

论文发现retrieval heads对于chain-of-thought reasoning任务非常重要，但对于直接生成答案的任务影响较小。论文还探讨了retrieval heads的动态激活和因果关系，并证明完全剪枝retrieval heads会导致信息检索失败和产生幻觉。

最近的相关研究包括《Attention is not Explanation》和《On the Relationship between Self-Attention and Convolutional Layers》等。