arXiv 2506.099442025-09-28cs.CL

聚焦查询的检索头改善长上下文推理与重排序

Query-Focused Retrieval Heads Improve Long-Context Reasoning and Re-ranking

Princeton Language and Intelligence, Princeton University（普林斯顿语言与智能研究所，普林斯顿大学）；The University of Texas at Austin（德克萨斯大学奥斯汀分校）

Wuwei Zhang, Fangcong Yin, Howard Yen, Danqi Chen, Xi Ye

备注：: EMNLP 2025; Code at https://github.com/princeton-pli/QRHead

阅读 PDF

arXiv 页面 ↗HTML 原文 ↗

AI总结：

本文提出通过查询相关注意力识别 QRHead，并基于其累积注意力质量构建 QRRetriever，在长上下文多跳推理和 BEIR 零样本重排序中取得显著提升。

摘要：

近期工作识别出了检索头（retrieval heads），即长上下文语言模型（LM）中负责检索显著信息的一小组注意力头，其判定依据是它们在 Needle-in-a-Haystack 任务中的复制粘贴行为。本文提出 QRHead（Query-Focused Retrieval Head，聚焦查询的检索头），这是一组经过改进、可增强长上下文检索能力的注意力头。我们利用来自真实任务（例如长上下文问答）的少量示例，通过聚合相对于输入查询的注意力分数来识别 QRHead。我们进一步提出 QRRetriever，这是一种高效且有效的检索器，将 QRHead 的累积注意力质量作为检索分数。我们通过选择检索分数最高、最相关的部分，将 QRRetriever 用于长上下文推理。在多跳推理任务 LongMemEval 和 CLIPPER 上，该方法相较使用完整上下文带来超过 10% 的性能提升，并且优于强稠密检索器。我们还在 BEIR 基准上将 QRRetriever 作为重排序器进行评估，发现其取得了强劲的零样本性能，优于 RankGPT 等其他基于 LLM 的重排序器。进一步分析表明，查询—上下文注意力评分和任务选择对于识别具有强大下游效用的 QRHead 都至关重要。总体而言，我们的工作贡献了一个通用检索器，并为理解 LM 的长上下文能力提供了可解释性洞见。

英文摘要：

Recent work has identified retrieval heads, a subset of attention heads responsible for retrieving salient information in long-context language models (LMs), as measured by their copy-paste behavior in Needlein-a-Haystack tasks. In this paper, we introduce QRHead (Query-Focused Retrieval Head), an improved set of attention heads that enhance retrieval from long context. We identify QRHead by aggregating attention scores with respect to the input query, using a handful of examples from real-world tasks (e.g., long-context QA). We further introduce QRRetriever, an efficient and effective retriever that uses the accumulated attention mass of QRHead as retrieval scores. We use QRRetriever for long-context reasoning by selecting the most relevant parts with the highest retrieval scores. On multi-hop reasoning tasks LongMemEval and CLIPPER, this yields over 10% performance gains over full context and outperforms strong dense retrievers. We also evaluate QRRetriever as a re-ranker on the BEIR benchmark and find that it achieves strong zero-shot performance, outperforming other LLM-based re-rankers such as RankGPT. Further analysis shows that both the query-context attention scoring and task selection are crucial for identifying QRHead with strong downstream utility. Overall, our work contributes a general-purpose retriever and offers interpretability insights into the long-context capabilities of LMs.

机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。