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

arXiv: https://arxiv.org/abs/2506.09944 · HTML: https://mlsanity.com/paper/2506.09944/

## Key facts

- Cited by 31 papers in the map (data through 14 Oct 2026).
- Placed in Reasoning benchmarks and evaluation, one of 17 directions in LLM reasoning.
- Builds on 23 of its 56 references that are in the map.
- First posted on arXiv 11 Jun 2025.

## Placed in

- LLM reasoning / [Reasoning benchmarks and evaluation](https://mlsanity.com/lr/reasoning-benchmarks-and-evaluation/)

## Builds on (most cited first, 23 of 23)

- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://mlsanity.com/paper/1908.10084/) (2019)
- [Unsupervised Dense Information Retrieval with Contrastive Learning](https://mlsanity.com/paper/2112.09118/) (2021)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [LongMemEval: Benchmarking Chat Assistants on Long-Term Interactive Memory](https://mlsanity.com/paper/2410.10813/) (2024)
- [Passage Re-ranking with BERT](https://mlsanity.com/paper/1901.04085/) (2019)
- [Is ChatGPT Good at Search? Investigating Large Language Models as Re-Ranking Agents](https://mlsanity.com/paper/2304.09542/) (2023)
- [Autoregressive Entity Retrieval](https://mlsanity.com/paper/2010.00904/) (2020)
- [Transformer Memory as a Differentiable Search Index](https://mlsanity.com/paper/2202.06991/) (2022)
- [SGPT: GPT Sentence Embeddings for Semantic Search](https://mlsanity.com/paper/2202.08904/) (2022)
- [Improving Passage Retrieval with Zero-Shot Question Generation](https://mlsanity.com/paper/2204.07496/) (2022)
- [RankZephyr: Effective and Robust Zero-Shot Listwise Reranking is a Breeze!](https://mlsanity.com/paper/2312.02724/) (2023)
- [Jasper and Stella: distillation of SOTA embedding models](https://mlsanity.com/paper/2412.19048/) (2024)
- [Zero-Shot Listwise Document Reranking with a Large Language Model](https://mlsanity.com/paper/2305.02156/) (2023)
- [Copy Suppression: Comprehensively Understanding an Attention Head](https://mlsanity.com/paper/2310.04625/) (2023)
- [Learning to Tokenize for Generative Retrieval](https://mlsanity.com/paper/2304.04171/) (2023)
- [FIRST: Faster Improved Listwise Reranking with Single Token Decoding](https://mlsanity.com/paper/2406.15657/) (2024)
- [Attention in Large Language Models Yields Efficient Zero-Shot Re-Rankers](https://mlsanity.com/paper/2410.02642/) (2024)
- [Phrase Retrieval Learns Passage Retrieval, Too](https://mlsanity.com/paper/2109.08133/) (2021)
- [Improving Contextual Faithfulness of Large Language Models via Retrieval Heads-Induced Optimization](https://mlsanity.com/paper/2501.13573/) (2025)
- [PaRaDe: Passage Ranking using Demonstrations with Large Language Models](https://mlsanity.com/paper/2310.14408/) (2023)
- [CLIPPER: Compression enables long-context synthetic data generation](https://mlsanity.com/paper/2502.14854/) (2025)
- [Sliding Windows Are Not the End: Exploring Full Ranking with Long-Context Large Language Models](https://mlsanity.com/paper/2412.14574/) (2024)

## Built on by (most cited first, 30 of 31)

- [MemVerse: Multimodal Memory for Lifelong Learning Agents](https://mlsanity.com/paper/2512.03627/) (2025)
- [Advancing General-Purpose Reasoning Models with Modular Gradient Surgery](https://mlsanity.com/paper/2602.02301/) (2026)
- [Query-focused and Memory-aware Reranker for Long Context Processing](https://mlsanity.com/paper/2602.12192/) (2026)
- [ReAttn: Improving Attention-based Re-ranking via Attention Re-weighting](https://mlsanity.com/paper/2602.19969/) (2026)
- [E2Rank: Unifying Text Embedding and Listwise Reranking for Effective and Efficient Search](https://mlsanity.com/paper/2510.22733/) (2025)
- [From Interpretability to Performance: Optimizing Retrieval Heads for Long-Context Language Models](https://mlsanity.com/paper/2601.11020/) (2026)
- [Hierarchical Memory Orchestration for Personalized Persistent Agents](https://mlsanity.com/paper/2604.01670/) (2026)
- [Mechanistic Insights into Functional Sparsity in Multimodal LLMs via CoRe Heads](https://mlsanity.com/paper/2606.05843/) (2026)
- [Modern Transformers Are Implicit Hybrids: From Functional Differentiation to Principled Hybrid Architecture Design](https://mlsanity.com/paper/2609.02986/) (2026)
- [Retrieval Heads are Dynamic](https://mlsanity.com/paper/2602.11162/) (2026)
- [DySCO: Dynamic Attention-Scaling Decoding for Long-Context Language Models](https://mlsanity.com/paper/2602.22175/) (2026)
- [HeadRank: Decoding-Free Passage Reranking via Preference-Aligned Attention Heads](https://mlsanity.com/paper/2604.17237/) (2026)
- [Preference Heads in Large Language Models: A Mechanistic Framework for Interpretable Personalization](https://mlsanity.com/paper/2604.22345/) (2026)
- [RL-Index: Reinforcement Learning for Retrieval Index Reasoning](https://mlsanity.com/paper/2606.16316/) (2026)
- [Can Language Models Actually Retrieve In-Context? Drowning in Documents at Million Token Scale](https://mlsanity.com/paper/2607.01538/) (2026)
- [S3-Attention:Attention-Aligned Endogenous Retrieval for Memory-Bounded Long-Context Inference](https://mlsanity.com/paper/2601.17702/) (2026)
- [Can Retrieval Heads See Images? Multimodal Retrieval Heads in Long-Context Vision-Language Models](https://mlsanity.com/paper/2605.27243/) (2026)
- [Self-Guided Test-Time Training for Long-Context LLMs](https://mlsanity.com/paper/2607.09415/) (2026)
- [Retrieval Heads Meet Vision: Uncovering How VLMs Locate and Extract Visual Information](https://mlsanity.com/paper/2608.27417/) (2026)
- [Citation Failure: Definition, Analysis and Efficient Mitigation](https://mlsanity.com/paper/2510.20303/) (2025)
- [Identifying Crucial Attention Heads for Multilingual Language Models: Retrieval and Retrieval-Transition Heads](https://mlsanity.com/paper/2602.22453/) (2026)
- [PRISM: Breaking the O(n) Memory Wall in Long-Context LLM Inference via O(1) Photonic Block Selection](https://mlsanity.com/paper/2603.21576/) (2026)
- [Learning to Route Queries to Heads for Attention-based Re-ranking with Large Language Models](https://mlsanity.com/paper/2604.24608/) (2026)
- [Cultural Binding Heads in Language Models](https://mlsanity.com/paper/2605.28543/) (2026)
- [CompRank: Efficient LLM Reranking via Token-Level Compression and Decoding-Free Scoring](https://mlsanity.com/paper/2606.11700/) (2026)
- [DREAM: Dense Retrieval Embeddings via Autoregressive Modeling](https://mlsanity.com/paper/2606.24667/) (2026)
- [NOWJ@COLIEE 2026: Adaptive Pipelines for Legal Retrieval and Reasoning](https://mlsanity.com/paper/2607.16603/) (2026)
- [Mechanistic Attention Guidance for Agent Memory Refinement](https://mlsanity.com/paper/2607.17621/) (2026)
- [AttnLink: Turning Attention into Schema Links for Text-to-SQL](https://mlsanity.com/paper/2608.00693/) (2026)
- [How Calibration Content Shapes Attention-Based Reranking](https://mlsanity.com/paper/2609.17764/) (2026)

## Abstract (arXiv)

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.