[Paper Review] Pre-training Tasks for Embedding-based Large-scale Retrieval

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

Topic ModelingComputer Science101 citations

TL;DR

The paper analyzes pre-training tasks for two-tower Transformer retrieval models and shows paragraph-level tasks (Inverse Cloze Task, Body First Selection, Wiki Link Prediction) substantially improve retrieval performance over BM-25, while token-level MLM provides limited gains.

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.

Motivation & Objective

• Motivate the large-scale query-document retrieval problem and the need for efficient retrieval in two-stage systems.
• Investigate how pre-training tasks influence performance of two-tower Transformer retrievers.
• Evaluate paragraph-level pre-training tasks ICT, BFS, WLP, and their combination, against token-level MLM and BM-25.
• Demonstrate that properly designed pre-training tasks enable two-tower models to outperform BM-25 and BoW baselines in retrieval settings.

Proposed method

• Define a two-tower retrieval model with query and document encoders producing embeddings via Transformer architectures.
• Train using Softmax over candidate documents with Sampled Softmax to approximate the fullSoftmax.
• Propose and evaluate three paragraph-level pre-training tasks: Inverse Cloze Task (ICT), Body First Selection (BFS), Wiki Link Prediction (WLP), and a combination ICT+BFS+WLP, compared to token-level MLM.
• Use Wikipedia-derived data to construct positive (q, d) pairs for pre-training, and fine-tune on downstream retrieval QA datasets (SQuAD, Natural Questions) and open-domain settings.
• Compare against BM-25 and BoW-MLP baselines to assess effectiveness in recall@k metrics.
• Experiment setup uses two-tower encoders with 512-dimensional embeddings, 64-token queries, 288-token documents, and 100K pre-training steps on 32 TPU v3s.

Experimental results

Research questions

• RQ1 How do different pre-training tasks impact the effectiveness of two-tower Transformer retrieval models in large-scale retrieval?
• RQ2 Do paragraph-level pre-training tasks outperform token-level MLM and traditional IR baselines like BM-25 for retrieval tasks?
• RQ3 Does combining ICT, BFS, and WLP yield additional gains, especially in low-data or open-domain settings?
• RQ4 How do model depth and embedding dimensionality interact with pre-training tasks to affect retrieval recall?

Key findings

• Two-tower Transformer models with properly designed paragraph-level pre-training tasks significantly outperform BM-25 and BoW baselines in retrieval tasks.
• Paragraph-level pre-training ICT, BFS, and WLP yield substantial gains, while token-level MLM provides only marginal improvements.
• The combination ICT+BFS+WLP consistently outperforms individual tasks on SQuAD and Natural Questions, especially in low-resource settings and open-domain scenarios.
• Transformer encoders benefit more from paragraph-level pre-training than shallow BoW-MLP encoders, and larger embedding dimensions improve performance.
• Open-domain retrieval experiments show ICT+BFS+WLP and ICT provide robust gains even with large candidate sets.