[论文解读] Efficiently Teaching an Effective Dense Retriever with Balanced Topic Aware Sampling

Sebastian Hofstätter, Sheng-Chieh Lin+3 more|arXiv (Cornell University)|Apr 14, 2021

Topic Modeling被引用 4

一句话总结

本文提出 TAS-Balanced，一种主题感知且边际平衡的采样方法，用于通过成对和批次内教师模型的双监督训练高效密集检索器。在单张消费级 GPU 上训练时间不足 48 小时，实现了最先进的低延迟性能（64ms/查询），并在两个 TREC-DL 数据集上相较 BM25 在 NDCG@10 上提升 44%。

ABSTRACT

A vital step towards the widespread adoption of neural retrieval models is their resource efficiency throughout the training, indexing and query workflows. The neural IR community made great advancements in training effective dual-encoder dense retrieval (DR) models recently. A dense text retrieval model uses a single vector representation per query and passage to score a match, which enables low-latency first stage retrieval with a nearest neighbor search. Increasingly common, training approaches require enormous compute power, as they either conduct negative passage sampling out of a continuously updating refreshing index or require very large batch sizes for in-batch negative sampling. Instead of relying on more compute capability, we introduce an efficient topic-aware query and balanced margin sampling technique, called TAS-Balanced. We cluster queries once before training and sample queries out of a cluster per batch. We train our lightweight 6-layer DR model with a novel dual-teacher supervision that combines pairwise and in-batch negative teachers. Our method is trainable on a single consumer-grade GPU in under 48 hours (as opposed to a common configuration of 8x V100s). We show that our TAS-Balanced training method achieves state-of-the-art low-latency (64ms per query) results on two TREC Deep Learning Track query sets. Evaluated on NDCG@10, we outperform BM25 by 44%, a plainly trained DR by 19%, docT5query by 11%, and the previous best DR model by 5%. Additionally, TAS-Balanced produces the first dense retriever that outperforms every other method on recall at any cutoff on TREC-DL and allows more resource intensive re-ranking models to operate on fewer passages to improve results further.

研究动机与目标

• 解决训练高效密集检索模型的高计算成本问题，此类模型通常需要大规模 GPU 集群或超大批次。
• 通过引入一种最大化每批次信息增益的采样策略，减少对昂贵基础设施的依赖。
• 在不增加硬件要求的前提下，提升训练效率和模型性能，使更多用户能够访问高性能密集检索模型。
• 通过增强第一阶段召回能力和与重排序器的兼容性，实现低延迟推理的高性能检索流水线。
• 设计一种资源高效且对不同随机种子和教师监督配置均具鲁棒性的训练方法。

提出的方法

• 在训练前一次性使用语义点积相似度将查询聚类为主题组，实现主题感知的批次构成。
• 通过每批次仅从单一主题聚类中采样查询来构建训练批次，以集中语义信息并提升批次内负样本教学效果。
• 通过选择边际值相近的文本对来平衡成对教师模型的得分边际，以稳定训练并提升泛化能力。
• 结合使用成对拼接 BERT（BERT-CAT）和基于 ColBERT 的批次内负样本教师模型，实施双监督蒸馏以提升信号质量。
• 使用 Margin-MSE 损失联合蒸馏教师信号，训练一个轻量级的 6 层双编码器 BERT-DOT 模型，以提升排序质量。
• 全程使用单张消费级 GPU（11GB 显存）进行端到端训练，避免对大批次或持续更新的索引策略的依赖。

实验结果

研究问题

• RQ1 与随机采样相比，主题感知采样（TAS）在提升训练效率和检索性能方面有多有效？
• RQ2 在不同训练配置下，平衡成对教师得分边际（TAS-Balanced）是否能进一步提升模型性能？
• RQ3 来自成对教师和批次内负样本教师的双监督如何影响模型的有效性和鲁棒性？
• RQ4 TAS-Balanced 是否能在单张消费级 GPU 上训练出高性能密集检索器，且优于在大规模 GPU 集群上训练的模型？
• RQ5 在与重排序器集成的端到端检索流水线中，TAS-Balanced 在低延迟和中等延迟场景下的表现如何？

主要发现

• TAS-Balanced 在两个 TREC-DL 数据集上实现了 SOTA 的 NDCG@10 性能，相较 BM25 提升 44%，相较普通微调的 DR 模型提升 19%，相较此前最佳 DR 模型提升 5%。
• 模型实现了 64ms/查询 的低延迟推理时间，适用于实时应用场景。
• TAS-Balanced 是首个在 TREC-DL 上所有截断点（cutoff）上均超越所有基线模型的密集检索器，显著提升了第一阶段召回率。
• 将 TAS-Balanced 与 docT5query 融合后，Recall@1K 显著提升，且性能优于延迟高 2–6 倍的系统，展现出强大的流水线兼容性。
• 模型可在单张消费级 GPU（11GB 显存）上于 48 小时内完成训练，无需多 GPU 或大批次训练设置。
• 双监督结合 Margin-MSE 损失取得最强性能，尤其在召回指标上，且在不同随机种子和教师配置下均表现出一致的性能增益。