Multi-task Retrieval for Knowledge-Intensive Tasks

Given a large collection of unstructured text passages, information retrieval (IR) can be broadly defined as finding a small set of passages that satisfies an information need, often presented in the form of a short-text query Manning et al. (2008). Traditional IR methods, such as tf-idf and BM25 Robertson and Zaragoza (2009), match keywords efficiently with an inverted index. Such methods can be seen as representing queries and passages in high-dimensional, sparse vectors, where each dimension corresponds to a term in the vocabulary and the weight indicates its importance.

In contrast to tf-idf and BM25, dense retrieval methods encode text as a latent semantic vector of a fixed, much smaller dimensionality. Whether a passage is relevant to a given query is determined by the distance of their vectors Deerwester et al. (1990). Although dense representations do not encode tokens explicitly and can potentially map paraphrases of completely different tokens to close vectors, performance of early dense retrieval methods was often inferior to term-matching approaches, except when large labelled data is available Yih et al. (2011); Gao et al. (2011); Huang et al. (2013). Thanks to success of large pre-trained models Devlin et al. (2019); Liu et al. (2019b), however, recent dense retrieval methods have shown to outperform the sparse counterparts, when fine-tuned on a small set of in-domain labelled data Karpukhin et al. (2020); Lewis et al. (2020b); Xiong et al. (2020). Efficient index and search of dense vectors are made possible by maximum inner product search (MIPS) algorithms (e.g., Shrivastava and Li, 2014; Guo et al., 2016), as well as tools like FAISS Johnson et al. (2019).

Our work is built upon the Dense Passage Retriever (DPR) architecture of Karpukhin et al. (2020), which was initially proposed for the task of open-domain question answering. DPR is a neural bi-encoder model which embeds queries with an encoder $\boldsymbol{f}\,(\cdot)$ and passages with a separate encoder $\boldsymbol{g}\,(\cdot)$. Given an input query $x$ and a target passage $y$, we have

$$\operatorname{p}\left(x\mid y\right)\propto\operatorname{sim}(x,y),$$

where the similarity score $\operatorname{sim}\left(x,y\right)$ is defined as the inner product of the embeddings of its arguments, $\boldsymbol{f}(x)\cdot\boldsymbol{g}(y)$. Given a query at inference time, calculating its similarity with every possible passage would be prohibitive for large knowledge sources. Therefore, DPR makes use of the FAISS library (Johnson et al., 2019) to perform fast approximate nearest neighbour search in sub-linear time.

Training is based on a contrastive loss. Given a query $x$, a relevant passage $y$, and a set of $n$ irrelevant passages $y^{-}_{i}$, we train the model by optimising the following negative log likelihood:

$$\mathcal{L}=-\log\frac{\exp(\operatorname{sim}(x,y))}{\exp(\operatorname{sim}(x,y))+\sum_{i=1}^{n}\exp(\operatorname{sim}\left(x,y^{-}_{i}\right))}.$$

As the set of irrelevant passages, we use the relevant passages for other queries within the same batch, as well as a specially selected “hard” confounder. This is a passage which has high lexical overlap with the query (high BM25 score), but is not among the set of relevant passages for the given data point. Karpukhin et al. (2020) have shown that the inclusion of such “hard” confounders leads to substantially improved training results. This training process is illustrated in Figure 1.

For the training and evaluation of all models in the paper we make use of KILT, a benchmark and library of datasets (Petroni et al., 2020). KILT consists of a selection of datasets spanning five varied classes of knowledge-intensive tasks (i.e., question answering, slot filling, fact checking, dialogue, entity linking), with the aim to cover many different ways of seeking knowledge. Input queries can vary wildly from one task to the other, and include classic examples of open-domain retrieval tasks such as natural language questions and claims to be verified, as well as more unusual examples like conversation fragments and long chunks of annotated text. Crucially, all datasets distributed in KILT have been re-aligned such that they are all grounded in the same snapshot of Wikipedia, which the authors distribute. The knowledge required to answer any of the queries in the library of tasks can thus be found within the same unified knowledge source.

To illustrate the variety of ways in which the input queries for different tasks can be formulated, we provide a few simple examples in Table 1. In spite of the differences between query formulations, all these tasks share one crucial aspect: they all require a retriever to fetch the relevant passages from the knowledge source, in order to support the final downstream task.

Using task-specific models to tackle our collection of retrieval tasks would involve completely separate models, one per dataset. As illustrated in Figure 2, this would lead to a proliferation of models and data, down to separate indexed copies of the knowledge source itself (Wikipedia). This setup will form one of our baselines.

Multi-task training has been successfully used to allow models to leverage cross-task data, as well as to provide a regularisation effect leading to better generalisation ability (Liu et al., 2019a). We apply this concept to neural retrievers, with the aim of improving performance by jointly leveraging multiple different retrieval datasets.

