SimKGC: Simple Contrastive Knowledge Graph Completion with Pre-trained Language Models

In plain words

A knowledge graph stores facts as triples: a head entity, a relation and a tail entity. Many facts are missing. Completing the graph means guessing the missing tail for a given head and relation. Older methods learn one vector per entity from the graph's links. They work well, but they cannot handle entities they never saw during training. Text-based methods read each entity's name and description with a language model, which lets them handle new entities, but they have been less accurate. The paper's diagnosis is that text-based models learn from too few wrong answers per example, because every text has to go through a large language model.

SimKGC encodes the head and relation once, encodes candidate tails separately, and compares the two with a simple similarity score. That makes it cheap to reuse the same vectors as many wrong answers. It adds three kinds of extra wrong answers: the other examples in the batch, the tails from the last batches, and the head entity itself, which is a tricky wrong answer because it shares much of its text with the query. A learned temperature and a small margin sharpen the training signal. At evaluation, entities close to the head in the graph get a small score boost. The authors report that this beats earlier text-based and embedding-based methods on several benchmarks.

TL;DR

• SimKGC scores tail candidates for a (head, relation, ?) query with two BERT encoders that have separate weights, using cosine similarity §3.2.
• Training uses InfoNCE with three negative sources: in-batch, pre-batch and self-negatives. Hyperparameters are shared across datasets except learning rate and epochs §3.3 §3.5 §4.1.
• Optional re-ranking adds α = 0.05 to entities within k hops of the head, with k = 5 for WN18RR and 2 otherwise §3.4 app. A.
• Author-reported gains: WN18RR MRR 66.6 against 47.6 for RotatE; Wikidata5M 35.8 transductive and 71.4 inductive. FB15k-237 is not state of the art (33.6 against TuckER's 35.8) tab. 2 tab. 3.
• The released scripts do not all reproduce the headline configuration. See Paper vs code.

Problem

Text-based knowledge graph completion (KGC) lagged embedding-based methods such as TransE and RotatE, even with pre-trained language models §1. The authors attribute this to inefficient contrastive learning. RotatE's default setting trains 1000 epochs with 64 negatives on Wikidata5M, while the text-based KEPLER trains 30 epochs with one negative, because of the cost of its RoBERTa encoder §1.

Cross-encoders such as KG-BERT need |E| × 2 × |T| forward passes for inference, where |E| is the number of entities and |T| the number of test triples §3.5. Wikidata5M has about 4.6 million entities tab. 1; the authors estimate about 3000 hours for KG-BERT on the transductive setting §4.2. The method assumes that every entity has a textual name and, where available, a description. WN18RR and FB15k-237 use the text supplied by KG-BERT; Wikidata5M ships descriptions for all entities and relations §4.1.

Method

Every training triple is also added with an inverse relation, so only tail prediction is needed §3.1. The inverse relation is the original relation with the prefix "inverse " app. A.

1. Two encoders. Both start from the same pre-trained model, bert-base-uncased §4.1. They do not share parameters §3.2. Encoder A reads the head entity's description and the relation as a pair, separated by [SEP]. Encoder B reads the tail entity's description alone §3.2.
2. Pooling. Use the mean of the last-layer hidden states, then L2 normalisation, rather than the first token §3.2. The authors cite prior work that found mean pooling gives better sentence embeddings §3.2.
3. Score. Both vectors are unit length, so the cosine similarity is a dot product §3.2, eq. 1. Prediction takes the candidate with the highest score over all entities §3.2, eq. 2.
4. Negatives §3.3:
   • In-batch: the other tails in the same batch are negatives.
   • Pre-batch: tail vectors cached from earlier batches are negatives. The cache holds the tails of one or two previous batches, and its logits are multiplied by 0.5 §3.3 §4.1.
   • Self-negative: the head entity's own vector is one extra negative per row, as a hard negative §3.3.
   • False-negative mask: any candidate that is a known true tail for the same head and relation in the training graph is masked out §3.3.
   • Count: with batch size 1024 and two pre-batches, there are 1023 in-batch, 2048 pre-batch and 1 self-negative, 3072 in total §3.3.
5. Loss. InfoNCE with an additive margin γ subtracted from the positive score and a temperature τ. The temperature is learned, parameterised as log(1/τ), and initialised to 0.05 §3.5 §4.1 eq. 5.
6. Re-ranking, at evaluation only §3.4 eq. 4. Candidates within k hops of the head, on the training graph, get +α. The paper says the re-ranking does not apply to the inductive setting §5.2.
7. Inference §3.5. Encode every entity once with encoder B. Encode each query once with encoder A. Score with a dot product and retrieve the top-k with a similarity-search tool such as Faiss §3.5. The cost is |E| + 2|T| forward passes §3.5.

# SimKGC training step. The paper has no algorithm box; this is written from
# §3.2-§3.5 and the training code.  [§3.2] [§3.3] [§3.5]
# Input:  batch of B triples (h, r, t) with text for h and t
# Output: scalar loss

def training_step(batch, queue):
    # Encode head text joined with the relation, using encoder A.
    # Mean pooling, then L2 normalisation.  [§3.2] [code: models.py#L66-L69]
    hr_vectors = encode(encoder_a, batch.head_text, batch.relation)  # (B, d)

    # Encode tail text with encoder B, same pooling.  [§3.2]
    # [code: models.py#L71-L74]
    tail_vectors = encode(encoder_b, batch.tail_text)                # (B, d)

    # Encode head text alone with encoder B, for self-negatives.
    # The paper does not say which encoder is used.  [§3.3]
    # [code: models.py#L76-L79]
    head_vectors = encode(encoder_b, batch.head_text)                # (B, d)

    # In-batch scores: cosine similarity between query and every tail.
    # [§3.2] [eq. 1]
    logits = hr_vectors @ tail_vectors.T                             # (B, B)

    # Subtract the additive margin from the positive pair only.  [eq. 5]
    # [code: models.py#L92-L93]
    logits = logits - margin * identity(B)                           # (B, B)

    # Divide by the temperature tau, i.e. multiply by 1/tau.  [§3.5]
    # [code: models.py#L94]
    logits = logits / tau                                            # (B, B)

    # Mask in-batch tails that are known true tails for (h, r).
    # [§3.3] [code: triplet_mask.py#L12-L42]
    logits = mask_known_true(logits, batch)                          # (B, B)

    # Pre-batch negatives from cached tails of earlier batches.
    # Weight 0.5 on their logits.  [§3.3] [§4.1] [code: models.py#L117-L124]
    past_logits = hr_vectors @ queue.vectors.T * 0.5 / tau           # (B, P)

    # Mask cached tails that are true tails; applied only once the
    # cache is full.  [code: models.py#L125-L127]
    past_logits = mask_known_true(past_logits, batch, queue)         # (B, P)
    logits = concat([logits, past_logits], dim=1)                    # (B, B+P)

    # Self-negative: one column per row, the dot product with the head.
    # Dropped when (h, r, h) is a training triple.  [§3.3]
    # [code: models.py#L104-L109] [code: triplet_mask.py#L45-L52]
    self_logit = sum(hr_vectors * head_vectors, dim=1) / tau         # (B,)
    self_logit = where(is_known(h, r, h), -1e4, self_logit)          # (B,)
    logits = concat([logits, self_logit[:, None]], dim=1)  # (B, B+P+1)

    # InfoNCE over tails: the correct tail for row i is column i.  [eq. 5]
    loss = cross_entropy(logits, labels=arange(B))

    # Code also adds the reverse direction, over in-batch columns only.
    # The paper's loss has no such term.  [code: trainer.py#L159-L161]
    loss = loss + cross_entropy(logits[:, :B].T, labels=arange(B))

    # Store this batch's tail vectors for later batches.
    # [code: models.py#L129-L131]
    queue.push(tail_vectors.detach())
    return loss

Key equations, as in the paper:

Cosine similarity of the query and a tail, equal to the dot product of unit vectors eq. 1:

\cos(\mathbf{e}_{hr}, \mathbf{e}_t) = \mathbf{e}_{hr} \cdot \mathbf{e}_t

Prediction over all entities eq. 2:

\underset{t_i}{\mathrm{argmax}}\ \cos(\mathbf{e}_{hr}, \mathbf{e}_{t_i}),\ t_i \in \mathcal{E}

Negative set for (h, r): all tails in the batch, the pre-batch caches and the self-negative, excluding known true triples eq. 3, unnumbered in the source:

\mathcal{N}(h,r) = \{ t' \mid t' \in \mathcal{N}_{IB} \cup \mathcal{N}_{PB} \cup \mathcal{N}_{SN},\ (h,r,t') \notin \mathcal{G} \}

Graph re-ranking eq. 4:

\underset{t_i}{\mathrm{argmax}}\ \cos(\mathbf{e}_{hr},\mathbf{e}_{t_i}) + \alpha\,\mathbb{1}(t_i \in \mathcal{E}_k(h))

InfoNCE with additive margin γ and temperature τ eq. 5:

\mathcal{L} = -\log \frac{e^{(\phi(h,r,t)-\gamma)/\tau}}{e^{(\phi(h,r,t)-\gamma)/\tau} + \sum_{i=1}^{|\mathcal{N}|} e^{\phi(h,r,t_i')/\tau}}

Margin-loss ablation, not the main method eq. 6:

\frac{1}{|\mathcal{N}|}\sum_{i=1}^{|\mathcal{N}|}\max(0, \lambda + \phi(h,r,t'_i) - \phi(h,r,t))

with λ = 0.8 §5.1.

Inference, one query per direction, using the same encoders §3.5:

# Filtered tail prediction with graph re-ranking  [§3.4] [§3.5] [eq. 4]
# Input:  head text, relation, entity vectors E (N, d), training graph G
# Output: score per entity (N,)

def score_tails(head_text, relation):
    # Encode the query with encoder A, pooled and normalised.  [§3.2]
    query = encode(encoder_a, head_text, relation)                   # (d,)

    # Cosine scores against every entity, all encoded once.  [§3.5]
    scores = E @ query                                               # (N,)

    # Boost entities within k hops of the head on the training graph.
    # The head itself is in this set.  [eq. 4] [code: triplet.py#L108-L121]
    for node in k_hop_nodes(G, head, k):
        scores[node] += alpha                                        # 0.05

    # Filtered setting: remove other known true tails for (h, r).
    # [§4.1] [code: evaluate.py#L69-L80]
    for other in known_tails(head, relation):
        if other != gold_tail:
            scores[other] = -1
    return scores

Backward direction, for head prediction, reuses the same code with the relation prefixed by "inverse " triplet.py.

Hyperparameters

Results

Author-reported. Setup: datasets tab. 1, metrics MRR and Hits@1/3/10 under the filtered setting, averaged over tail and head prediction §4.1. Baselines are reported as published, with TransE and DKRL numbers taken from KEPLER's evaluation and RotatE numbers from the GraphVite benchmark §4.2. Hardware: 4 V100 GPUs §4.1.

Dataset statistics tab. 1:

Wikidata5M tab. 2:

tab. 2. Embedding methods cannot do inductive KGC tab. 2 caption. The inductive setting ranks about 7,475 candidates and the transductive setting about 4.6 million, so the inductive metrics are much higher tab. 2 caption.

WN18RR and FB15k-237 tab. 3. Rows marked † are copied from StAR tab. 3 caption:

Other results:

• Loss and negatives, WN18RR, batch 256 baseline tab. 4: InfoNCE with 255 negatives 64.4 MRR; InfoNCE with 5 negatives 48.8; margin loss with 255 39.5; margin loss with 5 38.0; margin-τ with 255 57.8.
• Re-ranking, Wikidata5M transductive tab. 6: with re-ranking MRR 35.8, without 35.5.
• MRR by relation cardinality, Wikidata5M tab. 8: transductive 1-1 30.4, 1-n 8.3, n-1 71.1, n-n 10.6; inductive 83.5, 71.1, 80.0, 54.7.
• Human evaluation, Wikidata5M transductive tab. 9: 100 wrong predictions sampled by H@1 were judged by annotators. Averaged over both directions, 49% were correct, 34% wrong and 17% unknown.
• Batch size, Wikidata5M transductive tab. 11: 256 gives 33.8 MRR, 512 34.6, 1024 35.3 (IB only). FB15k-237 tab. 12: 32.4, 32.7, 33.3.
• Margin γ on FB15k-237 tab. 13: 0 gives 33.4 MRR, 0.02 gives 33.6, 0.05 gives 33.6.
• MRR against the number of negatives, WN18RR, batch 1024 fig. 2: the text says it rises from 48.8 to 67.1.

Headline comparisons §4.2 tab. 2 tab. 3: WN18RR MRR 47.6 to 66.6; Wikidata5M transductive 29.0 to 35.8; Wikidata5M inductive 49.3 to 71.4. The paper says SimKGC is the first text-based method to beat embedding-based methods on these benchmarks §4.2.

Paper vs code

The repository has no LICENSE file and the README gives no licence. Code is cited by path and line only, and no code is quoted.

• Loss. Paper: one-direction InfoNCE over tail candidates §3.5 eq. 5. Code: the total loss adds a second cross-entropy, tail to head, over in-batch columns only trainer.py. The paper does not mention the second term.
• Warm-up. Paper: 400 steps tab. 10. Code: capped at a tenth of total steps trainer.py. The cap does not bind for these four settings (our arithmetic from tab. 1, with inverse triples).
• Pre-batch per script. Paper: two pre-batches §4.1; the headline rows for WN18RR and Wikidata5M are IB+PB+SN tab. 2 tab. 3. Code: WN18RR and Wikidata5M scripts pass --pre-batch 0 train_wn.sh train_wiki.sh. That is the IB+SN row (WN18RR 66.7, Wikidata5M 35.6 and 71.3) rather than the headline row. FB15k-237 passes 2 train_fb.sh. The shipped predictions match the IB+PB+SN row on all four datasets within rounding; average MRR is 0.6655 for WN18RR, 0.3581 for Wikidata5M transductive, 0.7140 for inductive and 0.3355 for FB15k-237 metrics.json metrics.json metrics.json metrics.json. The predictions README does not state which script produced them.
• Self-negative encoder. Paper: the head is a hard negative, with no encoder named §3.3. Code: head text is encoded with the tail encoder, and the self-negative logit is a dot product with the query models.py models.py.
• Relation text. Paper: "textual descriptions" of head and relation §3.2. Code: the relation field of the dataset file is passed as the second text of the pair, with no description step doc.py.
• Token limit. Paper: entity descriptions truncated to 50 tokens §4.1. Code: 50 is the maximum length of the whole pair, which includes the entity name and the relation for encoder A doc.py doc.py.
• Neighbour text. Paper: neighbour names are appended for WN18RR and FB15k-237 app. A. Code: only when the description is under 20 words, with at most 10 neighbours in sorted entity-id order. The gold tail is removed after the cap, and only during training doc.py doc.py triplet.py. The paper states none of these thresholds.
• Test-time neighbour text. Training removes the gold tail from neighbour text. Evaluation does not doc.py. The paper says nothing.
• Dropout. Paper: 0.1 tab. 10. Code: --dropout defaults to 0.1 but no code reads it config.py. Effective dropout is whatever the pretrained model's config sets, which this repository does not state.
• Re-rank set. Paper: the k-hop neighbours of the head §3.4 eq. 4. Code: breadth-first search over an undirected training graph triplet.py triplet.py. The search starts from the head, so the head's own score also gets +α rerank.py. If the set exceeds 100,000 nodes, the boost is silently dropped triplet.py.
• Filtered setting. Paper: known true triples from train, valid and test are filtered §4.1. Code: every *.txt.json in the training directory is used dict_hub.py evaluate.py. This matches if those files are the three splits.
• Rank ties. Code: rank is the position after a descending sort evaluate.py. The paper does not say how ties are broken.
• Evaluation wrapper. The sharded Wikidata5M script reuses the same filtering and re-ranking through eval_single_direction eval_wiki5m_trans.py eval_wiki5m_trans.py.
• No discrepancy found for: learning rates, epochs, batch size, γ = 0.02, τ initial 0.05 with a learned temperature, AdamW with linear decay, weight decay 1e-4, gradient clip 10, mean pooling, pre-batch logit weight 0.5, inverse "inverse " prefix, α = 0.05, 5 hops for WN18RR and 2 otherwise §4.1 app. A train_wn.sh train_fb.sh train_wiki.sh eval.sh config.py.

Pitfalls

• Pick the configuration deliberately. The shipped scripts for WN18RR and Wikidata5M disable pre-batch negatives, which the headline rows include train_wn.sh. The IB+SN rows differ from the headline rows by 0.1 MRR on WN18RR, 0.2 on Wikidata5M transductive and 0.1 on inductive tab. 2 tab. 3. The gaps are small, but the scripts and the headline are different configurations.
• Averaging hides a large direction gap. The reported MRR averages tail and head prediction §4.1. In the shipped Wikidata5M transductive predictions, forward MRR is 0.6295 and backward 0.0867, and the average is 0.3581 metrics.json. On FB15k-237 the same pattern is 0.4251 and 0.2459 metrics.json.
• The abstract's "+19%" is absolute. It is 19.0 MRR points, 47.6 to 66.6, not a relative gain §4.2 tab. 3. Relative to 47.6, it is about 40%. The Wikidata5M figures (+6.8 and +22) are also absolute points §4.2 tab. 2.
• Figure 2 and Table 4 disagree on a starting point. The text describes Figure 2 at batch 1024 starting from 48.8 §5.1 fig. 2. Table 4's 48.8 is five negatives at batch 256 tab. 4. The paper does not reconcile them. Re-run the sweep rather than trusting either endpoint.
• Queue cold start and checkpoints. The pre-batch cache starts from random unit vectors. Its entries are not masked until the cache is full models.py models.py. The cache is a non-persistent buffer plus a Python list, and the checkpoint saves only the state dict models.py trainer.py. A saved checkpoint therefore does not restore the cache.
• Batch size is a real dependency. Smaller batches hurt: on Wikidata5M transductive, MRR goes from 33.8 at batch 256 to 35.3 at 1024 with in-batch negatives only tab. 11. Multi-GPU runs use DataParallel only; the README says DistributedDataParallel is not supported because masks need the whole batch [README].
• Compute. Wikidata5M takes about 12 hours per epoch on 4 V100 GPUs app. A. Embedding about 4.6 million entities takes about 40 minutes on 2 GPUs §4.2.
• Human evaluation. Annotators judged 49% of wrong top-1 predictions correct on average tab. 9. The authors say automatic metrics underestimate performance §5.4. Do not compare these MRRs with other papers' scores without this caveat.
• FB15k-237 is where SimKGC trails. The authors attribute it to a denser graph (average degree about 37 with about 15k entities) and to unpredictable links §4.2. They note an ensemble with embedding methods as future work §4.2.
• Inductive setting. Re-ranking is disabled. The code asserts it rerank.py and the evaluation script sets the weight to zero eval.sh.
• Numbers are author-reported. The shipped predictions reproduce the headline table, but the repository commit used here is 97cc43e; the paper's camera-ready version was not checked against it.

Lineage

• Builds on: KG-BERT, the cross-encoder text-based baseline that SimKGC replaces with a bi-encoder §1 §3.3. Mean pooling follows SimCSE and Sentence-BERT §3.2. In-batch and pre-batch negatives follow dense passage retrieval and the pre-batch scheme of Lee et al. §3.3. The margin-τ variant resembles RotatE's self-adversarial sampling §5.1.
• Compared against: TransE, RotatE, DistMult, TuckER, DKRL, KEPLER, BLP, KG-BERT, MTL-KGC and StAR tab. 2 tab. 3.
• Superseded by: not in this wiki.