K-BERT: Enabling Language Representation with Knowledge Graph

• Open-domain pre-trained models may fail on domain-specific tasks, while training a domain-focused model is computationally expensive.Source: Introduction · p. 1
• Knowledge graphs offer structured domain knowledge that can improve domain-task performance while reducing large-scale pre-training costs and enabling manual editing.Source: Introduction · p. 1
• K-BERT injects KG knowledge while addressing heterogeneous embedding spaces and knowledge noise, and can load pre-trained BERT parameters without additional pre-training.This compatibility makes it convenient for users with limited computing resources.Source: Introduction · pp. 1–2
• K-BERT significantly outperforms BERT on domain-specific tasks and many open-domain tasks across twelve Chinese NLP tasks.Source: Introduction · p. 2

• BERT optimization focused on pretraining, including whole-word and span masking and RoBERTa’s changes, or encoder changes such as XLNet’s Transformer-XL and THU-ERNIE’s word-entity aggregator.Source: Related Work · p. 2
• Fusion of knowledge graphs with pre-trained language representation was limited: THU-ERNIE (Zhang et al. 2019) fused entities but ignored relations, while COMET (Bosselut et al. 2019) trained GPT on KG triples inefficiently.Source: Related Work · p. 2
• Before pre-trained language models, KG–word-vector work jointly represented entities and words, textual relations, knowledge bases, or cross-lingual words and entities (Wang et al. 2014; Toutanova et al. 2015; Han, Liu, and Sun 2016; Cao et al. 2018).Source: Related Work · p. 2
• These methods relied on “word2vec + transE” rather than pre-trained language models, retained HES problems, and could require entity tables exceeding GPU memory for KGs with millions of entities.Source: Related Work · p. 2

• The section details K-BERT’s implementation and presents its overall framework in Figure 1.Source: Methodology · p. 2

Notation

• A sentence s = {w0, w1, ..., wn} contains vocabulary tokens, with English tokenized by word and Chinese by character; each KG triple ε = (wi, rj, wk) links two entity names via relation rj.Here, n is the sentence length, and the relation is also in the vocabulary V.Source: Notation · p. 2

Model architecture

• K-BERT comprises a knowledge layer, embedding layer, seeing layer, and mask-transformer; the knowledge layer injects relevant KG triples to form a knowledge-rich sentence tree.The tree is then fed to the embedding and seeing layers and converted into token-level embeddings and a visible matrix.Source: Model architecture · pp. 2–3
• Tokens are flattened by hard-position index, while soft-position indices provide position embeddings alongside token embeddings.Tokens in the first sentence receive segment label “A”.Source: Figure 2 · p. 3
• The visible matrix controls each token’s visible area, preventing excessive injected knowledge from changing the original sentence’s meaning.For example, the figure shows that Apple(4) cannot see China(9).Source: Model architecture · p. 3, Figure 2 · p. 3

• KL converts an input sentence into a sentence tree by querying a knowledge graph for entity-linked triples and injecting them at their corresponding sentence positions.The process consists of K-Query, which selects sentence entities and collects corresponding triples, and K-Inject, which stitches those triples into the sentence.Source: Knowledge layer · p. 3
• The sentence tree can have multiple branches, but its depth is fixed at one, preventing entities in injected triples from deriving branches iteratively.Source: Knowledge layer · p. 3

Embedding layer

• The embedding layer converts a sentence tree into a representation for the Mask-Transformer, combining token, position, and segment embeddings as in BERT.Unlike BERT, K-BERT’s input to this layer is a sentence tree rather than a token sequence.Source: Embedding layer · p. 3

Token embedding

• K-BERT uses BERT’s vocabulary and trainable token embeddings, but rearranges branch tokens after their corresponding nodes before embedding.Subsequent tokens are shifted back; the resulting sequence can become unreadable and lose structural information, which soft-position and visible matrix address.Source: Token embedding · pp. 3–4

Soft-position embedding

• Soft-position embeddings restore sentence order after tree rearrangement by assigning tokens positions that reflect their original relationships.This can give unrelated tokens the same position, a problem addressed by Mask-Self-Attention in the next subsection.Source: Soft-position embedding · p. 4
• For multiple input sentences, K-BERT joins them with [SEP] and uses segment tags A and B to distinguish them.The combined sequence begins with [CLS].Source: Segment embedding · p. 4

Seeing layer

