[论文解读] HLATR: Enhance Multi-stage Text Retrieval with Hybrid List Aware Transformer Reranking

Yanzhao Zhang, Dingkun Long+2 more|arXiv (Cornell University)|May 21, 2022

Topic Modeling被引用 12

一句话总结

HLATR 提出了一种轻量级、混合列表感知的 Transformer 重排序模型，通过在多阶段文本检索中融合来自独立阶段的检索与重排序特征，实现高效性能。该模型利用源自检索排名的位置感知嵌入以及基于 BERT 的相关性 logit，实现了最先进的性能，同时推理开销极低。与基线模型相比，其在 MS MARCO Passage 数据集上将 MRR@10 提升了最多 4.3 个百分点，在 Document 数据集上提升了 3.5 个百分点。

ABSTRACT

Deep pre-trained language models (e,g. BERT) are effective at large-scale text retrieval task. Existing text retrieval systems with state-of-the-art performance usually adopt a retrieve-then-reranking architecture due to the high computational cost of pre-trained language models and the large corpus size. Under such a multi-stage architecture, previous studies mainly focused on optimizing single stage of the framework thus improving the overall retrieval performance. However, how to directly couple multi-stage features for optimization has not been well studied. In this paper, we design Hybrid List Aware Transformer Reranking (HLATR) as a subsequent reranking module to incorporate both retrieval and reranking stage features. HLATR is lightweight and can be easily parallelized with existing text retrieval systems so that the reranking process can be performed in a single yet efficient processing. Empirical experiments on two large-scale text retrieval datasets show that HLATR can efficiently improve the ranking performance of existing multi-stage text retrieval methods.

研究动机与目标

• 解决多阶段文本检索系统中检索与重排序阶段联合优化的空白。
• 通过有效融合粗粒度检索与细粒度重排序特征，提升排序性能。
• 设计一种轻量级、可并行化的重排序模块，可无缝集成至现有检索流水线。
• 验证基于注意力机制的列表感知建模相较于简单模型，能更好地捕捉文档相关性依赖关系。
• 证明对比学习损失在处理不平衡文本排序数据时，优于二元交叉熵损失。

提出的方法

• HLATR 使用 Transformer 编码器，其中每个标记代表一个文档，输入嵌入由预训练重排序模型的相关性 logit 构成。
• 位置嵌入源自文档在初始检索列表中的排名，实现列表感知建模。
• 模型采用对比学习目标以优化相关性排序，提升在不平衡数据上的泛化能力。
• 设计为即插即用模块，可无缝集成至任意现有检索与重排序流水线，无需架构重构。
• 该架构轻量化，参数量极少，支持快速推理（每 1000 个查询仅需 2ms）。
• 消融研究对比了使用 MLP 替代 Transformer 以及使用 BCE 损失替代对比损失的变体，以验证设计选择。

实验结果

研究问题

• RQ1 融合检索阶段的排名与重排序阶段的相关性得分，能否提升整体文本检索性能？
• RQ2 相较于非列表感知模型，列表感知的 Transformer 架构在重排序任务中是否更有效？
• RQ3 在不平衡文本排序数据背景下，对比学习是否优于二元交叉熵损失？
• RQ4 HLATR 模型的架构设计（深度、宽度）如何影响其性能与效率？
• RQ5 HLATR 是否能在现实系统中高效部署，且不会显著增加推理成本？

主要发现

• 在 MS MARCO Passage 数据集上，HLATR 的 MRR@10 达到 42.6，比最佳基线模型（WCR）高出 0.8 个百分点。
• 在 MS MARCO Document 数据集上，HLATR 将 MRR@10 从检索阶段的 39.4 提升至 43.5，提升 4.1 个百分点。
• 消融研究显示，将 Transformer 编码器替换为两层 MLP 会降低性能，证实注意力机制在建模列表依赖关系中的重要性。
• 使用二元交叉熵损失替代对比损失导致 MRR@10 下降 0.9 个百分点，表明对比学习更适合处理不平衡排序数据。
• HLATR 的推理时间仅为每 1000 个查询 2ms，比 BERT-base 和 BERT-large 模型快 300 至 450 倍。
• 最优性能在较小的隐藏维度（d=128）和 3–4 个 Transformer 层时达到，表明高容量对特征融合并非必需。