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

Abstract: 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.

• The paper introduces RouteHead, a mechanism that dynamically selects optimal attention heads based on query characteristics for enhanced re-ranking.
• The paper employs a two-stage approach with offline pseudo-labeling and bilinear router training to capture query-specific signals while ensuring sparsity.
• The paper demonstrates superior nDCG@10 scores across diverse benchmarks, underscoring improved scalability and interpretability compared to static head selection methods.

Query-Dependent Attention Head Routing for LLM-Based Re-Ranking

Motivation and Problem Statement

The study "Learning to Route Queries to Heads for Attention-based Re-ranking with LLMs" (2604.24608) introduces RouteHead, a novel mechanism for fine-grained, query-dependent selection of attention heads within LLMs for zero-shot re-ranking tasks. Traditional neural re-rankers, including bi-encoders and cross-encoders, achieve strong results but require extensive supervised training and labeled data, limiting scalability in diverse IR domains. Recent advances leverage LLMs as zero-shot re-rankers, either via generation-based list ranking or by extracting relevance from internal attention metrics. However, current approaches aggregate signals across all heads or rely on static, heuristic-based head subsets, failing to capture the dynamic, query-specific utility of attention heads and often degrading performance due to mixing informative and redundant ranking signals.

RouteHead: Query-Dependent Head Routing

RouteHead addresses the static nature of previous attention-based re-ranking by introducing a lightweight router capable of selecting the optimal head set for each query in real time. The method consists of two stages:

1. Offline Pseudo-Label Construction: For each query, pseudo-optimal head sets are generated via a constrained combinatorial search in a high-quality head pool, constructed based on individual head nDCG@10 scores. Forward selection with early stopping and swap-based local search maximize ranking objectives, providing query-specific multi-hot labels for top head selections.
2. Router Architecture and Training: Drawing from matrix factorization principles, each head is parameterized by a learnable embedding and queries are embedded by mean pooling over frozen LLM hidden states. A bilinear-style interaction produces head activation scores which are independently converted to probabilities and regularized for sparsity. Training minimizes binary cross-entropy between router outputs and pseudo-labels, plus an $\ell_1$ penalty to avoid excessive head activation.

This flexible routing paradigm enables document relevance scores to be aggregated from query-adaptive head subsets, mitigating redundancy and enhancing signal diversity.

Empirical Evaluation and Comparative Performance

Experiments are conducted on two heterogeneous benchmarks: BEIR (eleven retrieval datasets, including NQ, COVID, SciFact, FEVER, Robust04) and the reasoning-focused BRIGHT benchmark. Baselines include both traditional dense and cross-encoder re-rankers (e.g., MSMARCO-MiniLM, GTR-T5-base, Contriever), as well as prior attention-based methods (ICR, QRhead, Corehead) and generation-based LLM approaches (RankGPT).

RouteHead consistently achieves superior performance among attention-based methods across all tested LLMs and datasets. For example, with Llama-3.2-3B, RouteHead attains an average nDCG@10 of 49.6 on BEIR, surpassing RankGPT (46.3/43.6 for Llama/Qwen) and even the strongest MSMARCO-MiniLM traditional re-ranker (48.0). On BRIGHT, RouteHead with Llama-3.2-3B achieves 18.9, exceeding Qwen3-235B-A22B RankGPT (18.8). Notably, RouteHead improves monotonically with larger backbone LLMs, indicating strong complementarity to backbone model scaling. However, for domain-specialized datasets (e.g., Climate, DBPedia), traditional cross-encoders still hold an advantage, reflecting the need for enhanced domain adaptation.

Figure 1: Performance comparison (nDCG@10) across three datasets for different attention-based re-ranking methods on Llama-3.2-1B. ICR(middle) denotes direct selection of all heads from the middle layers.

The comparative results Figure 1 underscore the dataset-dependence of static head selection and highlight RouteHead's advantage in adapting to query/domain-specific signal requirements.

Theoretical and Practical Implications

The introduction of query-dependent head routing has several implications:

• Signal Diversification: Aggregating relevance from optimal, query-specific head subsets avoids dilution from noisy/redundant heads, yielding higher quality ranking signals.
• Generalization Across Heterogeneous Collections: The dynamic routing mechanism is significantly more robust to distributional shifts and domain heterogeneity, unlike static head selection approaches.
• Efficiency in Zero-Shot Settings: By bypassing autoregressive decoding and reducing reliance on prompt engineering, RouteHead offers practical deployment opportunities for LLM-based IR systems.
• Mechanistic LLM Interpretability: The learnable head embeddings and the query-head interaction mapping provide opportunities for mechanistic analysis of attention specialization and adaptation, informing future LLM interpretability studies.
• Scalability Potential: The lightweight router architecture is amenable to integration across increasingly large LLMs and long-context retrieval applications.

Limitations and Future Directions

Despite RouteHead's strengths, several limitations persist:

• Performance on domain-focused datasets (entity-centric or terminology-heavy) remains inferior to end-to-end cross-encoders.
• Router training is limited by available pseudo-label coverage; training on datasets with broader domain representation (e.g., code/math reasoning) may further enhance generalization.
• The pseudo-label generation process is computationally intensive, though mitigated via head pool constraining.

Future work should explore enriching router input with additional query metadata (type, length), improving the quality and coverage of pseudo-labels (possibly via weak supervision or data augmentation), and investigating deeper mechanistic connections between head specialization and reasoning signal transfer in LLMs.

Conclusion

RouteHead establishes a principled framework for real-time, query-adaptive attention head routing in LLM-based zero-shot re-ranking. By combining pseudo-optimal head set construction and robust embedding-driven routing, it achieves consistently strong performance across multi-domain IR benchmarks and outperforms both static attention-based and generation-based LLM re-rankers. The method advances signal interpretability, generalization, and practical efficiency in retrieval system architectures. Enhancing domain adaptation and expanding interpretability analysis represent promising avenues for future research in attention-based LLM ranking.