Pre-Training Tasks For Embedding-Based Large-Scale Retrieval

Contributors

Pierre McWhannel wrote this summary with editorial and technical contributions from STAT 946 fall 2020 classmates. The summary is based on the paper "Pre-Training Tasks for Embedding-Based Large-Scale Retrieval" which was presented at ICLR 2020. The authors of this paper are Wei-Cheng Chang, Felix X. Yu, Yin-Wen Chang, Yiming Yang, Sanjiv Kumar [1].

Introduction

The problem domain which the paper is positioned within is large-scale query-document retrieval. This problem is: given a query/question to collect relevant documents which contain the answer(s) and identify the span of words which contain the answer(s). This problem is often separated into two steps 1) a retrieval phase which reduces a large corpus to a much smaller subset and 2) a reader which is applied to read the shortened list of documents and score them based on their ability to answer the query and identify the span containing the answer, these spans can be ranked according to some scoring function. In the setting of open-domain question answering the query, [math]\displaystyle{ q }[/math] is a question and the documents [math]\displaystyle{ d }[/math] represent a passage which may contain the answer(s). The focus of this paper is on the retriever in open-domain question answering, in this setting a scoring function [math]\displaystyle{ f: \mathcal{X} \times \mathcal{Y} \rightarrow \mathbb{R} }[/math] is utilized to map [math]\displaystyle{ (q,d) \in \mathcal{X} \times \mathcal{Y} }[/math] to a real value which represents the score of how relevant the passage is to the query. We desire high scores for relevant passages and and low scores otherwise and this is the job of the retriever.

Certain characteristics desired of the retriever are high recall since the reader will only be applied to a subset meaning there is an an opportunity to identify false positive but not false negatives. The other characteristic desired of a retriever is low latency. The popular approach to meet these heuristics has been the BM-25 [2] which is relies on token-based matching between two high-dimensional and sparse vectors representing [math]\displaystyle{ q }[/math] and [math]\displaystyle{ d }[/math] this approach performances well and retrieval can be performed in time sublinear to the number of passages with inverter indexing. The downfall of this type of algorithm is its inability to be optimized for a specific task. An alternate option is BERT or transformer based models with cross attention between query and passage pairs these can be optimized for a specific task. However, these models suffer from latency as the retriever needs to process all the pairwise combinations of a query with all passages. The next type of algorithm is an embedding-based model that jointly embeds the query and passage in the same embedding space and then uses measures such as cosine distance, inner product, or even the Euclidean distance in this space to get a score. The authors suggest the two-tower models can capture deeper semantic relationships in addition to being able to be optimized for specific tasks. This model is referred to as the "two-tower retrieval model", since it has two transformer based models where one embeds the query [math]\displaystyle{ \phi{(\cdot)} }[/math] and the other the passage [math]\displaystyle{ \psi{(\cdot)} }[/math], then the embeddings can be scored by [math]\displaystyle{ f(q,d) = \langle \phi{(q)},\psi{(d)} \rangle \in \mathbb{R} }[/math]. This model can avoid the latency problems of a cross-layer model by pre-computing [math]\displaystyle{ \psi{(d)} }[/math] before hand and then by utilizing efficient approximate nearest neighbor search algorithms in the embedding space to find the nearest documents.

These two-tower retrieval models often use BERT with pretrained weights, in BERT the pre-training tasks are masked-LM (MLM) and next sentence prediction (NSP). The authors note that the research of pre-training tasks that improve the performance of two-tower retrieval models is an unsolved research problem. This is exactly what the authors have done in this paper is develop pre-training tasks for the two-tower model based on heuristics and then validated by experimental results. The pre-training tasks suggested are Inverse Cloze Task (ICT), Body First Search (BFS), Wiki Link Prediction (WLP), and combining all three. The authors used the Retrieval Question-Answering (ReQA) benchmark and used two datasets SQuAD and Natural Questions for training and evaluating their models.

Contributions of this paper

• The two-tower transformer model with proper pre-training can significantly outperform the widely used BM-25 Algorithm.
• Paragraph-level pre-training tasks such as ICT, BFS, and WLP hugely improve the retrieval quality, whereas the most widely used pre-training task (the token-level masked-LM) gives only marginal gains.
• The two-tower models with deep transformer encoders benefit more from paragraph-level pre-training compared to its shallow bag-of-word counterpart (BoW-MLP)

Background on Two-Tower Retrieval Models

