LinkBERT: Pretraining Language Models with Document Links

In plain words

Language models such as BERT learn from one document at a time. Each training input is a chunk of a single article, so facts that only make sense across two articles are never seen together. A Wikipedia page about a park may mention a festival, while the linked page about the festival explains its history. Read separately, neither page gives the full answer.

LinkBERT changes what goes into a training input. It uses the links between documents, such as hyperlinks on Wikipedia or citations between biomedical papers, and sometimes places a linked document right after the first chunk. The model is trained to fill in hidden words, as before, and also to guess how the second chunk relates to the first: it follows on from it, is an unrelated document, or is a linked document.

The payoff shows up in question answering. Questions that need information from several documents benefit most, and the authors report the gains also hold when the model is later used with little labelled data. A biomedical variant trained on citation links reaches strong results on medical benchmarks.

TL;DR

• Builds a graph of documents from hyperlinks or citation links and fills each LM input with an anchor segment plus a contiguous, random or linked segment §4.1 §4.2.
• Trains with masked language modeling (MLM) plus a three-way document relation prediction (DRP) loss that replaces BERT's next sentence prediction §4.2 eq. 1 eq. 2.
• Author-reported: MRQA F1 gains over BERT on all six datasets, largest on multi-document QA, and BioLinkBERT sets new BLURB and MedQA-USMLE results §5.3 §6.3 tab. 1 tab. 7 tab. 8.

Problem

Standard masked language model pretraining, as in BERT, puts text from one document, or two unrelated documents, in each input context, so it does not model dependencies or knowledge that span documents §1 §3. Documents often have such dependencies, such as hyperlinks in Wikipedia or references in academic articles §1 §4.1. The paper assumes these links exist in the corpus, or can be built by a relevance measure such as TF-IDF similarity §4.1. Hyperlinks are assumed to be relevant and to add background that lexical similarity alone would miss §4.1 §4.3.

Method

Overview fig. 2. A corpus is treated as a directed graph of documents §4.1. Pretraining samples an anchor segment and a second segment, places them in one input, and optimizes two self-supervised losses §4.2.

1. Document graph §4.1. Add a directed edge from document i to document j if i contains a hyperlink to j. For the lexical comparison, take the top k = 5 documents by TF-IDF cosine similarity for each document §4.1.
2. Anchor segment §4.2. Sample a text segment X_A from a document.
3. Second segment §4.2. Choose one of three relations, each with equal probability in the main setting §5.1 §6.1:
   • contiguous: the next segment of the same document;
   • random: a segment from a different document;
   • linked: a segment from a document that the anchor document links to. The linked document is chosen with probability inversely proportional to its in-degree, so that heavily linked pages do not dominate §4.3.
4. Input §4.2. Join the segments as [CLS] X_A [SEP] X_B [SEP].
5. Masked language modeling §3 §4.2. Mask 15% of tokens; of those, 80% become [MASK], 10% become a random token and 10% stay unchanged. The model predicts the original tokens §3.
6. Document relation prediction §4.2 eq. 1. A three-class classifier on the [CLS] representation predicts contiguous, random or linked. It replaces NSP, which has two classes §3 §4.2.
7. Loss §4.2 eq. 1 eq. 2. The total loss is the MLM loss plus the DRP loss, with equal weight. Eq. 2 writes the MLM term as a sum of log-likelihoods over the masked tokens, and the DRP term as the log-likelihood of the relation label.

Relevance, salience and diversity are the three criteria the paper uses to choose linked documents §4.3. Diversity is handled by the in-degree rule in step 3 §4.3.

Equations (as written in the paper):

\mathcal{L} = \mathcal{L}_{\text{MLM}} + \mathcal{L}_{\text{DRP}} eq. 1

\mathcal{L} = -\sum_{i} \log p(x_i \mid \mathbf{h}_i) - \log p(r \mid \mathbf{h}_{\texttt{[CLS]}}) eq. 2

where x_i are the tokens of the input and h_i their representations eq. 2.

# Build one LM input from an anchor document  [§4.2] [§4.3]
# Output: tokens (n,), relation in {contiguous, random, linked}

