Pre-training Tasks for Embedding-based Large-scale Retrieval (2002.03932v1)

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.

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

The paper explores the large-scale query-document retrieval landscape, emphasizing the understudied retrieval phase compared to the scoring phase typically enhanced by BERT-style models. The research pivots around embedding-based retrieval models, specifically leveraging Transformer architectures to supplant traditional IR methods like BM25. A primary focus is placed on effective pre-training tasks for two-tower Transformer models, underscoring their impact on retrieval performance.

Key Findings and Contributions

• Pre-training Task Design: The research identifies the relevance of paragraph-level pre-training tasks for enhancing a two-tower Transformer model's performance in large-scale retrieval. These tasks—Inverse Cloze Task (ICT), Body First Selection (BFS), and Wiki Link Prediction (WLP)—outperform traditional token-level Masked LLMing (MLM) pre-training.
• Performance Comparison: Models pre-trained with paragraph-level tasks showed significant improvements over the conventional BM25 algorithm. This progression is attributed to the deeper semantic understanding and context captured between queries and documents by utilizing specific pre-training strategies.
• Task Effectiveness: Among the three tasks, ICT demonstrated superior performance in the tested settings, followed by BFS and WLP. This suggests that capturing local semantic context within passages is crucial for retrieval efficacy.
• Transformer Utilization: Two-tower Transformer models with appropriate pre-training considerably outperform MLP architectures, especially under conditions involving extensive pre-training. This highlights the capability of Transformers to capture complex semantic relations inherent in the retrieval tasks compared to shallow MLP models.
• Empirical Analysis: Comprehensive experiments using datasets such as ReQA demonstrated the robustness of Transformer-based retrieval models, specifically in open-domain settings. The superior performance of models employing paragraph-level pre-training tasks underpins the necessity of these tasks in large-scale retrieval problems.

Implications for Future Research

The paper's insights point towards several implications:

• Enhanced Model Architecture: There is a potential to further refine the two-tower model architectures by experimenting with additional pre-training data sources and examining the transferability to other NLP domains.
• Scaling Pre-trained Models: Future work could investigate the effect of progressively larger pre-training datasets or more diverse tasks beyond Wikipedia, aiming to build models capable of scaling efficiently with the rapid expansion of information in an interconnected digital ecosystem.
• Real-world Applications: The enhanced retrieval quality achieved could be translated into improved performance in real-world applications like recommendation systems, where understanding nuanced semantic connections between user queries and items is vital.

Conclusion

This paper makes a significant contribution by systematically evaluating the role of various pre-training tasks in two-tower Transformer retrieval models. The findings emphatically suggest that task selection is integral to embedding-based retrieval model success, prompting a paradigm shift in approaching large-scale retrieval nuances. This comprehensive examination offers a foundation for further explorations into efficient retrieval mechanisms, potentially transforming the landscape of information retrieval by infusing it with contextual intelligence powered by advanced deep learning methodologies.