Large Dual Encoders Are Generalizable Retrievers — part 1/2
Cleaned-text code points [0, 15694). Bibliography omitted here; source.md includes it.
# Large Dual Encoders Are Generalizable Retrievers
###### Abstract
It has been shown that dual encoders trained on one domain often fail to generalize to other domains for retrieval tasks. One widespread belief is that the bottleneck layer of a dual encoder, where the final score is simply a dot-product between a query vector and a passage vector, is too limited to make dual encoders an effective retrieval model for out-of-domain generalization. In this paper, we challenge this belief by scaling up the size of the dual encoder model while keeping the bottleneck embedding size fixed. With multi-stage training, surprisingly, scaling up the model size brings significant improvement on a variety of retrieval tasks, especially for out-of-domain generalization. Experimental results show that our dual encoders, G eneralizable T 5-based dense R etrievers (GTR), outperform existing sparse and dense retrievers on the BEIR dataset Thakur et al. (2021) significantly. Most surprisingly, our ablation study finds that GTR is very data efficient, as it only needs 10% of MS Marco supervised data to achieve the best out-of-domain performance. 1 1 1 All the GTR models are released at https://tfhub.dev/google/collections/gtr/1 .
## 1 Introduction
Typical neural retrieval models follow a dual encoder paradigm Gillick et al. (2018) ; Yang et al. (2020) ; Karpukhin et al. (2020) . In this setup, queries and documents are encoded separately into a shared fixed-dimensional embedding space where relevant queries and documents are represented in each other’s proximity. Then, approximated nearest neighbor search Vanderkam et al. (2013) ; Johnson et al. (2021) is applied to efficiently retrieve relevant documents given an encoded input query.
While dual encoders are popular neural retrievers, the expressiveness of the model is limited by a bottleneck layer consisting of only a simple dot-product between query embeddings and passage embeddings. Several papers Lu et al. (2021) ; Khattab and Zaharia (2020) have discussed that the simple dot-product (or cosine similarity) between the embeddings might not be powerful enough to capture semantic relevance. Thakur et al. (2021) studied whether the retriever models can generalize to other domains and conclude that dual encoder models have “issues for out-of-distribution data”, and showed that models with more interactions between queries and documents have better generalization ability.
In this paper, we challenge this belief by scaling up the dual encoder model size while keeping the bottleneck embedding size fixed. Note that scaling up a dual encoder is different from scaling up pretrained language models such as BERT Devlin et al. (2019) and T5 Raffel et al. (2020) because of the presence of the bottleneck layer. While increasing the model size can greatly increase the capacity of the model, for dual encoders, where the embedding size is fixed, the interactions between queries and documents are still limited by a simple dot-product.
In order to test this hypothesis, we take advantage of the existing T5 model architecture and checkpoints, which allows us to build encoders of up to 5 billion parameters while keeping the bottleneck embedding dimension of 768 in all configurations, as illustrated in Figure 2 . Following Ni et al. (2021) , we build dual encoders by taking the encoder part of T5. For effectively using the power of large models, we collect roughly two billion community question-answer pairs as generic pre-training data. By combining pre-training using generic training data and fine-tuning using MS Marco Nguyen et al. (2016) , we are able to train large-scale dual encoder retrieval models. We call the resulting models G eneralizable T 5-based dense R etrievers ( GTR ).
We evaluate the zero-shot performance of GTR on the BEIR benchmark Thakur et al. (2021) , which consists of 18 selected information retrieval tasks across 9 domains. 2 2 2 We focus on evaluating the performance on the 18 BEIR Thakur et al. (2021) tasks other than MS Marco and we did not use in-domain training data or question generation Ma et al. (2021) . Our results show that, surprisingly, scaling up of dual encoders leads to better generalizability despite the fixed bottleneck embedding dimension. Second, pre-training on community question-answer pairs and fine-tuning on human curated data are both important to fully utilize the power of the scaled up model. In addition, with scaling and pre-training, we found GTR to be highly data efficient in terms of human annotated queries, as it only needs to use 10% of MS Marco to match the overall out-of-domain performance.
## 2 Background
### 2.1 Dual Encoder and dense retrieval
Classic retrieval models such as BM25 Robertson and Zaragoza (2009) relies on lexical overlap, term frequency heuristics, inverse document frequency and document length. This type of retrieval models does not require any training and can generalize reasonably well due to its emphasis on lexical match. However, these retrieval models fall short of finding documents that are only semantically related to the query and have low lexical overlap.
This is where the dense retrieval models come into play. The retrieval process is dense because both queries and documents are embedded into low-dimensional dense representations, in contrast to the high-dimensional sparse representations used in lexical based retrieval functions. Such an encoding process is often accomplished by a dual encoder architecture, with one of the encoders tailored to the queries and the other to the documents. Matching using dense representations enables dense retrieval models to go beyond lexical overlap to retrieve semantically similar documents. This powerful capability of semantic matching, however, often requires relatively large amounts of training data.
Another critical challenge for dual encoder models is that the performance is possibly bounded by the dot-product similarity function. As such, there is growing interest in applying lightweight interaction layers to replace the single dot-product function. On the one hand, Luan et al. (2020) proposes a multi-vector encoding model, which represents each document as a fixed-size set of multiple vectors, and calculate the relevance scores as the maximum inner product over this set. This model combines the efficiency of dual encoders with some of the expressiveness of attention based architectures. On the other hand, ColBERT ( Khattab and Zaharia, 2020 ) proposes to learn embeddings for each token and then use a “MaxSim” operation to select the best candidate. These models achieve significant improvement but also introduce a large latency overhead. In this paper, we take a step back and aim to empirically study how to improve the performance of single dot-product based methods. Specifically, we study whether simply scaling up the model capacity can lead to better fixed embeddings to improve the performance of single dot-product retrievers.
### 2.2 BEIR generalization task
For evaluation in this paper we use BEIR, a heterogenous benchmark for zero-shot evaluation of information retrieval models. The BEIR zero-short evaluation suit contains 18 information retrieval datasets 3 3 3 MS Marco is excluded from the zero-shot comparison as many baseline model used it as training data. across 9 domains, including Bio-Medical , Finance , News , Twitter , Wikipedia , StackExchange , Quora , Scientific , and Misc . The majority of the datasets have binary relevancy labels indicating whether a document is relevant to a given query or not. A small part of the datasets have 3-level or 5-level relevancy judgements. We refer readers to the original BEIR paper Thakur et al. (2021) for more details.
## 3 Generalizable T5 Retriever
### 3.1 T5 dual encoder
We use the dual encoder framework to train dense retrieval models. We follow prior work ( Xiong et al., 2020 ; Hofstätter et al., 2021 ) to initialize dual encoders from pre-trained language models. In this work, we found convenient to use the pre-trained T5 model family as our backbone encoder because the T5 model family provides off-the-shelf pre-trained models (e.g. T5, mT5, byT5) with a wide range of model capacity from millions to billions of parameters ( Raffel et al., 2020 ; Xue et al., 2020 ; Xue et al., 2021 ) . The architectures of our models are illustrated in fig. 2 .
Let paired examples 𝒟 = { ( q i , p i + ) } \mathcal{D}=\{(q_{i},p_{i}^{+})\} be the training set, where q i q_{i} is an input question and p i + p_{i}^{+} is a related passage (e.g., a semantically relevant passage to the question). Following Ni et al. (2021) , we encode the question q i q_{i} and passage p i + p_{i}^{+} into embeddings by feeding them to the T5 encoder and taking the mean pooling of the encoder as output. In all outr experiments, we fix the output embeddings to be of size 768.
We train the model using an in-batch sampled softmax loss ( Henderson et al., 2017 ) : ℒ = e sim ( q i , p i + ) / τ ∑ j ∈ ℬ e sim ( q i , p j + ) / τ , \mathcal{L}=\frac{e^{\text{sim}(q_{i},p_{i}^{+})/\tau}}{\sum_{j\in\mathcal{B}}{e^{\text{sim}(q_{i},p_{j}^{+})/\tau}}}, (1) where the similarity scoring function sim is the cosine similarity between the embeddings of q i q_{i} and p i + p_{i}^{+} . ℬ \mathcal{B} is a mini-batch of examples and τ \tau is the softmax temperature.
Additional negatives p j − p_{j}^{-} can be given for input question q q . The loss is computed by including them in the denominator: ℒ = e sim ( q i , p i + ) / τ ∑ j ∈ ℬ e sim ( q i , p j + ) / τ + e sim ( q i , p j − ) / τ . \mathcal{L}=\frac{e^{\text{sim}(q_{i},p_{i}^{+})/\tau}}{\sum_{j\in\mathcal{B}}{e^{\text{sim}(q_{i},p_{j}^{+})/\tau}+e^{\text{sim}(q_{i},p_{j}^{-})/\tau}}}. (2)
We also apply a bi-directional in-batch sampled softmax loss, where we compute losses for both question to document matching and document to question matching.
### 3.2 Multi-stage training
As shown in fig. 3 , we use a multi-stage dual encoder training approach to achieve generalizable retrieval models.
The training process includes a pre-training stage on a web-mined corpus and a fine-tuning stage on search datasets. The web-mined corpus provides a large amount of semi-structured data pairs (such as question-answer pairs and conversations), which can provide rich semantic relevance information. It is easy to collect but it is often not well annotated, if at all. The search datasets are often annotated by humans, and the queries and documents are also authored by humans. These datasets are of high quality but costly to collect.
In this work, for dual encoder pre-training, we initialize the dual encoders from the T5 models and train on question-answer pairs collected from the Web. Recently, Sentence-T5 ( Ni et al., 2021 ) explored different ways to extract strong text embeddings and achieved remarkable performance on SentEval and Sentence Textual Similarity tasks. We follow that setting to encode queries and passages via mean pooling from the T5 encoders and focus on the dense retrieval tasks.
For fine-tuning, our aim is to adapt the model to retrieval using a high quality search corpus so the model can learn to better match generic queries to documents. In this paper, we consider two datasets for fine-tuning: MS Marco ( Nguyen et al., 2016 ) and Natural Questions ( Kwiatkowski et al., 2019 ) .
## 4 Experimental setup
### 4.1 Training Data
#### Community QA.
In order to leverage most of the power from the large scale models, we collect input-response pairs and question-answer pairs from online forums and QA websites including Reddit, Stack-Overflow, etc. This results in 2 billion question-answer pairs that we use to pre-train the dual encoder.
#### MS Marco.
We consider the MS Marco dataset Nguyen et al. (2016) , which includes 532K query and document pairs, as search data for fine-tuning. The dataset is sampled from Bing search logs, which covers a broad range of domains and concepts. Most of the neural models compared in Thakur et al. (2021) are trained on MS Marco, including DeepCT Dai and Callan (2020) , DocT5Query Nogueira (2019) , ANCE Xiong et al. (2020) and ColBERT Khattab and Zaharia (2020) . Some of these models have shown great generalization with comparable or even better performance relative to BM25.
#### Natural Questions.
In the fine-tuning stage, we also consider the Natural Questions dataset ( Kwiatkowski et al., 2019 ) , which has been widely used in the dense retrieval literature ( Karpukhin et al., 2020 ; Xiong et al., 2020 ) . It consists of 130k query and passage pairs which are also human-annotated.
### 4.2 Configurations
We implement GTR models in JAX 4 4 4 https://github.com/google/jax and train them on Cloud TPU-V8. We consider different sizes of the T5 transformer Vaswani et al. (2017) architecture including Base, Large, XL and XXL. Their number of parameters are listed in table 1 .
Note that we only use the encoder portion of the T5 models and thus the number of parameters are less than half of the full model size. We use the off-the-shelf checkpoints as the initial parameters and use the same sentencepiece vocabulary model. 5 5 5 https://github.com/google-research/text-to-text-transfer-transformer
During pre-training and fine-tuning, we set the batch size to 2048 and use a softmax temperature τ \tau of 0.01. We use Adafactor optimizer ( Shazeer and Stern, 2018 ) and set the initial learning rate to 1e-3 with a linear decay. We train the model for 800K steps and 20K steps for the pre-training and fine-tuning stages, respectively.
For fine-tuning, we use the hard negatives released by RocketQA ( Qu et al., 2021 ) when fine-tuning with MS Marco data and the hard negatives release by ( Lu et al., 2021 ) for Natural Questions, which were proven to lead to better retriever performance. By default, we use the complete MS Marco dataset and the NQ dataset for fine-tuning.
When evaluating on the BEIR benchmark, we use sequences of 64 tokens for the questions and 512 for the documents in all datasets except Trec-News, Robust-04 and ArguAna. In particular, we set the document length to 768 for Trec-News and Robust-04 while setting the question length to 512 for ArguAna, in accordance to the average query and document lengths in these datasets.
### 4.3 Models for comparison
We consider various baselines, including sparse retrieval models: BM25, DocT5Query, and dense retrieval models: DPR, ANCE, TAS-B, and GenQ ( Thakur et al., 2021 ) .
We conduct experiments on four different sizes of our GTR models (GTR-Base, GTR-Large, GTR-XL, and GTR-XXL). We consider three different settings for GTR to investigate the scaling up effect for different training stages: • GTR. This is the full GTR model that conducts both pre-training and fine-tuning.
• GTR-FT. This is a fine-tune only version of GTR where the T5 dual encoders on the MS Marco dataset.
• GTR-PT. This is a pre-training only version of GTR where the T5 dual encoders is only pre-trained on the CommunityQA dataset.
We evaluate baseline and our models on the BEIR generalization task Thakur et al. (2021) as discussed in section 2.2 ,. We consider two main retrieval metrics: NDCG@10 and Recall@100 following BEIR Thakur et al. (2021) . Due to space limitations, we report NDCG@10 in the main section of the paper and include Recall@100 results in appendix A .
## 5 Evaluation Results
We present three groups of experiments to study the a) in-domain performance of the GTR models on MS Marco, b) their out-of-domain generalization performance on BEIR, and c) their data efficiency.
### 5.1 Results on MS Marco