Deep Bidirectional Language-Knowledge Graph Pretraining

In plain words

A language model reads words well but has no built-in map of facts. A knowledge graph is that map: a web of entities and labelled links, such as "round brush" linked to "art supply". This paper builds one model that reads a passage and the small part of the map that matches it at the same time. Each side can then inform the other: the words help pick the right links, and the links help fill in the words.

The model is trained by playing two guessing games at once. In one, some words of the passage are hidden and must be guessed from the rest of the text and the map. In the other, some links of the map are hidden and must be guessed from the passage and the rest of the map. Training on both at once is meant to force a joint understanding of text and facts, rather than two separate ones glued together.

The model is an encoder that reads and scores answers; it does not write text. The authors test it on multiple-choice question answering in everyday commonsense and in medicine, and report gains over the same language model without the map and over an earlier model that uses the map only during fine-tuning.

TL;DR

• One encoder reads a text segment and a local knowledge-graph subgraph, fusing them in both directions after each layer §2.2 fig. 1.
• Pretrained with masked language modeling plus link prediction on held-out graph edges, summed with equal weight §2.3.
• Author-reported: best on nine commonsense QA tasks (average gain of about 4.4 points over the RoBERTa baseline by our reading of tab. 1), and new best on MedQA, PubMedQA and BioASQ tab. 1 tab. 6.
• The paper gives one training recipe in the text (BookCorpus and ConceptNet; PubMed and UMLS for biomedicine) §3.1 §4. The released script is a small demo on CommonsenseQA sentences, not that recipe run_pretrain.sh.

Problem

Language models pretrained on text learn strong contextual representations, but facts that appear rarely in text are weakly represented. Knowledge graphs hold such facts in structured form and help multi-step reasoning. Earlier systems fused text and graph in one direction only, or added the graph only during fine-tuning and never pretrained on it at scale §1. The authors argue that a deep, two-way fusion trained with self-supervision is missing §1.

Method

Overview

The model takes a pair: a text segment and a local subgraph of a large knowledge graph. The subgraph is found by linking entity mentions in the text to graph nodes, adding two-hop bridge nodes, and keeping the edges between them §2.1 app. B.1 fig. 1. Two self-supervised tasks run on the pair: masked language modeling (MLM) on the text, and link prediction (LinkPred) on the graph §2.3 fig. 1. The encoder is the GreaseLM architecture §2.2.

Step 1: build the input pair

# Build one (text, local KG) training pair  [§2.1] [app. B.1]
def build_pair(segment, kg):
    # Link entity mentions in the segment to KG nodes  [app. B.1]
    linked = entity_link(segment)
    # Add two-hop bridge nodes between pairs of linked nodes  [app. B.1]
    nodes = linked | bridge_nodes_2hop(kg, linked)
    # Cap the local graph at 200 nodes, sampled at random  [app. B.1]
    if len(nodes) > 200:
        nodes = random_sample(nodes, 200)
    # Keep every KG edge whose two ends are both retained  [app. B.1]
    edges = edges_between(kg, nodes)
    # Prepend the interaction token to the text  [§2.1]
    tokens = ["INT"] + tokenize(segment)
    # Add the interaction node, linked to the entity-linked nodes  [§2.1]
    graph = add_interaction_node(nodes, edges, link_to=linked)
    return tokens, graph

Notes on this step. The text segment can be up to 512 tokens and the local KG up to 200 nodes in the general domain §3.1. If no entity is linked, the graph is a single dummy node initialised to zero, and the model falls back to the text side app. B.2.

Step 2: the cross-modal encoder

The encoder maps text tokens and KG nodes into fused vectors. Its first stage runs unimodal layers, and its second stage runs fusion layers that pass information across the modalities §2.2.

(\mathbf{H}_{\rm int}, \mathbf{H}_1, \dots, \mathbf{H}_I), (\mathbf{V}_{\rm int}, \mathbf{V}_1, \dots, \mathbf{V}_J) = f_{\rm enc}\big((w_{\rm int}, w_1, \dots, w_I), (v_{\rm int}, v_1, \dots, v_J)\big) eq. 1

(\mathbf{H}^{(0)}_{\rm int}, \mathbf{H}^{(0)}_1, \dots, \mathbf{H}^{(0)}_I) = \text{LM-Layers}(w_{\rm int}, w_1, \dots, w_I) eq. 2

