Sequence-to-Sequence Knowledge Graph Completion and Question Answering

In plain words

A knowledge graph stores facts as triples, such as "Barack Obama, born in, United States". Most link-prediction models give every entity its own vector, so the model grows with the number of entities. This paper takes a different route: it writes a question such as "predict the tail of (Barack Obama, born in)" as plain text and trains a small text-to-text Transformer to write the answer's name. Ranking candidate answers is done by letting the model generate many answers and keeping the ones it produces, so no score has to be computed for every entity. The same model is then fine-tuned to answer natural-language questions over knowledge graphs that are missing facts.

The model is small and needs no per-entity table, and the authors report competitive link-prediction scores on large graphs. On question answering it does well on several datasets, but it is weak at recalling facts it saw during training, so the authors suggest pairing it with a traditional method.

TL;DR

KGT5 is an encoder-decoder with the T5-small architecture, trained from scratch on link prediction over verbalised triples, then fine-tuned for KGQA with a link-prediction regulariser §3 §4.3. Author-reported: on Wikidata5M it reaches MRR 0.300 with about 60M parameters, against 0.308 for ComplEx with 614M tab. 2. On incomplete-KG QA it beats the compared baselines on CWQ and WQSP, and matches or beats them on MetaQA 1-hop and 3-hop tab. 4 tab. 5 §4.5. Use it when a parameter-light, text-based link predictor or KGQA model is wanted and the KG has textual names for every entity. Do not use it as a memory of the KG.

Problem

Conventional knowledge-graph embedding (KGE) models learn one embedding per entity and relation. Their model size and inference cost grow linearly with the number of entities, which matters for graphs such as Wikidata with about 90 million entities §1. Such embeddings also need extra modules before they can serve question answering, so KGQA systems built on them have several stages §1 §2.1 §2.2. Earlier text-based approaches, KG-BERT's cross-encoder in particular, are not scalable §1.

The authors list four desiderata for a KGE model: scalability, quality, versatility and simplicity. They argue that existing methods satisfy at most some of them, and summarise this in a comparison table §1 tab. 9.

The setting for KGQA is an incomplete KG. The authors randomly drop 50% of the edges of each background KG, following PullNet, so the answer is not always reachable by graph traversal §4.1. The model is trained on (question, answer entity) pairs only, not on ground-truth SPARQL queries §4.1.

Method

Textual representations §3.1 app. A. Each entity and relation gets a unique text mention. For Wikidata-based KGs, the mention is the canonical Wikidata title. Identical titles are disambiguated by appending a one-line description or, failing that, the Wikidata id app. A. Freebase-based QA datasets (WQSP, CWQ) use the English name, disambiguated with an alias or a description where available app. A. For QA-only datasets the one-to-one mapping is not enforced, because surface forms are what the evaluation compares §3.1.

