Improving Multitask Retrieval by Promoting Task Specialization (2023)

Abstract: In multitask retrieval, a single retriever is trained to retrieve relevant contexts for multiple tasks. Despite its practical appeal, naive multitask retrieval lags behind task-specific retrieval, in which a separate retriever is trained for each task. We show that it is possible to train a multitask retriever that outperforms task-specific retrievers by promoting task specialization. The main ingredients are: (1) a better choice of pretrained model—one that is explicitly optimized for multitasking—along with compatible prompting, and (2) a novel adaptive learning method that encourages each parameter to specialize in a particular task. The resulting multitask retriever is highly performant on the KILT benchmark. Upon analysis, we find that the model indeed learns parameters that are more task-specialized compared to naive multitasking without prompting or adaptive learning.1

## introduction a standard approach to knowledge-intensive language tasks such as question answering (qa), entity disambigution, and fact verification is retrieval-based. given an query, a retriever is used to efficiently search a large knowledge base (kb) to retrieve relevant ''contexts'', typically in the form of short paragraphs. how these contexts are used is task-specific (e.g., entity disambiguation takes the title of the article in which the top retrieved context is found; qa predicts an answer from the contexts through a reader model). in this paper, we focus on the retrieval step. in particular, we focus on multitask retrieval. in this setting, there are @xmath0. a single retriever is then tasked with performing retrieval for k tasks. multitask retrieval contrasts with task-specific retrieval, in which a separate retriever is trained for each task, and has compelling advantages such as model simplicity (i.e., we can use the same model for all tasks rather than having to design potentially different models for different tasks) and memory efficiency at test time (k times smaller). despite the practical appeal, the performance of multitask retrieval has been underwhelming, severely limiting its real-world applicability. specifically, previous work by @xcite trains dpr @xcite on the union of all training datasets in the kilt benchmark @xcite , but the model is outperformed by task-specific retrievers in 5 out of 8 tasks (page-level r-precision, validation split). in our experiments, we find that it is in fact outperformed in all tasks (often by substantial margins) when a stronger task-specific baseline is used. this result is surprising as well as disappointing given the usual benefits of multitask learning (e.g., data efficiency, reduced overfitting) when properly done. we debunk the previous negative result by presenting a multitask retriever that outperforms task-specific retrievers. the main theme of our work is that it is beneficial to explicitly promote task specialization. a first important source of improvement is a better choice of pretrained model, one that is explicitly optimized for multitasking. specifically, instead of the standard retrieval encoder bert @xcite , we use t5 @xcite , which includes multitasking in its pretraining stage. importantly, we use the same prompting as in pretraining (i.e., task indicator) to reduce the gap between pretraining and finetuning for multitask retrieval. a second source of improvement is a novel adaptive learning method in which we adatively upweight the task gradients by the parameter's sensitivity to these tasks to encourage task specialization. the resulting multitask retriever is highly performant on the kilt benchmark. we achieve 73.74% average page-level r-precision on kilt validation data and 72.84% average page-level r-precision on kilt test data. upon analysis, we find that the model indeed learns parameters that are more task-specialized compared to naive multitasking without prompting or adaptive learning. we emphasize that achieving competitive multitask retrieval in practice is a highly difficult empirical problem. one might think that it is simply an application of multitask learning, which has no shortage of sophisticated techniques. these techniques typically modify the gradients during training, such as gradient surgery @xcite , gradient vaccine @xcite , common gradient descent @xcite , and gradnorm @xcite . we experiment with these techniques and find that they do not help, thus motivating us to develop one that works. our technical contribution is a new method for multitask learning based on the notion of task sensitivity. given a loss function j(θ), the sensitivity of the i-th parameter to the loss at θ is defined as the absolute change in the loss when θ i is set to zero, which can be approximated by a first-order taylor approximation as we briefly differentiate our work from other recent work on multitask retrieval. @xcite present corpusbrain, an autoregressive multitask retriever trained in largely the same style as genre @xcite with excellent performance. autoregressive retrieval has different pros and cons compared to dense retrieval which is our setting; it can be more memory and runtime efficient, but it does not ''read'' the description of the target and thus not suited for retrieval tasks that require involved reasoning over query-target pairs (e.g., zero-shot entity retrieval @xcite ). thus we consider the contribution of corpusbrain to be at least partially orthogonal to ours. nevertheless, we show that our model outperforms corpusbrain in a similar training setting in experiments. @xcite propose instruction-based retrieval in which the retriever is given an intent as well as a query to find the intended target. while this is a form of multitask retrieval, the problem formulation is different and it is evaluated on its own dataset benchmark. ## method we build on the well-established framework of dual encoder @xcite @xcite . let x denote the set of all queries and y the set of all targets (i.e., kb). first, we assume mappings text x : @xmath1: @xmath2. second, we assume encoders enc θ x , enc θ y : v + → r d with parameters θ defining the relevance score function n k ) ∼ pop k drawn iid from the k-th population distribution pop k . we use the kilt benchmark, which includes @xmath3, entity linking, fact checking, slot filling, and dialogue. 2 the per-task loss is previous work by @xcite uses the following setting. the kb y consists of 100token disjoint wikipedia passages. the text mappings text x , text y apply the bert tokenizer to unmodified queries and passages. the encoders enc θ x , enc θ y are initialized with independent pretrained bert-bases (uncased). the task-specific training datasets are downsampled to be of similar sizes. as in dpr, they train the model using hard negatives based on bm25, followed by one round of hard negative mining from the model (only on natural questions and triviaqa in which verifying if a candidate negative is indeed incorrect is expedient). we now describe the main sources of improvement that we achieve over the baseline multitask retriever: a better choice of the base model with appropriate prompting, and better optimization.