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

TL;DR

This paper introduces RouteHead, a method that dynamically selects the most relevant attention heads in LLMs for query re-ranking, improving relevance estimation by query-dependent head selection.

Contribution

It proposes a lightweight router that learns to select optimal attention heads per query, enhancing re-ranking performance over static or heuristic methods.

Findings

RouteHead outperforms strong baselines on diverse benchmarks.

Query-dependent head selection improves relevance estimation.

The method is effective across multiple LLM backbones.

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.