def make_instance(corpus, graph, in_degree):
    # Sample the anchor document and its segment X_A  [§4.2]
    anchor_doc = sample_document(corpus)
    segment_a = sample_segment(anchor_doc)
    # Pick the relation of the second segment, 1/3 each  [§5.1] [§6.1]
    relation = sample_uniform(["contiguous", "random", "linked"])
    if relation == "contiguous":
        # Take the next segment of the same document  [§4.2]
        segment_b = next_segment(anchor_doc, segment_a)
    elif relation == "random":
        # Take a segment from a different document  [§4.2]
        other_doc = sample_other_document(corpus, anchor_doc)
        segment_b = sample_segment(other_doc)
    else:
        # Take a document linked from the anchor, with probability
        # proportional to 1 / in-degree  [§4.2] [§4.3]
        linked_doc = sample_linked_document(graph, anchor_doc, in_degree)
        segment_b = sample_segment(linked_doc)
    # Join as [CLS] X_A [SEP] X_B [SEP]  [§4.2]
    tokens = join_with_special_tokens(segment_a, segment_b)
    return tokens, relation

# One pretraining loss: MLM plus DRP  [§4.2] [eq. 1] [eq. 2]
# Input:  tokens (B, n), relations (B,) with values 0, 1, 2
# Output: scalar loss

def pretraining_loss(model, tokens, relations):
    # Mask 15% of tokens: 80% [MASK], 10% random, 10% kept  [§3]
    masked_tokens, targets = apply_mlm_masking(tokens)
    # Encode the masked two-segment input  [§4.2]
    hidden = model.encode(masked_tokens)            # (B, n, d)
    # Token-level cross-entropy on masked positions  [§4.2] [eq. 2]
    mlm_loss = masked_cross_entropy(model.mlm_head(hidden), targets)
    # Three-way relation prediction from the [CLS] vector  [§4.2]
    relation_logits = model.drp_head(hidden[:, 0])  # (B, 3)
    drp_loss = cross_entropy(relation_logits, relations)
    # Equal-weight sum; the weighting is stated as 1:1  [eq. 1]
    return mlm_loss + drp_loss

Model sizes §5.1 §6.1: tiny (4.4M parameters), base (110M) and large (340M). The architectures follow BERT §5.1. The biomedical models are base (110M) and large (340M) §6.1.

Hyperparameters

Pretraining. "Same as tiny" means the values in the tiny row apply §5.1.

Fine-tuning. The same settings are applied to all models, including baselines app. B. The paper gives grids and the authors select one value per task; which value was selected is not stated.

Results

Author-reported. Setup: pretraining on Wikipedia with hyperlinks plus BookCorpus §5.1; fine-tuning on MRQA and GLUE §5.2; the baseline is BERT at the same scale, pretrained or continued with BERT's objectives on the same corpus for the same number of steps §5.1. Biomedical setup: PubMed with citation links against PubmedBERT-base §6.1 §6.2.

MRQA question answering, F1 (dev and test split of MRQA dev sets) tab. 1:

GLUE average score tab. 2: BERT-tiny 64.3, LinkBERT-tiny 64.6; BERT-base 79.2, LinkBERT-base 79.6; BERT-large 80.7, LinkBERT-large 81.1. The authors describe the GLUE gain as moderate §5.3 tab. 2.

Distracting documents, F1 on SQuAD tab. 3: BERT-base 88.7 without and 85.9 with distractors; LinkBERT-base 90.1 without and 89.6 with distractors.

Few-shot, 10% of fine-tuning data, F1 tab. 4: BERT-base vs LinkBERT-base: HotpotQA 64.8 vs 70.5; TriviaQA 59.2 vs 66.0; NaturalQ 64.8 vs 70.2; SQuAD 79.6 vs 82.8.

Ablation on linked documents, LinkBERT-tiny, F1 tab. 5:

The authors state that the random-link variant matches BERT-tiny §5.5. The TF-IDF variant is 2.3 points above BERT on average §5.5.

Ablation on DRP, LinkBERT-base, F1 tab. 6:

BLURB, test scores tab. 7. Columns: PubMedBERT-base, BioLinkBERT-base, BioLinkBERT-large.

MedQA-USMLE accuracy, 4 options tab. 8: BioBERT-large 36.7; QA-GNN 38.0; GreaseLM 38.5; PubmedBERT-base 38.1; BioLinkBERT-base 40.0; BioLinkBERT-large 44.6.

MMLU professional medicine accuracy, using the MedQA-fine-tuned model without further adaptation tab. 9 §6.2: GPT-3 (175B) 38.7; UnifiedQA (11B) 43.2; BioLinkBERT-large 50.7.

