G OING B EYOND T OKEN - LEVEL P RE - TRAINING FOR E MBEDDING - BASED L ARGE - SCALE R ETRIEVAL
Anonymous authors
Paper under double-blind review
A BSTRACT
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 ﬁrst reduces the solution space, returning a subset of candidate documents.
The scoring phase then scores and re-ranks the documents.
The algorithm used in the retrieval phase is critical.
On the one hand, it needs to have high recall – otherwise some relevant documents won’t even be considered in the scoring phase.
On the other hand, it needs to be highly efﬁcient, returning the candidate documents in time sublinear to the total number of documents.
Unlike the scoring phase which witnessed signiﬁcant advances recently due to the BERT-style cross-attention models, the retrieval phase remains less well studied: most previous works rely on the classic Information Retrieval (IR) methods such as BM-25 (token matching + TF-IDF weights).
In this paper, we conduct a comprehensive study on different retrieval algorithms and show that the two-tower Transformer models with properly designed pre-training tasks can largely improve over the widely used BM-25 algorithm.
The pre-training tasks we studied are Inverse Cloze Task ( ICT ), Body First Selection ( BFS ), Wiki Link
Prediction ( WLP ) and the combination of them.
I NTRODUCTION
We consider the large-scale retrieval problem: given a query, return the most relevant documents from a large corpus, where the size of the corpus can be hundreds of thousands or more.
One can view this problem as learning a scoring function f : X × Y → R , that maps a pair of query and document ( q , d
) ∈ X ×Y to a score f ( q , d ) .
The function should be designed such that the relevant ( q , d ) pairs have high scores, whereas the irrelevant ones have low scores.
Many real-world applications besides query-document retrieval can be cast in this form.
For example, in recommender systems, q represents a user query and d represents a candidate item to recommend (Krichene et al., 2019).
In extreme multi-label classiﬁcation, q represents a web-page document and d represents the categories or hashtags of interests (Jain et al., 2019;
Chang et al., 2019).
In open-domain question answering, q represents a question and d represents an evidence passage containing the answer (Chen et al., 2017; Hu et al., 2019; Lee et al., 2019).
Central to the above is designing the scoring function f .
Recently BERT (Devlin et al., 2019), along with its successors XLNet
(Yang et al., 2019b) and RoBERTa
(Liu et al., 2019), leads to signiﬁcant improvements to many NLP tasks such as sentence pairs classiﬁcation and questionanswering.
In BERT, the scoring function f is a pre-trained deep bidirectional Transformer model.
While BERT-style models are very successful, it cannot be directly applied to large-scale retrieval problems because computing f ( q , d ) for every possible document can be prohibitively expensive.
Thus, one typically ﬁrst uses a less powerful but more efﬁcient algorithm (another scoring function f ) to reduce the solution space (the “retrieval phase”), and then use the BERT-style model to re-rank the retrieved documents (the “scoring phase”).
The retrieval phase is critical.
Ideally speaking, the algorithm should have a high recall; otherwise, many relevance documents won’t even be considered in the scoring phase.
The algorithm also needs to be very fast: it should return a small subset of relevant documents in time sublinear to the number of all documents.
Although signiﬁcant developments are advancing the scoring algorithms, the retrieval algorithms remain to be less well studied, and this is the focus of this paper.
The retrieval algorithm can be put into two categories.
The ﬁrst type is classic information retrieval (IR) algorithms relying on token-based matching with TF-IDF weights.
One example is
BM-25 (Robertson et al., 2009), which remains to be the most widely used (Nguyen et al., 2016;
Yang et al., 2017; 2019a) and hard to beat (Chapelle & Chang, 2011; Lee et al., 2019) algorithm.
Here the scoring function f is based on token-matching between the two high-dimensional sparse vectors with TF-IDF token weights, and retrieval can be done in sublinear time using inverted index.
The second option is to jointly embed queries and documents in the same embedding space and use an inner product or cosine distance to measure the similarity between queries and documents.
Let the query embedding model be φ ( · ) and the doc embedding model be ψ ( · ) , The scoring function is
[Equation]
In the inference stage, retrieving relevant documents then becomes ﬁnding the nearest neighbors of a query in the embedding space.
Since the embeddings of all candidate documents can be precomputed and indexed, the inference can be done efﬁciently with approximate nearest neighbor search algorithms in the dense embedding space (Shrivastava & Li, 2014; Guo et al., 2016).
In this paper, we refer the above as the two-tower retrieval model , because the query and document embeddings are coming from two separate “towers” of neural networks.
In the literature, it is also known as the Siamese network (Das et al., 2016;
Triantaﬁllou et al., 2017) or dual-encoder model (Cer et al., 2018; Mazar ´ e et al., 2018).
Compared to the sparse token-based models such as BM-25 , the two-tower models capture deeper semantic relationships within queries and documents.
In the heart of two-tower models is to design the embedding function φ ( · ) and ψ ( · ) .
A modern choice is using Transformers to model the attention within queries and within documents, rather than the cross-attention between them as in the BERT model.
The token-level Masked-LM pre-training task is crucial to the success of BERT-style cross-attention models.
Nevertheless, what pre-training tasks are useful for improving two-tower Transformer models in large-scale retrieval, remains a crucial yet unsolved research problem.
In this paper, we aim to answer this question by studying different pre-training tasks for the two-tower Transformer models.
We contribute the following insight:
• The two-tower Transformer models with proper pre-training can signiﬁcantly outperform the widely used BM-25 algorithm; • Paragraph-level pre-training tasks such as Inverse Cloze Task ( ICT ), Body First Selection ( BFS ), and Wiki Link Prediction ( WLP ) hugely improve the retrieval quality, whereas the most widely used pre-training task (the token-level masked-LM) gives only marginal gains.
To the best of our knowledge, this is the ﬁrst comprehensive study on pre-training tasks for efﬁcient large-scale retrieval algorithms.
The rest of the paper is organized as follows.
We start by introducing the two-tower retrieval model in Section 2.
The pre-training tasks are presented in 3, and the experiments and analysis are presented in Section 4.
Finally, we conclude this work in Section 5.
T HE TWO - TOWER RETRIEVAL MODEL
We begin with introducing some backgrounds of the two-tower models.
Given a query q ∈ X and a document d ∈ Y , we consider two-tower retrieval models that consist of two encoder functions, φ : X → R k and ψ : Y → R k which map a sequence of tokens in X and Y to their associated embeddings φ ( q ) and ψ ( d ) , respectively.
The scoring function f : R k × R k
→ R is then deﬁned to be the inner product 1 of the embeddings
[Equation]
In this paper, we are interested in parameterizing the encoders φ, ψ as deep Transformer models (Vaswani et al., 2017) due to its expressive power in modeling natural language.
In the rest of this section, we illustrate the advantage of two-tower models in the inference phase, then discuss the pros and cons of two-tower models in comparison with BERT-like cross-attention models; present the learning procedure of estimating model parameters under maximum likelihood principle; and review the related works.
Inference The difference between two-tower models and cross-attention models is shown in Figure 1.
The advantage of the two-tower model is its efﬁciency in large-scale retrieval in the inference time.
First, all the document embeddings can be pre-computed.
Then, given an unseen query q , we only need to rank the document based on its inner product with the query embedding.
This is way more efﬁcient than running inference on a cross-attention BERT-style model (often used in the scoring stage).
To see this, the scoring function of BERT-style model is with the form
[Equation]
where ⊕ denotes the concatenate operation of the query and the document sequence and w ∈ R
k is an additional model parameters.
In BERT, for each query, one has to do the above expensive inference on all documents.
For example, with the 128-dimensional embedding space, inner product between 1000 query embeddings with 1 million document embeddings only takes hundreds of milliseconds on CPUs, while computing the same scores with cross-attention models takes hours if not more even on GPUs.
Furthermore, retrieving the documents with the maximum inner product (MIPS) is, in fact, a wellstudied problem: it can be further efﬁciently approximated by hashing and other indexing algorithms in sublinear time, with almost no loss in the recall (Shrivastava & Li, 2014; Guo et al., 2016).
Learning In this paper, we assume that the training data is presented as relevant “positive” querydocument pairs T = { ( q i , d i ) } |T | i =1 .
Let θ be the model parameters.
We estimate the model parameters by maximizing the log likelihood max θ (cid:80) ( q , d ) ∼T log p θ ( d | q ) where the conditional probability is deﬁned over the Softmax distribution:
[Equation]
The Softmax involves the expensive partition function in the denominator of equation 3 that scales linearly to the number of documents.
In practice, we consider Sampled Softmax, an approximation of the full-Softmax where pairs of query and document are uniformly sampled as mini-batches
[Equation]
[Equation]
Sampled Softmax has been widely used in language modeling in NLP (Chen et al., 2016;
Grave et al., 2017), recommendation system (Yu et al., 2017; Krichene et al., 2019) and extreme classiﬁcations (Blanc & Rendle, 2018; Reddi et al., 2019).
We then optimize the Sampled Softmax objective with variants of stochastic gradient methods such as Adam with weight decay (Kingma & Ba, 2014).
Here, the positive pairs T can be either from the downstream task or the pre-training tasks.
Since we often have a limited amount of supervised data from the downstream task, the model is ﬁrst trained with pre-training tasks and then ﬁne-tuned with the downstream task.
We will present the set of pre-training tasks we study in Section 3.
Related Works Cer et al.
(2018) study the two-tower Transformer model as a universal sentence encoder.
The model is learned with multiple tasks including the unsupervised Skip-Thought task, the supervised conversation input-response task (Henderson et al., 2017), and the supervised sentence classiﬁcation SNLI task (Bowman et al., 2015).
Humeau et al.
(2019) propose the Poly-encoders architecture to balance the computation/expressiveness tradeoff between two-tower models and crossattention models.
Reimers & Gurevych (2019) ﬁne-tune the deep two-tower models on two supervised datasets, SNLI and MNLI (Williams et al., 2018), then apply it in solving other downstream tasks.
Unlike all the above works that consider training the two-tower Transformer models on a limited amount of supervised corpus for the sentence classiﬁcation tasks, we study different pre-training tasks and their contributions in the large-scale retrieval settings.
Another closely related topic is open-domain question answering.
Previous works consider using
BM25 or other lexical matching methods to efﬁciently retrieve the top-k relevant passages and then deploy the more expensive cross-attention scoring function to ﬁnd the answer (Chen et al., 2017;
Yang et al., 2017; 2019a).
Das et al.
(2019) encode query and document separately with LSTM encoders.
They employ a training procedure different from ours and do not consider pre-training.
Very recently, Lee et al.
(2019) propose to pre-train two-tower Transformer models with the Inverse
Cloze Task ( ICT ) to replace BM25 in the passage retrieval phase.
The advantage is that the retriever can be trained jointly with the reader/scorer.
Nevertheless, their pre-trained two-tower models do not outperform BM25 on the SQuAD dataset, potentially because the ﬁne-tuning is only performed on the query-tower.
Model distillation (Hinton et al., 2015) can be used to compress expensive BERT-like cross-attention models into efﬁcient two-tower Transformer models for large-scale retrieval problems.
For example,
Tang et al.
(2019) demonstrate initial success in distilling the BERT model into a two-tower model with BiLSTM as encoders.
The pre-training tasks we study in this paper can be used as additional supervision in the distillation process, and therefore complementary to model distillation.
P RE - TRAINING TASKS OF DIFFERENT SEMANTIC GRANULARITIES
As mentioned in Section 2, due to the limited amount of supervised data from downstream tasks, a crucial step of learning deep retrieval models is to pre-train the model with a set of pre-training tasks (we will verify this in Section 4).
Sentence-level pre-training tasks have been studied before.
One example is reconstructing the surface form of surrounding sentences given the encoded sentence (Le & Mikolov, 2014; Kiros et al., 2015), and another one is discriminating the next sentence from random candidates (Jernite et al., 2017; Logeswaran & Lee, 2018).
In this paper, we assume that the pre-training data is deﬁned as positive query-document ( q , d ) pairs.
A good pre-training task should have the following two properties.
1) It should be relevant to the downstream task.
For example, when solving the question-answering retrieval problem, the model should capture different granularities of semantics between the 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) It should be cost-efﬁcient to collect the pre-training data, ideally not requiring additional human supervision.
Next, we present three pre-training tasks that emphasize on different aspects of semantics between queries and documents: Inverse Cloze Task ( ICT ), Body First Selection ( BFS ), and Wiki Link
Prediction ( WLP ).
In speciﬁc, BFS and WLP are newly proposed in this paper.
The training data of all these tasks can be freely obtained based on Wikipedia without additional manual labeling process.
Figure 2 provides illustrative examples of these tasks.
Inverse Cloze Task (ICT)
Given a passage p consisting of n sentences p = { s 1 , . . .
, s n } , the query q is a sentence randomly drawn from the passage q
= s i , i ∼ [1 , n ] and the document is others sentences d
= { s 1 , . . .
, s i − 1 , s i +1 , . . .
, s n } .
See the ( q 1 ,
) in Figure 2 for an example.
This is proposed
Body First Selection (BFS)
We propose BFS to capture semantic relationship outside of the local passage.
Here, the query q 2 is a random sentence in the ﬁrst section of a Wikipedia page, and the document d is a random passage from the same page (Figure 2).
Since the ﬁrst section of a Wikipedia article is often the description or summary of the whole page, we expect it contains information central to the topic.
Wiki Link Prediction (WLP)
We propose WLP to capture inter-page semantic relation.
The query q 3 is a random sentence in the ﬁrst section of a Wikipedia page, and the document d is a passage from another page where there is a hyperlink link to the page of q 3 (Figure 2).
Intuitively, a hyperlink link indicates relationship between the two Wikipedia pages.
Again, we take a sentence from the ﬁrst section because it is often the description or summary of the topic.
Masked LM (MLM)
In addition to the above tasks, we also consider the classic masked language model ( MLM ) pre-training task as a baseline: predict the randomly masked tokens in a sentence.
MLM is the primary pre-training task used in BERT (Devlin et al., 2019).
E XPERIMENTS 4. E XPERIMENTAL S ETTING
The two-tower retrieval model Each tower of the retrieval model follows the architecture and hyper-parameters of the 12 layers BERT-base model.
For both towers, the ﬁnal embedding is generated by applying a linear layer on the hidden state of the [ CLS ] token.
The embedding dimension is 512 .
The sequence length for the query encoder and document encoder are set to be 64 and 288, respectively.
We pre-train the model on 32 TPU v3 chips for 100K steps with an Adam optimizer and batch size of 8192.
This process takes about 2.5 days.
The Adam optimizer has an initial learning rate 1e-4 with the warm-up ratio 0 .
1 , followed by a linear learning rate decay.
For ﬁne-tuning, the learning rate of Adam is set to 3e-5 with 2000 training steps and batch size 512.
[Table] [Table]
Pre-training tasks We compare the token-level task MLM with the paragraph-level tasks including ICT , BFS and WLP .
The data of ICT , BFS and WLP are generated with the Wikipedia data.
The data statistics are reported in Table 1.
Note that #tokens represents number of sub-words tokenized by WordPiece (Wu et al., 2016).
The pre-training tasks deﬁne the positive ( q , d ) pair for learning the two-tower Transformer models.
For ICT , the d is a pair of article title and passage separated by [ SEP ] symbol as input to the doc-tower.
We propose to pre-train the two-tower Transformer models jointly with all three paragraph-level pre-training tasks, hence the name ICT+BFS+WLP .
Speciﬁcally, the model is pre-trained on one combined set of ( q , d ) pairs, where each pair is uniformly sampled from the three pre-training tasks in Table 1.
See Section 4.2 and 4.3 for its outstanding performance over other baselines.
Downstream tasks We consider the Retrieval Question-Answering ( ReQA ) benchmark, proposed by Ahmad et al.
2 Note that each entry of QA datasets is a tuple ( q , a , p ) , where q is the question, a is the answer span, and p is the evidence passage containing a .
Following Ahmad et al.
(2019), we split a passage into sentences, p = s 1 s 2 . . .
s n and transform the original entry ( q , a , p ) to a new tuple ( q , s i , p ) where s i is the sentence contains the answer span a .
The retrieval problem is that given a question q , retrieve the correct sentence evidence passage pair ( s , p ) from all candidates.
For each passage p , we create a set of candidate pairs ( s i , p ) where i = 1 . . .
n , and the retrieval candidate set is built by combining such pairs for all passages.
This problem is more challenging than retrieving the evidence passage only since the larger number of candidates to be retrieved.
The data statistics of the downstream ReQA benchmark are shown in
Table 2.
Note that, similar to Ahmad et al.
(2019), the ReQA benchmark is not entirely opendomain QA retrieval as the candidates ( s , p ) only cover the training set of QA dataset instead of entire Wikipedia articles.
For open-domain retrieval experiment, See details in Section 4.4.
Evaluation For each dataset, we consider different training/test split of the data ( 1% / 99% , 5%/95% , 10% / 90% and, 80% / 20% ) in the ﬁne-tuning stage and the 10% of training set is hold out as the validation set for hyper-parameter tuning.
The split is created assuming a cold-start retrieval scenario where the queries in the test (query, document) pairs are not seen in training.
For the evaluation metric, we focus on Recall@k 3 because the goal of the retrieval phase is to capture the positives in the top-k results.
The retrieval performance can be understood independently of the scoring model used by measuring recall at different k.
In fact, in the extreme cases when the scoring model is oracle or random, the ﬁnal precision metric is proportional to recall@k.
[Table]
4. M AIN R ESULTS
Table 3 and Table 4 compare the proposed combination of pre-training methods, namely
ICT+BFS+WLP , to various baselines on SQuAD and Natural Questions , respectively.
In both benchmarks, ICT+BFS+WLP notably outperforms all other methods.
This suggests that one should use a two-tower Transformer model with properly designed pre-training tasks in the retrieval stage to replace the widely used BM-25 algorithm.
We present some of the detailed ﬁndings below.
The BM-25 baseline BM-25 is the state-of-the-art unsupervised retrieval method based on tokenmatching with TF-IDF weights.
The performance of BM-25 is competitive on SQuAD because
SQuAD is biased towards question-answer pairs with overlapping tokens.
On both datasets, BMoutperforms the two-tower Transformer model without pre-training ( No-SSL ) and with the tokenlevel masked language model pre-training MLM in most cases.
This veriﬁes that BM-25 is a robust retrieval model and therefore widely used in recent works (Yang et al., 2017; Lee et al., 2019).
Encoder architecture We also justify the use of Transformer as encoders by comparing it with the simple MLP model.
MLP lookups the uni-gram representations from the embedding table (we empirically found that adding bi-grams does not further improve the performance on these tasks possibly due to over-ﬁtting), aggregate the embeddings with average pooling, and pass them through a shallow two-layer MLP network with tanh activation to generate the ﬁnal 512-dimensional query/document embeddings.
For the two-tower models, we compare the MLP encoder vs. the Transformer encoder ( No-SSL ).
Although without pre-training tasks both of them are not performing well, we see that the quality of MLP is much worse than Transformer in most cases, conﬁrming that modeling attentions within queries and documents is important.
Pre-training tasks When pre-training the two-tower Transformer model, we compare the pretraining tasks to two baselines: No-SSL and MLM .
No-SSL represents no pre-training, and MLM is the token-level Masked-LM task introduced in Section 3.
On both datasets, the token-level pre-training task MLM only marginally improves over the no-pretraining baseline No-SSL , whereas combining the paragraph-level pretraining tasks
ICT+BFS+WLP provides huge boost on the performance.
This veriﬁes our assumption that the design of task-related pre-training tasks is crucial.
The performance of adding individual pre-training tasks will be presented in the next section.
[Table] [Table]
4. A BLATION S TUDY
We conduct a more thorough ablation study on Natural Questions involving (1) the number of layers in Transformer; (2) different pre-training tasks; and (3) dimension of the embedding space.
The result is presented in Table 5.
Index 1, 2, and 3 show the individual performance of three pre-training tasks.
All of these tasks are much more effective than MLM .
Among them, ICT has the best performance, followed by BFS , and then WLP .
This suggests that the (query,document) pairs deﬁned by local context within passage are suitable for the ReQA task.
Also note from Index 6 and 7, ICT+BFS+WLP pre-training is better than ICT with 1.5% absolute improvement over ICT in the low-data regime.
This reﬂects that, when there’s no sufﬁcient downstream training data, more globally pre-training tasks is beneﬁcial as it encodes multi-hop reasoning priors such as different passages within the same article (BFS) or even going beyond different articles linked by the same entities (WLP).
Finally, The advantage of increasing number of layers is manifest at Index 7, 8 and 9, showing the beneﬁt of increasing the dimension of embedding space.
4. E VALUATION OF O PEN - DOMAIN R ETRIEVAL
We consider the open-domain retrieval setting by augmenting the candidate set of the ReQA benchmark with large-scale (sentence, evidence passage) pairs extracted from general Wikipedia articles.
In particular, we preprocess/sub-sample the open-domain Wikipedia retrieval set of the DrQA paper (Chen et al., 2017) into one million (sentence, evidence passage) pairs, and add this external 1M candidate pairs into the existing retrieval candidate set of the ReQA benchmark.
[Table]
The results of open-domain retrieval on Natural Questions are presented in Table 6.
Firstly, we see that the two-tower Transformer models pretrained with ICT+BFS+WLP and ICT substantially outperform the BM-25 baseline.
Secondly, ICT+BFS+WLP pre-training method consistently improves the ICT pre-training method in most cases.
Interestingly, the improvements are more noticeable at
R@50 and R@100, possibly due to that the distant multi-hop per-training supervision induces better retrieval quality at the latter part of the rank list.
Finally, we conclude that the evaluation results of the 1M open-domain retrieval are consistent with our previous empirical evaluation on the ReQA benchmark with smaller retrieval candidate sets (Section 4.2).
C ONCLUSION
We conducted a comprehensive study on how different pre-training tasks help in the large-scale retrieval problem such as evidence retrieval for question-answering.
We showed that the two-tower
Transformer models with random initialization ( No-SSL ) or the unsuitable token-level pre-training task ( MLM ) are no better than the robust IR baseline BM-25 in most cases.
With properly designed paragraph-level pre-training tasks inlcuding ICT , BFS and WLP , the two-tower Transformer models can considerably improve over the widely used BM-25 algorithm.
For future works we plan to study how the pre-training tasks apply to other types of encoders architectures, generating the pre-training data from corpora other than Wikipedia, and how pretraining compares with different types of regularizations.