This section is to present the two-tower model in more detail. As seen previously [math]\displaystyle{ q \in \mathcal{X} }[/math], [math]\displaystyle{ d \in \mathcal{y} }[/math] are the query and document or passage respectively. The two-tower retrieval model consists of two encoders: [math]\displaystyle{ \phi: \mathcal{X} \rightarrow \mathbb{R}^k }[/math] and [math]\displaystyle{ \psi: \mathcal{Y} \rightarrow \mathbb{R}^k }[/math]. The tokens in [math]\displaystyle{ \mathcal{X},\mathcal{Y} }[/math] are a sequence of tokens representing the query and document respectively. The scoring function is [math]\displaystyle{ f(q,d) = \langle \phi{(q)},\psi{(d)} \rangle \in \mathbb{R} }[/math]. The cross-attention BERT-style model would have [math]\displaystyle{ f_{\theta,w}(q,d) = \psi_{\theta}{(q \oplus d)^{T}}w }[/math] as a scoring function and [math]\displaystyle{ \oplus }[/math] signifies the concatenation of the sequences of [math]\displaystyle{ q }[/math] and [math]\displaystyle{ d }[/math], [math]\displaystyle{ \theta }[/math] are the parameters of the cross-attention model. The architectures of these two models can be seen below in figure 1.

Figure 1: end-to-end model

Comparisons

Inference (Latency focused)

In terms of inference the two-tower architecture has a distinct advantage as be easily seen from the architecture, since the query and document are embedded independently, [math]\displaystyle{ \psi{(d)} }[/math] can be pre computed. This is an important advantage as it makes the two-tower method feasible as an example the authors state the inner product between 100 query embeddings ad 1 million document embeddings only takes hundreds of milliseconds on CPUs, however, for the equivalent calculation for a cross-attention-model it takes hours on GPUs. In addition they remark the retrieval of the relevant documents can be done in sublinear time with respect [math]\displaystyle{ |\mathcal{Y}| }[/math] using a maximum inner product (MIPS) algorithm with little loss in recall.

Learning

In terms of learning the two-tower model compared to BM-25 has a unique advantage of being possible to fine tuned for downstream tasks. This problem is within the paradigm of metric learning as the scoring function can be passed used as input to a objective function of the log likelihood [math]\displaystyle{ max_{\theta} \space log(p_{\theta}(d|q)) }[/math]. The conditional probability is represented by a SoftMax and then can be rewritten as:

[math]\displaystyle{ p_{\theta}(d|q) = \frac{exp(f_{\theta}(q,d))}{\sum_{d' \in D}} exp(f_{\theta}(q,d')) }[/math]

Here [math]\displaystyle{ \mathcal{D} }[/math] is the set of all possible document, the denominator is a summation over [math]\displaystyle{ \mathcal{D} }[/math]. As is customary they utilized a sampled SoftMax technique to approximate the full-SoftMax. The technique they have employed is by using a small subset of documents in the current training batch, while also using a proper correcting term ensure the unbiasedness of the partition function.

Why do we need pre-training then?

Although the two-tower model can be fine-tuned for a downstream task getting sufficiently labelled data is not always possible. This is an opportunity to look for better pre-training tasks as the authors have done. The authors suggest getting a set of pre-training task [math]\displaystyle{ \mathcal{T} }[/math] with labelled positive pairs of queries of documents/passages and then afterward fine-tuning on the limited amount of data for the downstream task.

Pre-Training Tasks

The authors suggest three pre-training tasks for the encoders of the two-tower retriever, these are ICT, BFS, and WLP. Keep in mind that ICT is not a newly proposed task. These pre-training tasks are developed in accordance to two heuristics they state and our paragraph-level pre-training tasks in contrast to more common sentence-level or word-level tasks:

1. The tasks should capture a different resolutions of semantics between query and document. The semantics can be the local context within a paragraph, global consistency within a document, and even semantic relation between two documents.
2. Collecting the data for the pre-training tasks should require as little human intervention as possible and be cost-efficient.

All three of the following tasks can be constructed in a cost-efficient and with minimal human intervention by utilizing Wikipedia articles. In figure 2 below the dashed line surrounding the block of text is to be regarded as the document [math]\displaystyle{ d }[/math] and the rest of the circled portions of text make up the three queries [math]\displaystyle{ q_1, q_2, q_3 }[/math] selected for each task.

Figure 1: end-to-end model

Inverse Cloze Task (ICT)