Our base setup is illustrated in Figure 3(b) and involves using a shared passage encoder — so that a single index of encoded passages can be used — as well as a query encoder that is shared across all tasks. In essence, in this setup a single DPR model is used to perform all retrieval tasks.

Due to the complexity of training and evaluating retrieval models (which involves training the retriever, embedding all of Wikipedia, and building an index), our main set of experiments is all based on this configuration, which was found to work well in preliminary experiments. However, in order to report on the performance of alternative architectures, we also investigate the following additional variants in a restricted experimental setting, limited to a few tasks:

• •

  Task-specific query encoder. A different query encoder is used for each family of tasks, e.g. all question answering tasks use the same query encoder, but fact checking uses a different one. This is meant to allow for potentially different needs in processing queries, given the fundamentally diverse nature of the tasks at hand. This setup configuration is illustrated in Figure 3(a).

• •

  Task markers. This approach is similar to our base setup, where a single model performs all tasks. Additionally, we introduce specialised tokens which are inserted at the beginning of each query. Their aim is to help the model distinguish between the different tasks, by marking them. We use one task marker for each of the five task classes of KILT, such that all question answering tasks share the same marker.

We saw in § 2.1 how “hard” confounder passages are collected using a BM25 baseline, following the standard approach in DPR. However, any other retriever can be used to select such confounders, including the very retriever being trained, leading to an iterative, self-adversarial training. Concretely, this amounts to following steps: (1) a first version of the retriever is trained with BM25 confounders; (2) new confounders are selected with the trained model, by retrieving high-ranking passages which are not among the set of relevant ones; (3) a second version of the model is trained using the additional new confounders.

Intuitively, it is expected that this approach should lead to higher quality confounders compared to those selected by BM25 based on simple keyword matching. Based on our own experience as well as relevant literature (Khattab et al., 2020), this adversarial approach has been shown to work well for question answering.

As a way of further pushing the performance of the model, we experiment with this adversarial confounder selection on two datasets, Natural Questions (Kwiatkowski et al., 2019) and TriviaQA (Joshi et al., 2017). We selected these two datasets since, out of all of the tasks we are considering, they have an easy way of checking whether a certain passage is relevant or not for a given query – namely, by checking whether the answer is present in the passage. This enabled us to automatically build sets of confounders, ensuring relevant passages would be excluded.²²2Strictly speaking, assuming a passage to be irrelevant because of the absence of the answer span is not formally correct. However, experiments show a good correlation between this simple check and the overall model quality.

The results of the evaluations reported in Petroni et al. (2020) show that retrievers trained for question answering have poor performance outside of their domain. We would like to understand if it is possible to design a single model which can accurately satisfy the information needs of a wide variety of knowledge-intensive tasks. In short: Can a neural retriever be universal?

We perform a comprehensive evaluation of several models on the eight tasks of Table 2. The setups we evaluate include eight task-specific models (one trained on each of the eight datasets), for which we measure both in-domain and out-of-domain performance, and a BM25 baseline. Additionally, we include a multi-task trained model – as described in §3.1 – with the hope that it can learn to perform all tasks satisfyingly. This amounts to 10 models evaluated on eight tasks each, for a total of 80 evaluations.

To measure retrieval performance, we adopt the main metric used for the KILT benchmark, $R$-precision. This is calculated as $r/R$, where $R$ is the total number of relevant passages for a given query, and $r$ is the number of relevant passages returned among the top-$R$ retrieval results. For the case of $R=1$ this is therefore equivalent to precision@1. Table 3 shows retrieval performance on the validation data, with the best performance on a given dataset marked in bold, and the second best performance underlined.

While the KILT evaluation focuses on retrieval at the level of Wikipedia pages (thereby marking as “hits” any results that lie within the correct page), we are also interested in performing an evaluation at a more fine-grained level. We therefore also evaluate our models at the passage level, using a modified version of the official KILT evaluation scripts. These are shown as the second number in each column.

We straight away notice that the task-specific models tend to achieve high performance on their respective tasks, often taking one of the top two spots. Interestingly, we also note that these neural retrievers consistently outperform the BM25 baseline, showing that the result which Karpukhin et al. (2020) achieved for open-domain question answering also holds for other knowledge-intensive tasks.

The results reveal a strong performance for the multi-task model, confirming the hypothesis that a single model can be successfully trained to perform a wide variety of retrieval tasks. With the exception of one dataset, the shared model achieves the best retrieval performance or is within a few percentage points of the top score. We note that the one exception is the Zero-shot RE task (Levy et al., 2017), a trivial task in which the query will always contain the title of the page to be retrieved. Indeed, the model specific to this task manages to achieve a near-perfect score.

Another task which stands out for being markedly different in formulation is AIDA-YAGO 2 (Hoffart et al., 2011). As shown in Table 2, models that were not trained on this specific task perform it very poorly. Entity linking is a task that is normally better performed by models which are explicitly designed for it (Cao et al., 2020). We nevertheless include it to showcase the ability of neural retrievers to adapt to it, and note how well the multi-task retriever performs on it in spite of its unusual nature.

