Query Distillation: BERT -based Distillation for Ensemble Ranking (2020)

Abstract: Recent years have witnessed substantial progress in the development of neural ranking networks, but also an increasingly heavy computational burden due to growing numbers of parameters and the adoption of model ensembles. Knowledge Distillation (KD) is a common solution to balance the effectiveness and efficiency. However, it is not straightforward to apply KD to ranking problems. Ranking Distillation (RD) has been proposed to address this issue, but only shows effectiveness on recommendation tasks. We present a novel two-stage distillation method for ranking problems that allows a smaller student model to be trained while benefitting from the better performance of the teacher model, providing better control of the inference latency and computational burden. We design a novel BERT-based ranking model structure for list-wise ranking to serve as our student model. All ranking candidates are fed to the BERT model simultaneously, such that the self-attention mechanism can enable joint inference to rank the document list. Our experiments confirm the advantages of our method, not just with regard to the inference latency but also in terms of higher-quality rankings compared to the original teacher model.

## introduction the information retrieval (ir) community has witnessed the flourishing development of neural ranking models in the past several years, examples including drmm @xcite , @xcite , pacrr @xcite , and co-pacrr @xcite . recently, bert @xcite , the pre-trained deep bidirectional transformer, has shown strong performance on a broad range of language processing tasks and has wide application in ranking tasks as well @xcite @xcite . to further boost the results, it is common to adopt model ensembles, as modern neural ranking models provide a wealth of options for sub-models. the scores of all such sub-models with regard to the relevance of a candidate document to a query are collected and fed into a lambdamart @xcite or xgboost @xcite model to obtain the final ensemble relevance score. however, the computational burden is extremely heavy when drawing on such ensembles, and the prominence of increasingly deep and large neural networks such as bert exacerbates this problem even more. furthermore, it is rather inconvenient to update and maintain an ensemble of large models, which is concerning in real-world online deployments. knowledge distillation (kd) is a common approach to balance effectiveness and efficiency @xcite . a well-trained large model serves as a teacher for a smaller student model that is trained not only based on the ground truth labeled data, but also using the label distribution emitted by the teacher, such that the student is ultimately able to replace the teacher. however, it is not straightforward to apply kd to ranking problems. first of all, in ranking, we focus on the relative order of documents rather than the label distribution in classification problems, which kd is designed for. on top of this, the total number of documents is often so large that we retrieve only a subset of relevant documents to reduce the size of the ranking list, and computing the overall distribution over all documents is impractical. recently, @xcite proposed ranking distillation (rd) to address these issues in the context of recommender systems. based on the ground truth training data set, the teacher model makes predictions for additional unlabeled documents to obtain a top-k unlabeled document ranking. the student model is then trained to minimize not just the ranking loss on the training data set but also a distillation loss with the example top-k ranking of unlabeled documents generated by its teacher. this method is reminiscent of that of @xcite and has shown its effectiveness in recommender systems. however, this approach is beset by several problems when considering standard query document ranking tasks. unlike recommender systems, in the case of search engines, it is rare for the teacher to predict an unlabeled document as being a positive sample, since the set of unlabeled documents consists of all documents that are not in the recall set, and the vast majority of these documents bear little connection with the query. also, the weight of unlabeled samples needs to be set empirically, making it unclear whether this method can be effective for query document ranking. in this paper, we propose a substantially different distillation method for ranking tasks. instead of adding additional unlabeled documents for all the queries in the training data, our query distillation approach incorporates additional queries. it uses the teacher to predict top-ranked documents for these queries, enabling the teacher to guide the student's training. furthermore, we divide the student training into two phases: the student model is first trained using teacher-labeled ranking data as teacher guidance and subsequently is fine-tuned using the ground truth labeled training data so as to obtain higher-quality rankings. through this two-stage fine-tuned training regimen, we hope to obtain a student model that benefits from the strong retrieval effectiveness of the richer model but reduces the inference latency and computational burden. we adopt a bert-based model as the student in light of its outstanding performance in ranking tasks @xcite @xcite . we also propose a novel bert qd-list model structure for ranking, which, contrary to the common practice of treating a query document pair (bert qd-pair) as the input to bert, jointly considers a query and the entire candidate document list. the self-attention mechanism in bert allows the query to be evaluated with regard to all documents simultaneously. in addition, a custom attention mask is applied to the kqv self-attention layers to help boost the results. our experiments show that a single student model outperforms the original large ensemble model through the two-stage teacher-student training and our qd-list model saves substantial inference time while obtaining better results compared to the bert qd-pair model. the contributions of our work are two-fold: (1) we introduce a novel two-stage fine-tuned teacherstudent training method and obtain a single student model that benefits from the high quality of the ensemble model while reducing the inference latency and computational burden. (2) we propose a list-wise bert-based ranking model, which through self-attention, allows the model to observe more information while ranking. the experiments confirm the advantages of our contributions. ## conclusion in this paper, we introduce query distillation as a two-stage fine-tuned distillation training process for large ensemble ranking models. furthermore, we propose a novel list-wise bert model structure for ranking tasks (bert qd-list), which is used as our student model. the experiments confirm the advantages of our query document list model not just with regard to the inference latency but also with regard to the retrieval quality over regular query document pair bert (bert qd-pair) as well as the original teacher models. in the future, we will try to apply our ranking distillation method on further tasks, such as answer selection in machine reading. additionally, bert variants that can handle larger sequence lengths such as xlnet @xcite will be evaluated to process even more documents at a time.