This task is given a passage [math]\displaystyle{ p }[/math] consisting of [math]\displaystyle{ n }[/math] sentences [math]\displaystyle{ s_{i}, i=1,..,n }[/math]. [math]\displaystyle{ s_{i} }[/math] is randomly sampled and used as the query [math]\displaystyle{ q }[/math] and then the document is constructed to be based on the remaining [math]\displaystyle{ n-1 }[/math] sentences. This is utilized to meet the local context within a paragraph of the first heuristic. In figure 2 [math]\displaystyle{ q_1 }[/math] is an example of a potential query for ICT based on the selected document [math]\displaystyle{ d }[/math].

Body First Search (BFS)

BFS is [math]\displaystyle{ (q,d) }[/math] pairs are created by randomly selecting a sentence [math]\displaystyle{ q_2 }[/math] and setting it as the query from the first section of a Wikipedia page which contains the selected document [math]\displaystyle{ d }[/math]. This is utilized to meet the global context consistency within a document as the first section in Wikipedia generally is an overview of the whole page and they anticipate it to contain information central to the topic. In figure 2 [math]\displaystyle{ q_2 }[/math] is an example of a potential query for BFS based on the selected document [math]\displaystyle{ d }[/math].

Wiki Link Prediction (WLP)

The third task is selecting a hyperlink that is within the selected document [math]\displaystyle{ d }[/math], as can be observed in figure 2 "machine learning" is a valid option. The query [math]\displaystyle{ q_3 }[/math] will then be a random sentence on the first section of the page where the hyperlinks redirects to. This is utilized to provide the semantic relation between documents.

Additionally Masked LM (MLM) is considered during the experimentations which again is one of the primary pre-training tasks used in BERT.

Experiments and Results

Maybe provide a description.. we will see.

Training Specifications

• Each tower is constructed from the 12 layers BERT-base model.
• Embedding is achieved by applying a linear layer on the [CLS] token output to get an embedding dimension of 512.
• Sequence lengths of 64 and 288 respectively for the query and document encoders.
• Pre-trained on 32 TPU v3 chips for 100k step with Adam optimizer learning rate of [math]\displaystyle{ 1 \times 10^{-4} }[/math] with 0.1 warm-up ratio, followed by linear learning rate decay and batch size of 8192. (~2.5 days to complete)
• Fine tuning on downstream task [math]\displaystyle{ 5 \times 10^{-5} }[/math] with 2000 training steps and batch size 512.
• Tokenizer is WordPiece.

Pre-Training Tasks Setup

• MLM, ICT, BFS, and WLP are used.
• Various combinations of the aforementioned tasks are tested.
• Tasks define [math]\displaystyle{ (q,d) }[/math] pairings.
• ICT's document [math]\displaystyle{ d }[/math] has the article title and passage separated by a [SEP] symbol as input to the document encoder.

Below table 1 shows some statistics for the datasets constructed for the ICT, BFS, and WLP tasks. Token counts considered after the tokenizer is applied.

Downstream Tasks Setup

Retrieval Question-Answering (ReQA) benchmark used.

• Datasets: SQuAD and Natural Questions.
• Each entry of data from datasets is [math]\displaystyle{ (q,a,p) }[/math]. Where each element is as follows query [math]\displaystyle{ q }[/math], answer [math]\displaystyle{ a }[/math], passage [math]\displaystyle{ p }[/math] containing the answer.
• Authors split the passage to sentences [math]\displaystyle{ s_i }[/math] where [math]\displaystyle{ i=1,...n }[/math] and [math]\displaystyle{ n }[/math] is the number of sentences in the passage.
• The problem is then recast for a given query to retrieve the correct sentence-passage [math]\displaystyle{ (s,p) }[/math] from all candidates for all passages split in this fashion.
• They remark their reformulation makes it a more difficult problem.

Below table 2 show statistics of the datasets used for the downstream QA task.

Evaluation

• 3 train/test splits: 1%/99%, 5%/95%, and 80%/20%.
• 10% of training set is used as a validation set for hyper-parameter tuning.
• Training never has seen the test queries in the test pairs of [math]\displaystyle{ (q,d). }[/math]

References

[1] Wei-Cheng Chang, Felix X Yu, Yin-Wen Chang, Yiming Yang, and Sanjiv Kumar. Pre-training tasks for embedding-based large-scale retrieval. arXiv preprint arXiv:2002.03932, 2020.

[2] Stephen Robertson, Hugo Zaragoza, et al. The probabilistic relevance framework: BM25 and beyond. Foundations and Trends in Information Retrieval, 3(4):333–389, 2009.