检索头从机制上解释了长上下文事实性
Retrieval Head Mechanistically Explains Long-Context Factuality
AI总结：
本文揭示了Transformer模型中一种特殊的“检索头”负责从长上下文中检索信息，并验证了其普遍性、稀疏性、内在性、动态激活和因果性等特性，为减少幻觉、改进推理和压缩KV缓存提供了机制性解释。
摘要：
尽管长上下文语言模型近期取得了进展，但基于Transformer的模型如何具备从长上下文中任意位置检索相关信息的能力，仍然是一个未解之谜。本文旨在解答这一问题。我们对广泛的模型进行了系统调查，发现一种特殊类型的注意力头主要负责信息检索，我们将其称为检索头。我们发现了检索头几个引人注目的特性：(1) 普遍性：所有被探索的具备长上下文能力的模型都拥有一组检索头；(2) 稀疏性：只有一小部分（不到5%）的注意力头执行检索功能；(3) 内在性：检索头在以短上下文预训练的模型中已经存在。当通过持续预训练扩展上下文长度时，执行信息检索的仍是同一组头；(4) 动态激活：以Llama-2 7B为例，12个检索头无论上下文如何变化，始终关注所需信息。其余的检索头则在不同上下文中被激活；(5) 因果性：完全剪枝检索头会导致无法检索相关信息并产生幻觉，而剪枝随机的非检索头则不影响模型的检索能力。我们进一步表明，检索头强烈影响思维链推理，在此类任务中模型需要频繁回溯问题及先前生成的上下文。相反，模型直接利用其内在知识生成答案的任务受屏蔽检索头的影响较小。这些观察结果共同解释了模型的哪个内部部分从输入词元中寻找信息。我们相信，这些见解将促进未来在减少幻觉、改进推理和压缩KV缓存方面的研究。
英文摘要：
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.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
