Task-Aware Specialization for Efficient and Robust Dense Retrieval for Open-Domain Question Answering (2023)

Abstract: Given its effectiveness on knowledge-intensive natural language processing tasks, dense retrieval models have become increasingly popular. Specifically, the de-facto architecture for open-domain question answering uses two isomorphic encoders that are initialized from the same pretrained model but separately parameterized for questions and passages. This biencoder architecture is parameter-inefficient in that there is no parameter sharing between encoders. Further, recent studies show that such dense retrievers underperform BM25 in various settings. We thus propose a new architecture, Task-Aware Specialization for dEnse Retrieval (TASER), which enables parameter sharing by interleaving shared and specialized blocks in a single encoder. Our experiments on five question answering datasets show that TASER can achieve superior accuracy, surpassing BM25, while using about 60% of the parameters as bi-encoder dense retrievers. In out-of-domain evaluations, TASER is also empirically more robust than bi-encoder dense retrievers. Our code is available at https://github.com/microsoft/taser .

## introduction empowered by learnable neural representations built upon pretrained language models, the dense retrieval framework has become increasingly popular for fetching external knowledge in various natural language processing tasks @xcite @xcite . for opendomain question answering (odqa), the de-facto dense retriever is the bi-encoder architecture @xcite , consisting of a question encoder and a passage encoder. typically, the two encoders are isomorphic but separately parameterized, as they are initialized from the same pretrained model and then fine-tuned on the task. despite of its popularity, this bi-encoder architecture with fully decoupled parameterization has some open issues. first, from the efficiency perspective, the bi-encoder parameterization apparently results in scaling bottleneck for both training and inference. second, empirical results from recent studies show that such bi-encoder dense retrievers underperform its sparse counterpart bm25 @xcite in various settings. for example, both @xcite and @xcite suggest the inferior performance on squad @xcite is partially due to the high lexical overlap between questions and passages, which gives bm25 a clear advantage. @xcite also find that bi-encoder dense retrievers are more sensitive to distribution shift than bm25, resulting in poor generalization on questions with rare entities. in this paper, we develop task-aware specialization for dense retrieval, taser, as a more parameter-efficient and robust architecture. instead of using two isomorphic and fully decoupled transformer @xcite encoders, taser interleaves shared encoder blocks with specialized ones in a single encoder, motivated by recent success in using mixture-of-experts (moe) to scale up transformer @xcite . for the shared encoder block, the entire network is used to encode both questions and passages. for the specialized encoder block, some sub-networks are task-specific and activated only for certain encoding tasks. to choose among task-specific sub-networks, taser uses an input-dependent routing mechanism, i.e., questions and passages are passed through separate dedicated sub-networks. we carry out both in-domain and out-of-domain evaluation for taser. for the in-domain evaluation, we use five popular odqa datasets. our best model outperforms bm25 and existing biencoder dense retrievers, while using much less parameters. it is worth noting that taser can effectively close the performance gap on squad between dense retrievers and bm25. one interest-ing finding from our experiments is that excluding squad from the multi-set training is unnecessary, which was a suggestion made by @xcite and adopted by most follow-up work. our out-of-domain evaluation experiments use en-tityquestions @xcite and beir @xcite . consistent improvements over the doubly parameterized bi-encoder dense retriever are observed in these zero-shot evaluations as well. our code is available at https: //github.com/microsoft/taser. ## background in this section, we provide necessary background about the bi-encoder architecture for dense passage retrieval which is widely used in odqa @xcite and is the primary baseline model in our experiments. the bi-encoder architecture consists of a question encoder and a passage encoder, both of which are usually transformer encoders @xcite . a transformer encoder is built up with a stack of transformer blocks. each block consists of a multi-head self-attention (mha) sub-layer and a feed-forward network (ffn) sub-layer, with residual connections @xcite and layernormalization @xcite applied to both sublayers. given an input vector @xmath0, the ffn sub-layer produces an output vector as following the bi-encoder model independently encodes questions and passages into d-dimension vectors, using the final output vectors for [cls] from the corresponding encoders, denoted as @xmath1, respectively. the relevance between a question and a passage can then be measured in the vector space using dot product, i.e., sim(q, p) = q t p. during training, the model is optimized based on a contrastive learning objective, ## proposed model: taser as shown in fig. 1 , taser interleaves shared transformer blocks with specialized ones. the shared transformer block is identical to the transformer block used in the bi-encoder architecture, but the entire block is shared for both questions and passages. in the specialized block, we apply moe-style task-aware specialization to the ffn sub-layer, following @xcite , where the router always routes the input to a single expert ffn sub-layer. in our experiments, we use a simple yet effective routing mechanism, which uses an expert sub-layer (q-ffn) for questions and another (p-ffn) for passages. the router determines the expert ffn sub-layer based on whether the input is a question or a passage. other routing mechanisms are discussed in appendix a. taser uses one specialized transformer block after every t shared transformer blocks in the stack, starting with a shared one at the bottom. our preliminary study indicates that the model performance is not sensitive to the choice of t , so we use @xmath2. similar to the bi-encoder architecture, taser is trained using the contrastive learning objective l sim defined in equation 2. specifically, the objective needs to use a set of negative passages p for each question. following @xcite and @xcite , we construct p via hard negatives mining (appendix b). our experiments use the multi-set training paradigm, i.e., the model is trained by combining data from multiple datasets to obtain a model that works well across the board. ## related work recent seminal work on dense retrieval demonstrates its effectiveness using transformer-based bi-encoder models by either continual pre-training with an inverse cloze task @xcite or careful fine-tuning @xcite . one line of follow-up work improves dense retrieval models via various continual pre-training approaches @xcite @xcite @xcite . better contrastive learning objectives are also introduced @xcite @xcite . motivated by the success of augmenting dense models with sparse models, @xcite combine the dense retriever with a dense lexical model that mimics sparse retrievers. all above work focus on improving the accuracy of biencoder dense retrievers, whereas our work tackles the parameter efficiency issue. unlike most bi-encoder dense retrievers which measure the similarity between a question and a passage using their corresponding [cls]vectors, colbert @xcite ) develops a late-interaction paradigm and measures the similarity via a maxsim operator that computes the maximum similarity between a token in a sequence and all tokens in the other sequence. such architecture has shown promising results in odqa @xcite and the beir benchmark @xcite . our work instead focus on the improvement on the underlying text encoders, and the maxsim operator introduced by colbert can be applied on top of taser. xiong et al. ( 2021 ) use the bert-siamese architecture for dense retrieval, where all transformer blocks are shared. compared with this architecture, taser is a more effective and general way to increase the parameter efficiency, by interleaving specialized transformer blocks with shared ones. ## conclusion we propose a new parameterization framework, taser, for improving the efficiency and robustness of dense retrieval for odqa. it interleaves shared encoder blocks with specialized ones in a single encoder where some sub-networks are task-specific. as the specialized sub-networks are sparsely activated, taser can provide better parameter efficiency with almost no additional computation cost. experiments show that taser substantially outperforms existing fully supervised biencoder dense retrievers on both in-domain and out-of-domain generalization.