PaperlayerView exact revision on arXiv

Source-linked AI summary

LinkBERT: Pretraining Language Models with Document Links

Michihiro Yasunaga, Jure Leskovec, Percy Liang

arXiv:2203.15827v1cs.CLcs.LG

TL;DR

BERT-style pretraining does not capture dependencies and knowledge spanning documents. LinkBERT incorporates linked documents into shared contexts with masked language modeling and document relation prediction, consistently outperforming baseline LMs across domains and tasks, including +3% absolute on BLURB.

  • Problem

    Existing language-model pretraining methods model single documents and miss dependencies and knowledge spanning documents.

  • Method

    LinkBERT views corpora as document graphs, places linked documents in shared contexts, and jointly trains masked language modeling with document relation prediction.

  • Results

    Across domains and tasks, LinkBERT consistently outperforms baseline language models, including +3% absolute on BLURB and +7% absolute accuracy on MedQA-USMLE.

  • Takeaways & Limitations

    LinkBERT is especially effective for multi-hop reasoning, multi-document understanding, and few-shot question answering across general and biomedical domains.

  • Takeaways & Limitations

    LinkBERT depends on semantically relevant document links; randomly built links reduce it to BERT-like inputs and perform substantially worse.

Abstract

from arXiv · show

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. 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 at https://github.com/michiyasunaga/LinkBERT.

1 Introduction

LinkBERT addresses the limitation of single-document language-model pretraining by placing linked documents together and jointly optimizing masked language modeling with Document Relation Prediction. It improves baseline language models across general and biomedical domains, especially on multi-hop and biomedical reasoning tasks.

  • Motivation: Existing language-model pretraining typically uses a single document per input and misses dependencies and knowledge spanning linked documents.Hyperlinks can connect facts such as Tidal Basin hosting the National Cherry Blossom Festival and the festival celebrating Japanese cherry trees.
  • Method: LinkBERT views the corpus as a document graph and places linked documents in the same context alongside contiguous or random-document inputs.This incorporates link knowledge directly into language-model pretraining.
  • Method: Document Relation Prediction classifies whether the second segment is contiguous, random, or linked, while masked language modeling promotes learning multi-hop knowledge.Together, the objectives encourage relevance and bridging concepts between documents.
  • Results: +4% absolute in F1-score: LinkBERT outperforms BERT on the MRQA benchmark and also improves on GLUE.The general-domain model uses Wikipedia articles with hyperlinks.
  • Results: +3% absolute in BLURB score: LinkBERT surpasses PubmedBERT on the biomedical NLP benchmark.The biomedical model is pretrained on PubMed articles using citation links.
  • Results: +7% absolute in accuracy: LinkBERT improves on the MedQA-USMLE reasoning task.The paper reports notably large gains for multi-hop reasoning overall.

2 Related work

Prior work uses related documents, hyperlinks, citation links, and knowledge graphs to augment language models or retrieval, while LinkBERT focuses on document links for pretraining general-purpose LMs. Related-document pretraining methods provide precedents, but LinkBERT presents document links as a general mechanism for linking context across documents.

  • Retrieval-augmented LMs: Retrieval-augmented LMs add retrieved text to the same context at inference time, but do not study this strategy during pretraining.The cited retrieval works use anchor texts such as questions to improve inference, including answer prediction.
  • Pretrain LMs with related documents: Concurrent pretraining methods place topical or lexically similar documents or sentences in the same language-model context.Caciularu et al. use news articles about the same topic, whereas Levine et al. use sentences with high lexical similarity.
  • Pretrain LMs with related documents: LinkBERT generalizes related-document pretraining by incorporating document links, with lexical or topical similarity treated as possible link types alongside hyperlinks.The work focuses on hyperlinks because they can bring salient knowledge into context.
  • Hyperlinks and citation links for NLP; Graph-augmented LMs: Unlike hyperlink-based retriever training and graph-augmented LMs using knowledge graphs, LinkBERT uses hyperlinks to improve context for general-purpose LM learning and document graphs to span documents.Knowledge-graph approaches represent entities as nodes and their relations as edges, whereas LinkBERT works with document graphs.

3 Preliminaries

