Wikiformer: Pre-training with Structured Information of Wikipedia for Ad-hoc Retrieval

Pre-trained Language Models (PLMs) have achieved great success in recent years(Devlin et al. 2018; Vaswani et al. 2017; Yang et al. 2019; Liu et al. 2019; Yasunaga, Leskovec, and Liang 2022). These models are firstly trained on large-scale unlabeled text corpora and then fine-tuned on certain downstream tasks with labeled data. Benefiting from the self-supervised learning on a large-scale pre-training corpus, these models own a powerful ability on contextual text representation.

Among these PLMs, Transformer based models(Vaswani et al. 2017) show great performance in most downstream NLP tasks. One of the remarkable examples is the BERT model(Devlin et al. 2018), a bi-directional Transformer based pre-trained language model. BERT has two self-supervised tasks in the pre-training stage: Masked Language Modeling (MLM) and Next Sentence Prediction (NSP). Following BERT, researchers redesign and optimize the pre-training tasks of PLMs. For example, Roberta(Liu et al. 2019) uses a dynamic masking strategy and is trained on a larger text corpus. In addition, some researchers explore the integration of structured information into PLMs(Yasunaga, Leskovec, and Liang 2022; Colon-Hernandez et al. 2021; Kaur et al. 2022; Zhang et al. 2019). For example, LinkBERT(Yasunaga, Leskovec, and Liang 2022) replaces the NSP task of BERT with the Document Relation Prediction (DRP) task, which enables the model to learn cross-document knowledge from hyperlinks among web pages. ERNIE(Zhang et al. 2019) utilizes both textual corpora and Knowledge Graphs to train an enhanced PLM.

Considering the great success that PTMs have achieved in NLP tasks, the IR community begins to apply PLMs to solve IR tasks (Li et al. 2023d, c; Chen et al. 2023; Li et al. 2023b; Ye et al. 2023), and devise pre-training methods tailored for IR(Ma et al. 2023, 2021a, 2021d, 2021b; Chang et al. 2020; Fan et al. 2021; Guo et al. 2022; Chen et al. 2022; Su et al. 2023a, b; Li et al. 2023a). For example, HARP(Ma et al. 2021d) utilizes hyperlinks and anchor texts in the pre-training stage. As the anchor texts are edited by humans, constructing pseudo query-document pairs from them may be more reliable than an algorithm. Webformer(Guo et al. 2022) is a pre-trained language model based on large-scale web pages, HTML tags, and the DOM (Document Object Model) tree structures of web pages. Ma et al.(Ma et al. 2021a) devised a self-supervised learning task Representative Words Prediction (ROP) based on the Query Likelihood model and train the Transformer encoder with a self-supervised contrastive learning strategy. From another angle, ARES(Chen et al. 2022) propose several pre-training objectives based on Axiomatic Regularization. Experimental results on several IR benchmarks show that ARES, PROP, Webformer, and HARP perform significantly better than traditional methods such as BM25 after fine-tuning. Also, some researchers explore incorporating structure information for entity retrieval(Gerritse, Hasibi, and de Vries 2020; Nikolaev and Kotov 2020; Chatterjee and Dietz 2022; Gerritse, Hasibi, and de Vries 2022).

Different from the above approaches, we propose four new pre-training objectives using the titles, abstracts, hierarchical heading (multi-level title) structure, relationship between articles, references, hyperlink structures, and the writing organizations of Wikipedia to leverage the wisdom of crowds brought by Wikipedia editors. Compared to previous work, Wikiformer captures more internal relationships between the paragraph structure in Wikipedia web pages, which helps it better model relevance matching.

The SRR task is inspired by an important IR problem: document re-ranking. In general, the goal of the document re-ranking task is to sort a series of documents that are highly related to the query, and then select the ones that are most related to the query. According to the characteristics of this task, we aim to design a self-supervised learning task to select the most relevant document from a series of documents with similar contents. In the SRR task, we make full use of the hierarchical heading (multi-level title) structure of Wikipedia to achieve the above objective. Every article on Wikipedia is organized by the hierarchical heading (multi-level title) structure, the subtitle corresponding to a certain section tends to be the representative words or summarization of the text. Besides, different subsections of the same section share similar semantics. As a result, through this structure, we can obtain a series of texts that are highly similar but slightly different in content and generate the query through the multi-level titles as shown in Figure 2.

To be specific, we modeled each Wikipedia article into a tree structure namely Wiki Structure Tree (WST) based on the hierarchical heading structure. It can be defined as:

$WST=<D,R>$, where $D$ is a finite set containing $n$ nodes, and $R$ is the root node of $WST$. Each node in $D$ consists of two parts: the subtitle and its corresponding content. The root node $R$ contains the main title and the abstract of this article. Starting from the root node $R$, recursively take all the corresponding lower-level sections as its child nodes until every section in this article is added to the $WST$.

