[Paper Review] K-BERT: Enabling Language Representation with Knowledge Graph

Weijie Liu, Peng Zhou+5 more|arXiv (Cornell University)|Sep 17, 2019

Topic ModelingComputer Science21 references84 citations

TL;DR

K-BERT injects domain knowledge from knowledge graphs into BERT-like models using a knowledge-enabled sentence tree with a soft-position and visible matrix to control knowledge influence, improving domain-specific NLP tasks without extra pre-training.

ABSTRACT

Pre-trained language representation models, such as BERT, capture a general language representation from large-scale corpora, but lack domain-specific knowledge. When reading a domain text, experts make inferences with relevant knowledge. For machines to achieve this capability, we propose a knowledge-enabled language representation model (K-BERT) with knowledge graphs (KGs), in which triples are injected into the sentences as domain knowledge. However, too much knowledge incorporation may divert the sentence from its correct meaning, which is called knowledge noise (KN) issue. To overcome KN, K-BERT introduces soft-position and visible matrix to limit the impact of knowledge. K-BERT can easily inject domain knowledge into the models by equipped with a KG without pre-training by-self because it is capable of loading model parameters from the pre-trained BERT. Our investigation reveals promising results in twelve NLP tasks. Especially in domain-specific tasks (including finance, law, and medicine), K-BERT significantly outperforms BERT, which demonstrates that K-BERT is an excellent choice for solving the knowledge-driven problems that require experts.

Motivation & Objective

• Motivate the need for domain knowledge in language representations beyond open-domain pre-training.
• Propose a KG-enabled BERT variant that injects structured knowledge without requiring full pre-training on KG data.
• Offer a mechanism (soft-position and visible matrix) to mitigate knowledge noise and preserve sentence meaning.
• Demonstrate effectiveness across multiple open-domain and domain-specific Chinese NLP tasks.

Proposed method

• Introduce K-BERT, which loads a pre-trained BERT and augments inputs with KG triples to form a knowledge-rich sentence tree.
• Use a knowledge layer to query KG triples for sentence entities (K-Query) and inject them into the sentence (K-Inject).
• Transform the sentence tree into a token-level embedding via an embedding layer, adopting a soft-position embedding to preserve structure.
• Employ a seeing layer with a visible matrix to regulate which tokens can attend to which others, implemented as a Mask-Transformer that restricts self-attention according to visibility.
• Train with the same parameter budget as BERT (12 layers, 12 attention heads, 768 hidden units) and enable KG only during fine-tuning/inference (no KG pre-training).
• Evaluate on twelve Chinese NLP tasks spanning open-domain and domain-specific domains (finance, law, medicine).

Experimental results

Research questions

• RQ1 Can KG-enabled injection improve domain-specific language understanding without re-pretraining on KG data?
• RQ2 How does integrating KG knowledge affect performance on open-domain versus domain-specific tasks?
• RQ3 Do controllable mechanisms (soft-position and visible matrix) mitigate knowledge noise while preserving or enhancing sentence semantics?
• RQ4 Which types of KGs (encyclopedic vs. language-focused) are more beneficial for different task categories (NLI, QA, NER, sentiment)?
• RQ5 Is K-BERT compatible with standard BERT pre-trained parameters and practical for resource-constrained settings?

Key findings

• K-BERT shows superior performance on domain-specific tasks (finance, law, medicine) compared to BERT, validating the value of KG-enabled knowledge. CN-DBpedia (encyclopedic KG) improves Q&A and NER tasks, while HowNet (language-based KG) better supports semantic similarity tasks like XNLI and LCQMC; WebtextZh provides additional gains when used with KG.
• In specific-domain results, K-BERT with CN-DBpedia improves MSRA-NER F1 from 93.6% to 95.7% (approx. +2.1 percentage points) and similarly boosts other domain tasks; MedicalKG yields notable gains for Medicine_NER.
• Ablation studies show that removing soft-position or the visible matrix reduces performance, and omitting the visible matrix can even degrade performance below BERT in some cases, confirming KN mitigation is essential.
• K-BERT converges faster than BERT in Law_Q&A (peak at epoch 2 vs. epoch 4 for BERT) and remains robust to KN through the visible matrix and soft-position mechanisms.
• The approach allows leveraging existing BERT pre-trained parameters without KG-specific pre-training, facilitating practical adoption with limited resources.