PaperlayerView exact revision on arXiv

Source-linked AI summary

Pre-training Tasks for Embedding-based Large-scale Retrieval

Wei-Cheng Chang, Felix X. Yu, Yin-Wen Chang, Yiming Yang, Sanjiv Kumar

arXiv:2002.03932v1cs.LGcs.CLcs.IRstat.ML

TL;DR

Large-scale retrieval needs high-recall, sublinear-time candidate search, yet retrieval pre-training is less studied than cross-attention scoring. The paper comprehensively studies pre-training tasks for two-tower embedding models. It finds that paragraph-level tasks substantially improve Transformer retrieval, including over BM-25 and non-Transformer embedding models.

  • Problem

    Efficient retrieval requires high recall and sublinear-time candidate search, but useful pre-training tasks for two-tower Transformer retrieval remain insufficiently studied.

  • Method

    The paper conducts a comprehensive study of two-tower embedding-based retrieval using paragraph-level tasks ICT, BFS, WLP, and their combination.

  • Results

    Proper paragraph-level pre-training substantially improves two-tower Transformer retrieval, outperforming BM-25, while masked-LM yields only marginal gains.

  • Takeaways & Limitations

    Pre-training task design is central to learning strong embedding-based Transformer models for efficient large-scale retrieval.

  • Takeaways & Limitations

    The study assumes pre-training data consists of positive query-document pairs and reports that adding bi-grams does not further improve performance, possibly because of over-fitting.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Large-scale retrieval must efficiently find high-recall candidates before expensive re-ranking, but retrieval has received less study than BERT-style scoring. This paper examines two-tower Transformer retrieval and finds that paragraph-level pre-training tasks are central to its effectiveness.

  • Retrieval setting: Retrieval first narrows a large corpus to candidates, requiring both high recall and sublinear-time search before scoring re-ranks them.Computing cross-attention scores for every document is prohibitively expensive.
  • Retrieval models: BM-25 remains a widely used token-matching baseline, whereas embedding-based models jointly represent queries and documents in a shared space.Embedding retrieval can use precomputed document vectors and approximate nearest-neighbor search.
  • Retrieval models: Two-tower models use separate query and document networks, enabling semantic representations and task-specific optimization beyond sparse token-based features.Transformers model attention within each query and document rather than cross-attention between them.
  • Research question: The paper studies which pre-training tasks improve two-tower Transformer retrieval, an important question left unresolved despite masked-LM success for BERT-style models.The focus is efficient large-scale retrieval rather than the downstream scoring phase.
  • Findings: With proper pre-training, two-tower Transformer models significantly outperform BM-25 and improve over shallow bag-of-word MLP counterparts.The study presents itself as a comprehensive examination of pre-training tasks for efficient large-scale retrieval.
  • Findings: Paragraph-level tasks such as ICT, BFS, and WLP substantially improve retrieval, while token-level masked-LM provides only marginal gains.The studied paragraph-level tasks target semantic information at the paragraph scale.

2 THE TWO-TOWER RETRIEVAL MODEL

The two-tower retrieval model separately embeds queries and documents, enabling efficient large-scale inference while supporting task-specific training. The section contrasts this approach with cross-attention models and classic retrieval methods.

  • Model formulation: Two-tower models use separate encoder functions to map queries and documents into a shared embedding space, scoring pairs with an inner product.The encoders can be parameterized as deep Transformer models.
  • Inference: Document embeddings can be pre-computed, allowing query-time ranking through inner products and approximate nearest-neighbor or MIPS search.This supports efficient retrieval without evaluating every query-document pair with a cross-attention model.
  • Inference: For 1,000 query embeddings and 1 million document embeddings, 128-dimensional inner products take hundreds of milliseconds on CPUs, versus hours or more for cross-attention scoring on GPUs.The comparison illustrates the inference-efficiency advantage of two-tower scoring.
  • Model advantages: Two-tower models can capture deeper semantic relationships than sparse token-based models and can be optimized for specific downstream tasks.This task-specific trainability is presented as a key advantage over classic IR methods.
  • Learning: Because downstream supervised data is limited, the retrieval model is first trained on positive query-document pairs from pre-training tasks and then fine-tuned on downstream pairs.The paper studies the relevant pre-training tasks in the following section.