After building $WST$, we use a contrastive sampling strategy to construct pseudo query-document pairs based on the tree. For a non-leaf node $F$ in the $WST$, we add all its child nodes to the set $S$. A node $d_{i}$ is randomly selected from $S$. Traversing from the root node to node $d_{i}$, all the titles on the path are put together to form a query $q$. This process is shown in Figure 2. The content of the node $d_{i}$ is defined as $d^{+}$, and the content of the other nodes in $S$ is defined as $d^{-}$. We use a Transformer based PLM to compute the relevance score of a pseudo query-document pair:

RWI task is inspired by an IR axiom which assumes that the user’s query is the representative words extracted from the relevant documents. According to the Wikipedia structure, we regard the subtitle of each section as representative words, and then we sample pseudo query-document pair via a simple strategy based on the hierarchical heading (multi-level title) structure, as shown in Figure 3.

Specifically, pseudo query-document pairs are organized as follows: for each Wikipedia article, we first model it as the $WST$ structure. Then we add all nodes of $WST$ except the root node to the set $S$. A node $d_{i}$ is randomly selected from $S$, and we define the depth of this node in $WST$ as $n$. Traversing from the root node to node $d_{i}$, all the titles on the path are put together to form a query $q^{+}$. The content of the node $d_{i}$ is defined as $D$. For the negative queries, we randomly select $n-1$ nodes from $S$, and concatenate the main title and subtitles of the selected nodes to define it as $q^{-}$. The relevance score is defined in Equation2. The loss function of the RWI task is defined as:

In this task, although both positive and negative queries contain the subtitles of the document, the positive query is more representative compared to the negative query. The model gives higher scores to the positive query through contrastive learning, so that the model can recognize the representative words in the text, and assign higher weights to these words if they are matched to the query. Therefore, through the RWI task, the model can learn how to identify the keywords in the text, which further leads to a better performance in the IR downstream task.

In the ATI task, we utilize the abstract and the inner structure of Wikipedia. The abstract (the first section) of Wikipedia is regarded as the summarization of the whole article. Compared with other sections of the same article, the abstract is more likely to meet the user’s information needs when the query is the title. Therefore, we extract the title from the Wikipedia article as the query (denoted as $q$). Then the abstract of the same article is regarded as a positive document (denoted as $d^{+}$). For the negative ones, we use the other sections of the same article (denoted as $d^{-}$). The relevance score of a pseudo query-document pair is defined in Equation 2. The loss function of the ATI task is defined as:

After pre-training with RWI, ATI, and SRR tasks, Wikiformer acquires the ability to measure the relevance between a short text (query) and a long text. This can help the model better handle the vast majority of ad-hoc retrieval tasks. However, there are also scenarios involving ”long queries”, such as legal case retrieval and document-to-document search. In these scenarios, the model is required to match the relevance between two long texts. Fortunately, with the structured information of Wikipedia, especially hyperlinks, we can build a series of informative pseudo long query-document pairs. To be specific, we utilize the See Also section of Wikipedia which consists of hyperlinks that link to the other articles related to or comparable to this article. The See Also section is mainly written manually, based on the judgment and common sense of the authors and editors. Thus, we can obtain a series of reliable web pages that are highly related to the content of this page.

To this end, we design the Long Texts Matching (LTM) task to encourage the Wikiformer to learn the relevance matching ability between two long documents. Initially, we transformed the complete Wikipedia corpus into a graph structure by leveraging the interconnections provided by the ’See Also’ links. This graph is designated as the See Also Graph (SAG). Each hyperlink in the See Also section can be formally represented as $(v_{i},v_{j})$, which means that $v_{j}$ appears in the See Also section of $v_{i}$. Consequently, $SAG$ can be defined as a directed graph: $SAG=(V,E)$, where $E$ is the above-mentioned set of ordered pairs $(v_{i},v_{j})$ and $V$ is a set of Wikipedia articles. The order of an edge indicates the direction of hyperlinks. After building $SAG$, we use a contrastive sampling strategy based on the graph. For each node in $SAG$, we define its content as query $D$ and define all its adjacent nodes as positive documents $d^{+}$. We randomly select other documents as $d^{-}$. The relevance score of a pseudo query-document pair is defined in Equation 2. The loss function of the LTM task is defined as:

We add the loss of the proposed four tasks together as the overall loss of the model:

