Source-linked AI summary
K-BERT: Enabling Language Representation with Knowledge Graph
Weijie Liu, Peng Zhou, Zhe Zhao, Zhiruo Wang, Qi Ju, Haotang Deng, Ping Wang
TL;DR
Pre-trained language models lack domain-specific knowledge, while injecting too much knowledge can introduce semantic distortion. K-BERT injects knowledge-graph triples into BERT-compatible models and controls their influence with soft-position and visible-matrix mechanisms. Across twelve NLP tasks, it reports promising results, especially on domain-specific tasks, without requiring users to pre-train the model themselves.
Problem
Pre-trained language models lack domain-specific knowledge for knowledge-driven tasks, while excessive knowledge incorporation can divert sentences from their correct meaning.
Method
K-BERT injects knowledge-graph triples into sentences and uses soft-position and visible-matrix mechanisms to control knowledge influence while loading parameters from pre-trained BERT.
Results
Across twelve open- and specific-domain NLP tasks, K-BERT shows promising results and significantly outperforms BERT on domain-specific tasks.
Takeaways & Limitations
K-BERT can incorporate domain knowledge without user pre-training and supports knowledge-driven problems requiring domain expertise.
Takeaways & Limitations
Future work is needed to improve K-Query's filtering of unimportant triples and extend the approach to other language representation models.
Abstract
from arXiv · showhide
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.
Introduction
Pre-trained language models struggle with knowledge-driven domain tasks because open-domain pre-training leaves limited domain knowledge, while retraining for each domain is costly. K-BERT integrates knowledge graphs into BERT-compatible models without additional pre-training and reports stronger performance across domain-specific and open-domain tasks.
- Open-domain pre-training creates a domain discrepancy that limits models such as BERT on knowledge-driven tasks, including electronic medical record analysis.
- Experts use relevant domain knowledge for inference, whereas publicly provided language models rely primarily on contextual comprehension.
- Domain-specific pre-training could address this gap, but its time and computational costs make it unacceptable to most users.
- Knowledge graphs offer structured domain knowledge, avoiding the inefficient requirement for repeated textual co-occurrences during pre-training.
- K-BERT addresses heterogeneous embedding space and knowledge noise while loading parameters from pre-trained BERT and injecting domain knowledge without pre-training.
- Across twelve Chinese NLP tasks, K-BERT significantly outperforms BERT on domain-specific and many open-domain tasks.
Related Work
Prior work optimized BERT pre-training or combined knowledge graphs with word vectors and language models, but these approaches left important integration and scalability limitations. K-BERT is positioned against this background of limited KG fusion and inefficient or memory-intensive alternatives.
- BERT extensions such as ERNIE, BERT-WWM, SpanBERT, and RoBERTa primarily optimize pre-training, masking, or encoder objectives.
- Research on fusing pre-trained language models with knowledge graphs remains limited.
- THU-ERNIE fuses entity information but ignores relations between entities, while COMET trains GPT on knowledge-graph triples inefficiently.
- Earlier joint word-entity representation methods combine entities and words in a shared vector space using word2vec- and TransE-related approaches.
- For knowledge graphs with millions of entities, joint representation can create an entity table too large for GPU memory.
Methodology
K-BERT injects KG triples into a sentence tree, preserves structural information through soft positions, and restricts attention with a visible matrix. Its Mask-Transformer uses this structure to limit knowledge-driven semantic changes while retaining useful enrichment.
- Knowledge layer: K-BERT’s knowledge layer queries triples for sentence entities and stitches them into corresponding positions, producing a sentence tree.The tree can contain multiple branches, but its depth is fixed at 1, so injected entities do not recursively derive further branches.
- Embedding layer: The embedding layer flattens the sentence tree by hard-position order, then combines token, soft-position, and segment embeddings.Soft positions restore the original structural order after branch tokens are inserted into the token sequence.
- Seeing layer: The seeing layer constructs a visible matrix that prevents tokens in unrelated branches from directly affecting one another, limiting knowledge noise.For example, [Apple] and [China] are mutually invisible, while [Apple] can enrich [Cook] through their permitted visibility relationship.
- Mask-Transformer: Mask-self-attention applies the visible matrix to restrict each token’s attention region within the Transformer.When a token is invisible to another token, its attention score is masked so it contributes nothing to the target hidden state.
- Mask-Transformer: The Mask-Transformer stacks multiple mask-self-attention blocks and uses the visible matrix as structural information from the sentence tree.The attention computation includes the matrix M and the scaling factor √d_k.
Experiments
Experiments compare K-BERT with BERT across open-domain and specific-domain Chinese NLP tasks, using multiple corpora, knowledge graphs, and matched model settings. Results show that task-appropriate knowledge graphs improve domain-oriented performance, while soft-position and visible-matrix mechanisms reduce knowledge-noise effects.
- Experimental setup: Experiments evaluate K-BERT on twelve Chinese NLP tasks, including eight open-domain and four specific-domain tasks.The evaluation uses train, development, and test splits and compares K-BERT with Google BERT and a reimplemented BERT.
- Experimental setup: K-BERT and BERT use matched configurations with L = 12, A = 12, H = 768, and 110M trainable parameters.This controls parameter count when reflecting on the role of the knowledge graph.
- Open-domain tasks: Knowledge graphs have task-dependent effects: sentiment analysis gains little, HowNet suits semantic similarity, and CN-DBpedia suits Q&A and NER.The results support choosing the knowledge graph according to task type.
- Open-domain tasks: CN-DBpedia raises MSRA-NER F1 from 93.6% to 95.7%, whereas adding WebtextZh raises it only to 94.6%.The reported comparison attributes the larger improvement to the knowledge graph rather than the additional corpus.
- Specific-domain tasks: On specific-domain tasks, K-BERT with CN-DBpedia improves F1 by 1∼2% across all tasks compared with BERT.The experiments cover finance and law Q&A and finance and medicine NER; MedicalKG produces an especially obvious gain on Medicine NER.
- Ablation studies: Removing soft-position or the visible matrix lowers performance, while K-BERT without the visible matrix becomes worse than BERT on Law Q&A.The ablation results also report K-BERT peaking at epoch 2 versus epoch 4 for BERT, indicating faster convergence in that comparison.
Conclusions
K-BERT injects knowledge-graph information into sentences while using soft-position and a visible matrix to limit knowledge from changing the original meaning. It produces promising results across twelve open- and specific-domain NLP tasks, while future work targets context-based triple filtering and broader model compatibility.
- Conclusions: K-BERT injects knowledge from a KG into a sentence and controls its scope with soft-position and a visible matrix.The resulting knowledge-rich sentence tree is designed to prevent injected knowledge from deviating from the sentence’s original meaning.
- Conclusions: K-BERT shows promising results on twelve open- and specific-domain NLP tasks, with KG information especially helpful for knowledge-driven domain tasks.The authors position the approach for problems requiring domain expertise.
- Conclusions: Future work includes filtering unimportant triples based on context and extending K-BERT to other language-representation models.The proposed extensions include K-Query improvements and application to models such as ELMo and XLNet.