Title: Less is More: Contrastive Retrieval Heads Improve Attention-Based Re-Ranking
Primary area: foundation or frontier models, including LLMs
Keywords: Large Language Models, Information Retrieval, Reranking

Abstract: The strong zero-shot and long-context capabilities of recent Large Language Models (LLMs) have paved the way for highly effective re-ranking systems. Attention-based re-rankers leverage attention weights from transformer heads to produce relevance scores, but not all heads are created equally: many contribute noise and redundancy, thus limiting performance. To address this, we introduce $\textit{CoRe heads}$, a small set of retrieval heads identified via a contrastive scoring metric that explicitly rewards high attention heads that correlate with relevant documents, while downplaying nodes with higher attention that correlate with irrelevant documents. This relative ranking criterion isolates the most discriminative heads for re-ranking and yields a state-of-the-art list-wise re-ranker. Extensive experiments with three LLMs show that aggregated signals from CoRe heads, constituting less than 1% of all heads, substantially improve re-ranking accuracy over strong baselines. We further find that CoRe heads are concentrated in middle layers, and pruning the computation of final 50% of model layers preserves accuracy while significantly reducing inference time and memory usage.

TL;DR: We introduce contrastive retrieval heads that improves list-wise re-ranker and allows layer pruning to reduce inference latency and memory usage without sacrificing accuracy.