Verbalisation §3.1. A query (s, p, ?) becomes an input string made of a task prefix, the subject mention and the relation mention, separated by a bar. The example in the text is predict tail: barack obama | born in, whose target output is united states §3.1. The README describes the same shape as <prefix>:<head entity><sep><relation> with the tail as output [code: README.md#L9]. The head-prediction prefix and the separator are not stated in the compiled text.

Link-prediction training §3.2 app. B. Each training triple (s, p, o) gives two sequences: one for (s, p, ?) with target o, and one for (?, p, o) with target s §3.2. The model is trained with teacher forcing and token-level cross-entropy. There is no negative sampling: the loss pushes probability onto the true next token at each decoding step, and this is repeated for every token of the target mention §3.2 app. B. The per-sequence loss averages the token losses over the target length T app. B.

# Build (input, target) pairs from one training triple  [§3.1] [§3.2]
# Input:  triple (s, p, o); mention(x) gives the text mention of x
# Output: two (input text, target text) pairs

def build_pairs(s, p, o, mention):
    pairs = []
    # Tail query (s, p, ?): target is the mention of o  [§3.1] [§3.2]
    tail_input = "predict tail: " + mention(s) + " | " + mention(p)
    pairs.append((tail_input, mention(o)))
    # Head query (?, p, o): target is the mention of s  [§3.2]
    # The head prefix string is not stated in the paper; HEAD_PREFIX
    # is a placeholder for the implementer to choose.
    head_input = HEAD_PREFIX + mention(o) + " | " + mention(p)
    pairs.append((head_input, mention(s)))
    return pairs

# Teacher-forced token loss for one pair  [§3.2] [app. B]
# Input:  input_ids (N,); target_ids = [w_1 .. w_T] (T,); model
# Output: scalar loss for this pair

def pair_loss(model, input_ids, target_ids):
    # The decoder is fed the true previous target tokens  [app. B]
    logits = model(input_ids, shift_right(target_ids))  # (T, vocab)
    # Cross-entropy at each position, averaged over the T positions
    # [app. B]
    return mean(cross_entropy(logits, target_ids))

Link-prediction inference by sampling §3.3 app. C. For a query (s, p, ?), the model's verbalised input is decoded by sampling: at each step a token is drawn from the output distribution until a stop token is produced. This is repeated a fixed number of times, 500 for link prediction §4.3. Each sample is mapped back to an entity id. The score of an entity is the sum of the log-probabilities of its tokens app. C. Entities that were never sampled get score minus infinity, so they rank last §3.3. Decoded strings that are not entity mentions are not filtered; the authors found that more samples were enough, and they did not use constrained decoding §3.3, footnote.

# Sampled link prediction for one query (s, p, ?)  [§3.3] [app. C]
# Input:  model; s, p; num_samples (500 in the paper)  [§4.3]
# Output: dict entity -> log-probability; unsampled entities are absent

def predict_tails(model, s, p, num_samples):
    text = "predict tail: " + mention(s) + " | " + mention(p)  # [§3.1]
    scores = {}
    for _ in range(num_samples):
        # Draw tokens one at a time until the stop token  [app. C]
        tokens, logprob = sample_until_stop(model, text)
        entity = entity_for_mention(tokens)  # None if not an entity
        if entity is None:
            continue
        # Repeated samples: the paper does not say how to combine them.
        # This sketch keeps the first score. Choose and document a rule.
        scores.setdefault(entity, logprob)
    # Entities missing from scores rank with score minus infinity  [§3.3]
    return scores

KGQA training §3.4 §4.3. The model is first trained on link prediction over the background KG. It is then fine-tuned on questions: the input is the prefix predict answer: followed by the question, and the target is the mention of the answer entity §3.4. To reduce overfitting on small KGs, every batch holds equal numbers of QA sequences and link-prediction sequences drawn at random from the background KG §3.4. The QA batch size is half the link-prediction batch size, and the learning rate is fixed at 0.001 §4.3.

# One fine-tuning batch for KGQA  [§3.4] [§4.3]
# Input:  qa_pairs (question, answer mention); lp_pairs from the KG
# Output: shuffled batch with equal QA and link-prediction parts

def make_qa_batch(qa_pairs, lp_pairs, half_size):
    # Half the batch: QA pairs with the KGQA prefix  [§3.4]
    qa = [("predict answer: " + q, a)
          for q, a in sample(qa_pairs, half_size)]
    # Other half: random link-prediction sequences from the KG  [§3.4]
    lp = sample(lp_pairs, half_size)
    return shuffle(qa + lp)

KGQA inference §3.4 §4.3 eq. 1. Beam search with beam size 4 produces candidate answers with their probabilities. A neighbourhood bonus then re-ranks them: a candidate a inside the n-hop neighbourhood N(e) of the question's topic entity e gets alpha added to its probability eq. 1. The authors use n = 1, 2 or 3 but do not say which per dataset, and alpha is described only as a constant hyperparameter whose value is not stated §4.3. Re-ranking is applied only to datasets whose test questions carry topic-entity annotations §4.3.

score(a) = \begin{cases} p_a + \alpha & a \in \mathcal{N}(e) \\ p_a & \text{otherwise} \end{cases} eq. 1

# Neighbourhood re-ranking of beam candidates  [eq. 1] [§4.3]
# Input:  candidates = [(answer, p_answer)] from beam size 4  [§4.3]
#         topic entity e; neighbourhood(e) gives N(e); alpha (not stated)
# Output: predicted answer entity

def rerank_answer(candidates, e, alpha):
    best_answer = None
    best_score = None
    for answer, p_answer in candidates:
        # Bonus only inside the n-hop neighbourhood of e  [eq. 1]
        bonus = alpha if answer in neighbourhood(e) else 0.0
        score = p_answer + bonus
        if best_score is None or score > best_score:
            best_answer = answer
            best_score = score
    return best_answer

Architecture and output size §4.3 app. A. The model is a T5-small encoder-decoder (about 60M parameters) with the same architecture but no pretrained weights §4.3. It is trained on the verbalised KG, with a SentencePiece BPE vocabulary trained on those verbalised strings §4.3. Entity identity is the output string, so the model has no entity embedding table; the ranking is over sampled strings §3.3.

Hyperparameters

Sampling versus beam search for link prediction app. C fig. 3: sampling more often keeps raising MRR as the sample size grows. Beam search with beam size 5 or more does worse as the beam grows. Beam size 50 gives far worse results than 50 samples app. C. The beam-versus-sample comparison is shown in fig. 3, described in words here.

Results

Author-reported. Hardware and KG splits are described in §4.1 and §4.3. Scores are single-model unless marked as an ensemble.

Link prediction, Wikidata5M tab. 2. Filtering protocol is not stated in the compiled caption. Ensemble rule: a query with no answer in the training graph goes to KGT5, otherwise to ComplEx (614M). KGT5 answered 42% of the ensemble's queries §4.4 §4.4 footnote.

Table 2 sources: TransE to QuatE are the best pre-trained Graphvite models; ComplEx 14-dim and 26-dim come from a LibKGE hyperparameter search; ComplEx (614M) is from Kochsiek et al. 2021; KEPLER and DKRL are from Wang et al. 2021; MLMLM is from Clouâtre et al. 2021 tab. 2 caption. The text says KGT5 beats all low-parameter models on all metrics and is 0.008 MRR below ComplEx and 0.012 Hits@1 above it §4.4. The table's own values agree with that.

Parameter reduction, from the listed counts: 60M is 97.5% below 2,400M, 90% below 614M, and 99.7% below the 18.2B of the WikiKG90Mv2 baselines tab. 2 tab. 3. The abstract's "up to 98%" is closest to the first of these comparisons (97.5%).

Link prediction by query type, Wikidata5M tab. 13. Columns are the number of other correct answers in training: none, 1 to 10, more than 10.

The ComplEx checkpoint used here is slightly weaker than the SOTA one [tab. 13 caption].

Train versus test MRR, Wikidata5M tab. 8.

Link prediction, WikiKG90Mv2 (validation MRR) tab. 3. KGT5 has no test score: OGB-LSC did not evaluate it on the hidden test set because the entity mentions used were removed §4.4 footnote.

Baseline numbers come from the OGB-LSC leaderboard tab. 3 caption.

Link prediction on small KGs tab. 10. The paper says these are inferior to the baselines on FB15k-237 and YAGO3-10, and that WN18RR may be helped by the entity definitions tab. 10 caption §4.8.

KGQA, Freebase-based datasets, hits@1 with gain over ground-truth query (GT query) tab. 4. 50% KG setting. Baseline numbers come from LEGO [tab. 4 caption]. For WQSP the KG matches the baselines' KG; for CWQ it is somewhat harder tab. 4 caption.

KGQA, MetaQA, hits@1 with gain over GT query tab. 5. 50% KG setting. Baseline numbers come from LEGO. There are two GT query rows because the baseline KG differs from the authors' tab. 5 caption.

The KGT5 2-hop gain as printed (−8.2) does not match its own GT query value (36.2 − 48.7 = −12.5). The other rows agree with their GT values. See Pitfalls.

Neighbourhood re-ranking ablation, 50% KG tab. 6. The text says the gains are small but consistent on all datasets §4.5.

Full-KG KGQA, hits@1 tab. 7. Setting: the full KG rather than the 50% setting tab. 7 caption. The T5-small row is the comparison discussed in §4.7; the KGT5 rows are discussed in §4.8.

Dataset statistics. KG statistics tab. 1: WikiKG90Mv2 91M entities, 1,387 relations, 601M edges; Wikidata5M 4.8M, 828, 21M; MetaQA 43k, 9, 70k; WQSP subset 158k, 816, 376k; CWQ subset 3.9M, 326, 6.9M. QA question counts tab. 11: MetaQA 1-hop 96,106 / 9,992 / 9,947 (train / valid / test); 2-hop 118,980 / 14,872 / 14,872; 3-hop 114,196 / 14,274 / 14,274; WQSP 2,998 / 100 / 1,639; CWQ 27,639 / 3,519 / 3,531.

Paper vs code

The paper's footnote links https://github.com/apoorvumang/kgt5, so this is the official repository. It is cloned at commit c527e0d847c43cc0a3fa8d8304c774bbe0cc1b6a (default branch, shallow clone) under Apache-2.0 LICENSE. The README says the main branch covers link prediction on Wikidata5M alone and reports hits@1 without filtering, that a branch apoorv-dump holds the latest code, and that the data "is yet to be uploaded", while giving download links for the KGC and KGQA data README.md README.md. The main branch has no KGQA training or evaluation code, and the QA model import in dataset.py is commented out dataset.py. The apoorv-dump branch was not checked, so the KGQA details below come from the paper alone.

• Tokenizer. The paper trains a SentencePiece BPE on the verbalised KG §4.3. The code tokenises with the off-the-shelf t5-small tokenizer dataset.py. The vocabulary sizes in tab. 1 therefore may not describe the code's tokenizer.
• Model. The paper builds a T5-small-sized model with no pretrained weights §4.3. The code takes the t5-small configuration and initialises random weights main_accelerate.py. This matches the paper.
• Batch size. The paper states 320 §4.3. The code default is 64 main_accelerate.py, and the README's multi-GPU command passes 64 with six processes README.md. The global batch that the authors used is not stated.
• Learning rate and schedule. The paper gives no peak learning rate or warmup length for link prediction §4.3. With no --learning_rate, the code uses Adafactor with relative step size and warmup initialisation main_accelerate.py. With a learning rate given, it turns both off and uses that fixed rate main_accelerate.py. The Noam scheduler is imported but never instantiated in the training script main_accelerate.py. The operative_config.gin file is a Mesh-TensorFlow T5 span-corruption config with a Noam warmup of 10000 steps operative_config.gin. No Python file in the repository reads it.
• Stopping rule. The paper stops when validation accuracy has not risen for 10k steps §4.3. The training loop runs a fixed number of epochs (default 5) and saves a checkpoint every save_steps (default 5000), with no validation pass main_accelerate.py main_accelerate.py.
• Loss. The code uses the library's token-level loss on the labels, with padding set to −100 and ignored main_accelerate.py dataset.py. This matches the teacher-forced cross-entropy described in app. B up to the averaging convention, which the paper does not state for a batch.
• Inference. The paper ranks entities by sampling 500 times §4.3. The repository's evaluation does greedy decoding and compares the exact string with the target eval_accelerate.py eval_accelerate.py eval_accelerate.py. Beam search is available only through --beam_size and its default is 1 eval_accelerate.py. No sampling-based ranking is implemented here. The README says the evaluation is hits@1 unfiltered README.md. MRR and Hits@10 in tab. 2 are not computed by the main branch.
• Evaluation split. The default split is test eval_accelerate.py. Table 2's split is not stated in its caption.
• Dropout. The paper states 10% §4.3. No script sets a dropout value; the model uses whatever the Hugging Face configuration defaults to. This was not verified against the library version.
• Input and target length. The paper does not state these. The code uses at most 128 input and 32 target tokens dataset.py. Entity strings are padded to 32 tokens dataset.py.
• Entity text. The code loads entity mentions from data/<dataset>/entity_strings.txt dataset.py. That file is not in the repository, so the disambiguation rule in app. A cannot be checked here.
• Data split and KGQA. The paper says the authors' KG split is released with the implementation §4.1. The main branch does not include the split. The README points to download links for the KGC and KGQA data README.md.
• Successor. The README names KGT5-context, an extension with a new codebase that improves link prediction README.md. It is not compiled on this wiki.

Pitfalls

• Filtering is not stated for Table 2. The caption that said "filtered evaluation setting" is commented out in the source tab. 2. The public code computes unfiltered hits@1 only README.md. Do not compare the table's MRR with a filtered or unfiltered number without checking which one you have.
• Repeated samples. The paper samples 500 times and scores each sampled entity by its sequence log-probability §3.3 app. C, but does not say how repeated samples of the same entity are combined. The sketch in Method keeps the first one, which is a choice, not the paper's.
• Unmatched outputs. The sampled text is not always an entity mention. The authors rely on sampling more to avoid losing predictions rather than constraining decoding §3.3 footnote. Expect some samples to map to no entity, and count how many.
• Sampling and beam size. Beam search with larger beams hurts (MRR falls past beam size 5), and sampling is the chosen method app. C [fig. 3]. Changing the ranking method changes the results.
• Head prediction at inference. The compiled text describes tail queries at inference §3.3. It does not say how head queries are ranked. The head-prediction prefix is also not stated §3.1. Choose and document both.
• Ensemble results. The Wikidata5M and WikiKG90M ensemble numbers use ComplEx (614M) for queries with known answers §4.4 footnote. They are not single-model results and not comparable with the 60M model's scores.
• Fewer parameters, lower scores on small KGs. tab. 10 shows KGT5 behind ComplEx and NBFNet on FB15k-237 and YAGO3-10. On WN18RR it is ahead of ComplEx but behind NBFNet, and the authors partly attribute that margin to entity definitions §4.8 tab. 10 caption.
• Memorisation. Train MRR is 0.304 for KGT5 against 0.721 for ComplEx tab. 8. On the full KG it scores 56.1 on WQSP and 36.5 on CWQ tab. 7. The paper recommends against using it as a standalone KGQA method and suggests using it only when query parsing does not give good results §4.8.
• MetaQA 2-hop. KGT5 is weaker on 2-hop questions, and the printed gain of −8.2 does not agree with its own GT query value (36.2 − 48.7 = −12.5) tab. 5. Recompute from the absolute values before quoting it. The paper says the drop is worst for questions whose head and answer have the same type §4.5.
• KG splits. Baseline numbers come from other authors' KG splits, which were not published, so only gains over each method's own GT query are comparable across papers §4.1 tab. 5 caption.
• Textual names. KGT5 needs a text mention for every entity, so it cannot be applied directly to Freebase with CVT nodes. The paper flags possible definition leakage on WN18RR and text-processing shortcuts in WikiKG90Mv2 validation triples, and says comparisons must account for them §4.8.
• Unstated values. The alpha bonus, the neighbourhood size n per dataset, the peak learning rate, the warmup length and the QA batch size in absolute terms are not stated §4.3 eq. 1. Use the code only where it exists, and say so.
• Table and figure numbers. The source bundle has no compiled PDF. Numbers here follow the order of \begin{table} and \begin{figure} in the LaTeX source, and were not checked against a rendered PDF.
• Scale. Wikidata5M has about 21M edges and WikiKG90Mv2 about 601M [tab. 1]. Training runs were of 5M steps on large KGs §4.3 footnote. The paper reports 4 GTX 1080 Ti GPUs §4.3; no wall-clock time is given.

Lineage

• Builds on: the T5 encoder-decoder architecture, which the paper uses without its pretrained weights §1 §4.3. T5 is not compiled on this wiki.
• Compared against, link prediction: KGE models TransE, DistMult, SimplE, RotatE, QuatE and ComplEx tab. 2; the text-based KEPLER, DKRL and MLMLM §2.1; NBFNet on small KGs §4.2.
• Compared against, KGQA: PullNet, EmQL, EmbedKGQA and LEGO §4.2; CBR-KGQA in the full-KG table tab. 7. PathPred is the authors' own template-based traversal baseline app. D.
• Extended by: KGT5-context, named in the README README.md.
• Superseded by: not stated in the paper.