检索头机械解释长上下文事实性
计算与语言
2024-04-25 v1
摘要
尽管最近的长上下文语言模型取得了进展,但 transformer-based 模型如何表现出从长上下文中任意位置检索相关信息的能力仍不明了。本文旨在回答此问题。我们的系统性调查在广泛的模型范围内揭示了,检索信息的特殊注意力头类型在很大程度上负责检索,这些头我们称为检索头。我们识别了检索头的惊人特性:(1) 通用:所有探索的具备长上下文能力的模型都有一组检索头;(2) 稀疏:只有少量(不到 5%)的注意力头是检索头。(3) 内在:检索头已存在于使用短上下文预训练的模型中。通过持续预训练扩展上下文长度后,仍是同一组头执行信息检索。(4) 动态激活:以 Llama-2 7B 为例,12 个检索头无论上下文如何变化都始终注意所需信息。其余检索头在不同上下文中被激活。(5) 因果:完全剪枝检索头会导致检索相关信息失败并产生幻觉,而随机剪枝非检索头则不影响模型的检索能力。我们进一步表明,检索头强烈影响链式思维(CoT)推理,其中模型需要频繁引用问题和先前生成的上下文。相反,模型直接使用内在知识生成答案的任务受到屏蔽检索头影响较小。这些观察共同解释了模型内部寻找输入标记中信息的部分。我们相信这些见解将促进未来研究,减少幻觉、提高推理能力,并压缩 KV 缓存。
引用
@article{arxiv.2404.15574,
title = {Retrieval Head Mechanistically Explains Long-Context Factuality},
author = {Wenhao Wu and Yizhong Wang and Guangxuan Xiao and Hao Peng and Yao Fu},
journal= {arXiv preprint arXiv:2404.15574},
year = {2024}
}
备注
Preprint