(\mathbf{V}^{(0)}_{\rm int}, \mathbf{V}^{(0)}_1, \dots, \mathbf{V}^{(0)}_J) = \text{Node-Embedding}(v_{\rm int}, v_1, \dots, v_J) eq. 3

(\mathbf{H}_{\rm int}, \dots, \mathbf{H}_I), (\mathbf{V}_{\rm int}, \dots, \mathbf{V}_J) = \text{Fusion-Layers}\big((\mathbf{H}^{(0)}_{\rm int}, \dots), (\mathbf{V}^{(0)}_{\rm int}, \dots)\big) eq. 4

Each fusion layer \ell = 1, \dots, M runs three updates §2.2:

(\widetilde{\mathbf{H}}_{\rm int}^{(\ell)}, \mathbf{H}^{(\ell)}_1, \dots, \mathbf{H}^{(\ell)}_I) = \text{LM-Layer}(\mathbf{H}^{(\ell-1)}_{\rm int}, \dots, \mathbf{H}^{(\ell-1)}_I) eq. 5

(\widetilde{\mathbf{V}}_{\rm int}^{(\ell)}, \mathbf{V}^{(\ell)}_1, \dots, \mathbf{V}^{(\ell)}_J) = \text{GNN-Layer}(\mathbf{V}^{(\ell-1)}_{\rm int}, \dots, \mathbf{V}^{(\ell-1)}_J) eq. 6

[\mathbf{H}^{(\ell)}_{\rm int}; \mathbf{V}^{(\ell)}_{\rm int}] = \text{MInt}\big([\widetilde{\mathbf{H}}^{(\ell)}_{\rm int}; \widetilde{\mathbf{V}}^{(\ell)}_{\rm int}]\big) eq. 7

Only the interaction token and interaction node are mixed by MInt; all other token and node vectors pass through their own layer unchanged §2.2. The paper's equations are written with an align environment; the numbers above are counted from that source, so check them against the PDF (see Pitfalls).

# Fused encoder f_enc  [§2.2] [eq. 1]
def encode(tokens, graph):
    # 19 unimodal LM layers over the text  [§2.2] [app. B.3]
    text_vecs = lm_layers(tokens, depth=19)
    # Node embeddings for the local graph  [§2.2] [app. B.2]
    node_vecs = node_embedding(graph.nodes)
    # 5 fusion layers  [§2.2] [app. B.3]
    for layer in range(5):
        # Text side: one LM layer  [§2.2] [eq. 5]
        text_vecs = lm_layer(text_vecs)
        # KG side: one GNN layer over the local edges  [§2.2] [eq. 6]
        node_vecs = gnn_layer(node_vecs, graph.edges)
        # Mix the interaction token and interaction node  [§2.2] [eq. 7]
        text_int, node_int = mint(text_vecs[0], node_vecs[0])
        text_vecs[0] = text_int
        node_vecs[0] = node_int
    return text_vecs, node_vecs

The encoder is GreaseLM with 19 unimodal LM layers and 5 fusion layers, 360M parameters in total §3.1. The text side is initialised from RoBERTa-Large (general) or BioLinkBERT-Large (biomedical); node embeddings come from BERT-Large or BioLinkBERT name embeddings §3.1 app. B.2.

Step 3: pretraining objective

The two losses are summed with weight 1 each: the total is MLM plus LinkPred §2.3.

Masked language modeling. A subset of text tokens is replaced by a mask token. A linear head predicts the originals from the fused token vectors, with a cross-entropy loss over the masked positions §2.3. The paper's loss is:

\mathcal{L}_{\rm MLM} = -\sum_{i \in M} \log p(w_i \mid \mathbf{H}_i) eq. 8

Link prediction. Some edges of the input KG are held out as positive triples. A relation embedding is learned per relation type, and a scoring function rates each triple. The paper names three options; the main model uses DistMult §2.3 §3.1 tab. 5:

\text{DistMult: } \langle \mathbf{h}, \mathbf{r}, \mathbf{t} \rangle, \quad \text{TransE: } -\|\mathbf{h} + \mathbf{r} - \mathbf{t}\|, \quad \text{RotatE: } -\|\mathbf{h} \odot \mathbf{r} - \mathbf{t}\| eq. 9