Paper vs code

The official repository is https://github.com/michiyasunaga/LinkBERT. It is named in the abstract, and its owner is the first author. It was cloned at commit 210bc47a412dc3df3e09339a6cd76f8eac51123b. Its LICENSE file is Apache-2.0.

• Pretraining code is not in the repository. The repository holds the fine-tuning scripts, data preprocessing scripts and the pretrained models on Hugging Face README.md README.md. The document graph, the link sampling, the DRP head and the MLM pretraining loop cannot be checked against code §4 §5.1 §6.1.
• Fine-tuning values in the code mostly fall in the paper's grids. Example: MRQA-large uses learning rate 2e-5 and batch 16 [code: src/run_examples_mrqa_linkbert-large.sh#L15] app. B.
• MRQA-base SearchQA uses 1 epoch, but the paper's base grid is {2, 4} epochs run_examples_mrqa_linkbert-base.sh app. B. The other MRQA-base runs use 4 epochs (HotpotQA, NaturalQ, TriviaQA, NewsQA) or 2 epochs (SQuAD) run_examples_mrqa_linkbert-base.sh run_examples_mrqa_linkbert-base.sh.
• MedQA-USMLE, base: learning rate 2e-5, per-device batch 4 with 8 gradient accumulation steps (effective batch 32), 2 epochs, 100 warmup steps [code: src/run_examples_medqa_biolinkbert-base.sh#L13-L14]. These are inside the paper's grid, except that the paper does not state a warmup value.
• MedQA-USMLE, large: learning rate 3e-5, effective batch 64 (2 × 32), 6 epochs, 500 warmup steps run_examples_medqa_biolinkbert-large.sh. Warmup is not stated in the paper.
• GLUE: the repository has no GLUE run script. GLUE appears only in the README results table README.md, and the sequence-classification script refers to Transformers' GLUE example run_seqcls.py. The GLUE settings in the paper cannot be checked against code.

Pitfalls

• Pretraining is underspecified. The paper does not state the segment length in the link-based inputs, how a segment is chosen from a linked document, what happens when an anchor document has no outgoing links, or which graph the in-degree is counted on §4.2 §4.3. The 1/in-degree rule is "inversely proportional" §4.3; the normalization is inferred, not stated.
• Loss scale is ambiguous. Eq. 2 is written as a sum over tokens, but the paper does not say whether implementations average over masked tokens eq. 2. The two losses have equal weight eq. 1.
• DRP head initialization for base and large is not stated. These models continue from BERT checkpoints, whose head is NSP, not DRP §5.1. Whether the three-class head starts fresh is not stated.
• The MRQA evaluation split is not standard. The authors split the MRQA dev sets in half to get new dev and test sets §5.2. The numbers are not comparable to official leaderboard results.
• Ablations are on LinkBERT-tiny, except the DRP ablation, which is on base §5.5 tab. 5 tab. 6. Do not assume they transfer to other sizes.
• Compute is large and link extraction is not costed. Pretraining took 1 day on two 2080 Ti for tiny, 4 and 7 days on 4 and 8 A100 for base and large, and 7 and 21 days on 8 A100 for biomedical base and large §5.1 §6.1. Cost for building the graph is not stated.
• Text and tables differ by rounding. The text gives MMLU as "50%" and GPT-3 as "39%", where the table gives 50.7 and 38.7 §6.3 tab. 9. The text says MedQA-large adds "+5%"; the table gives 40.0 to 44.6 §6.3 tab. 8. Use the tables.
• Results have no stated variance. The paper does not report seeds, error bars or repeated runs for the reported results. Treat differences of about a point with care.
• Naming. The paper writes "DPR" once where it means DRP §4.2.
• Few-shot gains use 10% of fine-tuning data, not a different model §5.4 tab. 4.
• Ethics. The authors warn about biases from the same corpora used by other LMs, and do not encourage clinical use of the models app. A.

Lineage

• Builds on: BERT (arXiv:1810.04805), which uses MLM and next sentence prediction §3. LinkBERT keeps MLM and replaces NSP with DRP §4.2.
• Compared against: BERT at each scale tab. 1; PubmedBERT for biomedical tasks §6.1; BioBERT, QA-GNN and GreaseLM on MedQA-USMLE tab. 8; GPT-3 and UnifiedQA on MMLU tab. 9.
• Superseded by: not known from content in this wiki.