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

• URL: http://arxiv.org/abs/2604.24608v1
• Date: Mon, 27 Apr 2026 15:36:54 GMT
• Title: Learning to Route Queries to Heads for Attention-based Re-ranking with Large Language Models
• Abstract summary: We propose a query-dependent head selection method, RouteHead, for attention-based re-ranking with Large Language Models.<n>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.<n> Experiments on diverse benchmarks and multiple LLM backbones show that the proposed method consistently outperforms strong baselines.
• Score: 25.83277662723952
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Efficient and Trainable Language Model Test-Time Scaling via Local Branch Routing [65.59464371673623]
  Local Branch Routing (LBR) is a token-level test-time scaling framework.<n>It expands a small local lookahead tree, forwards all sampled branches through the language model, and uses a lightweight router to select the depth-1 subtree to commit.<n>LBR allows each token decision to use evidence beyond the root next-token distribution while avoiding full solution-level search.
  arXiv  Detail & Related papers  (2026-06-24T03:42:44Z)
• Diagnosing LLM Reranker Behavior Under Fixed Evidence Pools [0.17188280334580197]
  Reranking evaluations study how a reranker orders candidates returned by an upstream retriever.<n>This setup couples ranking behavior with retrieval quality, so differences in output cannot be attributed to the ranking policy alone.<n>We introduce a controlled diagnostic that isolates reranking by using Multi-News clusters as fixed evidence pools.
  arXiv  Detail & Related papers  (2026-02-20T21:07:32Z)
• Think When Needed: Model-Aware Reasoning Routing for LLM-based Ranking [25.69863022367215]
  Reasoning prompting can enhance ranking utility, but its benefits are inconsistent and come at a substantial computational cost.<n>We propose a reasoning routing framework that employs a lightweight, plug-and-play router head to decide whether to use direct inference (Non-Think) or reasoning (Think) for each instance before generation.
  arXiv  Detail & Related papers  (2026-01-26T05:09:07Z)
• Provably Learning Attention with Queries [15.606567893781367]
  We study the problem of learning Transformer-based sequence models with black-box access to their outputs.<n>In this setting, a learner may adaptively query the oracle with any sequence of vectors and observe the corresponding real-valued output.
  arXiv  Detail & Related papers  (2026-01-23T16:28:22Z)
• Lookahead Routing for Large Language Models [24.082620717301477]
  Lookahead is a routing framework that "foresees" potential model outputs and uses these predictions to guide model selection.<n> Empirical evaluations across seven public benchmarks show that Lookahead consistently outperforms existing routing baselines.
  arXiv  Detail & Related papers  (2025-10-22T12:00:21Z)
• Threading the Needle: Reweaving Chain-of-Thought Reasoning to Explain Human Label Variation [60.18907916989796]
  Large Language Models (LLMs) generate chains of thought (CoTs) before giving the final answer.<n>We propose a novel pipeline enriched with linguistically-grounded discourse segmenters to extract supporting and opposing statements for each answer option.<n>We also propose a rank-based HLV evaluation framework that prioritizes the ranking of answers over exact scores.
  arXiv  Detail & Related papers  (2025-05-29T11:47:18Z)
• RAG over Tables: Hierarchical Memory Index, Multi-Stage Retrieval, and Benchmarking [63.253294691180635]
  In real-world scenarios, beyond pure text, a substantial amount of knowledge is stored in tables.<n>We first propose a table-corpora-aware RAG framework, named T-RAG, which consists of the hierarchical memory index, multi-stage retrieval, and graph-aware prompting.
  arXiv  Detail & Related papers  (2025-04-02T04:24:41Z)
• Rank-R1: Enhancing Reasoning in LLM-based Document Rerankers via Reinforcement Learning [76.50690734636477]
  We introduce Rank-R1, a novel LLM-based reranker that performs reasoning over both the user query and candidate documents before performing the ranking task.<n>Our experiments on the TREC DL and BRIGHT datasets show that Rank-R1 is highly effective, especially for complex queries.
  arXiv  Detail & Related papers  (2025-03-08T03:14:26Z)
• Zero-Shot Listwise Document Reranking with a Large Language Model [58.64141622176841]
  We propose Listwise Reranker with a Large Language Model (LRL), which achieves strong reranking effectiveness without using any task-specific training data. Experiments on three TREC web search datasets demonstrate that LRL not only outperforms zero-shot pointwise methods when reranking first-stage retrieval results, but can also act as a final-stage reranker.
  arXiv  Detail & Related papers  (2023-05-03T14:45:34Z)
• LEAPS: End-to-End One-Step Person Search With Learnable Proposals [50.39493100627476]
  We propose an end-to-end one-step person search approach with learnable proposals, named LEAPS. Given a set of sparse and learnable proposals, LEAPS employs a dynamic person search head to directly perform person detection and corresponding re-id feature generation without non-maximum suppression post-processing.
  arXiv  Detail & Related papers  (2023-03-21T13:59:32Z)
• Online Learning of Optimally Diverse Rankings [63.62764375279861]
  We propose an algorithm that efficiently learns the optimal list based on users' feedback only. We show that after $T$ queries, the regret of LDR scales as $O((N-L)log(T))$ where $N$ is the number of all items.
  arXiv  Detail & Related papers  (2021-09-13T12:13:20Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.