[论文解读] Intra-Document Cascading: Learning to Select Passages for Neural Document Ranking

Sebastian Hofstätter, Bhaskar Mitra+3 more|arXiv (Cornell University)|May 20, 2021

Topic Modeling参考文献 40被引用 10

一句话总结

本文提出了一种名为文档内级联重排序（IDCM）的两阶段神经文档重排序模型，该模型利用快速的知识蒸馏学生模型（ESM）从文档中筛选出前k个段落，随后仅对这些段落应用高精度的基于BERT的教师模型（ETM）。通过早期剔除无关内容，IDCM在保持MS MARCO和TREC深度学习赛道基准测试最先进效果的同时，将查询延迟降低了400%以上。

ABSTRACT

An emerging recipe for achieving state-of-the-art effectiveness in neural document re-ranking involves utilizing large pre-trained language models - e.g., BERT - to evaluate all individual passages in the document and then aggregating the outputs by pooling or additional Transformer layers. A major drawback of this approach is high query latency due to the cost of evaluating every passage in the document with BERT. To make matters worse, this high inference cost and latency varies based on the length of the document, with longer documents requiring more time and computation. To address this challenge, we adopt an intra-document cascading strategy, which prunes passages of a candidate document using a less expensive model, called ESM, before running a scoring model that is more expensive and effective, called ETM. We found it best to train ESM (short for Efficient Student Model) via knowledge distillation from the ETM (short for Effective Teacher Model) e.g., BERT. This pruning allows us to only run the ETM model on a smaller set of passages whose size does not vary by document length. Our experiments on the MS MARCO and TREC Deep Learning Track benchmarks suggest that the proposed Intra-Document Cascaded Ranking Model (IDCM) leads to over 400% lower query latency by providing essentially the same effectiveness as the state-of-the-art BERT-based document ranking models.

研究动机与目标

• 为解决基于BERT的神经文档重排序模型推理延迟过高的问题，该问题由于逐段BERT打分机制而随文档长度线性增长。
• 通过引入级联的两阶段模型架构，在不牺牲排序效果的前提下提升效率。
• 通过降低计算成本，使强效神经重排序模型能够在大规模生产系统中实现实际部署。
• 探索知识蒸馏作为一种自监督方法，利用BERT自身输出的分数作为监督信号，训练高效的段落筛选模型。

提出的方法

• 该模型采用文档内级联策略：轻量级的知识蒸馏学生模型（ESM）首先对文档中的段落进行排序并选出前k个段落。
• ESM通过知识蒸馏进行训练，其监督信号来源于基于BERT的教师模型（ETM）生成的段落级分数，从而使其能够近似BERT的排序行为。
• 仅对选出的k个段落执行完整的BERT模型打分（ETM），从而大幅减少每个文档所需的昂贵BERT推理次数。
• ESM进一步通过文档内排序损失进行优化，以提升其选择质量，而不仅限于匹配BERT的分数。
• 框架采用对称填充以保留位置信号，使模型能够学习段落位置的提示信息。
• 整个流水线采用多阶段训练方式，以平衡选择模块与打分模块之间的协作。

实验结果

研究问题

• RQ1 RQ1：IDCM是否能在显著降低计算成本和查询延迟的同时，达到与完整BERT重排序器相当的排序效果？
• RQ2 RQ2：ESM模块选择的段落数量k如何影响效果与效率之间的权衡？
• RQ3 RQ3：ESM模块对BERT模型所选顶级段落的召回能力如何？
• RQ4 RQ4：在不同相关性等级和段落位置下，段落选择行为呈现出何种模式？

主要发现

• 与完整BERT重排序器相比，IDCM将中位查询延迟降低了400%以上，在MS MARCO和TREC DL 2019基准测试中保持了相同的先进效果。
• 当ESM通过知识蒸馏结合BERT段落分数与文档内排序损失进行训练，并选择k=40个段落供BERT打分时，性能达到最佳。
• ESM对顶级段落的召回率很高，尤其在高度相关文档中表现突出，表明其与BERT的选择行为高度一致。
• 模型学会了利用文档的开头和结尾，这可能是由于对称填充引入的位置信号，有助于识别文档边界处的相关内容。
• 单次BERT对一个段落的推理增加约25ms的中位延迟，而CK(S)处理40个段落的运行时间相近，充分体现了级联机制带来的效率优势。
• 该框架通过用自监督蒸馏替代人工标注，实现了强效神经重排序器的高效部署，显著降低了人力与成本。