The preliminaries define language-model pretraining as encoding token sequences into contextualized representations for self-supervised and downstream tasks. They review BERT’s MLM and NSP objectives before introducing document-link information as an extension.

  • Language-model pretraining: An LM encodes a token sequence into a contextualized vector for each token, which supports self-supervised pretraining and downstream fine-tuning tasks.The model is composed of an encoder and a task head.
  • BERT pretraining objectives: BERT’s masked language modeling objective masks 15% of input tokens and predicts their original values.Among selected tokens, 80% are replaced with [MASK], 10% with random tokens, and 10% remain unchanged.
  • BERT pretraining objectives: BERT’s next sentence prediction objective takes two text segments, (X_A,X_B), as input.The supplied passage introduces NSP as BERT’s second self-supervised task.
  • Transition to LinkBERT: LinkBERT incorporates document-link information into language-model pretraining while building on MLM and NSP.The approach is developed in Section 4.

4 LinkBERT

LinkBERT pretrains language models on a graph of linked documents rather than isolated documents, placing linked texts in shared contexts to capture cross-document knowledge. It combines masked language modeling with Document Relation Prediction and emphasizes relevance, salience, and diversity when selecting links.

  • 4.1 Document graph: LinkBERT represents the corpus as a document graph whose edges encode hyperlinks or other document-relevance links.The method replaces a set-of-documents view with G = (X, E), where E contains directed document links.
  • 4.2 Pretraining inputs: It creates language-model inputs by placing linked document segments in the same context window, alongside single-document and random-document alternatives.This is intended to expose dependencies and knowledge spanning multiple documents during pretraining.
  • 4.2 Training objectives: LinkBERT jointly trains masked language modeling and Document Relation Prediction to learn knowledge and relations introduced by document links.Document Relation Prediction classifies whether one segment is contiguous, random, or linked to another, encouraging recognition of bridging concepts and document relevance.
  • Link selection: Useful links require semantic relevance, because random links reduce LinkBERT to BERT-like contiguous or random inputs, while hyperlinks and lexical similarity perform substantially better.The passage identifies relevance as a prerequisite for effective document linking.
  • Link selection: Beyond relevance, link selection considers salience and diversity: hyperlinks can add non-obvious background knowledge, while sampling corrects overrepresentation of high-in-degree documents.The method notes that hyperlinks empirically produce a more performant LM than lexical-similarity links and that inverse-frequency sampling improves document coverage.

5 Experiments

Experiments compare LinkBERT with matched BERT baselines on Wikipedia-based pretraining and downstream QA and GLUE tasks. LinkBERT improves QA broadly, with especially large gains in multi-document reasoning, few-shot settings, and document-relation-sensitive conditions.

  • Experimental setup: The evaluation covers six MRQA extractive QA datasets and the GLUE sentence-level classification benchmark, with LinkBERT compared against BERT.The QA datasets are HotpotQA, TriviaQA, NaturalQ, SearchQA, NewsQA, and SQuAD; GLUE includes eight listed tasks.
  • Experimental setup: LinkBERT is pretrained on Wikipedia and BookCorpus, using contiguous, random, or hyperlink-linked segments for Wikipedia inputs and contiguous or random segments for BookCorpus.The three Wikipedia sampling options are used uniformly, while the two BookCorpus options are used equally.
  • Overall results: +4.1% absolute, +2.6% absolute, and +2.5% absolute are LinkBERT’s average MRQA gains over BERT at the tiny, base, and large scales, respectively.LinkBERT substantially outperforms BERT on all MRQA datasets, while performing moderately better on GLUE.
  • Improved multi-hop reasoning: +5% over BERTtiny on HotpotQA, +6% on TriviaQA, and +8% on SearchQA exceed the +1.4% gain on SQuAD, indicating stronger results on multi-document reasoning tasks.HotpotQA requires 2-hop reasoning, whereas SQuAD provides a single document per question.
  • Improved few-shot QA performance: With 10% of fine-tuning data, LinkBERT achieves a 5.4% versus 1.8% absolute F1 gain on NaturalQ and 15% versus 7% relative error reduction over BERT.The few-shot gains are larger than those in the full-resource regime.
  • Ablation studies: Removing DRP hurts downstream QA, especially on multi-document tasks, while replacing hyperlinks with random links yields -4.1% average performance and TF-IDF links cause a 1.8% drop.These ablations support the role of document relations and salient hyperlink connections in LinkBERT’s gains.

6 Biomedical LinkBERT (BioLinkBERT)

