[Paper Review] Integrating Graph Contextualized Knowledge into Pre-trained Language Models

Bin He, Di Zhou+5 more|arXiv (Cornell University)|Nov 30, 2019

Topic Modeling23 references41 citations

TL;DR

The paper introduces KG-Transformer based knowledge representation learning that uses arbitrary knowledge subgraphs to inject graph contextualized knowledge into a BERT-based medical language model, achieving state-of-the-art results on medical NLP tasks.

ABSTRACT

Complex node interactions are common in knowledge graphs, and these interactions also contain rich knowledge information. However, traditional methods usually treat a triple as a training unit during the knowledge representation learning (KRL) procedure, neglecting contextualized information of the nodes in knowledge graphs (KGs). We generalize the modeling object to a very general form, which theoretically supports any subgraph extracted from the knowledge graph, and these subgraphs are fed into a novel transformer-based model to learn the knowledge embeddings. To broaden usage scenarios of knowledge, pre-trained language models are utilized to build a model that incorporates the learned knowledge representations. Experimental results demonstrate that our model achieves the state-of-the-art performance on several medical NLP tasks, and improvement above TransE indicates that our KRL method captures the graph contextualized information effectively.

Motivation & Objective

• Motivate learning richer node representations by modeling arbitrary subgraphs from knowledge graphs rather than single triples.
• Develop a Transformer-based KG-Transformer to learn joint entity and relation embeddings from subgraphs.
• Integrate the learned graph knowledge into a pre-trained language model to enhance medical NLP tasks.
• Demonstrate improved performance on entity typing and relation classification in the medical domain.

Proposed method

• Represent subgraphs as node sequences and convert relations into graph nodes to enable joint learning of entities and relations.
• Encode the graph-structured input with a Transformer-based KG-Transformer using a masked attention mechanism that respects graph connectivity.
• Train with a margin-based translation energy loss on restored triples to obtain knowledge embeddings.
• Pre-train a medical language model (BERT-MK) by integrating KG-derived embeddings with a large PubMed/medical corpus.
• Evaluate on entity typing and relation classification datasets to compare with BERT-family baselines and TransE-based approaches.

Experimental results

Research questions

• RQ1 Does modeling arbitrary subgraphs with a Transformer improve knowledge representations over triple-based approaches like TransE?
• RQ2 Can graph-contextualized knowledge embedded via KG-Transformer improve performance of pre-trained language models on medical NLP tasks?
• RQ3 How does KG-Transformer compare to traditional KRL methods in relation classification and entity typing?
• RQ4 What is the impact of pre-training data size on the effectiveness of graph-augmented language models?
• RQ5 How do learned entity and relation embeddings interact with a BERT-based medical language model in downstream tasks?

Key findings

• BERT-MK outperforms BERT-Base, BioBERT, and SCIBERT on average entity typing accuracy (97.26 vs 97.24, 97.02, and 97.26 respectively).
• On relation classification, BERT-MK achieves higher average F1 across datasets (81.23 P, 85.32 R, 83.02 F1) compared with the baselines.
• KG-Transformer-based knowledge representations yield a 0.9 and 0.64 higher F1 on two relation classification datasets compared to MedERNIE (TransE-based).
• BERT-MK shows continued advantage over MedERNIE as pre-training data increases, indicating the benefit of integrating richer medical knowledge.
• Using as little as 10–20% of the medical pre-training data, BERT-MK can outperform BioBERT, suggesting knowledge integration can reduce computational cost while improving performance.