用对比方法筛选关键注意力头，显著提升检索排序效果
Contrastive Retrieval Heads Improve Attention-Based Re-Ranking
- 通过对比评分选出与相关文档高度相关的注意力头
- 仅用不到1%的注意力头就超越强基线模型
- 筛选出的头集中在中间层，可大幅降低推理开销
近期大语言模型（LLMs）具备强大的零样本和长上下文能力，推动了高效重排序系统的发展。基于注意力的重排序器利用Transformer头的注意力权重生成相关性分数，但并非所有头都有效：许多头引入噪声和冗余，限制性能。为此，本文提出CoRe头，通过对比评分指标识别出一组小型检索头，该指标明确奖励与相关文档高度相关的高注意力头，同时抑制与无关文档相关的高注意力头。这种相对排序准则能有效分离最具区分性的头用于重排序，实现当前最优的列表级重排序。在三个LLM上的大量实验表明，由CoRe头聚合的信号（占全部头数不足1%）显著优于强基线模型。进一步发现，CoRe头集中于模型中间层，剪枝最后50%层的计算可在保持准确率的同时显著降低推理时间和内存占用。
原文摘要 · Abstract (English)
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 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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。