Enhancing Document Ranking with Task-adaptive Training and Segmented Token Recovery Mechanism (2021)

Abstract: In this paper, we propose a new ranking model DR-BERT, which improves the Document Retrieval (DR) task by a task-adaptive training process and a Segmented Token Recovery Mechanism (STRM). In the task-adaptive training, we first pre-train DR-BERT to be domain-adaptive and then make the two-phase fine-tuning. In the first-phase fine-tuning, the model learns query-document matching patterns regarding different query types in a pointwise way. Next, in the second-phase fine-tuning, the model learns document-level ranking features and ranks documents with regard to a given query in a listwise manner. Such pointwise plus listwise fine-tuning enables the model to minimize errors in the document ranking by incorporating ranking-specific supervisions. Meanwhile, the model derived from pointwise fine-tuning is also used to reduce noise in the training data of the listwise fine-tuning. On the other hand, we present STRM which can compute OOV word representation and contextualization more precisely in BERT-based models. As an effective strategy in DR-BERT, STRM improves the matching perfromance of OOV words between a query and a document. Notably, our DR-BERT model keeps in the top three on the MS MARCO leaderboard since May 20, 2020.

## introduction document retrieval (dr) requires the machine to retrieve and rank documents according to their relevance with the query, which needs strong text understanding ability. as one basic and crucial task in nlp, it can aid several real applications, such as question answering systems and web-based search engines, e.g., google, yahoo, bing, etc. with the development of deep learning and the increasing emergence of large-scale datasets, e.g., ms † equal contribution. figure 1 : two cases, each listing an input query, a document candidate and the text tokenized by word-piece in bert. the oov words, e.g., "insouciantly" and "bogue", are tokenized into 2 or more sub-tokens. marco @xcite , dr has achieved remarkable advancements. lots of dr models have been studied over the last few years. traditional machine learning based dr models, like lambdarank @xcite , adarank @xcite , etc., rely heavily on manual feature engineering which is time-consuming and unsustainable. neural dr models, including dssm @xcite , knrm @xcite , etc., learn the query and document representation and ranking features in a continuous vector space, which obviate the need of manual feature design. formally, ranking models can be divided into three categories: pointwise, pairwise and listwise. it has been reported that listwise models perform comparatively better in ranking tasks @xcite . recently, the pre-trained language models have caused a stir in dr. taking bert @xcite as an example, it employs deep transformers to enhance language understanding from largescale texts, obtaining state-of-the-art results in a wide variety of nlp tasks, including dr. no matter whether bert is applied to dr in either a featurebased or a fine-tuning manner @xcite , substantial improvements in dr have been attained. the bert-based models, despite their powerfulness, could be further improved for dr in the following aspects: (1) current bert-based ranking models are either pointwise or pairwise style. their training targets are to optimize the relevance of query and document or the order of documents within pairs rather than minimizing errors in ranking of documents. (2) the desired model could be more task-adaptive by considering text domains and query types. (3) due to wordpiece @xcite segmentation method employed by bert, the matching of oov word might be miscalculated and further a document irrelevant to a query might be ranked high. figure 1 gives two examples to illustrate the problem induced by word-piece. in figure 1 (a), the query and the document are irrelevant, because the word "bogue" in the query and the word "bogus" in the document are unrelated. however, bert fails to distinguish the two words as a result of being misled by the separated tokens generated by wordpiece, i.e., "bog". in figure 1(b) , the tokens generated by wordpiece, i.e., "ins" in the query and "in" in the document, lead to an undesirable matching and further a high query-document relevance score. (4) these models do not have the ability to deal with noise in the training data, which is common in nlp corpora. in order to solve the above-mentioned problems of bert-based ranking models, we propose a new dr model named dr-bert. in the dr-bert model, we make the following improvements on the basis of bert: (1) as depicted in figure 2 , we construct a bert-based listwise method to learn document-level comparison with regard to a given query in fine-tuning. (2) we present a domainadaptive pre-training process and a query typeadaptive fine-tuning strategy to adapt the model to this dr task. (3) we add a segmented token recovery mechanism (strm) into dr-bert, which can effectively improve the matching accuracy of oov words. (4) we employ the model derived from pointwise fine-tuning to reduce noise in the training data of the listwise fine-tuning. we conduct extensive experiments on the ms marco dataset, which is a large-scale benchmark from search engine bing. in ms marco, all queries are sampled from real search queries and documents are real web documents. with one million queries, it is one of the most comprehensive realworld datasets of its kind in both quantity and quality. experimental results show that our dr-bert model showed excellent performance. our contributions can be summarized as follows. ## related work learning to rank refers to adopting machine learning algorithms to train models for ranking tasks. these ranking models can be employed in a wide variety of applications. while applied to the dr tasks, ranking models output a ranked document list for each query based on relevance scores computed by the models. depending on how the loss functions are defined, we can categorize learning to rank into three classes, namely pointwise, pairwise and listwise. pointwise approaches @xcite transform ranking tasks to classification or regression tasks by directly predicting the relevance score of a document with respect to a given query. although easy to implement, pointwise approaches do not make comparisons between documents which is the core of ranking. in pairwise approaches @xcite @xcite , ranking is transformed into classification on the order of document pairs. in particular, each document is compared to the other one at a time according to their relevance with the query. the final ranked list is arranged by their relative positions. however, the assumption that these approaches require, i.e., the documents are generated in pairs, is too strong. even worse, the uneven distribution of documents in different queries might cause the training bias. further, the listwise approaches @xcite @xcite @xcite train the models to learn the optimal ranking directly by taking the documents to be ranked as input. @xcite treats the sequence of documents ordered by the top-one probability distribution as the ranking list and optimizes the cross entropy. @xcite maximizes the likelihood of the golden ranking list. @xcite optimizes the similarity of the query and the documents. @xcite ; @xcite directly optimize the metrics like ndcg @xcite which measure the quality of a ranking list. on the other hand, most of the above methods rely heavily on handcrafted features which require much expertise. as the rapid developing of deep learning techniques, lots of work @xcite tempts to build neural ranking models which need no manual features and show impressive performance. recently, the pre-trained language models like bert @xcite have achieved the stateof-the-art results on several nlp tasks. leveraging their powerful language understanding abilities, several ranking models built on bert are proposed to improve the effectiveness and efficiency of dr, for effectiveness of dr, @xcite adopt a pointwise paradigm and nogueira et al. (2019a) adopt a pairwise one. further, @xcite append predicted queries to the document and rank the documents with bert as described in @xcite . however, their training targets are minimizing errors in classification of query-document relevance or the order of document pairs rather than minimizing errors in ranking of documents, which restricts their performance. for efficiency of dr, @xcite present a novel ranking model which encodes the query and the document using bert and employs an interaction mechanism to speed up the retrieval. ## our approach first of all, we formally describe the task definition as follows. given a query q and a very large set of documents ď, our goal is to produce a ranked list of documents, which is as close as possible to the oracle ranking of documents according to their relevance levels, i.e., @xmath0, 1}, where 1 indicates positive sample and 0 means negative one. basically, the pipeline of dr contains two stages: a retrieval stage and a re-ranking stage. the retrieval stage is to get a smaller set of document candidates for a query, which is an effective way to balance the performance and costs. in this work, we apply bm25 (deepct-index) @xcite to get the top-k document candi- readers can refer to @xcite for more details. in this work, we focus on the reranking stage. in the following subsections, we describe our ranking model dr-bert, whose training process comprises a domain-adaptive pre-training and a two-phase fine-tuning, i.e., type-adaptive pointwise fine-tuning and listwise fine-tuning, in which strm is proposed to solve the oov mismatching problem. ## conclusion we propose a model named dr-bert for dr task, in which we adopt a domain-adaptive pre-training and present a two-phase fine-tuning strategy, i.e., type-adaptive pointwise fine-tuning and listwise fine-tuning. besides, we present a very useful segmented token recovery mechanism to improve the matching performance of oov words, which can be also applicable to other bert-based models. experimental results show our model outperforms many strong baselines and keeps in the top three on the ms marco leaderboard since may 20, 2020.