查询聚焦检索头提升长上下文推理与重排序

近期工作识别出检索头——负责在长上下文语言模型（LM）中检索显著信息的一组注意力头——其衡量标准是在 Needlein-a-Haystack 任务中的复制粘贴行为。在本文中，我们引入了 QRHead（查询聚焦检索头），这是一组改进的注意力头，能够增强从长上下文中的检索。我们通过相对于输入查询聚合注意力分数来识别 QRHead，使用少量来自真实世界任务（如长上下文问答）的示例。我们进一步引入了 QRRetriever，一个高效且有效的检索器，利用 QRHead 累积的注意力质量作为检索分数。我们使用 QRRetriever 进行长上下文推理，通过选择具有最高检索分数的最相关部分。在多跳推理任务 LongMemEval 和 CLIPPER 上，这相比全上下文实现了超过 10% 的性能提升，并优于强力的密集检索器。我们还将 QRRetriever 评估为 BEIR 基准上的重排序器，发现其取得了强大的零样本性能，超越了其他基于 LLM 的重排序器如 RankGPT。进一步分析表明，查询-上下文注意力评分和任务选择对于识别具有强下游效用的 QRHead 均至关重要。总体而言，我们的工作贡献了一个通用检索器，并提供了对 LM 长上下文能力的解释性洞察。