Here \mathbf{h}, \mathbf{t} are the fused node vectors \mathbf{V}_h, \mathbf{V}_t, and \mathbf{r} is a learned relation vector §2.3. The loss uses n negative triples per positive, with margin \gamma:

\mathcal{L}_{\rm LinkPred} = \sum_{(h,r,t) \in S} \left( -\log \sigma\big(\phi_r(\mathbf{h}, \mathbf{t}) + \gamma\big) + \frac{1}{n} \sum_{(h', r, t')} \log \sigma\big(\phi_r(\mathbf{h}', \mathbf{t}') + \gamma\big) \right) eq. 10

The second term, as printed, has the sign of the first. The released code uses the negative-sample term -\log\sigma(-\phi) instead, so implement that form (see Pitfalls and Paper vs code).

# One pretraining step on a (text, local KG) pair  [§2.3]
def pretrain_loss(tokens, graph):
    # Mask 15% of the text tokens and keep the originals  [§2.3] [§3.1]
    masked, targets, masked_pos = mask_tokens(tokens, rate=0.15)
    # Hold out 15% of KG edges as positive triples  [§2.3] [§3.1]
    kept_edges, held_out = hold_out_edges(graph.edges, rate=0.15)
    text_vecs, node_vecs = encode(masked, graph.with_edges(kept_edges))
    # Cross-entropy over the masked positions  [§2.3] [eq. 8]
    logits = lm_head(text_vecs[masked_pos])
    loss_mlm = cross_entropy_sum(logits, targets)
    # Link prediction on the held-out triples  [§2.3] [eq. 10]
    loss_link = link_loss(node_vecs, rel_emb, held_out)
    # Unweighted sum of the two objectives  [§2.3]
    return loss_mlm + loss_link

# Score a triple with DistMult  [§2.3] [eq. 9]
def score(node_vecs, rel_emb, triple):
    head, rel, tail = triple
    return sum(node_vecs[head] * rel_emb[rel] * node_vecs[tail])

# Link loss with n corrupted triples per positive  [§2.3] [eq. 10]
def link_loss(node_vecs, rel_emb, held_out, n_neg=128):
    # Paper: n = 128 negatives; released script uses 64  [§3.1]
    total = 0.0
    for triple in held_out:
        pos = score(node_vecs, rel_emb, triple)
        # Positive term, margin gamma = 0  [§2.3] [§3.1] [eq. 10]
        total += -log_sigmoid(pos)
        # Corrupt the head or the tail with a random node  [§2.3]
        for corrupted in corruptions(triple, n_neg):
            neg = score(node_vecs, rel_emb, corrupted)
            # Code form of the negative term, weight 1/n  [eq. 10]
            total += -log_sigmoid(-neg) / n_neg
    return total

The paper's loss sums over held-out triples; the code averages them instead (Paper vs code).

Step 4: fine-tuning

For a downstream task, the same input construction and encoder are used. For multiple choice, the text is the question concatenated with one answer choice §2.4. The pooled vector combines three parts §2.4 modeling_dragon.py modeling_dragon.py modeling_dragon.py:

# Pooled input vector for a classifier head  [§2.4]
def pool(text_vecs, node_vecs):
    # Attention over KG nodes, query = interaction token  [§2.4]
    graph_vec = attention_pool(query=text_vecs[0], keys=node_vecs)
    # Pooler output of the interaction token (BERT-style)  [§2.4]
    sent_vec = pooler(text_vecs[0])
    # Vector of the interaction node  [§2.4]
    node_int_vec = node_vecs[0]
    return concat(graph_vec, sent_vec, node_int_vec)

The pooled vector goes to a classifier in the same way as the classification token of BERT or RoBERTa §2.4. Fine-tuning hyperparameters are in the table below.

Hyperparameters

Values from the paper's hyperparameter table tab. 7, with the text sources. Commonsense and biomedical columns are the paper's; "pre" is pretraining and "ft" is fine-tuning.

Other stated values:

• Pretraining in the general domain: BookCorpus, 6 GB, with ConceptNet (800K nodes, 2M edges) §3.1.
• Pretraining in the biomedical domain: PubMed, 21 GB, with UMLS (300K nodes, 1M edges) §4.
• Link prediction: DistMult, 128 negative triples, margin 0 §3.1.
• Masking: 15% of tokens for MLM and 15% of edges held out §3.1.
• Cost: 7 days on eight A100 GPUs in FP16 for the general model §3.1.
• Downstream sizes: CSQA 12,102; OBQA 5,957; RiddleSense 5,715; ARC 2,590; CosmosQA 35.6K; HellaSwag 70K; PIQA 20K; SIQA 37K; aNLI 170K; MedQA 12,723; PubMedQA 1,000; BioASQ 885 app. B.4.
• Commonsense splits: in-house splits for CSQA §3.2 app. B.4; OBQA uses only the question, not the extra science facts §3.2.

Results

Author-reported. Setup: nine commonsense QA datasets, RoBERTa-Large or GreaseLM as baselines, and the paper's own fine-tuning recipe §3.2 §3.3 tab. 1. "RoBERTa" in the table is RoBERTa-Large continued with the MLM objective on the same text for the same number of steps, not the released checkpoint §3.3.

Table 1: accuracy on commonsense QA (general domain) tab. 1

Averaged over the nine columns, our arithmetic on this table gives about +4.4 points over RoBERTa and +2.5 over GreaseLM. The paper does not state an average; see Pitfalls for the figures it quotes.

Table 2: accuracy on CSQA and OBQA dev sets, by complex-reasoning proxy tab. 2

Table 3: low-resource fine-tuning, 10% of training data tab. 3

Table 4: larger fusion depth (5 to 7 fusion layers, "-Ex") tab. 4

Table 5: ablations on CSQA and OBQA tab. 5

Table 6: biomedical QA tab. 6

Other reported results. On ConceptNet link prediction, Hit@3 is 61.3 for DistMult alone and 78.1 for DRAGON, which also uses retrieved text app. C tab. 9. The qualitative study on CSQA shows attention shifting between graph nodes when a question is negated or extended fig. 2 §3.4.1.

Paper vs code

Checked against the official repository at commit b21251e4482128cae9825e4c7b8ba771ab19ad5c, Apache-2.0 LICENSE. Each item names the paper's value and the code's value.

1. Pretraining recipe is a demo. The released pretraining script uses CommonsenseQA sentences as text (dataset="csqa") run_pretrain.sh, not BookCorpus or PubMed. The README says so README.md, section 3. The repository contains no biomedical pretraining script.
2. Learning rates and batch. Paper: LM 2e-5 and other 3e-4, batch 8,192, 20,000 steps tab. 7. Script: LM 1e-5 and other 1e-3 run_pretrain.sh, batch 128 with mini-batch 2 run_pretrain.sh, 30 epochs run_pretrain.sh, no step count.
3. Input length. Paper: up to 512 tokens tab. 7. Script: max_seq_len=100 run_pretrain.sh.
4. Warmup. Paper: warmup ratio 0.1 tab. 7. Script: 100 warmup steps run_pretrain.sh.
5. LM frozen epochs. Paper: 2 for pretraining tab. 7. Script: 3 run_pretrain.sh. The argparse default is 4 dragon.py.
6. Negatives per positive. Paper: 128 §3.1. Script: 64 run_pretrain.sh; the code splits them into 32 head and 32 tail corruptions modeling_dragon.py.
7. Negative weighting. Paper: uniform weight 1/n §2.3 eq. 10. Code default: self-adversarial softmax weights with temperature 1, detached from the gradient modeling_gnn.py dragon.py. The script does not override it.
8. Negative term form. Paper eq. 10 writes the negative term as log σ(φ′+γ) eq. 10. Code uses −log σ(−φ′) modeling_gnn.py.
9. Loss reduction. Paper: sum over held-out triples eq. 10. Code: the mean over positives and over negatives, averaged, plus a regulariser modeling_gnn.py.
10. Regulariser. Not in the paper. Code adds 0.01 times the mean square of the relation weights modeling_gnn.py dragon.py.
11. DistMult scaling. Paper: plain trilinear product eq. 9. Script sets scldstmlt=true run_pretrain.sh, which divides by the square root of the dimension modeling_gnn.py.
12. Margin γ. Paper: γ = 0 §3.1. Script: gamma=0 run_pretrain.sh. The code's DistMult scoring does not use γ; only RotatE reads link_gamma modeling_gnn.py. The argparse default is 12 dragon.py.
13. MLM masking. Paper: a fraction of tokens is masked §2.3. Code: whole-word masking by default dragon.py data_utils.py, with 80% replaced by the mask token, 10% randomised and 10% kept data_utils.py.
14. MLM loss reduction. Paper: sum over masked positions eq. 8. Code: mean cross-entropy over the masked positions modeling_dragon.py.
15. Edge holdout. Paper: 15% of edges held out §3.1. Code: per graph, 15% of non-special edges, capped at 100 data_utils.py run_pretrain.sh. Edges touching a type-3 node are never held out data_utils.py. The code also keeps 10% of the held-out triples in the input graph data_utils.py run_pretrain.sh. None of this is in the paper.
16. Negative sampling range. Paper: "random" triples §2.3. Code: uniform over the local graph's node count data_utils.py.
17. Pooling. Paper: MLP over the interaction token, interaction node and attention-pooled graph vector §2.4. Code: concatenation of the attention-pooled graph vector, the pooler output of the interaction token and the interaction-node vector modeling_dragon.py modeling_dragon.py modeling_dragon.py.
18. Not verified. I did not check the ConceptNet or UMLS node-embedding preprocessing against app. B.2.