For the pre-training dataset, we use the English Wikipedia (version 20220101). For the downstream datasets, we evaluate the performance of Wikiformer on five IR benchmarks. The basic statistics are shown in Table 1. MS MARCO Document Re-ranking (Nguyen et al. 2016) is a large-scale ad-hoc retrieval dataset with 0.37M queries and 3.2M documents. TREC DL 2019 (Craswell et al. 2020) shares the same document collection with MS MARCO but collects finer-grained human labels for 43 queries in the test set. TREC Covid Round2 (Roberts et al. 2021) is an ad-hoc retrieval dataset consisting of biomedical articles. It contains the May 1, 2020 version of the CORD-19 (Wang et al. 2020) document set and 35 queries written by biomedical professionals. LeCaRD (Ma et al. 2021c) is a legal case retrieval dataset, consisting of 107 query cases and 10700 candidate cases. The queries in the LeCaRD dataset are the factual description part of a legal case, while the candidate documents are complete legal cases. CAIL-LCR(Ma 2022) is a case retrieval dataset (document-to-document search) provided by CAIL 2022 consisting of 130 query cases and 100 candidate cases for each query case.

We consider three types of IR baselines for comparison, including traditional IR methods, Neural IR models, and pre-trained language models:

For the implementation of KNRM and Conv-KNRM, we use the OpenMatch²²2https://github.com/thunlp/OpenMatch toolkit, and the 300d GloVe (Pennington, Socher, and Manning 2014) vectors are used to initialize the word embeddings. For the implementation of BM25 and QL, we use the pyserini toolkit³³3https://github.com/castorini/pyserini. For the hyperparameter of BM25, we set $k1=3.8$ and $b=0.87$⁴⁴4This is the best hyperparameter we got after parameter searching.. Note that in our experiments, we use the scores of the BM25 and QL models to re-rank the candidate documents, rather than re-ranking the whole corpus. For the implementation of BERT, we use the Pytorch version BERT-base released by Google⁵⁵5https://github.com/google-research/bert. For the implementation of ARES, PROP_MS, and PROP_WIKI, we directly use the checkpoints released by the original paper. Since the original paper of Webformer and HARP did not release any checkpoints, we reproduce them on the same dataset based on their code and the details provided in their paper.

To facilitate comparison with previous baselines, we adopted the same architecture as BERT-base. This aligns with the settings of previous works such as ARES, HARP, PROP, B-PROP, and Webformer. To save computational resources during training, we initialized our model with BERT-base, following the same setting as previous works such as ARES and HARP. We use the AdamW optimizer with a learning rate of 1e-5 in the first 50k steps and 5e-6 in the following steps. We set the warm-up ratio to 0.1. In the RWI, ATI, and SRR tasks, we set the maximum length of the query as 30 and the maximum length of the documents as 480. In the LTM task, we set the maximum length of both documents as 255. We trained our model on four Nvidia GeForce RTX 3090 GPUs for 60 hours. After training for 50k steps, we save the checkpoint every 5k steps and evaluate the zero-shot performance of each checkpoint on a subset of the MS MARCO training set which has no overlap with our test set. We select the best zero-shot performance checkpoint as the final model.

For the two large-scale datasets MS MARCO and TREC DL 2019, we use Mean Reciprocal Rank at 10 and 100 (MRR@10 and MRR@100) for MS MARCO and normalized discounted cumulative gain at 10 and 100 (nDCG@10 and nDCG@100) for TREC DL 2019 as the evaluation metrics. For TREC Covid, we follow the setting of OpenMatch which re-ranks the top 60 candidates provided by the BM25-fusion method. We use precision at rank 5 (P@50) and nDCG@10 as the evaluation metrics for TREC Covid. For LeCaRD and CAIL-LCR datasets, we re-rank the candidate documents provided by the original dataset and use nDCG@5 and nDCG@15 as the evaluation metrics.

For the significance test, we adopt Fisher’s randomization test (Fisher 1936; Cohen 1995; Box et al. 1978) which is recommended for IR evaluation by previous work (Smucker, Allan, and Carterette 2007).

To investigate whether a larger training dataset enhances the performance of the pre-training phase, we evaluate the performance of Wikiformer on different sizes of training data varying from 100 to 1,000,000 pseudo query-document pairs. As shown in Figure 4, Wikiformer surpasses the Query Likelihood model by pre-training with only 100 pseudo query-document pairs in the SRR task. This experimental result shows the effectiveness of our pre-training task and our proposed pseudo query-document pair sampling strategy.

To further analyze the effectiveness of each pre-training task, we conduct ablation experiments on MS MARCO (zero-shot) and LeCaRD (fine-tuned). The experimental results in table 5 show that removing any pre-training tasks will lead to a drop in performance, indicating the effectiveness of each pre-training task on downstream IR tasks. On MS MARCO, among the four tasks, removing the SRR task leads to the largest performance degradation, which reveals that the hierarchical heading structure and the writing organization of Wikipedia contain valuable knowledge for ad-hoc retrieval which helps Wikiformer better at handling relevance matching. On LeCaRD, removing the LTM task leads to the largest performance degradation, which reveals that the LTM task is critical for improving the model’s ability on long text-matching tasks.