We saw that our proposed approach achieves strong performance across a variety of retrieval tasks. However, our interest in neural retrievers stems from their use as components within larger systems, to perform tasks such as question answering. Our next experimental question is therefore: Can a universal retriever lead to better downstream performance in knowledge-intensive tasks?

We perform a downstream evaluation of our approach used in conjunction with BART (Lewis et al., 2020a) as the generative component or classifier, adopting the same setup as Petroni et al. (2020). Results are reported in Table 4, with bold and underline marking the best and second best scores respectively.

The DPR + BART line refers to a setup similar to our own, but with the simpler retriever of Karpukhin et al. (2020). Therefore, comparing its performance to ours gives us a clear indication of the contribution of multi-task training on the overall performance on knowledge-intensive tasks. Our proposed model achieves significantly better performance than this baseline in AY2, zsRE and HoPo; while for the other tasks, the discrepancy is always below two points. This fact is reflected in the last column too, showing that on average multi-task training leads to better downstream performance. The model also compares favourably to RAG (Lewis et al., 2020b), a more advanced system in which the query encoder is fine-tuned on the end task.

Task-specific neural retrievers can achieve higher performance than IR-based methods, but they are not suitable for cases where no training data (or not enough) is available. In those cases, tf-idf and BM25 are the better choice. To evaluate the performance of a multi-task retriever as a suitable replacement for them in this scenario, we run a series of experiments in the low data regimes (few-shot and zero-shot).

We start by training a set of multi-task retrievers (using the base setup) in the leave-one-out setting for each of the datasets, in order to see how a neural retriever will perform when trained on all domains except for the one it is to be evaluated on. The results of these zero-shot experiments are reported in the second line of Table 5 (again, text here is in bold for the best overall performance, and underlined for second best). They show that, even in the zero-shot setting, the multi-task neural retriever achieves performance that is competitive to BM25, with retrieval being 10 points higher at the page level and 5 points lower at the passage level on average.

The advantage of neural retrievers over BM25 lies in their ability to improve with training. We therefore look at few-shot training for each task, and create two smaller copies for each of the original training sets with a random sample of 128 and 1,024 examples respectively. In order to evaluate the suitability of a multi-task trained retriever as a starting checkpoint for few-shot training, we take the various leave-one-out models and fine-tune them on our few-shot training sets. To check whether multi-task pre-training is effective, we also compare these to DPR models (which are just initialised with BERT weights) fine-tuned on the same data.

The bottom two sections of Table 5 report the results. The most dramatic gains from fine-tuning are seen for AY2, an “outlier” task whose formulation differs from that of the other tasks, and which seems to benefit the most from seeing in-domain data. The zsRE performance does not seem to improve from fine-tuning on the smaller dataset, but sees a very big jump when switching to the larger dataset. As a reminder, in this trivial task the title of the page to be retrieved always appears at the start of the query. It is therefore not surprising that models specifically fine-tuned on it can achieve near-perfect scores, as long as enough training data is provided.

In spite of the fine-tuning, we note that both DPR and the multi-task model fail to improve on their performance for T-REx, suggesting that large amounts of training data are required to learn this task. Nevertheless, the multi-task model proves itself more robust, and achieves the top performance on it.

Finally, we note for 2 out of 8 tasks, namely zsRE and WoW, DPR achieves lower page-level retrieval scores than the multi-task model, but performs better at the passage level. This shows that fine-grained and coarse-grained retrieval performance are not always perfectly correlated.

Overall, the experiments show strong results for the multi-task model, with the average zero-shot performance being competitive to BM25, and the average few-shot performance being markedly better than the alternatives. The discrepancy in performance between a vanilla DPR model and the leave-one-out multi-task model is especially noticeable when using the smaller of the two datasets, in which case average performance for the latter is more than double that of vanilla DPR.

In this set of experiments we compare our base multi-task model with the two variants described in § 3.1. Due to the high memory consumption of the “task-specific encoders” variant (requiring one full query encoder per task family, in addition to the passage encoder), it was only possible to perform these evaluations in a restricted setting of three datasets. The results in Table 6 do not reveal a clear winner, suggesting that the base architecture might be the better choice due to its simplicity and generally good performance.⁵⁵5Not included in this table, due to very poor performance in preliminary experiments, are two further variants: a base model with a single encoder for both queries and passages, and a base model trained from scratch without BERT pre-training.

Finally, we evaluate the adversarial confounder selection method described in § 3.2. This involves augmenting our regular training sets with additional confounders for TriviaQA and Natural Questions, selected using our top multi-task trained model. A new multi-task model is then trained from scratch on this augmented data. Its performance is reported in Table 7, showing an overall improvement across multiple tasks. While this approach is demonstrated here on our multi-task model, it is in fact orthogonal to it, and could be applied to any other neural retrievers trained with a contrastive loss.