Pitfalls

• Recipe is not in the repo. The paper's main pretraining run (BookCorpus and ConceptNet; PubMed and UMLS) is not reproducible from the released script, which is a CommonsenseQA demo run_pretrain.sh README.md, section 3. Reproducing the reported numbers needs the paper's values from tab. 7 and your own corpus preparation; the authors' checkpoints are linked from the README README.md, section 1.
• Negative-term sign. Eq. 10 as printed adds +log σ(φ′+γ) for negatives. As φ′ falls, that term goes to minus infinity, so the printed objective has no finite minimum. The code's −log σ(−φ′) form is bounded modeling_gnn.py. Implement the code's form and note the difference in your own docs.
• Average gain is not stated. The abstract says +5% on average and the introduction says +8% on average over RoBERTa §1 [abstract]. Table 1 gives about +4.4 points over RoBERTa and +2.5 over GreaseLM by our arithmetic tab. 1. Quote Table 1 directly rather than either summary figure.
• Baseline is continued pretraining. The RoBERTa row is RoBERTa-Large continued with MLM on the same data and steps §3.3. A stock RoBERTa-Large checkpoint is not the same baseline.
• Variance is not reported. The main text gives no seeds, error bars or repeated runs for Tables 1 to 6 §3. Treat single-point differences of 1 to 2 points with caution, especially on small dev sets such as the in-house splits app. B.4.
• Fine-tuning grid selection is not stated. Learning rates, epochs and maximum tokens are given as grids, but the selection rule is not stated tab. 7 app. B.3.
• Equation and section numbers are from source. Equations were numbered by counting align rows in the LaTeX source, and sections by source order. Related work is a subsection of the introduction in the source, which would make the method §2 in the rendered paper. Check against the PDF before quoting a number.
• Table numbering is from source order. Tables 3 to 5 come from one table float with three captions, and tables 7 to 9 are in the appendix. Check against the PDF.
• Compute and inference cost are not in the abstract. Training took 7 days on eight A100 GPUs §3.1. Each example needs entity linking and subgraph retrieval, with at most 200 nodes and 512 tokens §3.1 app. B.1. The retrieval cost is not reported.
• Encoder only. The model does not generate text; the authors name generation as future work §5.
• Entity linking can fail silently. Fewer than 1% of inputs link no entity, and those fall back to a zero dummy node app. B.2.
• Biomedical use. The authors caution against clinical use of the MedQA-trained models app. A.
• Ethical risk of the knowledge source. ConceptNet is reported to encode stereotypes, and the paper warns that using it outside benchmarks can propagate them app. A.

Lineage

• Builds on: GreaseLM (Zhang et al., 2022) for the encoder; DistMult (Yang et al., 2015) for link prediction; RoBERTa-Large and BioLinkBERT-Large as initial weights §2.2 §2.3 §3.1. None of these has a card in this wiki.
• Compared against: continued-MLM RoBERTa §3.3; QA-GNN (Yasunaga et al., 2021) and GreaseLM as fine-tuning-only KG models §3.3; BioBERT, PubmedBERT, BioLinkBERT and their KG variants tab. 6.
• Superseded by: not known from the sources used for this card.