BioLinkBERT pretrains LinkBERT on PubMed abstracts with citation links and evaluates it against PubmedBERT on biomedical benchmarks. It improves performance across BLURB, MedQA-USMLE, and MMLU-professional medicine, with especially large gains on document-level and multi-hop question answering.

  • Evaluation: BioLinkBERT is evaluated on BLURB, MedQA-USMLE, and MMLU-professional medicine against PubmedBERT and prior biomedical or general-domain models.MedQA-USMLE is a 4-way multiple-choice task requiring biomedical and clinical knowledge, often through multi-hop reasoning.
  • Pretraining: BioLinkBERT pretrains LinkBERT on PubMed using citation links between articles, with contiguous, random, and linked segments sampled uniformly.Its pretraining corpus matches PubmedBERT’s 21GB of PubMed abstracts, except for the added citation links.
  • BLURB: +3% absolute: BioLinkBERT outperforms the previous best on BLURB, establishing a new state of the art.BioLinkBERTbase improves +2% absolute over PubmedBERTbase on average, while BioLinkBERTlarge adds a further +1%.
  • MedQA-USMLE: +7% absolute: BioLinkBERT outperforms the previous best on MedQA-USMLE and sets a new state of the art.BioLinkBERTbase obtains a 2% accuracy boost over PubmedBERTbase, and BioLinkBERTlarge provides an additional +5% boost.
  • Question answering: 50% accuracy: BioLinkBERTlarge on MMLU-professional medicine outperforms GPT-3 175B params at 39% accuracy and UnifiedQA 11B params at 43% accuracy.A cited example illustrates two-hop reasoning enabled by linked documents: pancreatic cancer → deep vein thrombosis → compression ultrasonography.

7 Conclusion

LinkBERT incorporates document-link knowledge into language-model pretraining and outperforms previous BERT models across downstream tasks in both general and biomedical domains. Its gains are notably large for multi-hop reasoning, multi-document understanding, and few-shot question answering.

  • Method: LinkBERT is a language-model pretraining method that incorporates document-link knowledge, including hyperlinks and citation links.It is pretrained on Wikipedia with hyperlinks in the general domain and PubMed with citation links in the biomedical domain.
  • Results: LinkBERT outperforms previous BERT models across a wide range of downstream tasks in both general and biomedical domains.The reported domains are Wikipedia-based general-domain pretraining and PubMed-based biomedical pretraining.
  • Results: The gains are notably large for multi-hop reasoning, multi-document understanding, and few-shot question answering.

A Ethics, limitations and risks

The paper identifies ethical risks from LinkBERT’s training data and cautions against using its current models for real-world clinical prediction. Because it uses corpora also used by existing language models, LinkBERT may reproduce demographic biases and toxic behaviors.

  • Ethical risks: LinkBERT may reflect biases and toxic behaviors present in its training corpora, including biases concerning race, gender, and other demographic attributes.The cited corpora include Wikipedia, Books, and PubMed, which are also used by existing language models.
  • Ethical risks: The authors identify MedQA-USMLE as an additional ethical concern despite viewing it as an interesting testbed for LinkBERT and multi-hop reasoning.The evaluation is described as a clinical reasoning task.
  • Ethical risks: The authors do not encourage users to use the current models for real-world clinical prediction.This caution accompanies their discussion of MedQA-USMLE as an evaluation setting.

B Fine-tuning details

Fine-tuning uses task-specific sequence lengths and hyperparameter search spaces, with separate settings for MRQA, GLUE, BLURB, and MedQA-USMLE. MRQA also uses a sliding window for sequences exceeding its maximum length.

  • MRQA: MRQA uses max_seq_length = 384 and a sliding window of size 128 for longer extractive question-answering inputs.
  • MRQA: MRQA selects learning rates, batch sizes, and epochs by model scale: tiny uses {5e-5, 1e-4, 3e-4}, {16, 32, 64}, and {5, 10}; base and large use smaller ranges.Base uses learning rates {2e-5, 3e-5}, batch sizes {12, 24}, and epochs {2, 4}; large uses {1e-5, 2e-5}, {16, 32}, and {2, 4}.
  • GLUE: GLUE uses max_seq_length = 128, with tiny models using learning rates {5e-5, 1e-4, 3e-4}, batch sizes {16, 32, 64}, and 5 or 10 epochs.
  • GLUE: GLUE base and large models use learning rates {5e-6, 1e-5, 2e-5, 3e-5, 5e-5}, batch sizes {16, 32, 64}, and 3–10 epochs.
  • BLURB: BLURB uses max_seq_length = 512, learning rates {1e-5, 2e-5, 3e-5, 5e-5, 6e-5}, batch sizes {16, 32, 64}, and 1–120 epochs.
  • MedQA-USMLE: MedQA-USMLE uses max_seq_length = 512, learning rates {1e-5, 2e-5, 3e-5}, batch sizes {16, 32, 64}, and 1–6 epochs.
Loading 2203.15827v1…