从一次失败的问答说起
假设你在一家大型企业做知识库问答系统，用户问：“根据公司 2023 年年报，研发投入占比是多少？”系统把一份 100 页的 PDF 塞进上下文，模型却答非所问，甚至编造了一个数字。你检查了检索模块，发现相关片段确实被召回了；你检查了提示词，格式也正确。问题出在哪里？
一种常见的直觉是：模型既然能处理长文本，就应该能“看到”所有信息。但事实并非如此。Transformer 的注意力机制虽然理论上可以关注任意位置，但实际中，模型内部只有少数特定的注意力头真正负责从上下文中“找出”相关信息。这些头就是本文要讨论的“检索头”（Retrieval Heads）。
检索头是什么：从注意力头到信息检索
Transformer 的每一层都有多个注意力头，每个头独立计算查询（Query）、键（Key）、值（Value）之间的注意力分数。通俗地说，每个头都在回答一个问题：“当前这个 token 应该关注上下文中的哪些 token？”不同的头可能关注语法关系、指代、或者语义相似性。
检索头是一类特殊的注意力头，它们的功能是：当模型需要回答某个问题时，检索头会从上下文中定位包含答案的 token，并通过注意力机制将这些 token 的信息“复制”到当前生成的位置。这种“复制粘贴”行为是检索头的核心特征。
2024 年，北京大学、华盛顿大学、MIT 等机构的研究者在论文《Retrieval Head Mechanistically Explains Long-Context Factuality》中系统研究了这一现象。他们发现，在具有长上下文能力的模型中，存在少量检索头，它们负责从输入中检索相关信息。
检索头的五个关键特性
该论文通过对 Llama、Qwen 等 4 个模型家族、6 种模型规模和 3 种微调类型的实验，总结了检索头的五个特性：
- 普遍性：所有具备长上下文能力的模型都有一组检索头。
- 稀疏性：检索头只占注意力头总数的不到 5%。
- 固有性：检索头在预训练阶段就已存在，后续通过继续预训练扩展上下文长度时，仍然是同一组头在执行检索。
- 动态激活：最强的检索头无论上下文如何变化都会激活，而较弱的检索头则根据上下文内容选择性激活。
- 因果性：如果完全剪除检索头，模型将无法从上下文中检索信息，导致幻觉；而剪除随机的非检索头则不影响检索能力。
这些特性说明，检索头是模型内部专门负责信息检索的组件，而不是偶然出现的模式。
如何检测检索头：检索分数与大海捞针实验
要验证某个头是否是检索头，研究者引入了“检索分数”（Retrieval Score）。检索分数本质上是一个 token 级别的召回率：在生成答案时，该头成功从关键信息（“针”）中复制了多少比例的 token。
具体来说，给定一个关键信息 P（例如“在旧金山最好的活动是在公园吃三明治”），在模型生成每个 token 时，记录该注意力头最关注的输入 token。如果最关注的 token 恰好是 P 中的 token，并且生成的 token 也属于 P，则视为一次“复制”。检索分数就是复制成功的 token 数占 P 总 token 数的比例。
为了系统评估检索头，研究者使用“大海捞针”（Needle-in-a-Haystack）实验：构造一段长上下文（“海”），在其中随机位置插入一个短句（“针”），然后提问要求模型检索该短句。通过反复测试不同位置和长度，可以评估模型的检索能力。
如果某个头的平均检索分数 ≥ 0.1，即它在 10% 以上的时间里执行了复制粘贴操作，就判定为检索头。
头消融实验：验证检索头的因果作用
仅仅观察到相关性还不够，研究者通过头消融实验（Ablation）来证明因果性。消融实验的做法是：将某个注意力头的输出全部置零，然后观察模型行为的变化。
论文中的实验显示，在 Llama-2-7B-80K 模型上，如果屏蔽前 20 个最高得分的检索头，模型在“大海捞针”测试中的准确率从接近 100% 下降到 63.6%；而屏蔽 20 个随机的非检索头，准确率仍保持在 94.7%。这说明检索头是模型检索能力的必要组件。
更细致的观察发现：
- 完全屏蔽检索头，模型会产生幻觉，例如把“在公园吃三明治”答成“参观金门大桥”。
- 部分屏蔽检索头，模型只能检索到部分信息，例如记得“三明治”但忘了“公园”。
- 屏蔽随机非检索头，模型仍能完整找到“针”。
这些结果清晰地展示了检索头在信息检索中的因果作用。
检索头如何影响推理：从事实性到思维链
检索头不仅影响直接的事实检索，还影响需要多步推理的任务。论文发现，检索头对思维链（Chain-of-Thought, CoT）推理有强烈影响，因为模型在推理过程中需要频繁回看问题和之前生成的上下文。相反，对于直接利用内部知识生成答案的任务，屏蔽检索头的影响较小。
这意味着，检索头是模型“从上下文中找信息”这一能力的核心，而不仅仅是简单的复制粘贴。在长上下文问答中，如果检索头失效，模型就会退化为依赖内部知识，从而产生幻觉。
从检索头到 InfiniRetri：利用注意力进行检索
检索头的发现启发了新的应用。2025 年，研究者提出了 InfiniRetri 方法，它利用模型自身的注意力信息，实现对任意长度输入的精确检索，无需额外训练。
InfiniRetri 的核心思想是：既然模型在有限上下文窗口内能够通过注意力定位相关信息，那么我们可以将长文本切分成多个块，逐块处理，并利用注意力得分来识别与问题最相关的 token，将这些 token 所在的句子缓存下来，供后续推理使用。
具体流程如下：
flowchart TD
A[输入长文本] --> B[切分为文本块]
B --> C[逐块与缓存合并]
C --> D[模型推理，计算注意力得分]
D --> E[聚合最后一层注意力头得分]
E --> F[计算每个token的重要性分数]
F --> G[选取Top-K token所在句子存入缓存]
G --> H{是否还有文本块?}
H -- 是 --> C
H -- 否 --> I[基于缓存生成最终答案]
InfiniRetri 与传统 KV 缓存压缩方法（如 StreamingLLM、H2O、SnapKV）不同，它不是在每一层存储历史键值，而是在模型外部缓存 token ID，并在每次推理前将缓存与当前文本块合并。这样，模型始终使用标准注意力机制，而不是滑动窗口注意力。
实验表明，InfiniRetri 在“大海捞针”测试中，将 Llama3-8B 的有效上下文长度从 32K 扩展到 1M token，准确率保持 100%。在 LongBench 的真实数据集上，Qwen2-7B 在 HotpotQA 任务上的得分提升了 288%。
检索头与 RAG：互补还是替代？
检索增强生成（RAG）是目前工业界处理长上下文的主流方案。RAG 依赖外部嵌入模型从文档中检索相关片段，然后将其拼接到上下文中。然而，RAG 系统普遍难以建立检索信息之间的关联，因为检索模块和生成模块是分离的。
检索头的研究提供了一种新的视角：模型自身就具备检索能力，只是受限于上下文窗口。InfiniRetri 正是利用这一点，通过注意力机制在长文本中定位信息，而不需要外部检索器。
下表对比了 RAG、InfiniRetri 和直接长上下文模型三种方案：
| 方案 | 检索方式 | 训练成本 | 推理开销 | 适用场景 | 局限性 |
|---|---|---|---|---|---|
| RAG | 外部嵌入模型 | 低（无需训练） | 低（仅处理检索片段） | 知识库问答、开放域问答 | 检索质量依赖嵌入模型，难以建立多跳关联 |
| InfiniRetri | 模型自身注意力 | 低（无需训练） | 中（需处理所有块） | 超长文本检索、多文档问答 | 依赖模型内部检索能力，短上下文模型可能失效 |
| 直接长上下文 | 模型自身注意力 | 高（需长上下文预训练） | 高（计算量随长度平方增长） | 需要全局理解的场景 | 成本高，实际有效长度可能低于宣称值 |
从表中可以看出，RAG 和 InfiniRetri 各有优势。RAG 适合已有外部知识库的场景，而 InfiniRetri 适合需要从单一长文档中检索信息的场景。两者也可以结合：先用 RAG 粗筛，再用 InfiniRetri 精确定位。
检索头的局限与未解问题
尽管检索头的研究为理解 Transformer 提供了重要视角，但仍存在一些局限：
- 检索头的识别依赖特定任务：检索分数是基于“大海捞针”这类任务定义的，可能无法覆盖所有类型的检索行为。
- 跨模型泛化：不同模型家族的检索头分布差异较大，同一家族内相关性高，但跨家族相似度低。这意味着针对一个模型发现的检索头，不能直接迁移到另一个模型。
- 动态激活的机制尚不明确：为什么有些检索头始终激活，而有些则根据上下文选择性激活？背后的机制仍未完全解释。
- 与长上下文训练的交互：检索头在短上下文预训练中已经存在，但长上下文继续预训练如何影响它们的功能，仍需进一步研究。
此外，检索头的发现也启发了减少幻觉的新方法。例如，DeCoRe（Decoding by Contrasting Retrieval Heads）通过对比原始模型和屏蔽检索头模型的输出，来放大正确预测，从而减少幻觉。这种方法无需训练，但需要知道哪些是检索头。
结论
检索头机制揭示了 Transformer 模型内部信息检索的专门化组件。它们稀疏、固有、动态激活，并且对长上下文事实性至关重要。通过头消融实验，研究者证明了检索头的因果作用。这一发现不仅加深了我们对模型内部机制的理解，还催生了 InfiniRetri 等无需训练的长上下文处理方法。
对于工程师来说，理解检索头有助于诊断模型在长上下文任务中的失败原因：如果模型产生幻觉，可以检查检索头是否被抑制或未激活。对于研究者来说，检索头为压缩 KV 缓存、改进推理能力提供了新方向。然而，检索头的跨模型泛化和动态激活机制仍是未解之谜，需要更多研究。