# Learning to Route Queries to Heads for Attention-based Re-ranking with Large Language Models

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 43 of its 54 references that are in the map.
- First posted on arXiv 27 Apr 2026.

## Builds on (most cited first, 30 of 43)

- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://mlsanity.com/paper/1908.10084/) (2019)
- [Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2004.04906/) (2020)
- [FEVER: a large-scale dataset for Fact Extraction and VERification](https://mlsanity.com/paper/1803.05355/) (2018)
- [BEIR: A Heterogenous Benchmark for Zero-shot Evaluation of Information Retrieval Models](https://mlsanity.com/paper/2104.08663/) (2021)
- [ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT](https://mlsanity.com/paper/2004.12832/) (2020)
- [Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned](https://mlsanity.com/paper/1905.09418/) (2019)
- [Are Sixteen Heads Really Better than One?](https://mlsanity.com/paper/1905.10650/) (2019)
- [Extending Context Window of Large Language Models via Positional Interpolation](https://mlsanity.com/paper/2306.15595/) (2023)
- [C-Pack: Packed Resources For General Chinese Embeddings](https://mlsanity.com/paper/2309.07597/) (2023)
- [Fact or Fiction: Verifying Scientific Claims](https://mlsanity.com/paper/2004.14974/) (2020)
- [Large Dual Encoders Are Generalizable Retrievers](https://mlsanity.com/paper/2112.07899/) (2021)
- [SPECTER: Document-level Representation Learning using Citation-informed Transformers](https://mlsanity.com/paper/2004.07180/) (2020)
- [Document Ranking with a Pretrained Sequence-to-Sequence Model](https://mlsanity.com/paper/2003.06713/) (2020)
- [Is ChatGPT Good at Search? Investigating Large Language Models as Re-Ranking Agents](https://mlsanity.com/paper/2304.09542/) (2023)
- [Large Language Models are Effective Text Rankers with Pairwise Ranking Prompting](https://mlsanity.com/paper/2306.17563/) (2023)
- [Retrieval Head Mechanistically Explains Long-Context Factuality](https://mlsanity.com/paper/2404.15574/) (2024)
- [Distilling Knowledge from Reader to Retriever for Question Answering](https://mlsanity.com/paper/2012.04584/) (2020)
- [Multi-Stage Document Ranking with BERT](https://mlsanity.com/paper/1910.14424/) (2019)
- [CLIMATE-FEVER: A Dataset for Verification of Real-World Climate Claims](https://mlsanity.com/paper/2012.00614/) (2020)
- [Data Engineering for Scaling Language Models to 128K Context](https://mlsanity.com/paper/2402.10171/) (2024)
- [Improving Passage Retrieval with Zero-Shot Question Generation](https://mlsanity.com/paper/2204.07496/) (2022)
- [LLM Maybe LongLM: Self-Extend LLM Context Window Without Tuning](https://mlsanity.com/paper/2401.01325/) (2024)
- [ConvGQR: Generative Query Reformulation for Conversational Search](https://mlsanity.com/paper/2305.15645/) (2023)
- [Attention in Large Language Models Yields Efficient Zero-Shot Re-Rankers](https://mlsanity.com/paper/2410.02642/) (2024)
- [RATT: A Thought Structure for Coherent and Correct LLM Reasoning](https://mlsanity.com/paper/2406.02746/) (2024)
- [Unsupervised Dense Retrieval with Relevance-Aware Contrastive Pre-Training](https://mlsanity.com/paper/2306.03166/) (2023)
- [Query-Focused Retrieval Heads Improve Long-Context Reasoning and Re-ranking](https://mlsanity.com/paper/2506.09944/) (2025)
- [A Survey of Conversational Search](https://mlsanity.com/paper/2410.15576/) (2024)
- [Ranked List Truncation for Large Language Model-based Re-Ranking](https://mlsanity.com/paper/2404.18185/) (2024)

## Abstract (arXiv)

Large Language Models (LLMs) have recently been explored as fine-grained zero-shot re-rankers by leveraging attention signals to estimate document relevance. However, existing methods either aggregate attention signals across all heads or rely on a statically selected subset identified by heuristic rules. This solution can be suboptimal because the informative heads can vary across queries or domains. Moreover, naively combining multiple heads can degrade performance due to redundancy or conflicting ranking signals. In this paper, we propose a query-dependent head selection method, RouteHead, for attention-based re-ranking with LLMs. Specifically, we learn a lightweight router that can map each query to an optimal head set, and relevance scores are computed by aggregating attention signals only from these heads. Since query-to-head optimal labels are unavailable, we first construct pseudo labels via an offline search. The router represents each head with a learnable embedding and represents each query using an embedding extracted from the hidden states of the frozen LLM. Then it is trained on the pseudo labels with a sparsity regularizer. Experiments on diverse benchmarks and multiple LLM backbones show that the proposed method consistently outperforms strong baselines.