[论文解读] Ultron: An Ultimate Retriever on Corpus with a Model-based Indexer

Yujia Zhou, Jing Yao+4 more|arXiv (Cornell University)|Aug 19, 2022

Topic Modeling被引用 15

一句话总结

Ultron 提出了一种以模型为中心的文档检索系统，将文档知识直接编码到生成式 Transformer 模型中，从而无需使用独立的索引。通过使用语义文档ID（URL/标题或产品量化编码）和三阶段训练策略，Ultron 在相比 DPR 等双编码器基线方法显著降低内存占用和延迟的前提下，实现了最先进水平的检索性能。

ABSTRACT

Document retrieval has been extensively studied within the index-retrieve framework for decades, which has withstood the test of time. Unfortunately, such a pipelined framework limits the optimization of the final retrieval quality, because indexing and retrieving are separated stages that can not be jointly optimized in an end-to-end manner. In order to unify these two stages, we explore a model-based indexer for document retrieval. Concretely, we propose Ultron, which encodes the knowledge of all documents into the model and aims to directly retrieve relevant documents end-to-end. For the model-based indexer, how to represent docids and how to train the model are two main issues to be explored. Existing solutions suffer from semantically deficient docids and limited supervised data. To tackle these two problems, first, we devise two types of docids that are richer in semantics and easier for model inference. In addition, we propose a three-stage training workflow to capture more knowledge contained in the corpus and associations between queries and docids. Experiments on two public datasets demonstrate the superiority of Ultron over advanced baselines for document retrieval.

研究动机与目标

• 为克服传统流水线索引-检索框架在文档检索中的局限性，此类框架将索引与检索分离，阻碍端到端优化。
• 设计更丰富、语义意义明确的文档ID，以增强模型学习与推理效率。
• 通过引入三阶段训练工作流，在检索任务中缓解标注数据稀缺问题，即先在通用数据和搜索导向数据上进行预训练，再进行监督微调。
• 通过序列到序列生成实现端到端文档检索，直接从查询生成排序后的文档ID列表，无需外部索引。

提出的方法

• 使用基于 T5 的编码器-解码器模型，将文档检索视为从查询到文档ID列表的序列到序列生成任务。
• 引入两种语义文档ID类型：基于关键词的（URL 和标题）与基于语义的（文档嵌入的产物量化编码）。
• 采用三阶段训练流程：在文档内容与文档ID对上进行通用预训练，利用伪查询在搜索导向数据上进行预训练，最后在相关性标注数据上进行监督微调。
• 在推理阶段应用约束束搜索（constrained beam search），直接从模型的自回归概率中生成前k名的文档ID排序结果。
• 仅存储文档ID序列的前缀树，与向量化索引相比，大幅降低内存开销。
• 利用产品量化（PQ）生成紧凑且语义一致的标识符，使具有相似语义的文档共享相同前缀。

实验结果

研究问题

• RQ1 基于模型的索引器是否能在文档检索准确率上超越传统的索引-检索流水线？
• RQ2 不同类型的文档ID表示（关键词 vs. 语义）如何影响检索性能与模型收敛？
• RQ3 在标注数据有限的情况下，三阶段训练策略在多大程度上提升了模型泛化能力与检索性能？
• RQ4 通过序列生成实现的端到端检索，是否能在内存占用与延迟方面优于 DPR 等双编码器方法？

主要发现

• Ultron-PQ 在 MS MARCO 数据集上实现了 MRR@10 为 0.3155，优于 DPR（0.2908）及其他基线模型，且仅使用 69.6MB 内存。
• 在包含 320K 篇文档的语料上，推理延迟从 DPR 的 18.87ms 降低至 8.90ms，展现出更优的效率。
• 与 DPR 的向量化索引相比，内存成本降低了 90%，因为 Ultron 仅存储文档ID的前缀树。
• 三阶段训练策略显著提升了性能，其中监督微调对最终性能提升贡献最大。
• 基于语义的文档ID（PQ）在共享前缀的文档间表现出更高的相似性一致性，有助于模型收敛。
• 消融实验表明，若移除监督微调，检索性能将出现显著下降，验证了其必要性。