mlsanityData to 14 Oct 2026

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

arXiv 2506.09944 ↗First posted 11 Jun 2025cs.CLcited by 31

On the map: LLM reasoning → Reasoning benchmarks and evaluation

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.

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.

Builds on

23 of 56 references.

References · 23

  1. 2024The Llama 3 Herd of Modelscited by 14,715
  2. 2019Sentence-BERT: Sentence Embeddings using Siamese BERT-Networkscited by 9,170
  3. 2021Unsupervised Dense Information Retrieval with Contrastive Learningcited by 1,216
  4. 2022In-context Learning and Induction Headscited by 964
  5. 2024LongMemEval: Benchmarking Chat Assistants on Long-Term Interactive Memorycited by 707
  6. 2019Passage Re-ranking with BERTcited by 645
  7. 2023Is ChatGPT Good at Search? Investigating Large Language Models as Re-Ranking Agentscited by 360
  8. 2020Autoregressive Entity Retrievalcited by 334
  9. 2022Transformer Memory as a Differentiable Search Indexcited by 247
  10. 2022SGPT: GPT Sentence Embeddings for Semantic Searchcited by 204
  11. 2022Improving Passage Retrieval with Zero-Shot Question Generationcited by 165
  12. 2023RankZephyr: Effective and Robust Zero-Shot Listwise Reranking is a Breeze!cited by 155
11 more
  1. 2024Jasper and Stella: distillation of SOTA embedding modelscited by 112
  2. 2023Zero-Shot Listwise Document Reranking with a Large Language Modelcited by 82
  3. 2023Copy Suppression: Comprehensively Understanding an Attention Headcited by 76
  4. 2023Learning to Tokenize for Generative Retrievalcited by 68
  5. 2024FIRST: Faster Improved Listwise Reranking with Single Token Decodingcited by 54
  6. 2024Attention in Large Language Models Yields Efficient Zero-Shot Re-Rankerscited by 43
  7. 2021Phrase Retrieval Learns Passage Retrieval, Toocited by 33
  8. 2025Improving Contextual Faithfulness of Large Language Models via Retrieval Heads-Induced Optimizationcited by 17
  9. 2023PaRaDe: Passage Ranking using Demonstrations with Large Language Modelscited by 10
  10. 2025CLIPPER: Compression enables long-context synthetic data generationcited by 10
  11. 2024Sliding Windows Are Not the End: Exploring Full Ranking with Long-Context Large Language Modelscited by 8

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2025: 3; 2026: 28.

31 papers.

Sort:

By area:

  1. 2025MemVerse: Multimodal Memory for Lifelong Learning Agentscited by 35
  2. 2026Advancing General-Purpose Reasoning Models with Modular Gradient Surgerycited by 9
  3. 2026Query-focused and Memory-aware Reranker for Long Context Processingcited by 6
  4. 2026ReAttn: Improving Attention-based Re-ranking via Attention Re-weightingcited by 5
  5. 2025E2Rank: Unifying Text Embedding and Listwise Reranking for Effective and Efficient Searchcited by 4
  6. 2026From Interpretability to Performance: Optimizing Retrieval Heads for Long-Context Language Modelscited by 4
  7. 2026Hierarchical Memory Orchestration for Personalized Persistent Agentscited by 4
  8. 2026Mechanistic Insights into Functional Sparsity in Multimodal LLMs via CoRe Headscited by 4
  9. 2026Modern Transformers Are Implicit Hybrids: From Functional Differentiation to Principled Hybrid Architecture Designcited by 3
  10. 2026Retrieval Heads are Dynamiccited by 2
  11. 2026DySCO: Dynamic Attention-Scaling Decoding for Long-Context Language Modelscited by 2
  12. 2026HeadRank: Decoding-Free Passage Reranking via Preference-Aligned Attention Headscited by 2
19 more
  1. 2026Preference Heads in Large Language Models: A Mechanistic Framework for Interpretable Personalizationcited by 2
  2. 2026RL-Index: Reinforcement Learning for Retrieval Index Reasoningcited by 2
  3. 2026Can Language Models Actually Retrieve In-Context? Drowning in Documents at Million Token Scalecited by 2
  4. 2026S3-Attention:Attention-Aligned Endogenous Retrieval for Memory-Bounded Long-Context Inferencecited by 1
  5. 2026Can Retrieval Heads See Images? Multimodal Retrieval Heads in Long-Context Vision-Language Modelscited by 1
  6. 2026Self-Guided Test-Time Training for Long-Context LLMscited by 1
  7. 2026Retrieval Heads Meet Vision: Uncovering How VLMs Locate and Extract Visual Informationcited by 1
  8. 2025Citation Failure: Definition, Analysis and Efficient Mitigationcited by 0
  9. 2026Identifying Crucial Attention Heads for Multilingual Language Models: Retrieval and Retrieval-Transition Headscited by 0
  10. 2026PRISM: Breaking the O(n) Memory Wall in Long-Context LLM Inference via O(1) Photonic Block Selectioncited by 0
  11. 2026Learning to Route Queries to Heads for Attention-based Re-ranking with Large Language Modelscited by 0
  12. 2026Cultural Binding Heads in Language Modelscited by 0
  13. 2026CompRank: Efficient LLM Reranking via Token-Level Compression and Decoding-Free Scoringcited by 0
  14. 2026DREAM: Dense Retrieval Embeddings via Autoregressive Modelingcited by 0
  15. 2026NOWJ@COLIEE 2026: Adaptive Pipelines for Legal Retrieval and Reasoningcited by 0
  16. 2026Mechanistic Attention Guidance for Agent Memory Refinementcited by 0
  17. 2026AttnLink: Turning Attention into Schema Links for Text-to-SQLcited by 0
  18. 2026How Calibration Content Shapes Attention-Based Rerankingcited by 0
  19. 2026Lend Me Your Eyes: Instruction-Aware Text Embeddings via Attention Relaycited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2025Infinite Retrieval: Attention Enhanced LLMs in Long-Context Processing
  2. 2025Eliciting In-context Retrieval and Reasoning for Long-context Large Language Models
  3. 2026AttentionRetriever: Attention Layers are Secretly Long Document Retrievers
  4. 2026Retrieval and Multi-Hop Reasoning in 1M-Token Context Windows: Evaluating LLMs on Classical Chinese Text
  5. 2024Retrieval Meets Reasoning: Dynamic In-Context Editing for Long-Text Understanding
  6. 2025TongSearch-QR: Reinforced Query Reasoning for Retrieval
  7. 2026QReason: Query-Focused Decoupled Chain-of-Thought for Efficient Passage Reranking
  8. 2023Making Retrieval-Augmented Language Models Robust to Irrelevant Context