Large Dual Encoders Are Generalizable Retrievers

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.

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³³3MS 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.

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, 2021). The architectures of our models are illustrated in fig. 2.

Let paired examples $\mathcal{D}=\{(q_{i},p_{i}^{+})\}$ be the training set, where $q_{i}$ is an input question and $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}$ and passage $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):

$$\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}$ and $p_{i}^{+}$. $\mathcal{B}$ is a mini-batch of examples and $\tau$ is the softmax temperature.

Additional negatives $p_{j}^{-}$ can be given for input question $q$. The loss is computed by including them in the denominator:

$$\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.

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).

We implement GTR models in JAX⁴⁴4https://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.⁵⁵5https://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.

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.

We first analyze in-domain performance based on the evaluation results on MS Marco. As show in table 3, with scaling up, the models achieve consistent improvement on NDCG@10. We observe similar improvements on other evaluation metrics including MRR@10 and Recall@1000 and reported the numbers in table 7 of appendix A. This shows that increasing model capacity leads to better in-domain performance.

The next set of experiments investigates the effect of increasing model capacity on out-of-domain (OOD) performance.

As shown in table 3, we observe a clear gain on out-of-domain performance in terms of NDCG@10 when the model size increases. The GTR-Large model already outperforms the previous best dense retrieval model TAS-B as well as the best sparse model DocT5Query. Scaling up to GTR-XXL leads to another jump in retrieval performance. Similar improvements are found on Recall@100 as shown in the Appendix’s table 8. On average, the scaling up process demonstrates an encouraging ascending trend that eventually outperforms all baseline methods on all evaluation metrics. This confirms that scaling up is a valid path towards generalizability.

Previously, dual encoders failed to match the performance of BM25 for tasks that require better lexical matching capabilities. Thus, we wanted to investigate what kind of tasks can get improved by scaling up the model size. Figure 4 presents a detailed comparison of all sizes of GTR models against the BM25 baseline.

For tasks like NQ where dual encoders have been previously shown to be more effective than BM25, increasing the model size continues to advance the performance of dual encoders. This suggests scaling up can further boost the head start of dense models over sparse models on these datasets.

For tasks like BioASQ and NFCorpus, where dual encoders previously struggled to match the performance of BM25 for inherent reasons, we discovered that scaling up consistently improves the retrieval performance. In particular, for NFCorpus, our Base model under-performs BM25 but the XL model outperforms BM25 by 5.5% (0.343 vs. 0.325). This exciting finding verifies our assumption that scaling up can further exploit the powerful semantic matching capabilities of the dual encoder models and enable them to ultimately outperform BM25.

To better understand the data efficiency for large dual encoders, we trained models using different portions of the MS Marco dataset during fine-tuning. In particular, we sampled a subset of the training data by keeping only 10% of the training queries as well as their relevant (positive) passages and irrelevant (hard negative) passages.

As shown in table 4, using 10% of training data reduces the in-domain performance of the GTR models on MS Marco. For the GTR-FT (fine-tuning only) models, using 10% of the data leads to a mixed result of out-of-domain performance.

On the other hand, for full GTR models, using 10% of the MS Marco dataset is sufficient for fine-tuning. In particular, the GTR-Large, XL and XXL models achieve comparable or even better OOD performance than fine-tuning on the complete MS Marco dataset. This might suggest that GTR models have the benefit of data efficiency and could use less training data for domain adaptation.

The first ablation study aims to investigate how scaling up effects dual encoder pre-training and fine-tuning. Results are listed in table 5.

For fine-tuning only models, scaling up benefits both in-domain and out-of-domain performance. For pre-training only models, the improvement on in-domain performance is not obvious; meanwhile for out-of-domain tasks, scaling up also improves the generalization. Finally with both pre-training and fine-tuning, GTR models consistently improve over GTR-FT models of all sizes. This shows the power of combining scaling up and a generic pre-training stage.

In table 5, we compare GTR and GTR-PT on the BEIR benchmark to understand the importance of fine-tuning on MS Marco. The table shows that there is a clear gap between GTR models before and after fine-tuning. The result shows the necessity of leveraging a high quality dataset (e.g. search data) to fine-tune the dual encoders.

In table 6, we compare fine-tuning GTR on NQ instead of MS Marco. Compared to MS Marco, NQ only covers Wikipedia documents and is much smaller in size, which allows us to investigate the performance of GTR when fine-tuned on a less generalizable dataset. In addition, fine-tuning on NQ can give us a fair comparison with DPR.

As shown in table 6, the GTR-base model fine-tuned on NQ outperforms the original DPR model, which uses a BERT-Base model as the encoder backbone. This demonstrates the effectiveness of our pre-training on the Web dataset as well as the hard negatives introduced from Lu et al. (2021) for NQ. Fine-tuning on NQ leads to inferior performance compared to fine-tuning on MS Marco, which is consistent with prior work Thakur et al. (2021). However, importantly, scaling up GTR size improves zero-shot performance on BEIR when fine-tuning on NQ. This shows that the benefit of scaling up holds for different fine-tuning datasets. Furthermore, when scaling from Large to XL, we observe a larger gain when fine-tuning with NQ than with MS Marco, indicating that scaling up helps more when using weaker fine-tuning data.

Previously, BEIR has shown that models trained with cosine similarity prefer short documents while those trained with dot-product prefer long documents Thakur et al. (2021). We investigate whether scaling up affect this observation. Specifically, we compute the median lengths (in words) of the top-10 retrieved documents for all queries. Results are shown in fig. 5.

Though all GTR models are trained using cosine similarity, we found that scaling up the model size has influence over the lengths of retrieved documents. We observe an increasing trend of document length for DB-Pedia, Fever, HotpotQA, Signal-1M, Trec-News, and Web-Touche2020 with scaling up. In particular, for Web-Touche2020, the lengths of the retrieved documents grow drastically as the models scale up: The largest GTR-XXL retrieves documents that are on average twice as long compared with the smallest GTR-Base. This plays in our favor since Thakur et al. (2021) show that the majority of relevant documents in Web-Touche2020 are longer.

On the other hand, the only exception we observe is the Trec-Covid dataset, where GTR-XXL model retrieves much shorter documents than those retrieved by the smaller size counterparts. This may explain the inferior performance of GTR-XXL on Trec-Covid shown in table 3 and table 8. We leave it as future work to explore the effects of using the dot-product as similarity function for large dual encoders.