LinkBERT: Language Model Pretraining with Document Link Knowledge
Abstract
Language model (LM) pretraining can learn 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, citation links.
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 diverse downstream tasks across both general domain (pretrained on Wikipedia with hyperlinks) and biomedical domain (pretrained on PubMed with citation links).
In particular, LinkBERT is effective for knowledge- and reasoning-intensive tasks such as multi-hop reasoning and few-shot inference (+7% absolute gain on BioASQ and MedQA), and achieves new state-of-the-art results on various biomedicalNLPtasksincludingrelationextractionandliteratureclassification.
Ourresultssuggestthe promise of LinkBERT for scientific applications.
1. Introduction
Pretrainedlanguagemodels(LMs), likeBERTandGPTs(Devlin et al., 2019; Brown et al., 2020), have shown remarkable performance on many natural language processing (NLP) tasks, such as text classification and question answering, becoming the foundation of modern NLP systems (Bommasani et al., 2021).
By performing self-supervised learning, such as masked language modeling (Devlin et al., 2019), LMs learn to encode various knowledge from text corpora and produce informative representations for downstream tasks (Petroni et al., 2019; Bosselut et al., 2019; Raffel et al., 2020).
However, existing LM pretraining methods typically consider text from a single document in each input context [Table] [Table]
[Table]
Language Model [CLS]
The
Tidal Basin [SEP] [MASK] [MASK] trees ...
[SEP]
Segment A
Segment B
Japanese cherry Masked language modeling (MLM) Document relation prediction (DRP) ●
Linked ● Random ● Contiguous
Corpus with Document Links
[Table]
Doc
Doc
Doc Doc
Create
LM inputs
Train
LMs
[Table]
Doc 1 . Doc 1 .
Contiguous Random
Linked or or [
[Table]
[Table]
Cherry Blossom
Festival held
[Table]
Martin
Luther King Jr. Memorial
[Table]
Franklin Delano Roosevelt
Memorial,
[Table]
George Mason Memorial are situated adjacent
[Table]
[The National Cherry Blossom Festival] ...
It is a spring celebration commemorating the March 27, 1912, gift of Japanese cherry trees from Mayor of Tokyo City Yukio Ozaki to the city of Washington,
D.C. Mayor Ozaki gifted the trees to enhance the growing friendship between the United States and Japan.
Of the initial gift of 12 varieties of 3,020 trees, the Yoshino Cherry [Table] of [Table]
and Kwanzan Cherry
[Table]
) now dominate. ... Document Linked document
[Table]
reservoir located between the Potomac River and the
Washington Channel in Washington, D.C.
It is part of West Potomac Park, is near the
National Mall and is a focal point of the National Cherry Blossom
Festival held each spring.
The Jeﬀerson Memorial, the Martin
Luther King Jr. Memorial, the Franklin Delano Roosevelt
Memorial, and the George Mason Memorial are situated adjacent to the Tidal Basin.
[The National Cherry Blossom Festival] ...
It is a spring celebration commemorating the March 27, 1912, gift of Japanese cherry trees from Mayor of Tokyo City Yukio Ozaki to the city of Washington,
D.C. Mayor Ozaki gifted the trees to enhance the growing friendship between the United States and Japan.
Of the initial gift of 12 varieties of 3,020 trees, the Yoshino Cherry (70% of total) and Kwanzan Cherry (13% of total) now dominate. ...
Document
Linked document (
[Table]
hyperlink,
[Table]
(Liu et al., 2019; Joshi et al., 2020) and do not model links between documents. This can pose limitations because documents often have rich dependencies (e.g. hyperlinks, citation links), and knowledge can span across documents. As an example, in Figure 1, the Wikipedia article “Tidal
Basin, Washington D.C.” (left) describes that the basin hosts “National Cherry Blossom Festival”, and the hyperlinked article (right) reveals the background that the festival celebrates “Japanese cherry trees”.
Taken together, the hyperlink offers new, multi-hop knowledge “Tidal Basin has Japanese cherry trees”, which is not available in the single article “Tidal Basin” alone.
Acquiring such multi-hop knowledge in pretraining could be useful for various applications including question answering.
In fact, document links like hyperlinks and citation links are ubiquitous (e.g. web, books, scientific literature), and guide how we humans acquire knowledge and make discoveries too (Margolis et al., 1999).
In this work, we propose LinkBERT , an effective language model pretraining method that incorporates document link knowledge.
Given a text corpus, we obtain links between documents such as hyperlinks and citation links, and create LM
[Table]
(cid:17)(cid:45)(cid:32)(cid:47)(cid:45)(cid:28)(cid:36)(cid:41)(cid:1)(cid:47)(cid:35)(cid:32)(cid:1)(cid:13)(cid:14) (cid:35)(cid:52)(cid:43)(cid:32)(cid:45)(cid:39)(cid:36)(cid:41)(cid:38)(cid:1058)(cid:1)(cid:45)(cid:32)(cid:33)(cid:32)(cid:45)(cid:32)(cid:41)(cid:30)(cid:32)(cid:1058)(cid:1)(cid:32)(cid:47)(cid:30)(cid:1057) inputs by placing linked documents in the same context, besides the existing option of placing a single document or random documents as in BERT.
Specifically, as in Figure 2, after sampling an anchor text segment, we place either (1) the contiguous segment from the same document, (2) a random document, or (3) a document linked from anchor segment, as the next segment in the input.
We then train the LM with two joint objectives: We use masked language modeling (MLM) to encourage learning multi-hop knowledge of concepts brought into the same context by document links (e.g. “Tidal Basin” and “Japanese cherry” in Figure 1).
Simultaneously, we propose a Document Relation Prediction (DRP) objective, which classifies the relation of the second segment to the first segment ( contiguous , random , or linked ).
DRP encourages learning the relevance and bridging concepts (e.g. “National
Cherry Blossom Festival”) between documents.
We experiment and train LinkBERT in two domains: a general domain, using Wikipedia articles with hyperlinks (§4), and a biomedical domain, using PubMed articles with citation links (§6).
We then evaluate the pretrained models on a wide range of downstream tasks such as question answering, in both domains.
LinkBERT consistently improves on baseline LMs across domains and tasks.
For the general domain, LinkBERT outperforms BERT on MRQA benchmark (+4% absolute in F1-score) as well as GLUE benchmark.
For the biomedical domain, LinkBERT exceeds
PubmedBERT (Gu et al., 2020) and sets new states of the art on biomedical reasoning tasks such as MedQA-USMLE (+7% absolute in accuracy), and on the BLURB benchmark, a diverse suite of BioNLP tasks including biomedical relation extraction, literature classification and question answering (+3% absolute in the average BLURB score).
Overall, our finding is that by pretraining with document links, LinkBERT better captures knowledge about document and concept relations, and achieves notable improvements on multi-document and multi-hop reasoning as well as few-shot inference.
As scientific discovery involves knowledge- and reasoning-intensive processes (e.g., extracting insights from multiple prior publications, designing new experiments in a zero- or few-shot manner), our results suggest the promise of LinkBERT to be applied to various key tasks in science.
2. Related work
Retrieval-augmented LMs.
Several works (Lewis et al., 2020b; Karpukhin et al., 2020; Oguz et al., 2020; Xie et al., 2022) 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 the effect of doing so in pretraining.
Guu et al.
(2020) 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 (Asai et al., 2020).
Pretrain LMs with related documents.
Several concurrent works use multiple related documents to pretrain LMs.
Caciularu et al.
(2021) place documents (news articles) about the same topic into the same LM context, and Levine et al.
(2021) 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 (§D.1).
Additionally, we propose the DRP objective, which improves modeling multiple documents and relations between them in LMs (§D.1).
Hyperlinks and citation links for NLP.
Hyperlinks are often used to learn better retrieval models.
Chang et al.
Asai et al.
(2020); Seonwoo et al.
(2021) use Wikipedia hyperlinks to train retrievers for open-domain question answering.
Ma et al.
(2021) 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, Calixto et al.
(2021) use Wikipedia hyperlinks to learn multilingual
LMs.
Citation links are often used to improve summarization and recommendation of academic papers (Qazvinian &
Radev, 2008; Yasunaga et al., 2019; Bhagavatula et al., 2018;
Khadka et al., 2020; Cohan et al., 2020).
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.
Zhang et al.
; He et al.
; Wang et al.
(2021b)
combine LM training with KG embeddings.
Sun et al.
; Yasunaga et al.
; 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.
3. Preliminaries
A language model (LM) can be pretrained from a corpus of documents, X =
{ X ( i ) } .
An LM is a composition of two functions, f head ( f enc ( X )) , where the encoder f enc takes in a sequence of tokens X
=
( x 1 ,x 2 ,...,x n ) and produces a contextualized vector representation for each token, ( h 1 , h 2 ,..., h n ) .
The head f head uses these representations to perform self-supervised tasks in the pretraining step and to perform downstream tasks in the fine-tuning step.
We build on BERT (Devlin et al., 2019), which pretrains an LM with the masked language modeling task (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.
4. 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 X
=
X ( i ) } , we view it as a graph of documents,
G =( X , E ) , where E = { ( X ( i ) ,X ( j ) ) } denotes links between documents (§4.1).
The links can be existing hyperlinks and citation links, 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.
4.1. Document graph
Given a pretraining corpus, we link related documents so that the links can bring together knowledge that is not available in single documents.
We focus on hyperlinks of Wikipedia articles (§5) and citation links of scientific articles (§6).
Hyperlinks and citation links have a number of advantages.
They provide background knowledge about concepts that the document writers deemed useful—the links are likely to have high precision of relevance, and can also bring in relevant documents that may not be obvious via lexical similarity alone (e.g., in Figure 1, while the hyperlinked article mentions “Japanese” and “Yoshino” cherry trees, these words do not appear in the anchor article).
Hyperlinks and citation links are also ubiquitous on the web and easily gathered at scale (Aghajanyan et al., 2021).
To construct the document graph, we simply make a directed edge ( X ( i ) ,X ( j ) ) if there is a hyperlink from document X ( i ) to document X ( j ) .
4.2. Pretraining tasks
Creating input instances.
Several works (Gao et al., 2021; Levine et al., 2021) find that LMs can learn stronger dependencies between words that were shown together in the same context during training, than words that were not.
To effectively learn knowledge that spans across documents, we create LM inputs by placing linked documents in the same context window, besides the existing option of a single document or random documents.
Specifically, we first sample an anchor text segment from the corpus (Segment
A; X A ⊆ X ( i ) ).
For the next segment (Segment B; X B ), we either (1) use the contiguous segment from the same document ( X B ⊆
X ( i ) ), (2) sample a segment from a random document ( X B
⊆
X ( j ) where j ̸ = i ), or (3) sample a segment from one of the documents linked from Segment
A ( X B ⊆
X ( j ) where ( X ( i ) ,X ( j ) ) ∈ E ).
We then join the two segments via special tokens to form an input instance:
[CLS] X A [SEP] X B
[SEP] .
Training objectives.
To train the LM, we use two objectives.
The first is the MLM objective to encourage the LM to learn multi-hop knowledge of concepts brought into the same context by document links.
The second objective, which we propose, is Document Relation Prediction (DRP), which classifies the relation r of segment X B to segment
X A ( r ∈ { contiguous , random , linked } ).
By distinguishing linked from contiguous and random , DRP encourages the
LM to learn the relevance and existence of bridging concepts between documents, besides the capability learned in the vanilla NSP objective.
To predict r , we use the representation of [CLS] token, as in NSP.
Taken together, we optimize:
[Equation]
where x i is each token of the input instance, [CLS] X A [SEP] X B
[SEP] , and h i is its representation.
5. Experiments: General domain
We experiment in a general domain first, where we pretrain
LinkBERT on Wikipedia articles with hyperlinks and evaluate on diverse downstream tasks.
We compare with
BERT (Devlin et al., 2019) as our baseline.
We then experiment in a biomedical domain in §6.
5.1. Setup
Pretraining data.
We use the same pretraining corpus used by BERT: Wikipedia and BookCorpus (Zhu et al., 2015).
For
Wikipedia, we use the WikiExtractor 1 to extract hyperlinks between Wiki articles.
We then create training instances by sampling contiguous , random , or linked segments as described in §4, with the three options appearing uniformly (33%, 33%, 33%).
For BookCorpus, we create training instance by sampling contiguous or random segments (50%, 50%) as in BERT.
We then combine the training instances from Wikipedia and BookCorpus to train LinkBERT.
In summary, our pretraining data is the same as BERT, except that we have hyperlinks between Wikipedia articles.
Implementation.
We pretrain LinkBERT of two sizes, -base and -large, following the configurations of BERT base (110M params) and BERT large (340M params).
Pretraining hyperparameters can be found in Appendix B.
Baselines.
We compare LinkBERT with BERT, for which we pretrain with the vanilla BERT objectives on the same corpus for the same number of steps as LinkBERT.
Downstream evaluation tasks.
We evaluate on MRQA (Fisch et al., 2019), an extractive question answering benchmark containing 6 tasks:
HotpotQA , TriviaQA ,
NaturalQ , SearchQA , NewsQA , SQuAD .
We also evaluate on GLUE (Wang et al., 2018), a suite of 8 sentence-level classification tasks testing general language understanding.
Appendix C provides more details.
5.2. Results
Table 1 shows the performance (F1 score) on MRQA datasets.
LinkBERT substantially outperforms BERT on all datasets.
On average, the gain is +2.6% for the BERT base scale and +2.5% for the BERT large scale.
Table 2 shows the results on GLUE, where LinkBERT performs moderately better than BERT.
These results suggest that LinkBERT is especially effective at learning knowledge useful for QA tasks (e.g. world knowledge), while keeping performance on sentence-level language understanding.
5.3. Analysis
We study when LinkBERT is especially useful.
Improved multi-hop reasoning.
In Table 1, we find that LinkBERT obtains notably large gains on QA datasets that require reasoning with multiple documents, such as
HotpotQA (+3% over BERT large ), TriviaQA (+4%) and
SearchQA (+2%), as opposed to SQuAD (+1%) which just has a single document per question.
Our intuition is that because LinkBERT is pretrained with pairs of linked documents rather than purely single documents, it better learns how to flow information (e.g., do attention) across tokens when multiple related documents are given in the context.
Improved understanding of document relations.
While the MRQA datasets typically use ground-truth documents as context for answering questions, in open-domain QA,
QA systems need to use documents obtained by a retriever, which may include noisy documents besides gold ones (Chen et al., 2017; Dunn et al., 2017).
In such cases, QA systems need to understand the document relations to perform well (Yang et al., 2018).
To simulate this setting, we modify the SQuAD dataset by prepending or appending 1–2 distracting documents to the original document given to each question.
Table 3 shows the result.
While BERT incurs a large performance drop (-2.8%), LinkBERT is robust to distracting documents (-0.5%).
This result suggests that pretraining with document links improves the ability to understand document relations and relevance.
In particular, our intuition is that the DRP objective helps the LM to better recognize document relations like (anchor document, linked document) in pretraining, which helps to recognize relations like (question, right document) in downstream QA tasks.
The strength of understanding document relations also suggests the promise of applying LinkBERT to various retrievalaugmented methods and tasks (e.g. Lewis et al. 2020b), either as the main LM or the dense retriever component.
Improved few-shot QA performance.
We also find that
LinkBERT is notably good at few-shot learning.
Concretely, for each MRQA dataset, we fine-tune with only 10% of the available training data, and report the performance in Table 4.
In this few-shot regime, LinkBERT attains more significant gains over BERT, compared to the full-resource regime in
Table 1 (on NaturalQ, 5.4% vs 1.8% absolute in F1, or 15% vs 7% in relative error reduction).
This result suggests that
LinkBERT internalizes more knowledge than BERT during pretraining, which supports our core idea that document links can bring in new, useful knowledge for LMs.
6. Experiments: Biomedical domain
Pretraining LMs on biomedical text is shown to boost performance on biomedical NLP tasks (Beltagy et al., 2019;
Lee et al., 2020; Lewis et al., 2020a; Gu et al., 2020).
Biomedical LMs are typically trained on PubMed, which contains abstracts and citations of biomedical papers.
While prior works only use their raw text for pretraining, academic papers have rich dependencies with each other via citations (references).
We hypothesize that incorporating citation links can help LMs learn dependencies between papers and knowledge that spans across them.
[Table] [Table]
[Table] [Table]
With this motivation
[Table]
we pretrain LinkBERT on PubMed with citation links (§6.1), which we term BioLinkBERT , and evaluate on biomedical downstream tasks (§6.2).
As our baseline, we follow and compare with the state-of-the-art biomedical LM, PubmedBERT (Gu et al., 2020), which has the same architecture as BERT and is trained on PubMed.
6.1. Setup
Pretraining data.
We use the same pretraining corpus used by PubmedBERT: PubMed abstracts (21GB).
We use the
Pubmed Parser 3 to extract citation links between articles.
We then create training instances by sampling contiguous , random , or linked segments as described in §4, with the three options appearing uniformly (33%, 33%, 33%).
In summary, our pretraining data is the same as PubmedBERT, except that we have citation links between PubMed articles.
Implementation.
We pretrain BioLinkBERT of -base (110M params) and -large (340M) sizes from scratch, following the same hyperparamters as the PubmedBERT (Gu et al., 2020).
Appendix B provides more details.
Baselines.
We compare BioLinkBERT with PubmedBERT released by Gu et al.
6.2. Downstream evaluation tasks.
We evaluate on the BLURB benchmark (Gu et al., 2020), a diverse set of biomedical NLP datasets, and MedQA-USMLE (Jin et al., 2021), a challenging biomedical QA benchmark.
BLURB consists of five named entity recognition tasks,
We use papers published before Feb. 2020 as in PubmedBERT.
[Table] a PICO (population, intervention, comparison, and outcome) extraction task, three relation extraction tasks, a sentence similarity task, a document classification task [Table] and two question answering tasks, as summarized in Table 5.
These tasks are based on text from biomedical publications
[Table]
and hence directly useful for scientific discovery and applications.
We follow the same finetuning method and evaluation metric used by PubmedBERT (Gu et al., 2020).
MedQA-USMLE is a 4-way multi-choice QA task that tests applications of biomedical and clinical knowledge.
The questions are from practice tests for the US Medical
License Exams (USMLE).
The questions typically require multi-hop reasoning, e.g., given patient symptoms, infer the likely cause, and then answer the appropriate diagnosis procedure (Figure 3).
We follow the finetuning method in
Jin et al.
More details are provided in Appendix C.
MMLU-professional medicine is a multi-choice QA task that tests biomedical knowledge and reasoning, and is part of the popular MMLU benchmark (Hendrycks et al., 2021) that is used to evaluate massive language models.
We take the BioLinkBERT finetuned on the above MedQA-USMLE task, and evaluate on this task without further adaptation.
Three days
[Table]
undergoing a laparoscopic Whipple's procedure, a [Table] year-old woman has swelling of her right leg .
She was diagnosed with pancreatic cancer 1 month ago. ...
Her temperature is 38°C ( [Table] F), pulse is 90/min
[Table]
and blood pressure is 118/
[Table]
mm Hg.
Examination shows mild swelling of the right thigh to the ankle; there is no erythema or pitting edema. ...
Which of the following is the most appropriate
[Table]
in management?
[Table]
Knowledge learned via
document links
Reference
Question : Roden Brothers were taken over in 1953 by a group headquartered in which Canadian city?
Doc A : Roden Brothers was founded June 1, 1891 in Toronto, Ontario, Canada by Thomas and Frank Roden.
In the 1910s the firm became known as Roden Bros. Ltd. and were later taken over by Henry Birks and Sons in 1953. ...
In 1974 Roden Bros. Ltd. published the book, "Rich Cut Glass" with Clock House Publications in Peterborough, Ontario, which was a reprint of the 1917 edition published by Roden
Bros., Toronto.
Doc B : Birks Group (formerly Birks & Mayors) is a designer, manufacturer and retailer of jewellery, timepieces, silverware and gifts, with stores and manufacturing facilities located in Canada and the United States.
As of June 30, 2015, it operates stores under three diﬀerent retail banners: … The company is headquartered in Montreal , Quebec, with American corporate oﬀices located in Tamarac, Florida.
LinkBERT prediction: “Montreal”
[Figure]
BERT
[Figure] Question : Roden Brothers were taken over in 1953 by a group headquartered in which Canadian city?
Doc A : Roden Brothers was founded June 1, 1891 in Toronto, Ontario, Canada by Thomas and Frank Roden.
In the 1910s the firm became known as Roden Bros. Ltd. and were later taken over by Henry Birks and Sons in 1953. ...
In 1974 Roden Bros. Ltd. published the book, "Rich Cut
Glass" with Clock House Publications in Peterborough, Ontario, which was a reprint of the 1917 edition published by Roden Bros., Toronto.
Doc B : Birks Group (formerly Birks & Mayors) is a designer, manufacturer and retailer of jewellery, timepieces, silverware and gifts, with stores and manufacturing facilities located in Canada and the United States.
As of June 30, 2015, it operates stores under three different retail banners: ...
The company is headquartered in Montreal , Quebec, with American corporate offices located in Tamarac, Florida.
LinkBERT prediction: “Montreal” ( ✓ ) BERT
[Figure]
: “Toronto” (
[Figure]
HotpotQA example
LinkBERT predicts: “Montreal” (
[Figure]
BERT
[Figure] “Toronto” ( [Figure] in the question ( leg swelling , pancreatic cancer ), infer the cause ( deep vein thrombosis ), and then infer the appropriate diagnosis procedure ( compression ultrasonography ).
While the existing PubmedBERT tends to simply predict a choice that contains a word appearing in the question (“blood” for choice D), BioLinkBERT correctly predicts the answer (B). Our intuition is that citation links bring relevant documents together in the same context in pretraining (right), which readily provides the multi-hop knowledge needed for the reasoning (center). [Table] [Table]
[Table]
oLinkBERT significantly outperforms the largest general-domain
LM or QA model, despite having just 340M parameters.
6.3. Results
BLURB. Table 5 shows the results on BLURB. BioLinkBERT base outperforms PubmedBERT base on all task categories, attaining a performance boost of +2% absolute on average. Moreover, BioLinkBERT large provides a further boost of +1%.
In total, BioLinkBERT outperforms the previous best by +3% absolute, establishing a new state of the art on the BLURB leaderboard.
We see a trend that gains are notably large on document-level tasks such as question answering (+7% on BioASQ and PubMedQA).
This result is consistent with the general domain (§5.2) and confirms that LinkBERT helps to learn document dependencies better.
MedQA-USMLE.
Table 6 shows the results.
BioLinkBERT base obtains a 2% accuracy boost over
PubmedBERT base , and BioLinkBERT large provides an additional +5% boost.
In total, BioLinkBERT outperforms the previous best by +7% absolute, setting a new state of the art.
To further gain qualitative insights, we studied in what
QA examples BioLinkBERT succeeds but the baseline PubmedBERT fails.
Figure 3 shows a representative example.
Answering the question (left) needs 2-hop reasoning (center): from the patient symptoms described in the question ( leg swelling , pancreatic cancer ), infer the cause ( deep vein thrombosis ), and then infer the appropriate diagnosis procedure ( compression ultrasonography ).
We find that while the existing PubmedBERT tends to simply predict a choice that contains a word appearing in the question (“blood” for choice
D), BioLinkBERT correctly predicts the answer (B).
Our intuition is that citation links bring relevant documents and concepts together in the same context in pretraining (right),which readily provides the multi-hop knowledge needed for the reasoning (center).
Combined with the analysis on
HotpotQA (§5.3), our results suggest that pretraining with document links consistently helps for multi-hop reasoning across domains (e.g., general documents with hyperlinks and biomedical articles with citation links).
MMLU-professional medicine.
Table 7 shows the performance.
Despite having just 340M parameters,
BioLinkBERT large achieves 50% accuracy on this QA task, significantly outperforming the largest general-domain LM or QA models such as GPT-3 175B params (39% accuracy) and UnifiedQA 11B params (43% accuracy).
This result shows that with an effective pretraining approach, a small domain-specialized LM can outperform orders of magnitude larger language models on QA tasks.
7. Conclusion
We presented LinkBERT, a new language model (LM) pretraining method that incorporates document link knowledge For instance in Figure 3 (
), Ansari (2015) in PubMed mention that pancreatic cancer can induce deep vein thrombosis in leg , and it cites a paper in PubMed, Piovella et al.
(2002), which mention that deep vein thrombosis is tested by compression ultrasonography .
Placing these two documents in the same context yields the complete multi-hop knowledge needed to answerthequestion(“ pancreaticcancer ” → “ deepveinthrombosis ” → “ compression ultrasonography ”). such as hyperlinks and citation links.
In both general and biomedical domains, LinkBERT outperforms previous
BERT models across diverse downstream tasks.
The gains are notably large for knowledge- and reasoning-intensive tasks such as multi-hop and multi-document reasoning as well as few-shot inference, which are crucial in scientific discovery.
Our results suggest that LinkBERT can be a strong pretrained LM to be applied to various tasks in science.
In Empirical Methods in Natural Language Processing
Appendix A. Additional methodological details A.1. Graph machine learning perspective
Our two pretraining tasks, masked language modeling (MLM) and document relation prediction (DRP), are also motivated as graph self-supervised learning on the document graph.
In graph self-supervised learning, two types of tasks, node feature prediction and link prediction, are commonly used to learn the content and structure of a graph.
In node feature prediction (Hu et al., 2020), some features of a node are masked, and the task is to predict them using neighbor nodes.
This corresponds to our MLM task, where masked tokens in Segment A can be predicted using Segment B (a linked document on the graph), and vice versa.
In link prediction (Bordes et al., 2013; Wang et al., 2021a), the task is to predict the existence or type of an edge between two nodes.
This corresponds to our DRP task, where we predict if the given pair of text segments are linked (edge), contiguous (self-loop edge), or random (no edge).
LinkBERT can be viewed as a natural fusion of language-based (e.g. BERT) and graph-based self-supervised learning.
A.2. Strategy to obtain linked documents
As described in §4.1, §4.2, our method builds links between documents, and for each anchor segment, samples a linked document to put together in the LM input.
Here we discuss three key axes to consider to obtain useful linked documents in this process.
Relevance.
Semantic relevance is a requisite when building links between documents.
If links were randomly built without relevance, LinkBERT would be same as BERT, with simply two options of LM inputs ( contiguous or random ).
Relevance can be achieved by using hyperlinks or lexical similarity metrics, and both methods yield substantially better performance than using random links (§D.1).
Salience.
Besides relevance, another factor to consider ( salience ) is whether the linked document can offer new, useful knowledge that may not be obvious to the current LM.
Hyperlinks are potentially more advantageous than lexical similarity links in this regard: LMs are shown to be good at recognizing lexical similarity (Zhang et al., 2020), and hyperlinks can bring in useful background knowledge that may not be obvious via lexical similarity alone (Asai et al., 2020).
Indeed, we empirically find that using hyperlinks yields a more performant LM (§D.1).
Diversity.
In the document graph, some documents may have a very high in-degree (e.g., many incoming hyperlinks, like the “United States” page of Wikipedia), and others a low in-degree.
Ifweuniformlysamplefromthelinkeddocuments for each anchor segment, we may include documents of high in-degree too often in the overall training data, losing diversity.
Toadjustsothatalldocumentsappearwithasimilar frequency in training, we sample a linked document with probability inversely proportional to its in-degree, as done in graph data mining literature (Henzinger et al., 2000).
We find that this technique yields a better LM performance (§D.1).
B. Pretraining details B.1. General domain
We pretrain LinkBERT of two sizes, -base and -large, following the configurations of BERT base (110M params) and BERT large (340M params).
For -base, we initialize LinkBERT with the BERT base checkpoint released by Devlin et al.
and continue pretraining.
We use the AdamW (Loshchilov & Hutter, 2019) optimizer with ( β 1 ,β 2 )=(0 . 9 , 0 . 98) , warm up the learning rate for the first 5,000 steps and then linearly decay it.
We train for 40,000 steps with a peak learning rate 3e-4, weight decay 0.01, and batch size of 2,048 sequences withtokens.
Training took 4 days on four A100 GPUs with fp16.
For -large, we follow the same procedure as -base, except that we use a peak learning rate of 2e-4.
Training took 7 days on eight A100 GPUs with fp16.
B.2. Biomedical domain
We pretrain BioLinkBERT of -base (110M params) from scratch, following the same hyperparamters as the PubmedBERT (Gu et al., 2020).
Specifically, we use a peak learning rate 6e-4, batch size 8,192, and train for 62,500 steps.
We warm up the learning rate in the first 10% of steps and then linearly decay it.
Training took 7 days on eight A100 GPUs with fp16.
Additionally, while the original PubmedBERT release did not include the -large size, we pretrain BioLinkBERT of the -large size (340M params) from scratch, following the same procedure as -base, except that we use a peak learning rate of 4e-4 and warm up steps of 20%.
Training took 21 days on eight A100 GPUs with fp16.
C. Downstream task details
Here we provide the details of our downstream evaluation tasks and finetuning hyperparameters.
MRQA.
Given a document (or set of documents) and a question as input, the task is to identify an answer span from the document.
We evaluate on six popular datasets from the MRQA shared task (Fisch et al., 2019):
HotpotQA (Yang et al., 2018), TriviaQA (Joshi et al., 2017), NaturalQ (Kwiatkowski et al., 2019), SearchQA (Dunn et al., 2017),
NewsQA (Trischler et al., 2017), and SQuAD (Rajpurkar et al., 2016).
As the MRQA shared task does not have a public test set, we split the dev set in half to make new dev and test sets.
For all the MRQA datasets, we use max seq length =and a sliding window of size 128 if the lengths are longer than max seq length .
For -base (BERT base , LinkBERT base ), we choose learning rates from { 2e-5, 3e-5 } , batch sizes from { 12, 24 } , and fine-tuning epochs from { 2, 4 } .
For -large (BERT large , LinkBERT large ), we choose learning rates from { 1e-5, 2e-5 } , batch sizes from { 16, 32 } , and fine-tuning epochs from { 2, 4 } .
GLUE.
The General Language Understanding Evaluation (GLUE) benchmark (Wang et al., 2018) is a popular suite of sentence-level classification tasks.
Following BERT, we evaluate on CoLA (Warstadt et al., 2019), SST-2 (Socher et al., 2013), MRPC (Dolan & Brockett, 2005), QQP , STS-B (Cer et al., 2017), MNLI (Williams et al., 2017), QNLI (Rajpurkar et al., 2016), and RTE (Dagan et al., 2005; Haim et al., 2006; Giampiccolo et al., 2007).
For finetuning, we use max seq length = 128.
For -base and -large (BERT base , LinkBERT base , BERT large ,
LinkBERT large ), we choose learning rates from { 5e-6, 1e-5, 2e-5, 3e-5, 5e-5 } , batch sizes from { 16, 32, 64 } and fine-tuning epochs from 3–10.
BLURB.
BLURB consists of five named entity recognition tasks, a PICO (population, intervention, comparison, and outcome) extraction task, three relation extraction tasks, a sentence similarity task, a document classification task, and two question answering tasks, as summarized in Table 5.
We follow the same fine-tuning method and evaluation metric used by PubmedBERT (Gu et al., 2020).
We use max seq length = 512 and choose learning rates from { 1e-5, 2e-5, 3e-5, 5e-5, 6e-5 } , batch sizes from { 16, 32, 64 } and fine-tuning epochs from 1–120.
MedQA-USMLE.
This is a 4-way multi-choice QA task that tests biomedical and clinical knowledge.
The questions are from practice tests for the US Medical License
Exams (USMLE).
The questions typically require multi-hop reasoning, e.g., given patient symptoms, infer the likely cause, and then answer the appropriate diagnosis procedure (Figure 3).
For finetuning, we use max seq length =and choose learning rates from { 1e-5, 2e-5, 3e-5 } , batch sizes from { 16, 32, 64 } and fine-tuning epochs from 1–6.
MMLU-professional medicine.
This is a multi-choice
QA task that tests biomedical knowledge and reasoning, and is part of the popular MMLU benchmark (Hendrycks et al., 2021) that is used to evaluate massive language models.
We take the BioLinkBERT finetuned on the above
MedQA-USMLE task, and evaluate on this task without further adaptation.
D. Additional experimental results D.1. Ablation study
We conduct ablation studies on the key design choices of
LinkBERT.
What linked documents to feed into LMs?
We study the strategies discussed in §A.2 for obtaining linked documents: relevance, salience, and diversity.
Table 8 shows the ablation
[Table]
[Table]
result on MRQA datasets.
First, if we ignore relevance and use random document links instead of hyperlinks, we get the same performance as BERT (-4.1% on average; “random” in Table 8).
Second, using lexical similarity links instead of hyperlinks leads to 1.8% performance drop (“TF-IDF”).
Our intuition is that hyperlinks can provide more salient knowledge that may not be obvious from lexical similarity alone.
Nevertheless, using lexical similarity links is substantially better than BERT (+2.3%), confirming the efficacy of placing relevant documents together in the input for LM pretraining.
Finally, removing the diversity adjustment in document sampling leads to 1% performance drop (“No diversity”).
In summary, our insight is that to create informative inputs for
LM pretraining, the linked documents must be semantically relevant and ideally be salient and diverse.
Effect of the DRP objective.
Table 9 shows the ablation result on the DRP objective (§4.2).
Removing DRP in pretraining hurts downstream QA performance.
The drop is large on tasks with multiple documents (HotpotQA, TriviaQA, and SQuAD with distracting documents).
This suggests that
DRP facilitates LMs to learn document relations.
D.2. Qualitative example
To gain qualitative insights, we studied in what QA examples
LinkBERT succeeds but BERT fails.
Figure 4 shows a representative example from HotpotQA.
Answering the question needs 2-hop reasoning: identify “Roden Brothers were taken over by Birks Group” from the first document, and then “Birks Group is headquartered in Montreal” from the second document.
While BERT tends to simply predict an entity near the question entity (“Toronto” in the first document, which is just 1-hop), LinkBERT correctly predicts the answer in the second document (“Montreal”).
A representative example from the biomedical domain is discussed in §6.3.
E. Ethics, limitations and risks
We outline potential ethical issues with our work below.
First, LinkBERT is trained on the same text corpora (e.g.,
Wikipedia, Books, PubMed) as in existing language models.
Consequently, LinkBERT could reflect the same biases and
[Table]
[Table]
[Table]
[Table]
[Table]
Question : Roden Brothers were taken over in 1953 by a group headquartered in which Canadian city?
Doc A : Roden Brothers was founded June 1, 1891 in Toronto, Ontario, Canada by Thomas and Frank Roden.
In the 1910s the firm became known as Roden Bros. Ltd. and were later taken over by Henry Birks and Sons in 1953. ...
In 1974 Roden Bros. Ltd. published the book, "Rich Cut Glass" with Clock House Publications in Peterborough, Ontario, which was a reprint of the 1917 edition published by Roden
Bros., Toronto.
Doc B : Birks Group (formerly Birks & Mayors) is a designer, manufacturer and retailer of jewellery, timepieces, silverware and gifts, with stores and manufacturing facilities located in Canada and the United States.
As of June 30, 2015, it operates stores under three diﬀerent retail banners: … The company is headquartered in Montreal , Quebec, with American corporate oﬀices located in Tamarac, Florida.
LinkBERT prediction: “Montreal” ( ✓ ) BERT prediction: “Toronto” ( ✗ )
[Figure]
Question : Roden Brothers were taken over in 1953 by a group headquartered in which Canadian city?
Doc A : Roden Brothers was founded June 1, 1891 in Toronto, Ontario, Canada by Thomas and Frank Roden.
In the 1910s the firm became known as Roden Bros. Ltd. and were later taken over by Henry Birks and Sons in 1953. ...
In 1974 Roden Bros. Ltd. published the book, "Rich Cut
Glass" with Clock House Publications in Peterborough, Ontario, which was a reprint of the 1917 edition published by Roden Bros., Toronto.
Doc B : Birks Group (formerly Birks & Mayors) is a designer, manufacturer and retailer of jewellery, timepieces, silverware and gifts, with stores and manufacturing facilities located in Canada and the United States.
As of June 30, 2015, it operates stores under three different retail banners: ...
The company is headquartered in Montreal , Quebec, with American corporate offices located in Tamarac, Florida.
LinkBERT prediction: “Montreal” ( ✓ ) BERT prediction: “Toronto” ( ✗ )
[Figure]
Answering the question needs to identify “Roden Brothers were taken over by Birks Group” from the first document, and then “Birks Group is headquartered in Montreal” from the second document.
toxic behaviors exhibited by language models, such as biases about race, gender, and other demographic attributes (Sheng et al., 2020).
Another source of ethical concern is the use of the MedQAUSMLE evaluation (Jin et al., 2021).
While we find this clinical reasoning task to be an interesting testbed for
LinkBERT and for multi-hop reasoning in general, we do not encourage users to use the current models for real world clinical prediction.