L ink BERT : Pretraining Language Models with Document Links (2022)

Abstract: Language model (LM) pretraining captures various knowledge from text corpora, helping downstream tasks. However, existing methods such as BERT model a single document, and do not capture dependencies or knowledge that span across documents. In this work, we propose LinkBERT, an LM pretraining method that leverages links between documents, e.g., hyperlinks. Given a text corpus, we view it as a graph of documents and create LM inputs by placing linked documents in the same context. We then pretrain the LM with two joint self-supervised objectives: masked language modeling and our new proposal, document relation prediction. We show that LinkBERT outperforms BERT on various downstream tasks across two domains: the general domain (pretrained on Wikipedia with hyperlinks) and biomedical domain (pretrained on PubMed with citation links). LinkBERT is especially effective for multi-hop reasoning and few-shot QA (+5% absolute improvement on HotpotQA and TriviaQA), and our biomedical LinkBERT sets new states of the art on various BioNLP tasks (+7% on BioASQ and USMLE). We release our pretrained models, LinkBERT and BioLinkBERT, as well as code and data.

## introduction pretrained language models (lms), like bert and gpts @xcite , have shown remarkable performance on many natural language processing (nlp) tasks, such as text classification and question answering @xcite , becoming the foundation of modern nlp systems. by performing self-supervised learning on text, such as masked language modeling @xcite , lms learn to encode various knowledge from text corpora and produce informative language representations for downstream tasks @xcite . ## related work retrieval-augmented lms. several works @xcite @xcite introduce a retrieval module for lms, where given an anchor text (e.g. question), retrieved text is added to the same lm context to improve model inference (e.g. answer prediction). these works show the promise of placing related documents in the same lm context at inference time, but they do not study pretraining. @xcite pretrain an lm with a retriever that learns to retrieve text for answering masked tokens in the anchor text. in contrast, our focus is not on retrieval, but on pretraining a general-purpose lm that internalizes knowledge that spans across documents, which is orthogonal to the above works (e.g., our pretrained lm could be used to initialize the lm component of these works). additionally, we focus on incorporating document links such as hyperlinks, which can offer salient knowledge that common lexical retrieval methods may not provide @xcite . pretrain lms with related documents. several concurrent works use multiple related documents to pretrain lms. @xcite place documents (news articles) about the same topic into the same lm context, and @xcite place sentences of high lexical similarity into the same context. our work provides a general method to incorporate document links into lm pretraining, where lexical or topical similarity can be one instance of document links, besides hyperlinks. we focus on hyperlinks in this work, because we find they can bring in salient knowledge that may not be obvious via lexical similarity, and yield a more performant lm ( §5.5). additionally, we propose the drp objective, which improves modeling multiple documents and relations between them in lms ( §5.5). hyperlinks and citation links for nlp. hyperlinks are often used to learn better retrieval models. chang et al. (2020); asai et al. (2020); seonwoo et al. (2021) use wikipedia hyperlinks to train retrievers for open-domain question answering. @xcite study various hyperlink-aware pretraining tasks for retrieval. while these works use hyperlinks to learn retrievers, we focus on using hyperlinks to create better context for learning general-purpose lms. separately, @xcite use wikipedia hyperlinks to learn multilingual lms. citation links are often used to improve summarization and recommendation of academic papers @xcite @xcite @xcite . here we leverage citation networks to improve pretraining general-purpose lms. graph-augmented lms. several works augment lms with graphs, typically, knowledge graphs (kgs) where the nodes capture entities and edges their relations. @xcite al. (2019); he et al. (2020 ); wang et al. (2021b) combine lm training with kg embeddings. sun et al. (2020); yasunaga et al. (2021); zhang et al. (2022) combine lms and graph neural networks (gnns) to jointly train on text and kgs. different from kgs, we use document graphs to learn knowledge that spans across documents. ## preliminaries a language model (lm) can be pretrained from a corpus of documents, @xmath0. an lm is a composition of two functions, f head (f enc (x)), where the encoder f enc takes in a sequence of tokens @xmath1,x 2 ,...,x n ) and produces a contextualized vector representation for each token, (h 1 ,h 2 ,...,h n ). the head f head typically uses these representations to perform self-supervised tasks in the pretraining step, and perform particular downstream tasks in the fine-tuning step. we build on bert @xcite , which pretrains an lm with the following two self-supervised tasks. masked language modeling (mlm). given a sequence of tokens x, a subset of tokens y ⊆ x is masked, and the task is to predict the original tokens from the modified input. y accounts for 15% of the tokens in x; of those, 80% are replaced with [mask], 10% with a random token, and 10% are kept unchanged. next sentence prediction (nsp). the nsp task takes two text segments foot_0 (x a ,x b ) as input, and predicts whether x b is the direct continuation of x a . specifically, bert first samples x a from the corpus, and then either (1) takes the next segment x b from the same document, or (2) samples x b from a random document in the corpus. the two segments are joined via special tokens to form an input instance, [cls] x a [sep] x b [sep], where the prediction target of [cls] is whether x b indeed follows x a (contiguous or random). in this work, we will further incorporate document link information into lm pretraining. our approach ( §4) will build on mlm and nsp. ## linkbert we present linkbert, a self-supervised pretraining approach that aims to internalize more knowledge into lms using document link information. specifically, as shown in figure 2 , instead of viewing the pretraining corpus as a set of documents @xmath2, we view it as a graph of documents, @xmath3, e), where @xmath4, x (j) )} denotes links between documents ( §4.1). the links can be existing hyperlinks, or could be built by other methods that capture document relevance. we then consider pretraining tasks for learning from document links ( §4.2): we create lm inputs by placing linked documents in the same context window, besides the existing options of a single document or random documents. we use the mlm task to learn concepts brought together in the context by document links, and we also introduce the document relation prediction (drp) task to learn relations between documents. finally, we discuss strategies for obtaining informative pairs of linked documents to feed into lm pretraining ( §4.3). ## experiments we experiment with our proposed approach in the general domain first, where we pretrain linkbert on wikipedia articles with hyperlinks ( §5.1) and evaluate on a suite of downstream tasks ( §5.2). we compare with bert @xcite as our baseline. we experiment in the biomedical domain in §6. ## conclusion we presented linkbert, a new language model (lm) pretraining method that incorporates document link knowledge such as hyperlinks. in both the general domain (pretrained on wikipedia with hyperlinks) and biomedical domain (pretrained on pubmed with citation links), linkbert outperforms previous bert models across a wide range of downstream tasks. the gains are notably large for multi-hop reasoning, multi-document understanding and few-shot question answering, suggesting that linkbert effectively internalizes salient knowledge through document links. our results suggest that linkbert can be a strong pretrained lm to be applied to various knowledge-intensive tasks.