Paper deep dive
Learning to Route Queries to Heads for Attention-based Re-ranking with Large Language Models
Yuxing Tian, Fengran Mo, Zhiqi Huang, Weixu Zhang, Jian-Yun Nie
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 97%
Last extracted: 6/21/2026, 8:40:55 AM
Summary
The paper introduces RouteHead, a query-dependent head selection method designed to improve attention-based re-ranking using Large Language Models (LLMs). Unlike existing methods that use static or heuristic-based attention head selection, RouteHead employs a lightweight router to dynamically map each query to an optimal set of attention heads. The method uses a two-stage pipeline: first, constructing pseudo labels via an offline search (forward selection and swap-based local refinement), and second, training a router that uses learnable head embeddings and query embeddings from frozen LLM hidden states. Experimental results on BEIR and BRIGHT benchmarks demonstrate that RouteHead consistently outperforms static head selection baselines (like QRhead and CoRehead) and remains competitive with traditional re-rankers and generation-based methods like RankGPT.
Entities (10)
Relation Signals (5)
RouteHead → evaluatedon → BEIR
confidence 100% · We evaluate on two benchmarks. The first is BEIR...
RouteHead → evaluatedon → BRIGHT
confidence 100% · The second is BRIGHT (SU et al., 2025)...
RouteHead → outperforms → QRhead
confidence 100% · RouteHead consistently outperforms strong baselines... RouteHead consistently outperforms the strongest static head baseline (QRhead)
RouteHead → outperforms → CoRehead
confidence 100% · RouteHead achieves the best average performance among attention-based methods
RouteHead → improves → Large Language Model
confidence 90% · RouteHead... enables real-time, query-dependent head selection for attention-based re-ranking with LLMs.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract: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.
Tags
Links
- Source: https://arxiv.org/abs/2604.24608v1
- Canonical: https://arxiv.org/abs/2604.24608v1
Trouble viewing inline? Open PDF directly →
Full Text
38,499 characters extracted from source content.
Expand or collapse full text
Learning to Route Queries to Heads for Attention-based Re-ranking with Large Language Models Yuxing Tian Université de MontréalMontrealCanada yuxing.tian@umontreal.ca , Fengran Mo Université de MontréalMontrealCanada fengran.mo@umontreal.ca , Zhiqi Huang Capital OneBostonUSA zhiqi.huang@capitalone.com , Weixu Zhang McGill University & MILAMontrealCanada weixu.zhang@mail.mcgill.ca and Jian-Yun Nie Université de MontréalMontrealCanada jian-yun.nie@umontreal.ca (2026) Abstract. 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. Re-ranking; Large Language Model; Attention Head; Router †copyright: acmlicensed†journalyear: 2018†doi: X.X†conference: Make sure to enter the correct conference title from your rights confirmation email; June 03–05, 2018; Woodstock, NY†isbn: 978-1-4503-X-X/2018/06†ccs: Information systems Language models†ccs: Information systems Retrieval models and ranking†journalyear: 2026†copyright: c†conference: Proceedings of the 49th International ACM SIGIR Conference on Research and Development in Information Retrieval; July 20–24, 2026; Melbourne, VIC, Australia†booktitle: Proceedings of the 49th International ACM SIGIR Conference on Research and Development in Information Retrieval (SIGIR ’26), July 20–24, 2026, Melbourne, VIC, Australia†doi: 10.1145/3805712.3809945†isbn: 979-8-4007-2599-9/2026/07 by 1. Introduction Re-ranking is a critical stage in modern retrieval systems for applications such as search and recommendation (Lewis et al., 2020; Gao et al., 2023; Meng et al., 2025b; Mo et al., 2025b; Wang et al., 2026). Traditional neural re-rankers, such as bi-encoders (Karpukhin et al., 2020; Khattab and Zaharia, 2020; Mo et al., 2025a) and cross-encoders (Nogueira et al., 2019, 2020; Meng et al., 2024), can be effective, but they typically require substantial relevance labels and task-specific training. Recent studies show that large language models (LLMs) can perform re-ranking in a zero-shot setting by leveraging their strong ability (Sachan et al., 2022; Sun et al., 2023; Qin et al., 2024; Meng et al., 2026). Most of them follow a generation-based paradigm, prompting an LLM to output a ranked list of candidate identifiers. However, this paradigm is costly due to autoregressive decoding, and it is sensitive to prompt design and instruction-following errors, which can lead to invalid outputs, missing candidates, or unstable rankings. To avoid the brittleness and overhead of generation-based paradigm, an alternative line of work estimates relevance from internal LLM signals. Chen et al. (2025) propose In-Context Re-ranking (ICR), which feeds the documents and query into an LLM and compute document relevance scores by aggregating query-to-document attention across all attention heads. However, aggregating over all heads can mix informative signals with redundant or noisy patterns, which can degrade ranking quality. QRhead (Zhang et al., 2025b) shows that a small set of heads can be sufficient. Specifically, it constructs a head-detection dataset that contains queries, candidate documents, and a labeled gold document for each query. It then scores each head by its attention mass on the gold document, averages the scores across queries, and selects the top-k heads (e.g., k=16k=16) for re-ranking. However, this scoring rule does not penalize heads that also attend strongly to irrelevant documents. CoRehead (Tran et al., 2025) addresses this with a contrastive criterion that prefers heads that allocate more attention to the gold document while allocating less attention to irrelevant documents. Figure 1. Performance comparison (nDCG@10) across three datasets for different attention-based re-ranking methods on Llama-3.2-1B. ICR(middle) denotes the ICR variant that directly selects all attention heads from the middle layers (7-12). Despite recent progress, head selection is often treated as a static decision: a single set of heads, identified on a head-detection dataset, is reused across queries and domains. Figure 1 shows that the impact of head selection varies substantially across datasets. First, specialized head selection does not consistently outperform a simple layer-based strategy, ICR (middle), which directly selects all heads from the middle layers. Second, the performance of head selection methods differs across datasets: CoreHead performs best on LeetCode and Psychology, whereas QRhead performs best on Scifact. These results indicate that head utility is strongly dataset-dependent, and that a static head selection rule does not generalize reliably. To achieve better fine-grained re-ranking, we propose a query-dependent head routing method that dynamically selects attention heads based on query features, called RouteHead. Training such a router requires supervision that maps each query to an optimal head set, but the annotations are unavailable in practice. We therefore adopt a two-stage pipeline. First, we construct pseudo labels via offline search. The search uses forward selection with early stopping, followed by swap-based local refinement, and produces a multi-hot label for each query. Second, we train a router for query-dependent head selection. Specifically, the router assigns each attention head a learnable embedding and represents each query using an embedding extracted from frozen LLM hidden states. It then scores each query–head pair and applies an independent sigmoid to obtain activation probabilities. It is trained by minimizing binary cross-entropy against the multi-hot pseudo labels, together with a sparsity regularizer. At inference time, the router dynamically selects a small set of heads, and relevance scores are computed using only these heads. Our contributions are as follows: 1. We introduce a lightweight query-to-head router that enables real-time, query-dependent head selection for attention-based re-ranking with LLMs. 2. We propose an offline search procedure to build pseudo labels that approximate optimal query-to-head labels for router training. 3. We evaluate the proposed method against strong baselines across multiple LLM backbones on diverse benchmarks, and it consistently achieves the best performance. 2. Preliminary and Problem Definition Given a query q and a set of N retrieved candidate documents D=d1,…,dND=\d_1,…,d_N\, we feed the concatenated sequence X=[d1,d2,…,dN,q]X=[d_1,d_2,…,d_N,q] into the LLM with M=h1,h2,…,hMM=\h_1,h_2,…,h_M\ attention heads. For a document did_i, we define its relevance score under head hmh_m as: (1) shm(di,q)=1|ℐq|∑j∈ℐdi∑z∈ℐqaz→jhms_h_m(d_i,q)= 1|I_q| _j _d_i _z _qa^h_m_z→ j where ℐqI_q and ℐdiI_d_i are the index sets of query tokens and tokens in did_i, respectively, and az→jhma^h_m_z→ j denotes the attention weight from the z-th query token to the j-th token of did_i under head hmh_m. Problem definition. Our goal is to learn a router ϕφ that can route each query q to an optimal head set ϕ(q)φ(q) and aggregates the document relevance scores from the selected heads for re-ranking: (2) s(di,q)=∑hm∈ϕ(q)shm(di,q)s(d_i,q)= _h_m∈φ(q)s_h_m(d_i,q) 3. Methodology Our method follows a two-stage pipeline, which consists of an offline search procedure for pseudo-label construction and a training procedure for the router. We describe each stage in detail below. 3.1. Offline Pseudo-Label Construction Search space. Searching over all heads is impractical for modern LLMs (e.g., 1024 heads in Llama-3.1-8B), since it makes head-set search either too expensive or too crude to yield high-quality sets. We therefore restrict the search space. Specifically, we score each head by the nDCG@10 it achieves when used alone for re-ranking and retain only the top-K heads (e.g., K=64K=64) as a compact head pool ℋH . This focuses the subsequent search on a smaller, higher-quality pool, improving both efficiency and the resulting head sets. Search with an adaptive budget. For each query q, the goal is to identify a head set (q)⊂ℋS(q) to approximate ϕ(q)φ(q) that maximizes ranking quality, measured by nDCG@10 and denoted by ℒrank(q,)L_rank(q,S). Although exhaustive search can yield the optimal solution, it is computationally infeasible in practice, requiring evaluation of 2K2^K subsets per query. We therefore approximate the query-wise optimum with a two-phase search procedure. We set a maximum set size P (e.g., P=8P=8), but allow |(q)|<P|S(q)|<P because the best performance can occur with fewer heads. In the first phase, we perform forward selection with early stopping. Starting from 0=∅S_0= , at step t≥1t≥ 1 we add the head hmh_m that maximizes the objective after inclusion, (3) ht=argmaxℒrank(q,t−1∪hm),t=t−1∪hth_t= _rank (q,S_t-1∪\h_m\ ), _t=S_t-1∪\h_t\ and stop at the first T such that either |T|=P|S_T|=P or no remaining head improves the objective. This stage produces an initial set TS_T with |T|≤p|S_T|≤ p. The initial set can depend on the order in which the heads are considered. To reduce the dependency and explore alternative combinations of heads, in the second phase, the set is refined using swap-based local search. Let 0=TS^0=S_T and let rS^r denote the current set at iteration r. The search considers one-swap moves that replace a selected head hu∈rh_u ^r with an unselected head hv∈ℋ∖rh_v ^r, and selects the best swap by (4) (hu,hv)=argmaxℒrank(q,(r∖hu)∪hv)(h_u,h_v)= _rank\! (q,\;(S^r \h_u\)∪\h_v\ ) Let ′=(r∖hu)∪hvS =(S^r \h_u\)∪\h_v\. The swap is accepted if it improves the objective by at least a tolerance ϵ≥0ε≥ 0, in which case the update (r+1)←′S^(r+1) is applied; otherwise, the procedure terminates and (q)=rS(q)=S^r. The final head set defines a multi-hot label for query q: (5) yhm(q)=1,hm∈(q),0,otherwise.y_h_m(q)= cases1,&h_m (q),\\ 0,&otherwise. cases NQ COVID NFCorpus FiQA Scifact Scidocs FEVER Climate DBPedia Robust04 News Avg BM25 † 30.5 59.5 32.2 23.6 67.9 14.9 65.1 16.5 31.8 40.7 39.5 38.4 RankGPTLlama-3.1-8BRankGPT_Llama-3.1-8B 53.7 75.5 34.3 31.4 69.3 17.4 67.5 23.8 42.9 47.8 46.2 46.3 RankGPTQwen-2.5-7BRankGPT_Qwen-2.5-7B 42.7 70.5 34.1 29.5 69.3 16.6 70.5 19.7 37.1 46.4 43.6 43.6 Base LLM: Llama-3.2-1B-Instruct ICR 35.6 57.2 28.8 22.3 64.8 13.7 75.7 16.8 25.9 36.0 37.8 37.7 QRhead 41.4 67.3 31.6 25.1 69.5 15.5 77.4 23.3 31.8 42.2 42.2 42.5 Corehead 35.6 56.9 29.1 24.8 65.3 14.3 70.6 17.4 21.8 35.1 34.9 36.9 RouteHead 42.4 67.7 32.0 25.9 70.1 15.8 78.0 23.4 32.4 42.8 42.3 43.0 Base LLM: Llama-3.2-3B-Instruct ICR 49.3 72.5 33.8 31.7 73.3 17.4 82.7 24.2 34.7 47.1 44.6 46.5 QRhead 54.9 77.3 35.1 35.0 74.7 18.3 83.6 24.5 36.0 49.7 45.2 48.6 Corehead 51.4 77.3 34.2 35.0 73.8 18.4 82.0 23.2 32.0 48.7 43.2 47.2 RouteHead 56.3 78.6 35.6 36.4 75.1 18.9 84.4 25.3 37.6 50.9 46.4 49.6 Base LLM: Llama-3.1-8B-Instruct ICR 54.0 73.3 34.8 35.6 75.5 19.0 85.8 24.8 36.9 49.0 44.5 48.5 QRhead 58.6 77.5 35.3 39.1 76.2 19.4 85.3 23.8 37.2 51.4 46.1 50.0 Corehead 54.0 75.6 34.1 39.4 74.7 18.6 81.5 21.6 32.6 48.9 42.7 47.5 RouteHead 59.8 78.2 36.0 39.6 76.7 19.9 85.8 24.0 37.6 51.6 46.8 50.5 Base LLM: Qwen-2.5-7B-Instruct ICR 43.1 66.1 32.7 27.0 71.1 16.4 79.2 19.6 35.3 43.0 40.0 43.0 QRhead 49.9 67.7 33.1 29.2 71.0 15.3 80.7 20.1 35.7 43.7 39.8 44.2 Corehead 50.1 67.7 32.5 30.6 70.0 14.5 79.8 19.3 32.4 45.0 42.1 44.0 RouteHead 51.4 68.0 33.4 31.2 71.2 16.0 81.2 20.4 36.3 44.5 40.2 44.9 Traditional re-rankers Contriever 44.6 67.5 32.8 28.4 67.1 18.9 64.2 28.0 39.5 45.7 41.7 43.5 GTR-T5-base 51.4 74.8 32.5 34.7 62.1 15.8 72.9 26.8 37.1 46.1 42.8 45.2 BGE-Reranker-base 55.2 66.4 31.0 31.7 70.8 15.7 88.6 36.5 42.5 39.9 37.0 46.8 MSMARCO-MiniLM 55.8 74.3 35.2 35.1 68.5 17.5 80.4 25.5 45.3 47.9 43.0 48.0 Table 1. Performance comparison (nDCG@10) on BEIR. † represent initial retrieval. Best results are highlighted in Bold. 3.2. Router Training Router Architecture Design. Drawing inspiration from matrix factorization models in recommendation systems, we route query to heads using low-dimensional embeddings. Since head indices are discrete and non-differentiable, we assign each attention head hmh_m a learnable embedding ehm∈ℝdhe_h_m ^d_h to represent head identity for routing. This continuous parameterization allows the router to organize heads into functional regions in a shared embedding space and to share selection patterns across heads with similar behavior. For query representation, we extract a query embedding eq∈ℝdqe_q ^d_q from the frozen LLM hidden states using mean pooling over token representations. Given (eq,ehm)(e_q,e_h_m), the router computes a query-dependent activation score for each head via a bilinear-style interaction: (6) αhm(q)=W2⊤(eh⊙(W1⊤eq+b)) _h_m(q)=W _2 (e_h (W _1e_q+b) ) where ⊙ denotes the Hadamard product, W1⊤∈ℝdq×dhW _1 ^d_q× d_h and b∈ℝdhb ^d_h are parameters of a projection layer that align the dimension of eqe_q with ehme_h_m, and W2⊤∈ℝdhW _2 ^d_h is the weight vector of a linear layer that produces a scalar score. The score is then converted into an independent activation probability using the sigmoid function σ: (7) phm(q)=σ(αhm(q))p_h_m(q)=σ ( _h_m(q) ) Training Objective. Let yhm(q)∈0,1y_h_m(q)∈\0,1\ denote the set indicator of whether head hmh_m is selected for query q by the offline labeling stage. We train the router by minimizing a set-level binary cross-entropy loss: (8) ℒroute=∑m=1KBCE(phm(q),yhm(q))L_route= _m=1^KBCE (p_h_m(q),\,y_h_m(q) ) Because the router outputs independent, non-normalized activation probabilities across heads, an explicit constraint is required to control the number of activated heads. To mitigate over-activation caused by noisy or redundant pseudo labels, an ℓ1 _1-style penalty is added on the activation probabilities: (9) ℒsparse=λ∑m=1Kphm(q)L_sparse=λ _m=1^Kp_h_m(q) This regularizer discourages uniformly high activation probabilities across heads. Thus the total objective is (10) ℒ=ℒroute+ℒsparseL=L_route+L_sparse Model StackExchange Coding Theorem-based Avg. Bio. Earth. Econ. Psy. Rob. Stack. Sus. Pony LC. AoPS TheoT. TheoQ. BM25† 18.2 27.9 16.4 13.4 10.9 16.3 16.1 4.3 24.7 6.5 2.1 7.3 13.7 RankGPTGPT4RankGPT_GPT4 33.8 34.2 16.7 27.0 22.3 27.7 11.1 15.6 3.4 1.2 8.6 0.2 16.8 RankGPTQwen3-32BRankGPT_Qwen3-32B 24.9 29.4 20.9 25.7 18.3 16.0 23.2 7.6 27.6 7.8 8.9 8.4 18.2 RankGPTQwen3-235B-A22BRankGPT_Qwen3-235B-A22B 26.4 26.7 22.1 26.3 18.8 17.0 24.9 8.2 27.2 7.7 11.7 8.6 18.8 Base LLM: Llama-3.2-1B-Instruct ICR 26.2 28.9 11.8 16.5 8.7 6.1 11.5 2.7 7.8 6.0 2.6 4.0 11.1 QRhead 28.4 32.7 15.3 18.1 11.0 6.8 12.0 3.0 10.1 6.3 3.4 5.0 12.6 Corehead 28.6 38.2 19.2 22.6 13.8 11.8 16.9 3.5 20.8 10.7 4.6 10.1 16.7 RouteHead 29.6 39.0 20.0 23.5 14.6 11.3 17.8 3.3 20.3 11.4 5.2 10.0 17.2 Base LLM: Llama-3.2-3B-Instruct ICR 33.4 37.7 18.5 21.2 13.8 10.0 17.4 2.4 13.3 7.4 2.9 7.9 15.5 QRhead 31.5 40.5 22.3 27.9 16.5 12.7 20.0 2.8 18.8 7.8 5.5 9.7 18.0 Corehead 30.7 40.5 22.7 27.2 16.4 14.6 20.1 3.4 22.4 9.5 5.6 10.8 18.6 RouteHead 33.9 41.2 23.1 28.8 17.3 13.5 20.8 3.0 19.5 8.4 6.2 10.2 18.9 Base LLM: Llama-3.1-8B-Instruct ICR 31.1 33.3 18.2 20.5 15.7 13.5 17.8 2.4 16.8 8.3 5.2 11.3 16.2 QRhead 30.8 35.5 20.2 25.1 18.9 15.2 20.3 3.6 20.9 9.8 5.5 11.8 18.1 Corehead 31.2 37.9 22.3 26.9 18.2 16.7 22.1 3.7 24.6 11.0 6.7 12.8 19.5 RouteHead 32.0 38.4 22.8 27.6 18.8 16.2 22.8 3.5 24.0 11.4 7.3 12.5 19.8 Base LLM: Qwen-2.5-7B-Instruct ICR 27.6 32.6 14.5 19.2 14.0 9.4 16.5 2.0 15.7 7.2 3.3 6.7 14.0 QRhead 26.8 34.7 15.1 21.6 13.9 8.9 18.5 3.3 16.8 6.3 4.3 8.0 14.8 Corehead 28.3 35.1 16.7 23.8 15.1 8.7 19.0 4.9 18.6 7.6 4.7 8.8 15.9 RouteHead 29.3 36.5 17.3 23.8 15.5 10.4 20.0 4.1 18.2 7.3 5.5 9.2 16.5 Table 2. Performance comparison (nDCG@10) on BRIGHT . † represent initial retrieval. Best results are highlighted in Bold. 4. Experiments 4.1. Experimental Setup Datasets. We evaluate on two benchmarks. The first is BEIR (Thakur et al., 2021) with eleven sub-datasets across diverse domains: NQ (Kwiatkowski et al., 2019), COVID (Voorhees et al., 2021), NFCorpus (Boteva et al., 2016), FiQA (Maia et al., 2018), SciFact (Wadden et al., 2020), SciDocs (Cohan et al., 2020), FEVER (Thorne et al., 2018), Climate (Diggelmann et al., 2020), DBPedia (Hasibi et al., 2017), Robust04 (Jeronymo et al., 2022) and News. The second is BRIGHT (SU et al., 2025), a reasoning-intensive retrieval benchmark with 1,385 real-world queries from multiple domains (e.g., StackExchange, LeetCode, and math competitions). Baselines. We compare our approach with baselines from two paradigms: the generation-based method RankGPT (Sun et al., 2023) and the attention-based methods ICR (Chen et al., 2025), QRhead (Zhang et al., 2025b), and CoRehead (Tran et al., 2025). We also include several traditional re-rankers: Contriever (Lei et al., 2023), GTR-T5-base (Ni et al., 2022), BGE-Reranker-base (Xiao et al., 2024) and MSMARCO-MiniLM (Reimers and Gurevych, 2019). Base LLMs. We evaluate instruction-tuned LLMs from two families and multiple scales, including Llama-3.2-1B, Llama-3.2-3B and Llama-3.1-8B from the Llama series, as well as Qwen2.5-7B from the Qwen series. We also integrate the open-source reasoning models Qwen3-32B and Qwen3-235B-A22B into RankGPT to evaluate performance on the reasoning-intensive BRIGHT benchmark. Setting. We re-rank the top-200 BM25 candidates and randomly sample 512 queries for evaluation on BEIR. On BRIGHT, we re-rank the top-100 BM25 candidates and evaluate on the full query set. RankGPT uses bubble-sort with a sliding window of size 20 and stride 10, while others re-rank the full candidate list. 4.2. Experiment Results We report results on BEIR in Table 1 and on BRIGHT in Table 2. Across both benchmarks and all base LLMs, RouteHead achieves the best average performance among attention-based methods, and performance increases monotonically within each model family as the backbone size increases. This suggests that our method provides complementary gains on top of stronger LLM backbones. Specifically, on BEIR, RouteHead consistently outperforms the strongest static head baseline (QRhead), indicating that a single globally selected head set does not transfer reliably across heterogeneous IR collections. Under the same evaluation setting, RouteHead with Llama-3.2-3B reaches an average of 49.6, exceeding the strongest traditional re-ranker, MSMARCO-MiniLM (48.0). It also outperforms the generation-based method RankGPT while using a smaller backbone, surpassing RankGPTLlama-3.1-8BRankGPT_Llama-3.1-8B (46.3) and RankGPTQwen-2.5-7BRankGPT_Qwen-2.5-7B (43.6) with Llama-3.2-3B. Although RouteHead underperforms traditional re-rankers on Climate and DBPedia, this gap is not specific to RouteHead: the LLM-based re-ranking methods are generally weaker than strong end-to-end cross-encoders on these two datasets, which emphasize domain terminology and entity-centric matching. RouteHead remains competitive even under reasoning-intensive benchmark BRIGHT. With Llama-3.2-3B, it reaches 18.9, exceeding the RankGPT result with Qwen3-235B-A22B (18.8). Notably, RouteHead is often below CoreHead on coding- and math-oriented domains, whereas it is consistently strong on the StackExchange subsets and often improves multiple domains. One plausible explanation is a training–evaluation distribution mismatch. The router is trained on MSMARCO, whose queries primarily reflect general web information needs and may provide limited coverage of code- and math-oriented questions. Consequently, the learned routing patterns may be less aligned with the highly structured signals required by coding- and theorem-style queries. This observation motivates improving the coverage of router training data, for example by incorporating more code and math queries or by using domain-balanced sampling during pseudo-label construction and router training. Together, these results show that RouteHead delivers strong and consistent improvements within the attention-based re-ranking, , while remaining competitive against stronger traditional and generation-based baselines. 5. Related Work Our work is related to both zero-shot re-ranking with LLMs and the attention mechanism. Zero-shot re-ranking with LLMs is typically formulated in point-wise (Sachan et al., 2022; Liang et al., 2022; Mo et al., 2023), pair-wise (Qin et al., 2024; Meng et al., 2025a; Xu et al., 2025), or list-wise (Ma et al., 2023; Sun et al., 2023; Chen et al., 2025, 2023; Jin et al., 2024; Fu et al., 2024) manners. However, most of them adopt a generation-based paradigm, which incurs high cost and latency due to autoregressive decoding (Zhang et al., 2026b) and is sensitive to prompt design and instruction-following errors. Attention is a central mechanism in transformer LLMs and provides signals that are informative for retrieval (Izacard and Grave, 2021; Tian et al., 2026; Mo et al., 2026; Zhang et al., 2026a). Chen et al. (2025) treat attention as an implicit relevance signal and compute document relevance scores for re-ranking by aggregating query-to-document attention weights across all heads. Earlier mechanistic interpretability studies suggest that only a subset of attention heads is functionally critical (Michel et al., 2019; Voita et al., 2019; Zhang et al., 2024; Wu et al., 2025; Su et al., 2026; Zhang et al., 2025a, 2026c). Motivated by these findings, recent work proposes heuristic rules to identify a small set of important heads for attention-based re-ranking (Zhang et al., 2025b; Tran et al., 2025). However, these methods typically select a fixed head set, which can be suboptimal since the useful heads can differ across queries. This limitation motivated our exploration of real-time, query-dependent head selection. 6. Conclusion We introduce a lightweight router that learns to route queries to attention heads for attention-based re-ranking with LLMs, where relevance scores are computed by aggregating attention signals from the selected heads. The router is trained using pseudo labels constructed by an offline search procedure. Extensive experiments demonstrate the effectiveness of our method. Future work can explore improving the quality and consistency of pseudo labels, enriching the router inputs with additional information (e.g., query type and length). References V. Boteva, D. Gholipour, A. Sokolov, and S. Riezler (2016) A full-text learning to rank dataset for medical information retrieval. In Advances in Information Retrieval: 38th European Conference on IR Research, ECIR 2016, Padua, Italy, March 20–23, 2016. Proceedings 38, p. 716–722. Cited by: §4.1. S. Chen, B. J. Gutierrez, and Y. Su (2025) Attention in large language models yields efficient zero-shot re-rankers. In The Thirteenth International Conference on Learning Representations, Cited by: §1, §4.1, §5. S. Chen, S. Wong, L. Chen, and Y. Tian (2023) Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2306.15595. Cited by: §5. A. Cohan, S. Feldman, I. Beltagy, D. Downey, and D. S. Weld (2020) Specter: document-level representation learning using citation-informed transformers. arXiv preprint arXiv:2004.07180. Cited by: §4.1. T. Diggelmann, J. Boyd-Graber, J. Bulian, M. Ciaramita, and M. Leippold (2020) Climate-fever: a dataset for verification of real-world climate claims. arXiv preprint arXiv:2012.00614. Cited by: §4.1. Y. Fu, R. Panda, X. Niu, X. Yue, H. Hajishirzi, Y. Kim, and H. Peng (2024) Data engineering for scaling language models to 128k context. arXiv preprint arXiv:2402.10171. Cited by: §5. Y. Gao, Y. Xiong, X. Gao, K. Jia, J. Pan, Y. Bi, Y. Dai, J. Sun, H. Wang, and H. Wang (2023) Retrieval-augmented generation for large language models: a survey. arXiv preprint arXiv:2312.10997 2 (1). Cited by: §1. F. Hasibi, F. Nikolaev, C. Xiong, K. Balog, S. E. Bratsberg, A. Kotov, and J. Callan (2017) DBpedia-entity v2: a test collection for entity search. In Proceedings of the 40th international ACM SIGIR conference on research and development in information retrieval, p. 1265–1268. Cited by: §4.1. G. Izacard and E. Grave (2021) Distilling knowledge from reader to retriever for question answering. In International Conference on Learning Representations, External Links: Link Cited by: §5. V. Jeronymo, M. Nascimento, R. Lotufo, and R. Nogueira (2022) MRobust04: a multilingual version of the trec robust 2004 benchmark. External Links: 2209.13738, Link Cited by: §4.1. H. Jin, X. Han, J. Yang, Z. Jiang, Z. Liu, C. Chang, H. Chen, and X. Hu (2024) Llm maybe longlm: self-extend llm context window without tuning. arXiv preprint arXiv:2401.01325. Cited by: §5. V. Karpukhin, B. Oguz, S. Min, P. Lewis, L. Wu, S. Edunov, D. Chen, and W. Yih (2020) Dense passage retrieval for open-domain question answering. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), B. Webber, T. Cohn, Y. He, and Y. Liu (Eds.), Online, p. 6769–6781. External Links: Link, Document Cited by: §1. O. Khattab and M. Zaharia (2020) ColBERT: efficient and effective passage search via contextualized late interaction over bert. In Proceedings of the 43rd International ACM SIGIR Conference on Research and Development in Information Retrieval, SIGIR ’20, New York, NY, USA, p. 39–48. External Links: ISBN 9781450380164, Link, Document Cited by: §1. T. Kwiatkowski, J. Palomaki, O. Redfield, M. Collins, A. Parikh, C. Alberti, D. Epstein, I. Polosukhin, M. Kelcey, J. Devlin, K. Lee, K. N. Toutanova, L. Jones, M. Chang, A. Dai, J. Uszkoreit, Q. Le, and S. Petrov (2019) Natural questions: a benchmark for question answering research. Transactions of the Association of Computational Linguistics. Cited by: §4.1. Y. Lei, L. Ding, Y. Cao, C. Zan, A. Yates, and D. Tao (2023) Unsupervised dense retrieval with relevance-aware contrastive pre-training. In Findings of the Association for Computational Linguistics: ACL 2023, A. Rogers, J. Boyd-Graber, and N. Okazaki (Eds.), Toronto, Canada, p. 10932–10940. External Links: Link, Document Cited by: §4.1. P. Lewis, E. Perez, A. Piktus, F. Petroni, V. Karpukhin, N. Goyal, H. Küttler, M. Lewis, W. Yih, T. Rocktäschel, et al. (2020) Retrieval-augmented generation for knowledge-intensive nlp tasks. Advances in Neural Information Processing Systems 33, p. 9459–9474. Cited by: §1. P. Liang, R. Bommasani, T. Lee, D. Tsipras, D. Soylu, M. Yasunaga, Y. Zhang, D. Narayanan, Y. Wu, A. Kumar, et al. (2022) Holistic evaluation of language models. arXiv preprint arXiv:2211.09110. Cited by: §5. X. Ma, X. Zhang, R. Pradeep, and J. Lin (2023) Zero-shot listwise document reranking with a large language model. arXiv preprint arXiv:2305.02156. Cited by: §5. M. Maia, S. Handschuh, A. Freitas, B. Davis, R. McDermott, M. Zarrouk, and A. Balahur (2018) W’18 open challenge: financial opinion mining and question answering. In Companion Proceedings of the The Web Conference 2018, W ’18, Republic and Canton of Geneva, CHE, p. 1941–1942. External Links: ISBN 9781450356404, Link, Document Cited by: §4.1. C. Meng, N. Arabzadeh, A. Askari, M. Aliannejadi, and M. de Rijke (2024) Ranked list truncation for large language model-based re-ranking. In Proceedings of the 47th International ACM SIGIR Conference on Research and Development in Information Retrieval, p. 141–151. Cited by: §1. C. Meng, N. Arabzadeh, A. Askari, M. Aliannejadi, and M. d. Rijke (2025a) Query performance prediction using relevance judgments generated by large language models. ACM Transactions on Information Systems 43 (4), p. 1–35. Cited by: §5. C. Meng, L. Ou, S. MacAvaney, and J. Dalton (2026) Revisiting text ranking in deep research. In Proceedings of the 49th International ACM SIGIR Conference on Research and Development in Information Retrieval, Cited by: §1. C. Meng, F. Tonolini, F. Mo, N. Aletras, E. Yilmaz, and G. Kazai (2025b) Bridging the gap: from ad-hoc to proactive search in conversations. In Proceedings of the 48th International ACM SIGIR Conference on Research and Development in Information Retrieval, p. 64–74. Cited by: §1. P. Michel, O. Levy, and G. Neubig (2019) Are sixteen heads really better than one?. Advances in neural information processing systems 32. Cited by: §5. F. Mo, Y. Gao, C. Meng, X. Liu, Z. Wu, K. Mao, Z. Wang, P. Chen, Z. Li, X. Li, et al. (2025a) Uniconv: unifying retrieval and response generation for large language models in conversations. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), p. 6936–6949. Cited by: §1. F. Mo, K. Mao, Z. Zhao, H. Qian, H. Chen, Y. Cheng, X. Li, Y. Zhu, Z. Dou, and J. Nie (2025b) A survey of conversational search. ACM Transactions on Information Systems 43 (6), p. 1–50. Cited by: §1. F. Mo, K. Mao, Y. Zhu, Y. Wu, K. Huang, and J. Nie (2023) Convgqr: generative query reformulation for conversational search. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), p. 4998–5012. Cited by: §5. F. Mo, Z. Su, Y. Hui, J. Zhang, J. A. Sun, Z. Liu, C. Zhang, T. Sakai, and J. Nie (2026) Opendecoder: open large language model decoding to incorporate document quality in rag. arXiv preprint arXiv:2601.09028. Cited by: §5. J. Ni, C. Qu, J. Lu, Z. Dai, G. Hernandez Abrego, J. Ma, V. Zhao, Y. Luan, K. Hall, M. Chang, and Y. Yang (2022) Large dual encoders are generalizable retrievers. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, Y. Goldberg, Z. Kozareva, and Y. Zhang (Eds.), Abu Dhabi, United Arab Emirates, p. 9844–9855. External Links: Link, Document Cited by: §4.1. R. Nogueira, Z. Jiang, and J. Lin (2020) Document ranking with a pretrained sequence-to-sequence model. arXiv preprint arXiv:2003.06713. Cited by: §1. R. Nogueira, W. Yang, K. Cho, and J. Lin (2019) Multi-stage document ranking with bert. arXiv preprint arXiv:1910.14424. Cited by: §1. Z. Qin, R. Jagerman, K. Hui, H. Zhuang, J. Wu, L. Yan, J. Shen, T. Liu, J. Liu, D. Metzler, et al. (2024) Large language models are effective text rankers with pairwise ranking prompting. In Findings of the Association for Computational Linguistics: NAACL 2024, p. 1504–1518. Cited by: §1, §5. N. Reimers and I. Gurevych (2019) Sentence-BERT: sentence embeddings using Siamese BERT-networks. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), K. Inui, J. Jiang, V. Ng, and X. Wan (Eds.), Hong Kong, China, p. 3982–3992. External Links: Link, Document Cited by: §4.1. D. Sachan, M. Lewis, M. Joshi, A. Aghajanyan, W. Yih, J. Pineau, and L. Zettlemoyer (2022) Improving passage retrieval with zero-shot question generation. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, p. 3781–3797. Cited by: §1, §5. H. SU, H. Yen, M. Xia, W. Shi, N. Muennighoff, H. Wang, L. Haisu, Q. Shi, Z. S. Siegel, M. Tang, R. Sun, J. Yoon, S. O. Arik, D. Chen, and T. Yu (2025) BRIGHT: a realistic and challenging benchmark for reasoning-intensive retrieval. In The Thirteenth International Conference on Learning Representations, External Links: Link Cited by: §4.1. Z. Su, F. Mo, J. Zhang, Y. Hui, J. Sun, and J. Nie (2026) Parametric retrieval-augmented generation using latent routing of lora adapters. External Links: 2511.17044, Link Cited by: §5. W. Sun, L. Yan, X. Ma, S. Wang, P. Ren, Z. Chen, D. Yin, and Z. Ren (2023) Is ChatGPT good at search? investigating large language models as re-ranking agents. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, H. Bouamor, J. Pino, and K. Bali (Eds.), Singapore, p. 14918–14937. External Links: Link, Document Cited by: §1, §4.1, §5. N. Thakur, N. Reimers, A. Rücklé, A. Srivastava, and I. Gurevych (2021) BEIR: a heterogeneous benchmark for zero-shot evaluation of information retrieval models. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), External Links: Link Cited by: §4.1. J. Thorne, A. Vlachos, C. Christodoulopoulos, and A. Mittal (2018) FEVER: a large-scale dataset for fact extraction and verification. arXiv preprint arXiv:1803.05355. Cited by: §4.1. Y. Tian, F. Mo, W. Zhang, Y. Qi, and J. Nie (2026) ReAttn: improving attention-based re-ranking via attention re-weighting. In Findings of the Association for Computational Linguistics: EACL 2026, V. Demberg, K. Inui, and L. Marquez (Eds.), Rabat, Morocco, p. 1282–1295. External Links: Link, Document, ISBN 979-8-89176-386-9 Cited by: §5. L. Tran, Y. Li, R. Florian, and W. Sun (2025) Contrastive retrieval heads improve attention-based re-ranking. External Links: 2510.02219, Link Cited by: §1, §4.1, §5. E. Voita, D. Talbot, F. Moiseev, R. Sennrich, and I. Titov (2019) Analyzing multi-head self-attention: specialized heads do the heavy lifting, the rest can be pruned. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, p. 5797–5808. Cited by: §5. E. Voorhees, T. Alam, S. Bedrick, D. Demner-Fushman, W. R. Hersh, K. Lo, K. Roberts, I. Soboroff, and L. L. Wang (2021) TREC-covid: constructing a pandemic information retrieval test collection. In ACM SIGIR Forum, Vol. 54, p. 1–12. Cited by: §4.1. D. Wadden, S. Lin, K. Lo, L. L. Wang, M. van Zuylen, A. Cohan, and H. Hajishirzi (2020) Fact or fiction: verifying scientific claims. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), B. Webber, T. Cohn, Y. He, and Y. Liu (Eds.), Online, p. 7534–7550. External Links: Link, Document Cited by: §4.1. Y. Wang, Y. Han, L. Qian, Y. He, X. Peng, D. Feng, Z. Xie, V. J. Zhang, R. Guo, F. Mo, et al. (2026) Conv-finre: a conversational and longitudinal benchmark for utility-grounded financial recommendation. arXiv preprint arXiv:2602.16990. Cited by: §1. W. Wu, Y. Wang, G. Xiao, H. Peng, and Y. Fu (2025) Retrieval head mechanistically explains long-context factuality. In The Thirteenth International Conference on Learning Representations, Cited by: §5. S. Xiao, Z. Liu, P. Zhang, N. Muennighoff, D. Lian, and J. Nie (2024) C-pack: packed resources for general chinese embeddings. In Proceedings of the 47th International ACM SIGIR Conference on Research and Development in Information Retrieval, SIGIR ’24, New York, NY, USA, p. 641–649. External Links: ISBN 9798400704314, Link, Document Cited by: §4.1. Z. Xu, F. Mo, Z. Huang, C. Zhang, P. Yu, B. Wang, J. Lin, and V. Srikumar (2025) A survey of model architectures in information retrieval. arXiv preprint arXiv:2502.14822. Cited by: §5. J. Zhang, F. Mo, X. Wang, and K. Liu (2024) Blind spot navigation in llm reasoning with thought space explorer. arXiv preprint arXiv:2410.24155. Cited by: §5. J. Zhang, F. Mo, T. C. Weerasooriya, R. Dai, X. Han, Y. Fu, D. Wang, and K. Liu (2026a) StaRPO: stability-augmented reinforcement policy optimization. arXiv preprint arXiv:2604.08905. Cited by: §5. J. Zhang, X. Wang, W. Ren, L. Jiang, D. Wang, and K. Liu (2025a) Ratt: a thought structure for coherent and correct llm reasoning. In Proceedings of the AAAI Conference on Artificial Intelligence, Vol. 39, p. 26733–26741. Cited by: §5. W. Zhang, F. Ye, Q. Gao, J. Li, H. Wu, Y. Tian, S. Duan, N. Du, X. Li, and X. Liu (2026b) Context-fidelity boosting: enhancing faithful generation through watermark-inspired decoding. External Links: 2604.22335, Link Cited by: §5. W. Zhang, Y. Yuan, C. Han, Y. Tian, Z. Sun, L. Du, J. Kang, H. Kang, X. Liu, and H. Wu (2026c) Preference heads in large language models: a mechanistic framework for interpretable personalization. External Links: 2604.22345, Link Cited by: §5. W. Zhang, F. Yin, H. Yen, D. Chen, and X. Ye (2025b) Query-focused retrieval heads improve long-context reasoning and re-ranking. In Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing, p. 23791–23805. External Links: Link, Document, ISBN 979-8-89176-332-6 Cited by: §1, §4.1, §5.