• The visible matrix prevents unrelated tokens from interacting, so [China] cannot affect [Apple] and [CLS] cannot bypass [Cook] to access [Apple].These restrictions are intended to reduce semantic changes caused by injected knowledge.Source: Seeing layer · p. 4
• Matrix M marks tokens as visible to one another when they are in the same branch; i and j are hard-position indices.Source: Seeing layer · p. 4

Mask-Transformer

• Mask-Transformer modifies BERT’s encoder, which cannot accept M, to limit self-attention regions according to the visible matrix and stacks multiple mask-self-attention blocks.The number of layers is denoted L, hidden size H, and number of attention heads A.Source: Mask-Transformer · p. 4

Mask-Self-Attention:

• Invisible tokens receive zero attention contribution, while visible knowledge can enrich an intermediate representation without directly affecting the original sentence’s meaning.In the example, [Apple] enriches [Cook]’s representation but does not directly affect [CLS] because [Apple] is invisible to it.Source: Mask-Self-Attention · pp. 4–5

Pre-training corpora

• Pre-training uses WikiZh, with 1 million entries and 120 million sentences, and WebtextZh, with 4.1 million Q&A entries across 28,000 topics.WikiZh is the Chinese Wikipedia corpus; WebtextZh is a Chinese question-and-answer corpus.Source: Pre-training corpora · p. 5

Knowledge graph

• The experiments use CN-DBpedia, HowNet, and MedicalKG; the refined KGs contain 5.17 million, 52,576, and 13,864 triples, respectively.MedicalKG is a self-developed Chinese medical concept KG covering symptoms, diseases, parts, and treatments, and is open source as part of K-BERT.Source: Knowledge graph · p. 5

Parameter settings and training details

• K-BERT and the BERT baselines share 110M parameters and settings of L = 12, A = 12, and H = 768; K-BERT uses BERT’s pre-trained parameters and enables KG only during fine-tuning and inference.The baselines are Google BERT pretrained on WikiZh and the authors’ BERT reimplementation pretrained on WikiZh and WebtextZh.Source: Parameter settings and training details · p. 5, Baselines · p. 5

Domain Q&A

• K-BERT is evaluated against BERT on eight Chinese open-domain tasks, including single-sentence classification, two-sentence classification, Q&A matching, and NER, alongside specific-domain finance and law Q&A and finance and medicine NER.Finance and Law Q&A use about 770,000 and 36,000 Baidu Zhidao samples, respectively. Finance NER comprises 3,000 financial news articles with over 65,000 labeled entities; Medicine NER extracts medical-related entity names from electronic medical records.Source: Open-domain tasks · p. 5, Experiments · p. 5, Domain Q&A · p. 6

Open-domain tasks

• Open-domain results show no significant KG effect on sentiment tasks, HowNet performs better for XNLI and LCQMC, and CN-DBpedia suits NLPCC-DBQA and MSRA-NER better.The authors conclude that KG choice should depend on task type.Source: Table 1 · p. 6, Table 2 · p. 6, Open-domain tasks · p. 6
• On MSRA-NER, CN-DBpedia raises F1 from 93.6% to 95.7%, while WebtextZh raises it to 94.6%.The additional corpus also boosts performance, but less than the KG.Source: Table 2 · p. 6, Open-domain tasks · p. 6

Domain NER

• K-BERT with CN-DBpedia improves F1 by 1–2% across all specific-domain tasks; MedicalKG yields an especially obvious gain on Medicine NER.The paper attributes the improvements to domain knowledge in the KG.Source: Domain NER · pp. 6–7, Table 3 · p. 7

Ablation studies

• Removing soft-position or the visible matrix lowers K-BERT performance; without the visible matrix, Law Q&A falls below BERT, while K-BERT peaks at epoch 2 versus BERT at epoch 4.The authors conclude that both mechanisms improve robustness to knowledge noise and enable more efficient use of knowledge.Source: Figure 5 · p. 7, Ablation studies · p. 7

Conclusions

• K-BERT injects KG knowledge into a sentence to form a knowledge-rich sentence tree, then uses soft-position and a visible matrix to limit knowledge scope and preserve the sentence’s original meaning.Source: Conclusions · p. 7
• K-BERT shows promising results on twelve open- and specific-domain NLP tasks, with KG especially helpful for knowledge-driven specific-domain problems requiring domain expertise.K-BERT can use available pre-trained BERT parameters without users pre-training the model themselves.Source: Conclusions · p. 7
• Future work includes improving K-Query to filter unimportant triples based on context and extending the approach to other language representation models, including ELMo and XLNet.Source: Conclusions · p. 7