让大模型自动选最合适的注意力头来重排序,提升搜索精准度
Learning to Route Queries to Heads for Attention-based Re-ranking with Large Language Models

- 根据查询动态选择最优注意力头组合,避免固定选法的局限
- 在多个数据集和模型上均超越现有方法,提升重排序效果
- 适合需要高精度搜索的场景,如信息检索与问答系统
大语言模型(LLMs)最近被用于基于注意力信号的细粒度零样本重排序,以估计文档相关性。然而,现有方法要么对所有注意力头的信号进行聚合,要么依赖启发式规则静态选择子集,这可能导致次优结果,因为有效头会随查询或领域变化。此外,简单合并多个头可能因冗余或冲突信号导致性能下降。本文提出一种查询相关的头选择方法 RouteHead,通过学习一个轻量级路由器,将每个查询映射到最优头集合,并仅聚合这些头的注意力信号计算相关性得分。由于缺乏查询-头最优标签,我们先通过离线搜索构建伪标签。路由器用可学习嵌入表示每个头,用冻结的 LLM 隐状态提取查询嵌入,并在伪标签上结合稀疏正则化进行训练。在多种基准和多个 LLM 主干上的实验表明,该方法始终优于强基线。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。