基于BERT的列表感知架构IRAG：通过多阶段集体段落注意力实现高效段落重排序

《Array》：IRAG: A BERT-based list-aware architecture for efficient passage reranking via multi-stage collective passage attention

　　

在大规模检索系统中，重排序器需在严格延迟预算内对候选段落列表进行重新排序。基于BERT的交叉编码器虽成本低但独立评分每个段落，忽略了跨候选证据；而基于大语言模型的列表级重排序器精度更高，但推理成本显著增加。本研究提出IRAG（隐式关系增强生成），一种完全基于BERT-base的列表感知重排序架构。IRAG是一个确定性的三阶段流水线：首先，评分阶段使用BERT交叉编码器对每个查询-段落对进行编码并分配相关性得分；其次，选择阶段保留前k个（实验中k=16）高分候选作为有界工作集；第三，交互阶段采用集体段落注意力（CPA），让每个保留段落通过注意力机制更新其表征，其中成对注意力权重基于学习到的全局工作集摘要进行条件化。有界工作集将列表级交互成本O(k2)与候选列表规模n解耦。训练时采用标签感知的工作集构建（始终包含正例），而推理时则基于模型得分的无标签top-k选择。两个阶段通过共享编码器在联合目标下训练，消融实验表明联合训练比单独训练各阶段提升2.7个MAP点。研究人员在WIKIQA和MS MARCO 2.0段落重排序任务上进行评估，与同等条件下训练的BERT-base交叉编码器相比，列表感知阶段仅增加约13%推理延迟。本研究的贡献包括：(1) 描述了一种列表感知重排序架构，其中有界学习工作集将列表级交互成本与候选列表规模解耦；(2) 提出了CPA机制，其成对权重基于全局工作集摘要进行条件化；(3) 提供了控制实验和组件消融分析。

论文解读文章

研究背景与问题

在大规模检索系统中——无论是网络搜索引擎、企业知识库还是问答服务——重排序器从第一阶段检索器接收候选段落列表，必须在严格的延迟预算内重新排序。可接受的预算因应用而异，从亚秒级交互式搜索到多秒级离线或批处理设置不等，但在所有情况下都排除了每查询耗时数十秒的方法或依赖外部推理服务的方案。

这一约束迫使研究人员在两类方法间做出选择。基于BERT的交叉编码器实现了高精度和低成本，但独立评分每个段落：段落i的得分仅取决于查询和段落i，因此那些仅跨候选可见的证据——一个段落定义的术语被另一段落使用、两个段落从独立角度相互印证——被丢弃了。基于大语言模型的列表级重排序器联合处理候选列表并捕获此类列表级证据，但推理成本显著更高：例如，一个3B参数的MonoT5重排序器在V100 GPU上处理1000个候选查询需18.4秒，而RankGPT依赖专有云API，无法保证本地延迟。

研究内容与结论

本研究提出IRAG（隐式关系增强生成），一种完全基于BERT-base的列表感知重排序架构。IRAG是一个确定性的三阶段流水线：评分阶段、选择阶段和交互阶段。通过引入有界工作集，将列表级交互成本与候选列表规模解耦，使得列表级交互变得经济可行。实验表明，与同等条件下训练的BERT-base交叉编码器相比，IRAG的列表感知阶段仅增加约13%推理延迟，同时在WIKIQA和MS MARCO 2.0上取得显著性能提升。消融实验证实联合训练两阶段比单独训练提升2.7个MAP点，CPA机制相比普通自注意力提升1.5个MAP点。

关键技术方法

研究人员采用了三项核心技术方法：第一，评分阶段使用标准BERT-base交叉编码器对每个查询-段落对进行编码，输出独立相关性得分和段落表征；第二，选择阶段通过维护大小为k（实验中设为16）的有界优先级队列，基于粗粒度得分保留最高分候选形成工作集，从而将交互复杂度限制为O(k2)；第三，交互阶段实施集体段落注意力（CPA），先计算全局工作集摘要向量，再基于此摘要条件化计算工作集内各段落间的成对注意力权重，使每个段落能吸收其他段落的信息更新自身表征。训练过程中采用标签感知的工作集构建策略（强制包含正例），而推理时使用纯模型驱动的top-k选择。

研究结果

评分与选择阶段效果：评分阶段提供粗粒度排序信号，选择阶段通过有界工作集有效过滤低质量候选，将后续交互的计算量从O(n2)降低至O(k2)，同时通过联合训练确保选择模块与交互模块协同优化。

CPA交互机制优势：集体段落注意力通过两步实现：子步骤A计算加权全局摘要V_l = Σα_iP_i^(c)；子步骤B计算条件化成对权重w_ij = O_i^(c)?O_j^(c) + V_l^?O_j^(c)。这种设计使注意力权重不仅依赖于两两相似度，还受全局集合上下文调节。消融实验显示，相比普通点积自注意力，CPA带来1.5个MAP点的提升。

训练-推理选择分离策略：训练时采用标签感知工作集构建（始终包含正例并用高分负例填充剩余槽位），推理时使用纯模型得分top-k选择。这种分离避免了早期训练中工作集不含正例导致交互损失无效的问题。

联合训练增益：通过共享编码器和联合损失函数耦合两个阶段，无需可微松弛即可端到端训练。消融实验表明联合训练比分别训练各阶段提升2.7个MAP点。

效率测量：在匹配输入长度条件下，列表感知阶段仅增加约13%推理延迟，验证了有界工作集设计的有效性。

工作集大小影响：通过扫描不同k值发现，适度大小的有界工作集（k=16）优于全列表交互，表明选择性关注最有信息的候选子集比平等处理所有候选更有效。

与其他方法对比：在WIKIQA和MS MARCO 2.0上，IRAG超越同等规模的BERT交叉编码器基线，并与大型LLM基方法（如MonoT5、RankGPT）相比，在保持竞争性精度的同时大幅降低推理成本和基础设施需求。

讨论与结论

讨论部分指出，IRAG属于交互式重排序家族，不同于平面自注意力方法之处在于两点：交互限制在有界学习工作集而非全列表，且成对注意力权重基于全局工作集摘要条件化。与Set-Encoder等排列不变性方法的区别在于交互发生的位置和对象：交互是事后对缓存候选表征的一次性操作，而非贯穿骨干网络的注意力层；每个保留段落直接关注所有工作集同伴，而非通过逐段落的令牌桥接。与GraphRAG等图基检索系统的比较表明，这些系统解决的是不同的语料级离线索引问题，不在效率比较范围内。

研究结论部分总结了三方面贡献：第一，描述了列表感知重排序架构，其中有界学习工作集将列表级交互成本与候选列表规模解耦，列表感知阶段在匹配输入长度下仅增加约13%推理延迟；第二，描述了CPA机制，其成对权重基于全局工作集摘要条件化，消融实验表明相比普通自注意力提升1.5个MAP点；第三，在WIKIQA和MS MARCO 2.0上提供了控制实验和组件消融分析，量化了每个阶段的贡献，并给出了多种子变异性估计以及训练-推理选择差异和联合训练梯度路径的明确说明。