3 PRE-TRAINING TASKS OF DIFFERENT SEMANTIC GRANULARITIES

The paper studies paragraph-level pre-training tasks designed to capture semantic relationships at local, document-wide, and inter-document scales. ICT, BFS, and WLP construct positive query-document pairs from Wikipedia without additional manual labeling.

  • Design criteria: A useful pre-training task should match downstream semantics and be cost-efficient to collect without additional human supervision.The targeted semantic scales include local paragraph context, global document consistency, and relations between documents.
  • Task overview: The paper introduces ICT, BFS, and WLP to emphasize different semantic relationships between queries and documents, using Wikipedia-derived training data.BFS and WLP are newly proposed in this paper.
  • Inverse Cloze Task (ICT): ICT samples one sentence from a paragraph as the query and uses the remaining sentences as the document, capturing local semantic context.The task defines a positive pair within a single paragraph.
  • Body First Selection (BFS): BFS pairs a sentence from a Wikipedia page’s first section with a passage from the same page to capture semantic relationships beyond a local paragraph.The first section is treated as a description or summary containing information central to the page topic.
  • Wiki Link Prediction (WLP): WLP pairs a sentence from one Wikipedia page’s first section with a passage from another page linked to it, capturing inter-page semantic relations.The hyperlink provides the relationship between the two pages.
  • Baseline and training data: MLM is included as a baseline that predicts randomly masked tokens, while ICT, BFS, and WLP define paragraph-level positive pairs for two-tower Transformer pre-training.The task data for ICT, BFS, and WLP is generated from Wikipedia.

4 EXPERIMENTS

Experiments evaluate two-tower retrieval models, pre-training tasks, encoder architectures, and open-domain retrieval on ReQA benchmarks. Proper paragraph-level pre-training, especially ICT+BFS+WLP, consistently improves retrieval over weaker pre-training and BM-25 baselines.

  • Experimental Setting: The experiments evaluate SQuAD and Natural Questions using Recall@k under 1%, 5%, and 80% fine-tuning-data regimes.The test queries follow a cold-start split, and retrieval correctness requires recovering the gold sentence and evidence-passage pair.
  • Main Results: ICT+BFS+WLP notably outperforms all other methods on both SQuAD and Natural Questions, including the widely used BM-25 baseline.The combined model uses positive pairs uniformly sampled from ICT, BFS, and WLP.
  • Main Results: In limited-data SQuAD settings, BM-25 outperforms two-tower Transformers with no pre-training or MLM, confirming BM-25 as a robust baseline.BM-25 uses token matching with TF-IDF weights and benefits from token overlap in SQuAD’s data collection and annotations.
  • Encoder Architecture: With ICT+BFS+WLP, Transformer encoders considerably outperform comparable-size BoW-MLP encoders, while unpretrained Transformers perform worse than BoW-MLP.The authors attribute the latter result possibly to overfitting on limited labeled fine-tuning data.
  • Main Results: MLM provides only marginal gains over no pre-training on both datasets, whereas ICT+BFS+WLP produces a large performance boost.The comparison supports the importance of designing pre-training tasks related to retrieval rather than relying only on token-level masked language modeling.
  • Ablation Study: ICT is the strongest individual task, followed by BFS and WLP; combining all three improves over ICT by 1.5% absolute in the low-data regime.The combined tasks are associated with local, within-article, and cross-article semantic relationships, respectively.

5 CONCLUSION

The study finds that properly designed paragraph-level pre-training substantially improves two-tower Transformer retrieval over BM-25. Future work will test these tasks with other encoder architectures, non-Wikipedia corpora, and alternative regularization methods.

  • Paragraph-level pre-training with ICT, BFS, and WLP enables two-tower Transformer models to considerably improve over BM-25.Without pre-training or with token-level MLM pre-training, the models are no better than BM-25 in most cases.
  • The paper provides a comprehensive study of how pre-training tasks help large-scale retrieval, including evidence retrieval for question answering.
  • Future work will examine other encoder architectures, corpora beyond Wikipedia, and comparisons with different regularization methods.
Loading 2002.03932v1…