[论文解读] Pre-training Tasks for Embedding-based Large-scale Retrieval

Wei-Cheng Chang|arXiv (Cornell University)|Feb 10, 2020

Topic Modeling被引用 101

一句话总结

该论文分析两塔 Transformer 检索模型的预训练任务，并展示段落级任务（Inverse Cloze Task、Body First Selection、Wiki Link Prediction）在检索性能上显著优于 BM-25，而基于 token 的 MLM 提供有限的增益。

ABSTRACT

We consider the large-scale query-document retrieval problem: given a query (e.g., a question), return the set of relevant documents (e.g., paragraphs containing the answer) from a large document corpus. This problem is often solved in two steps. The retrieval phase first reduces the solution space, returning a subset of candidate documents. The scoring phase then re-ranks the documents. Critically, the retrieval algorithm not only desires high recall but also requires to be highly efficient, returning candidates in time sublinear to the number of documents. Unlike the scoring phase witnessing significant advances recently due to the BERT-style pre-training tasks on cross-attention models, the retrieval phase remains less well studied. Most previous works rely on classic Information Retrieval (IR) methods such as BM-25 (token matching + TF-IDF weights). These models only accept sparse handcrafted features and can not be optimized for different downstream tasks of interest. In this paper, we conduct a comprehensive study on the embedding-based retrieval models. We show that the key ingredient of learning a strong embedding-based Transformer model is the set of pre-training tasks. With adequately designed paragraph-level pre-training tasks, the Transformer models can remarkably improve over the widely-used BM-25 as well as embedding models without Transformers. The paragraph-level pre-training tasks we studied are Inverse Cloze Task (ICT), Body First Selection (BFS), Wiki Link Prediction (WLP), and the combination of all three.

研究动机与目标

• 推动大规模查询-文档检索问题及在两阶段系统中实现高效检索的需求。
• 研究预训练任务如何影响两塔 Transformer 检索器的性能。
• 评估段落级预训练任务 ICT、BFS、WLP 及其组合，相对于 token-level MLM 与 BM-25 的表现。
• 证明设计良好的预训练任务使两塔模型在检索设置中超越 BM-25 与 BoW 基线。

提出的方法

• 定义一个两塔检索模型，查询编码器和文档编码器通过 Transformer 架构生成嵌入。
• 使用对候选文档的 Softmax，通过 Sampled Softmax 近似 fullSoftmax 来训练。
• 提出并评估三种段落级预训练任务：Inverse Cloze Task (ICT)、Body First Selection (BFS)、Wiki Link Prediction (WLP)，以及 ICT+BFS+WLP 的组合，相较于 token-level MLM。
• 使用来源于维基百科的数据来构建用于预训练的正向 (q, d) 对，并在下游检索问答数据集（SQuAD、Natural Questions）以及开放域设置上进行微调。
• 与 BM-25 和 BoW-MLP 基线进行比较，以评估 recall@k 指标的有效性。
• 实验设置使用具有 512 维嵌入、64-token 查询、288-token 文档的两塔编码器，并在 32 TPU v3 上进行 100K 的预训练步骤。

实验结果

研究问题

• RQ1 不同的预训练任务如何影响大规模检索中两塔 Transformer 检索模型的有效性？
• RQ2 段落级预训练任务是否优于 token-level MLM 及诸如 BM-25 的传统 IR 基线在检索任务中的表现？
• RQ3 将 ICT、BFS 与 WLP 相结合是否在数据较少或开放域设置下带来额外增益？
• RQ4 模型深度和嵌入维度如何与预训练任务交互，影响检索召回率？

主要发现

• 正确设计的段落级预训练任务使两塔 Transformer 模型在检索任务中显著优于 BM-25 与 BoW 基线。
• 段落级预训练 ICT、BFS、WLP 带来显著增益，而 token-level MLM 仅提供边际改进。
• ICT+BFS+WLP 的组合在 SQuAD 与 Natural Questions 上始终优于单一任务，特别在低资源和开放域场景中。
• Transformer 编码器比浅层 BoW-MLP 编码器从段落级预训练中获益更多，且更大的嵌入维度提升了性能。
• 开放域检索实验表明 ICT+BFS+WLP 与 ICT 在候选集合很大时也能带来稳健的增益。