SlowGuess/NAACL_2022_OCR huggingface dataset:
Xiaozhi Zhu $^{1}$ , Tianyong Hao $^{1}$ , Sijie Cheng $^{3}$ , Fu Lee Wang $^{4}$ , and Hai Liu $^{1,2*}$
$^{1}$ School of Computer Science, South China Normal University, China
$^{2}$ Guangzhou Key Laboratory of Big Data and Intelligent Education, Guangzhou 510000, China
$^{3}$ Fudan University, China,  $^{4}$ Hong Kong Metropolitan University, China
{2020022975,haoty}@m.scnu.edu.cn
815106263@163.com, pwang@hkmu.edu.hk, namelh@gmail.com
Pretrained language models such as BERT have been successfully applied to a wide range of natural language processing tasks and also achieved impressive performance in document reranking tasks. Recent works indicate that further pretraining the language models on the task-specific datasets before fine-tuning helps improve reranking performance. However, the pre-training tasks like masked language model and next sentence prediction were based on the context of documents instead of encouraging the model to understand the content of queries in document reranking task. In this paper, we propose a new self-supervised joint training framework (SJTF) with a self-supervised method called Masked Query Prediction (MQP) to establish semantic relations between given queries and positive documents. The framework randomly masks a token of query and encode the masked query paired with positive documents, and use a linear layer as a decoder to predict the masked token. In addition, the MQP is used to jointly optimize the models with supervised ranking objective during fine-tuning stage without an extra further pre-training stage. Extensive experiments on the MS MARCO passage ranking and TREC Robust datasets show that models trained with our framework obtain significant improvements compared to original models.
The document ranking task is to generate a ranked list of candidate documents based on their relevance scores to a given query posed in natural language, which has been a longstanding problem that has been widely studied over natural language processing (NLP) and question answering. Pre-trained language models (PLMs) such as BERT (Devlin et al., 2019), RoBERTa (Liu et al., 2019) and ELECTRA (Clark et al., 2020),
have achieved impressive results on various NLP tasks and have outperformed conventional document ranking methods (Hui et al., 2018; Mitra and Craswell, 2019) for powerful contextual representation capability. In recent years, several studies (Karpukhin et al., 2020; Qu et al., 2021) have used pre-trained language models as dual-encoder to separately encode queries and documents for dense document retrieval. One of the most common approaches (Nogueira and Cho, 2019) uses a PLMs as an interaction-based reranker for passage ranking, which fine-tunes BERT simply with an extra linear layer on the top of BERT and using a special vector [CLS] to produce relevance score for each query document pair. Inspired by the fact that contextualized embeddings produced by PLMs are essential for the success of pre-trained models, CEDR (MacAvaney et al., 2019) employed the classification vector into existing neural ranking models and PARADE (Li et al., 2020) used a transformer module for passage-level representation aggregation to obtain performance improvements.
In contrast to the approaches of using contextual representation for reranking, prior works (Sun et al., 2019; Gururangan et al., 2020) suggest that further pre-train the PLMs on within-task training unsupervised data is able to learn domain-specific and task-specific language patterns effectively. To better understand the complex sentence relations, UED (Yan et al., 2021) transformed original next sentence prediction (NSP) task in BERT to a new sentence relation prediction (NSR) task. Gu et al. (2020) proposed a novel selective masking strategy to focus on masking the important tokens and then train a model to reconstruct input for further pretraining the PLMs to learn task-specific patterns. However, these approaches typically perform the pre-training task on task-specific corpus to understand the context of passages, while fail to consider the passages as the context of the given query to capture semantic consistency. In addition, further
* Corresponding author: Hai Liu.
Findings of the Association for Computational Linguistics: NAACL 2022, pages 1056 - 1065 July 10-15, 2022 ©2022 Association for Computational Linguistics
pre-training on task-specific domain datasets entails additional time cost and computational cost.
Therefore, this paper proposes a self-supervised joint training framework for encouraging the model to understand the content of queries based on the context of passages, where the auxiliary self-supervised method is combined with the ranking task to fine-tune pre-trained model. Specifically, the self-supervised joint training framework (SJTF) extends the typical reranking pipeline with the auxiliary self-supervised method MQP and a decoder for predicting the masked token after the pretrained model. On the one hand, the self-supervised approach enables the model to establish semantic relations between queries and positive passages to better identify relevant passages from a large number of candidate passages. On the other hand, the proposed training framework reduces the training time by simultaneously performing the self-supervised method and the ranking task in the finetuning stage, while the time of the ranking model is not increased in the inference stage. We evaluate the proposed training framework on two widely used document ranking datasets MS MARCO and Robust04. The experimental results indicate that the models trained with the proposed SJTF framework obtain a performance improvement against original models.
In summary, the contributions of our paper are as follows:
- A self-supervised joint training framework (SJTF) is proposed to improve the representation learning without additional pre-training.
- A strategy to integrate the SJTF to existing passage reranking methods is proposed without architecture modification and inference time increasing.
- Experiments on standard datasets show that reranking models with SJTF integration achieve significant performance improvement.
In recent years, pretrained language models such as BERT (Devlin et al., 2019), RoBERTa (Liu et al., 2019) and ELECTRA (Clark et al., 2020) had substantially outperformed the traditional neural ranking models like DRMM (Guo et al., 2016), Co-PACRR (Hui et al., 2018) and Conv-KNRM (Dai et al., 2018). Nogueira and Cho (2019) first
employed the pretrained language model BERT to passage reranking tasks using the classification vector [CLS]. CEDR (MacAvaney et al., 2019) incorporated the contextualized embeddings of BERT into existing IR models for document ranking. PARADE (Li et al., 2020) utilized Transformer (Vaswani et al., 2017) blocks to aggregate relevance passage-level representations to predict a document ranking. These approaches provided different perspectives on score prediction using the relevance representation vectors produced by pretrained language models.
To better improve the representation learning of pre-trained language models in the target domain, Sun et al. (2019) further pre-trained BERT with masked language model (MLM) and next sentence prediction (NSP). Gururangan et al. (2020) observed that the less relevant the pre-trained corpus was to the target corpus, the more the pre-trained language model would benefit from further pretraining. In this paper, the proposed joint training framework could be used in conjunction with those reranking models in the fine-tuning phase, which reduced the time spent in further pre-training.
Representation learning has been shown to be critical on natural language tasks and has a significant influence on downstream tasks (Devlin et al., 2019; Peters et al., 2018; Yan et al., 2021). Devlin et al. (2019) adopted the self-supervised MLM task to encourage the model to learn contextual representations by predicting the masked token from the context. From the perspective of optimizing the masking strategy, RoBERTa (Liu et al., 2019) modified the static masking strategy by dynamically masking the input examples during the training stage, while (Gu et al., 2020) proposed a selective masking strategy that masked important words rather than any word in the sentence.
For understanding the content of documents, Cross-Thought (Wang et al., 2020) proposed a method for recovering masked words from documents that contain the most important information in a nearby sequence. However, understanding the content of both the query and the document is crucial in the question answering tasks (Zhang et al., 2019; Mudrakarta et al., 2018; Nogueira et al., 2019b), while these methods focused on representation learning for understanding the content of the document. The aim of our self-supervised task MQP is to predict the masked query word based on the semantic consistency of the query and the
Figure 1: The architecture of joint training with Masked Query Prediction (MQP) task during fine-tuning stage.
positive passages, which forces the model to consider the positive passages as the context of the given query. The MQP differs from the standard self-supervised approaches of using unsupervised data (He et al., 2021; Hendrycks et al., 2019; You et al., 2021) in that the supervised signal is used to select relevant passage as the context of the query.
Multi-task learning (MTL) is an effective training setting for allowing model to obtain shared knowledge from several related supervised tasks in document ranking. Fun et al. (2021) enhanced the common representation learning using a retrieval optimized multi-task framework (ROM) for jointly training the retriever, reader and self-supervised tasks with a single encoder. UED (Yan et al., 2021) jointly trained both the ranking and query generation tasks to exploit the task relationships for enhancing the neural re-ranker. Liu et al. (2019) and Maillard et al. (2021) leveraged supervised data from related tasks to enhance the robustness of the model and generic knowledge representation learning. Although the datasets of related tasks are available, there may be differences in language patterns and data distribution between datasets. Our proposed joint training framework uses relevance labels to construct data on the target domain of ranking tasks without the requirement of external data, which is achievable on any of the question answering datasets.
3 The Approach
This section describes the proposed self-supervised joint training framework (SJTF), which employs a self-supervised method (MQP) to jointly fine-tune reranking models with the ranking task. In
general, the whole model consists of three parts: a pretrained encoder for producing interaction-based representation of the given query and passages, a scorer for calculating a precise relevance score for each query-passage pair and an extra decoder to predict the masked query token based on the interaction-based representation. The overall framework of our approach is shown in Figure 1.
In the task of passage reranking, a natural language question and a list of candidate passages retrieved by traditional methods or dense retriever are provided. The question is denoted as a  $k$ -length sequence of tokens  $Q = < q_{1}, q_{2}, \dots, q_{k} >$ , while each candidate passage can be denoted by an  $m$ -length sequence of tokens  $P = < p_{1}, p_{2}, \dots, p_{m} >$ . The passage reranking task requires the model to learn informative representation and produce a precise relevance score for each query-passage pairs to return the best permutation of candidate passages.
3.1 Radom Masking
In contrast to the MLM task that random masks the tokens of the passage, we assume that understanding the content of the given query is necessary and the semantic information of the positive passage can be used to infer the masked token in the query, which finally allows the model to learn the semantic relation between the query and the positive passage. The tokens to be masked in the query are selected randomly following a uniform distribution and replaced with special token [MASK], and for simplicity, each query token is considered here to have the same importance. Then the masked query and positive passage are concatenated
and preprocessed into a sequence  $I_{mask} = < [CLS], q_1, \dots, [MASK], \dots, q_k, [SEP], p_1, \dots, p_m, [SEP] >$ , where special token [CLS] indicates the start of a sentence, token [MASK] indicates the masked token and token [SEP] is a separator symbol.
The reason for masking the query is that the original query can be reconstructed by understanding the content of the processed query and the positive passage. However, if the goal is to predict the masked tokens in the passage, the semantic information of the query is redundant as it can be inferred from the context of the passages alone. Without requiring extra data augmentation, the framework SJTF only utilizes the relevance label and a simple masking strategy to construct the masked input sequence  $I_{mask}$ , which can be easily implemented in question-answering tasks.
3.2 Masked Query Prediction
The pre-trained language model BERT (Devlin et al., 2019) uses the MLM task to learn contextual language representations of individual texts in a large corpus. However, establishing contextual semantic relations between queries and candidate passages is critical in the question-answering domain. To achieve this, we propose a self-supervised auxiliary approach called Masked Query Prediction (MQP) so that the model uses positive passages as the context of the query and predicts masked token in conjunction with visible query tokens, which allows the model to extract semantic relation between the query and positive passage, thus pick relevant passages out of a large number of candidate passages.
After passing the masked input sequence  $I_{mask}$  constructed in section 3.1 through the BERT-like encoder which is shared with a passage reranking task, the representation vector  $T_{mask} \in \mathbb{R}^d$  of the masked query token in the last layer is obtained as:
$$
T _ {\text {m a s k e d}} = \operatorname {E n c o d e r} \left(q _ {\text {m a s k e d}}, p _ {\text {p o s}}\right) \tag {1}
$$
Finally, the masked token representation vector  $T_{\text{masked}}$  is fed into a decoder implemented by a neural linear layer to predict the original token. The parameters of MQP module are optimized by the cross-entropy loss function  $\mathcal{L}_{MQP}$  which is defined as:
$$
\mathcal {L} _ {M Q P} = - \sum_ {m = 1} ^ {M} \log P \left(t _ {m}\right), \tag {2}
$$
where  $M$  is the number of masked token and  $P(t_{m})$  denotes the probability that the token  $t_{m}$  is predicted over the whole vocabulary. For each query and positive passage pair, a single query token is replaced with the special token [MASK] for the self-supervised method. Note that the MQP decoder is only used for predicting the masked token during the fine-tuning stage, while in the inference phrase only the encoder and scorer are used for reranking passages. The decoder for the MQP task is implemented with one linear layer, meaning that this requires only a small number of additional neural network parameters to be trained, which makes the MQP task easy to be extended to the existing passage reranking methods without significant modifications to the model architecture.
3.3 Passage Ranking
Following the settings of Nogueira and Cho (2019), a pair of query  $q$  and candidate passage  $p_i$  is packed as an input sequence  $I = < [CLS], q_1, \dots, q_k, [SEP], p_1, \dots, p_m, [SEP]>$ , and the BERT-like pretrained language model is employed as a passage encoder  $E$  that produces a relevance representation vector for each QA pair. During fine-tuning the model, the  $[CLS]$  vector  $T_{cls} \in \mathbb{R}^d$  from the last layer of encoder is regarded as the final interaction-based representation:
$$
T _ {c l s} ^ {i} = \operatorname {E n c o d e r} (q, p _ {i}) \tag {3}
$$
For passage reranking task, the representation vector  $T_{cls}$  of query-passage pairs is calculated by a scorer  $S_{ranker}$  which generates a relevance score to quantify their relevance. The relevance score of  $i$ -th pairs of query  $q$  and candidate passage  $P_i$  is denoted as:
$$
S c o r e _ {c l s} ^ {i} = S _ {r a n k e r} \left(T _ {c l s} ^ {i}; \theta\right), \tag {4}
$$
where the scorer  $S_{\text{ranker}}$  can be implemented by a linear layer at top of the BERT or by an elaborate scoring module such as KNRM (Dai et al., 2018), SAN (Kingma and Ba, 2015) or PARADE (Li et al., 2020), and  $\theta$  contains the set of parameters of the scorer module.
Compared with the point-wise ranking loss used in (Nogueira and Cho, 2019), the pair-wise margin ranking loss discriminates the positive and negative examples by relative distance, allowing the model to learn the margins between the positive and negative examples to give an appropriate relevance score. Therefore, the reranking module is
optimized by the ranking loss  $\mathcal{L}_{\text{rank}}$  as Equation (5):
$$
\begin{array}{l} \mathcal {L} _ {r a n k} \\ = \max  (0, y \cdot \left(S c o r e _ {c l s} ^ {k} - S c o r e _ {c l s} ^ {k + 1}\right) + \gamma), \tag {5} \\ \end{array}
$$
where  $y = -1$  if the Score $_{cls}^k$  is higher than Score $_{cls}^{k+1}$ , and vice-versa for  $y = 1$ .  $\gamma$  is a hyperparameter that controls the margin of positive and negative examples.
Figure 2: The overall self-supervised jointly training framework SJTF.
3.4 Joint Training
Different from the self-supervised tasks that are used for further pre-training on the downstream dataset (Fun et al., 2021; Liu et al., 2019), our proposed self-supervised method MQP is combined with the ranking task in the fine-tuning phase, since it utilizes the relevance label information to construct masked input sequences. As shown in Figure 2, the joint training strategy simplifies the training procedure without further pre-training phases and training resources required. In the fine-tuning stage, the loss is defined as a linear combination of passage reranking loss and masked query prediction loss as:
$$
\mathcal {L} = \mathcal {L} _ {\text {r a n k}} + \alpha \cdot \mathcal {L} _ {M Q P} \tag {6}
$$
The hyperparameter  $\alpha$  is assigned by different values to tradeoff between passage reranking and masked query prediction. As the MQP is a secondary task used to help the model understand the query content, where the loss weight  $\alpha$  is usually set to a lower value, resulting in the parameters of the reranking model still being optimized primarily by the passage reranking task.
The proposed method is extended to existing ranking models and evaluated on two widely used datasets: MS MARCO Passage Ranking (Nguyen et al., 2016) and TREC Robust 2004 (Voorhees). The statistics of these two datasets are shown in Table 1.
MS MARCO Passage Ranking dataset is a large-scale dataset consisting of real anonymous questions from the Bing search engine and 8.8 million candidate passages for passage reranking task. The training set contains about 500 thousand positive query-passage pairs and each query has one relevant passage on average. The development set and evaluation set contains 6980 queries and 6837 queries respectively, where the relevance labels are provided for the development set only.
Robust04 dataset is a newswire collection used by TREC 2004 Robust track, which comprises 250 queries and 0.5 million documents (TREC Disks 4 and 5). Following the setting of CEDR (MacAvaney et al., 2019), we use the same five folds cross-validation with three folds for training, one fold for validation and one fold for testing.
Following the previous works, three widely used evaluation metrics are adopted to measure the performance of the proposed approach, including MRR@10 for the MS MARCO Passage Ranking dataset, P@20 and NDCG@20 evaluated by trec_eval<sup>1</sup> for the Robust04 dataset. The result reported for model performance are averaged over all test folds on the Robust04 dataset.
MRR@10(Mean Reciprocal Rank) This metric considers the reciprocal rank of the first relevant passage in ranked list to a given query as the precision. For MS MARCO passage ranking task only provides binary label and does not specify relative ranking order between passages, thus the MRR metric is used for evaluation.
P@20 The top-20 precision is defined as the proportion of relevant documents which are ranked in the top 20 candidate documents.
NDCG@20 NDCG is used to measure the discrepancy between the ranked list and the correct ranking list, which evaluates the ranking performance of models.
1https://trec.nist.gov/trec_eval
Table 1: The statistics of datasets MS MARCO Passage Ranking and TREC Robust04.
Datasets	Queries	Avg. of word length	Passages	Avg. of word length
MS MARCO	516,756	5.97	8,841,823	56.58
Robust04	250	2.65	528,155	577.82
The proposed training framework is integrated with the following methods and compared the reranking results with the original methods:
BM25+BERT (Nogueira and Cho, 2019) utilizes the traditional unsupervised ranking method BM25 as a first-stage retriever to generate a ranked list of candidate passages, and the relevance scores between query and candidate passages are produced by the BERT-base model with a linear layer.
BM25+BERT+FP (Gururangan et al., 2020) further pretrains the language model BERT with MLM objective on the target datasets before finetuning to learn the specific domain language patterns for document reranking task.
BERT+PACRR (MacAvaney et al., 2019) constructs query-document term similarity matrix in PACRR (Hui et al., 2018) using the interaction-based representation vectors generated by BERT.
PARADE (Li et al., 2020) aggregates passage-level relevance representations to predict a document relevance score, where the long document is split into several passages and each passage is encoded with a given query.
CEDR-KNRM (MacAvaney et al., 2019) incorporates the classification vector of a fine-tuned BERT into existing neural models KNRM (Dai et al., 2018) and leverage contextual information to improve ad-hoc document ranking.
4.4 Implementation Detail
For both datasets, following the setting of the first retrieval stage (Nogueira et al., 2019a), we employ BM25 (Robertson et al., 2009) as the retriever at the first stage to obtain a list of top-k candidate documents/passages for next reranking stage. The reranking models are optimized by Adam (Kingma and Ba, 2015) with a learning rate of 3e-6 and a batch size of 8. The maximum sequence length is limited to 128 tokens for MS MARCO passage ranking dataset, while the documents are truncated to 800 tokens for Robust04 dataset. The encoder of CEDR and PARADE are the base-size pre-trained language model which consists of 12 transformer blocks, the hidden size as 768, and 12 self-attention
heads. For Robust04 dataset, the dropout function is used with the rate of 0.1 for improving model robustness. We set the value of margin hyperparameter  $\gamma$  of pairwise ranking loss to 1. A uniform distribution is used to random mask a token in a given query, where all query tokens are assumed to be of equal importance. Since the MQP loss serves as an auxiliary loss to enhance the interactive-based representation learning for encoder, the weight hyperparameter  $\alpha$  is assigned to 0.2 for MS MARCO and 0.05 for Robust04 respectively. The experiments are conducted on a single GeForce RTX 3090 GPU.
The experimental results of model+SJTF on the MS MARCO passage ranking and Robust04 datasets are presented in Table 2, where model+SJTF refers to the joint training of different reranking models applying to SJTF during the fine-tuning stage. From Table 2, the reranking models adopting SJTF have obtained significant improvements using paired t-test  $(p < 0.05)$  compared to original models.