Multi-Task Dense Retrieval via Model Uncertainty Fusion for Open-Domain Question Answering (2021)

Abstract: Multi-task dense retrieval models can be used to retrieve documents from a common corpus (e.g., Wikipedia) for different open-domain question-answering (QA) tasks. However, Karpukhin et al. (2020) shows that jointly learning different QA tasks with one dense model is not always beneficial due to corpus inconsistency. For example, SQuAD only focuses on a small set of Wikipedia articles while datasets like NQ and Trivia cover more entries, and joint training on their union can cause performance degradation. To solve this problem, we propose to train individual dense passage retrievers (DPR) for different tasks and aggregate their predictions during test time, where we use uncertainty estimation as weights to indicate how probable a specific query belongs to each expert’s expertise. Our method reaches state-of-the-art performance on 5 benchmark QA datasets, with up to 10% improvement in top-100 accuracy compared to a joint-training multi-task DPR on SQuAD. We also show that our method handles corpus inconsistency better than the joint-training DPR on a mixed subset of different QA datasets. Code and data are available at https://github.com/alexlimh/DPR_MUF .

## introduction open-domain question-answering requires finding answers to given questions from a large collection of documents @xcite . therefore, a first-stage retrieval component that selects a set of potentially answer-containing documents is often involved for the second-stage reading comprehension model @xcite . traditional term-matching methods such as tf-idf and bm25 @xcite that leverage an inverted index to construct sparse textual representations have built strong baselines in the first-stage retrieval. ## dense passage retrieval retrieval/inference given a collection of documents {d 1 , d 2 , • • • , d n } and a question answering task, dpr @xcite encodes the questions and documents using a bi-encoder structure where encoders f q (•) and f d (•) are independent functions that map a question/document into a low-dimensional, real-value vector. specifically, the similarity s between the question q and document d is defined by the dot product between their encoded vectors v @xmath0: training as pointed out by @xcite , training the encoders such that eq. ( 1 ) becomes a good ranking function is essentially a metric learning problem @xcite . formally, let d be the random variable of documents, q be the r.v. of questions, and c be the r.v. of the set of retrieved documents. given a specific question q, let d + be the positive context that contains answers for q and d - 1 , d - 2 , ...d - k be the negative contexts. the collection of contexts {d the single dpr expert and the joint-training dpr model follow the same training scheme. in the next section, we describe how the second optionmodel uncertainty fusion-is implemented. ## multi-task model fusion given m question-answering tasks and m independent experts, the goal of multi-task model fusion is to find the optimal set of weights {w (i) } m @xmath1. we use dpr as the expert model. ## experimental setup we follow the dpr paper @xcite to train and evaluate our dense retrievers. we replicate their results on all benchmark datasets, with a maximum score difference between ours and their numbers of 1%. this work only focuses on retrieval accuracy as we only improve the retriever. we finally perform sensitivity analysis for the ensemble and uncertainty visualization of our fusion approach. more details are provided in appendix a. ## conclusions in this paper, we propose a model fusion approach for multi-task dense retrieval. instead of training a single dpr model on the union of datasets from different distributions, we leverage model uncertainty to merge different dpr expert's predictions during test time. for each expert, we train an ensemble of small neural networks on top of the pre-trained expert's dense representations and use the mutual information between the ensemble parameters and predictions as the weight, which can be interpreted as the "disagreement" among the ensemble. we compare our model uncertainty fusion approach with single specialists and the multi-task, joint-training dpr model on 5 benchmark qa datasets, as well as their dataset random mixes to test out-of-distribution performance. extensive experiments show that our method manages to outperform these approaches in terms of top-20/100 accuracy on most datasets, while it can also be combined with sparse retrieval methods such as bm25 for further performance gains. our proposed method is simple to implement and effective while enjoying the benefits of continual learning, faster training speed as the experts can be trained in parallel, as well as the flexibility to combine experts from different domains. for future research directions, one could leverage model compression techniques to reduce the index size, or knowledge distillation to learn a single student model from the experts. finally, learning a question-agnostic document index can further save storage space and enhance inference speed for this model fusion method.