Abstract
Recent work has identified retrieval heads (Wu et al., 2025b), a subset of attention heads responsible for retrieving salient information in long-context language models (LMs), as measured by their copy-paste behavior in Needle-in-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 QR- RETRIEVER, an efficient and effective retriever that uses the accumulated attention mass of QRHEAD as retrieval scores. We use QR- RETRIEVER 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 querycontext 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.
Abstract (translated)
最近的研究（Wu et al., 2025b）发现了一类称为检索头（retrieval heads）的注意力机制子集，这些头部在长上下文语言模型中负责从大量文本信息中提取关键内容，这一结论是通过它们在针刺干草堆任务中的复制粘贴行为得出的。在此论文中，我们提出了QRHEAD（查询聚焦检索头），这是一种改进后的注意力机制集合，旨在增强从长文本中进行有效检索的能力。我们通过聚集与输入查询相关的注意力得分，并结合真实世界任务（如长上下文问答）示例来识别QRHEAD。 此外，我们引入了QR- RETRIEVER，这是一个高效且有效的检索器，它使用QRHEAD累积的注意力质量作为检索分数。我们在多跳推理任务LongMemEval和CLIPPER中利用QR-RETRIEVER进行长文本推理，通过选择具有最高检索分数的相关部分来实现这一目标。相比全面考虑上下文的方法，这种方法在这些任务上带来了超过10%的性能提升，并且优于其他密集型检索器。 我们还在BEIR基准测试中将QR-RETRIEVER作为重排序器进行了评估，并发现它实现了强大的零样本性能，超过了其他基于大语言模型（LLM）的重排序器，如RankGPT。进一步分析表明，查询上下文注意力评分以及任务选择对于识别具有强大下游效用的QRHEAD至关重要。 总体而言，我们的工作为长文本推理贡献了一种通用检索方法，并提供了关于LMs在处理长上下文时能力的理解和解释性见解。
URL
https://arxiv.org/abs/2506.09944
