[Paper Review] KG-BERT: BERT for Knowledge Graph Completion

Liang Yao, Chengsheng Mao+1 more|arXiv (Cornell University)|Sep 7, 2019

Advanced Graph Neural Networks43 references377 citations

TL;DR

KG-BERT fine-tunes BERT on textual representations of entities and relations to score triples, achieving state-of-the-art results on triple classification, link prediction, and relation prediction across multiple benchmarks.

ABSTRACT

Knowledge graphs are important resources for many artificial intelligence tasks but often suffer from incompleteness. In this work, we propose to use pre-trained language models for knowledge graph completion. We treat triples in knowledge graphs as textual sequences and propose a novel framework named Knowledge Graph Bidirectional Encoder Representations from Transformer (KG-BERT) to model these triples. Our method takes entity and relation descriptions of a triple as input and computes scoring function of the triple with the KG-BERT language model. Experimental results on multiple benchmark knowledge graphs show that our method can achieve state-of-the-art performance in triple classification, link prediction and relation prediction tasks.

Motivation & Objective

• Motivate knowledge graph completion as a task suffering from incompleteness and sparsity.
• Introduce KG-BERT, a framework that treats (h, r, t) triples as textual sequences fed to BERT.
• Demonstrate state-of-the-art performance on triple classification, link prediction, and relation prediction across benchmarks.
• Show that contextual, text-enhanced representations improve reasoning over KGs.

Proposed method

• Represent head, relation, and tail as text descriptions or names and pack them into a BERT input sequence.
• Fine-tune a pre-trained BERT model to produce a 2-class (positive/negative) score for triple plausibility (KG-BERTa).
• Fine-tune a variant that predicts the relation given two entities using a softmax over relations (KG-BERTb).
• Construct negative samples by replacing heads or tails to train a binary cross-entropy loss for triple classification.
• Use standard Adam optimization with BERT-based fine-tuning hyperparameters (batch size 32, lr 5e-5, dropout 0.1).
• Evaluate under three tasks: triple classification, link prediction, and relation prediction, using real-valued scores and standard metrics.

Experimental results

Research questions

• RQ1 Can KG-BERT accurately determine the plausibility of unseen triples (h, r, t)?
• RQ2 Can KG-BERT predict the correct relation between two given entities?
• RQ3 How does KG-BERT perform on link prediction and relation prediction compared to state-of-the-art KG embeddings?
• RQ4 What is the impact of using textual descriptions and contextualized embeddings on KG completion tasks?

Key findings

• KG-BERT(a) achieves 93.5% on WN11, 90.4% on FB13, and 91.9% avg. in triple classification (outperforming baselines).
• KG-BERT(a) attains the best mean ranks in link prediction on WN18RR and FB15k-237 among tested methods.
• KG-BERT(b) yields strong relation prediction performance, achieving the highest Hits@1 on FB15K in the reported results.
• KG-BERT demonstrates substantial improvements in triple classification due to context-aware, text-rich representations and attention over tokens.
• Attention visualizations show multi-head attention focusing on informative words linking entities and relations.