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

TL;DR

This paper introduces QRHead and QRRetriever, novel methods that improve long-context reasoning and re-ranking in language models by leveraging query-focused attention, leading to significant performance gains and interpretability insights.

Contribution

The paper presents QRHead and QRRetriever, new techniques for enhancing retrieval and reasoning in long-context language models, with demonstrated improvements over existing methods.

Findings

Over 10% performance improvement on LongMemEval and CLIPPER

Outperforms strong dense retrievers in long-context reasoning

Achieves strong zero-shot re-ranking performance on BEIR

Abstract

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.