KGLM: Integrating Knowledge Graph Structure in Language Models for Link Prediction
Year: 2023

Abstract:
The ability of knowledge graphs to represent complex relationships at scale has led to their adoption for various needs including knowledge representation, question-answering, and recommendation systems. Knowledge graphs are often incomplete in the information they represent, necessitating the need for knowledge graph completion tasks. Pre-trained and finetuned language models have shown promise in these tasks although these models ignore the intrinsic information encoded in the knowledge graph, namely the entity and relation types. In this work, we propose the Knowledge Graph Language Model (KGLM) architecture, where we introduce a new entity/relation embedding layer that learns to differentiate distinctive entity and relation types, therefore allowing the model to learn the structure of the knowledge graph. In this work, we show that further pretraining the language models with this additional embedding layer using the triples extracted from the knowledge graph, followed by the standard fine-tuning phase sets a new state-of-the-art performance for the link prediction task on the benchmark datasets.

heading: 1 Introduction
text: Knowledge graph (KG) is defined as a directed, multi-relational graph where entities (nodes) are connected with one or more relations (edges) (Wang et al., 2017). It is represented with a set of triples, where a triple consists of (head entity, relation, tail entity) or (h, r, t) for short, for example (Bill Gates, founderOf, Microsoft) as shown in Figure 1. Due to their effectiveness in identifying patterns among data and gaining insights into the mechanisms of action, associations, and testable hypotheses (Li and Chen, 2014; Silvescu et al., 2012), both manually curated KGs like DBpedia (Auer et al., 2007), WordNet (Miller, 1998), KIDS (Youn et al., 2022), and CARD (Alcock et al., 2020), and automatically curated ones like FreeBase (Bollacker et al., 2008), Knowledge Vault
Bill Gates
Microsoft
Washington
Melinda French
Jennifer Gates
founder Of
locatedIn
bornIn
dau ght
erO f
daughterOf
di vo
rc ed
W ith
Figure 1: Sample knowledge graph with 6 triples. The graph contains three unique entity types (circle for person, triangle for company, and square for location) and 5 unique relation types or 10 if considering both the forward and inverse relations. The task of the knowledge graph completion is to complete the missing links in the graph, e.g., (Bill Gates, bornIn?, Washington) using the existing knowledge graph.
(Dong et al., 2014), and NELL (Carlson et al., 2010) exist. However, these KGs often suffer from incompleteness. For example, 71% of the people in FreeBase have no known place of birth (West et al., 2014). To address this issue, knowledge graph completion (KGC) methods aim at connecting the missing links in the KG.
Graph feature models like path ranking algorithm (PRA) (Lao and Cohen, 2010; Lao et al., 2011) attempt to solve the KGC tasks by extracting the features from the observed edges over the KG to predict the existence of a new edge (Nickel et al., 2015). For example, the existence of the path Jennifer Gates daughterOf−−−−−−−→ Melinda French divorcedWith←−−−−−−−− Bill Gates in Figure 1 can be used as a clue to infer the triple (Jennifer Gates, daughterOf, Bill Gates). Other popular types of models are latent feature models such as TransE (Bordes et al., 2013), TransH (Wang et al., 2014), and RotatE (Sun et al., 2019) where entities and relations are converted into a latent space using embeddings.
217
TransE, a representative latent feature model, models the relationship between the entities by interpreting them as a translational operation. That is, the model optimizes the embeddings by enforcing the vector operation of head entity embedding h plus the relation embedding r to be close to the tail entity embedding t for a given fact in the KG, or simply h + r ≈ t.
Recently, pre-trained language models like BERT (Devlin et al., 2018) and RoBERTa (Liu et al., 2019) have shown state-of-the-art performance in all of the natural language processing (NLP) tasks. As a natural extension, models like KG-BERT (Yao et al., 2019) and BERTRL (Zha et al., 2021) that utilize these pre-trained language models by treating a triple in the KG as a textual sequence, e.g., (Bill Gates, founderOf, Microsoft) as ‘Bill Gates founder of Microsoft’, have also shown state-of-the-art results on the downstream KGC tasks. Although such textual encoding (Wang et al., 2021) models are generalizable to unseen entities or relations (Zha et al., 2021), they still fail to learn the intrinsic structure of the KG as the models are only trained on the textual sequence. To solve this issue, a hybrid approach like StAR (Wang et al., 2021) has recently been proposed to take advantage of both latent feature models and textual encoding models by enforcing a translation-based graph embedding approach to train the textual encoders. Yet, current textual encoding models still suffer from entity ambiguation problems (Cucerzan, 2007) where an entity Apple, for example, can refer to either the company Apple Inc. or the fruit. Moreover, there are no ways to distinguish forward relation (Jennifer Gates, daughterOf, Melinda French) from
inverse relation (Melinda French, daughterOf -1, Jennifer Gates).
In this paper, we propose the Knowledge Graph Language Model (KGLM) (Figure 2), a simple yet effective language model pre-training approach that learns from both the textual and structural information of the knowledge graph. We continue pre-training the language model that has already been pre-trained on other large natural language corpora using the corpus generated by converting the triples in the knowledge graphs as textual sequences, while enforcing the model to better understand the underlying graph structure and by adding an additional entity/relation-type embedding layer. Testing our model on the WN18RR dataset for the link prediction task shows that our model improved the mean rank by 21.2% compared to the previous state-of-the-art method (51 vs. 40.18, respectively). All code and instructions on how to reproduce the results are available online.1

heading: 2 Background
text: Link Prediction. The link prediction (LP) task, one of the commonly researched knowledge graph completion tasks, attempts to predict the missing head entity (h) or tail entity (t) of a triple (h, r, t) given a KG G = (E,R), where {h, t} ∈ E is the set of all entities and r ∈ R is the set of all relations. Specifically, given a single test positive triple (h, r, t), its corresponding link prediction test dataset can be constructed by corrupting either the head or the tail entity in the filtered setting (Bordes et al., 2013) as
1https://github.com/ibpa/KGLM
D(h,r,t)LP = {(h, r, t’) | t′ ∈ (E − {h, t}) ∧ (h, r, t′) /∈ D} ∪{(h’, r, t) | h′ ∈ (E − {h, t}) ∧ (h′, r, t) /∈ D}
∪{(h, r, t)}, (1)
whereD = Dtrain∪Dval∪Dtest is the complete dataset. Evaluation of the link prediction task is measured with mean rank (MR), mean reciprocal rank (MRR), and hits@N (Rossi et al., 2021). MR is defined as
MR =
∑ (h,r,t)∈Dtest rank((h, r, t) | D(h,r,t)LP )
|Dtest| ,
(2) where rank(·|·) is the rank of the positive triple among its corrupted versions and |Dtest| is the number of positive test triples. MRR is the same as MR except that the reciprocal rank 1/rank(·|·) is used. Hits@N is defined as
hits@N =
∑ (h,r,t)∈Dtest { 1, if rank((h, r, t) | D(h,r,t)LP ) < N 0, otherwise
|Dtest| ,
(3) where N ∈ {1, 3, 10} is commonly reported. Higher MRR and hits@N values are better, whereas, for MR, lower values denote higher performance.

heading: 3 Proposed Approach
text: In this work, we propose to continue pre-training, instead of pre-training from scratch, the language model RoBERTaLARGE (Liu et al., 2019) that has already been trained on English-language corpora of varying sizes and domains, using both the forward and inverse knowledge graph textual sequences (Figure 2). Following the convention used in the KG-BERT and StAR (see Appendix A), we use a textual representation of a given triple, e.g., (Bill Gates, founderOf, Microsoft) as ‘Bill Gates founder of Microsoft’, to generate the pre-training corpus. However, instead of extracting only the forward triple as done in the previous work, we extract both the forward and inverse versions of the triple, e.g., (Jennifer Gates, daughterOf, Bill Gates) and (Bill Gates, daughterOf -1, Jennifer Gates), where the -1
notation denotes the inverse direction of the corresponding relation.
To enforce the model to learn the knowledge graph structure, we introduce a new embedding layer entity/relation-type embedding (ER-type embedding) in addition to the pre-existing token and position embeddings of RoBERTa as shown in Figure 2. This additional layer aims to embed the tokens in the input sequence with its corresponding entity/relation-type, where the set of entities E in the knowledge graph can have tE different entity types depending on the schema of the knowledge graph, (e.g., tE = 3 for person, company, and location in Figure 1). Note that many knowledge graphs do not specify the entity types, in which case tE = 1. For the set of relations R, there exist tR = 2nR, where nR is the number of unique relations in the knowledge graph and the multiplier of 2 comes from forward and inverse directions (e.g., tR = 10 for the sample knowledge graph in Figure 1).
In this work, we propose three different variations of ER-type embeddings. KGLMBase is the simplified version where all entities are assigned a single entity type and relations are assigned either forward or inverse relation type regardless of their unique relation types, resulting in a total of 3 ER-type embeddings. The KGLMGR is a version with granular relation types with tR + 1 ER-type embeddings. The KGLMGER is the most granular version where we utilize all tE + tR ER-type embeddings. In other words, all entity types as well as all relation types including both directions are considered.
To be specific, we convert a triple (h, r, t) to a sequence of tokens w(h,r,t) = ⟨[s]whawrbwtc[/s] : a ∈ {1..|h|} & b ∈ {1..|r|} & c ∈ {1..|t|}⟩ ∈ R(|h|+|r|+|t|+2), where [s] and [/s] are special tokens denoting beginning and end of the sequence, respectively. The input to the RoBERTa model is then constructed by adding the ER-type embedding t(h,r,t) and the p(h,r,t) position embeddings to the
w(h,r,t) token embeddings, as
X(h,r,t) = w(h,r,t) + p(h,r,t) + t(h,r,t). (4)
Unlike the segment embeddings in the KG-BERT and StAR that were used to mark the input tokens with either the entity (se) or relation (sr), the ERtype embedding now replaces its functionality. Finally, we pre-train the model using the masked language model (MLM) training objective (Liu et al., 2019).
For fine-tuning, we extend the idea of how the KG-BERT scores a triple (see Equation 6 in Appendix A) to take advantage of the ER-type embeddings learned in our pre-training stage. For a given target triple, we calculate the weighted average score of both directions as
scoreKGLM (h, r, t) = αSeqCls(X(h,r,t))+
(1− α)SeqCls(X(t,r−1,h)), (5)
where SeqCls(·) is a RoBERTa model transformer with a sequence classification head on top of the pooled output (last layer hidden-state of the [CLS] token followed by dense layer and tanh activation function), (t, r−1, h) denotes the inverse version of (h, r, t), and 0 ≤ α ≤ 1 denotes the weight used for balancing the scores from forward and inverse scores. For example, α = 1.0 considers only the forward direction score.

heading: 4 Experiments and Results

heading: 4.1 Datasets
text: We tested our proposed method on three benchmark datasets WN18RR, FB15k-237, and UMLS as shown in Table 1. WN18RR (Dettmers et al., 2018) is derived from WordNet (Miller, 1998), a large English lexical database of semantic relationships between words, FB15k-237 (Toutanova and Chen, 2015) is extracted from Freebase (Bollacker et al., 2008), a large community-drive KG of general facts about the world, and UMLS contains biomedical relationships. WN18RR and FB15k237 are subsets of WN18 (Bordes et al., 2013) and FB15k (Bordes et al., 2013), respectively, where the inverse relation test leakage problem, i.e. the problem of inverted test triples appearing in the training set, has been corrected.

heading: 4.2 Settings
text: We used RoBERTaLARGE (Liu et al., 2019), a BERTLARGE-based architecture with 24 layers, 1024 hidden size, 16 self-attention heads, and 355M parameters, for the pre-trained language model as it has been shown in a previous study to perform better than BERT (hits@1 0.243 vs. 0.222 and MR 51 vs. 99, link prediction on WN18RR) (Wang et al., 2021). For pre-training, we used learning rate = 5e-05, batch size = 32, epoch = 20 (WN18RR), 10 (FB15k-237), and 1,000 (UMLS),
and AdamW optimizer (Loshchilov and Hutter, 2017). For fine-tuning training data, we sampled 10 negative triples for a positive triple by corrupting both the head and tail entity 5 times each. We used the validation set to find the optimal learning rates = {1e− 06, 5e− 07}, batch size = {16, 32}, epochs = {1, 2, 3, 4, 5} for WN18RR and FB15k237 and 25, 50, 75, 100 for UMLS, and α from 0.0 to 1.0 with an increment of 0.1. For all experiments, we set α = 0.5 based on the WN18RR validation set performance. Both pre-training and fine-tuning were performed on 3 × Nvidia Quadro RTX 6000 GPUs in a distributed manner using the 16-bit mixed precision and DeepSpeed (Rasley et al., 2020; Rajbhandari et al., 2020) library in the stage-2 setting. We used the Transformers library (Wolf et al., 2019).

heading: 4.3 Link Prediction Results
text: The hypothesis behind the KGLM was that learning the ER-type embedding layers in the pre-training stage using the corpus generated by the knowledge graph, followed by fine-tuning has the best performance. To test our hypothesis, we broke down the hypothesis into two separate claims. For the first claim, we only continued pre-training RoBERTaLARGE followed by fine-tuning without the ER-type embeddings. This test removes the contribution from the ER-type embeddings and solely tests the performance gained by further pretraining the model with the knowledge graph as input. Table 3 shows that claim 1 falls behind the KGLMGR in all metrics except for hits @1 (0.331 vs. 0.330, respectively). For the second claim, we did not continue pre-training and instead used the RoBERTaLARGE pre-trained weights as-is. We then learned the ER-type embeddings in the fine-tuning stage. This test shows if the ER-type embeddings can be learned only during the fine-tuning stage. Table 3 shows that KGLMGR outperforms all of the metrics obtained using the second claim. This re-
sult shows that the combination of these two claims works in a non-linear fashion to maximize performance.
The results of performing link prediction on the benchmark datasets are shown in Table 2. Compared to StAR, which had the best performance on MR and hits@10 on WN18RR, KGLMGR outperformed all the metrics with 21.2% improved MR (40.18 vs. 51, respectively) and 4.5% increased hits@10 (0.709 vs. 0.741, respectively). Although still inferior compared to the graph embedding approaches, KGLMGR has 35.8% improved hits@1 compared to the best language model-based approach StAR (0.243 vs. 0.330, respectively). Across all model types, KGLMGR has the best performance on all metrics for WN18RR except for hits@1. Although we did not observe any improvement compared to StAR for the FB15k-237 dataset, we had the best performance on all metrics for UMLS with 21.2% improved MR than ComplEx (1.19 vs. 1.51, respectively). KGLMGR outperformed KGLMBase in all metrics.

heading: 5 Conclusion
text: In this work, we presented KGLM, which introduces a new entity/relation (ER)-type embedding layer for learning the structure of the knowledge graph. Compared to the previous language modelbased methods that only fine-tune for a given task, we found that learning the ER-type embeddings in the pre-training stage followed by fine-tuning resulted in better performance. In future work, we plan to further test the version of KGLM that takes into account entity types, KGLMGER, on domainspecific knowledge graphs like KIDS (Youn et al., 2022) with entity types in their schema.
Limitations
Although KGLM outperforms state-of-the-art models when the training set includes full sentences
(e.g., UMLS and WN18RR), the model performed similarly to the state-of-the-art in cases where the training dataset had only ontological relationships, such as the /music/artist/origin relation present in the FB15k-237 dataset. One major limitation of the proposed method is the long training and inference time, which we plan to alleviate by adopting Siamese-style textual encoders (Wang et al., 2021; Li et al., 2022) in future work.
Ethics Statement
The authors declare no competing interests.

heading: Acknowledgements
text: We would like to thank the members of the Tagkopoulos lab for their suggestions. This work was supported by the USDA-NIFA AI Institute for Next Generation Food Systems (AIFS), USDANIFA award number 2020-67021-32855 and the NIEHS grant P42ES004699 to I.T. J.Y. conceived the project and performed all experiments. Both J.Y. and I.T. wrote the manuscript. I.T. supervised all aspects of the project.

heading: A Previous Work
text: A.1 KG-BERT
KG-BERT (Yao et al., 2019) is a finetuning method that utilizes the base version of the pre-trained language model BERT (BERTBASE) (Devlin et al., 2018) as an encoder for entities and relations of the knowledge graph. Specifically, KG-BERT first converts a triple (h, r, t) to a sequence of tokens w(h,r,t) = ⟨[CLS]wha[SEP]wrb[SEP]wtc[SEP] : a ∈ {1..|h|} & b ∈ {1..|r|} & c ∈ {1..|t|}⟩, where wn denotes the nth token of either entity or relation, [CLS] and [SEP] are the special tokens, while |h|, |r|, and |t| denote the number of tokens in the head entity, relation, and tail entity, respectively. This textual token sequence is then converted to a sequence of token embeddings w(h,r,t) ∈ Rd×(|h|+|r|+|t|+4), where d is the dimension of the embeddings and 4 is from the special tokens. Then the segment embeddings s(h,r,t) = ⟨(se)×(|h|+2)(sr)×(|r|+1)(se)×(|t|+1)⟩, where se and sr are used to differentiate entities from relations, respectively, as well as the position embeddings p(h,r,t) = ⟨pi : i ∈ {1..(|h|+|r|+|t|+4)}⟩ are added to the token embeddings w(h,r,t) to form a final input representation X(h,r,t) ∈ Rd×(|h|+|r|+|t|+4) that is fed to BERT as input. Then, the score of how likely a given triple (h, r, t) is to be true is computed by
scoreKG-BERT(h, r, t) = SeqCls(X(h,r,t)). (6)
KG-BERT significantly improved the MR of the link prediction task compared to the previous stateof-the-art approach CapsE (Vu et al., 2019) (97 compared to 719, an 86.5% decrease), but suffered from poor hits@1 of 0.041 due to the entity ambiguation problem and lack of structural learning (Wang et al., 2021; Cucerzan, 2007).
A.2 StAR StAR (Wang et al., 2021) is a hybrid model that learns both the contextual and structural information of the knowledge graph by augmenting the structured knowledge in the encoder. It divides a triple into two parts, (h, r) and (t), and applies a Siamese-style transformer with a sequence classification head to generate u = Pool(X(h,r)) ∈ Rd×(|h|+|r|+3) and v = Pool(X(t)) ∈ Rd×(|t|+2), respectively, where Pool(·) is the output of the RoBERTa’s pooling layer. The first scoring module focuses on classifying the triple by applying a
scorecStAR(h, r, t) = Cls([u;u× v;u− u;v]), (7)
where Cls(·) is a neural binary classifier with a dense layer followed by a softmax activation function. The second scoring module then adopts the idea of how translation-based graph embedding methods like TransE learns the graph structure by minimizing the distance between u and v as
scoredStAR(h, r, t) = −||u− v||, (8) where || · || is the L2-normalization. During the training, StAR uses a weighted average of the binary cross entropy loss computed using scorecStAR(h, r, t) and the margin-based hinge loss computed using scoredStAR(h, r, t), whereas only the scorecStAR(h, r, t) is used for inference. This approach shows a new state-of-the performance over the metrics MR (51) and hits@10 (0.709), as well as significantly improving the hits@1 compared to the KG-BERT (0.041 to 0.243, a 492.7% increase).

heading: Future Work
text: In future work, we plan to further test the version of KGLM that takes into account entity types, KGLMGER, on domainspecific knowledge graphs like KIDS (Youn et al. , 2022) with entity types in their schema. Limitations
Although KGLM outperforms state-of-the-art models when the training set includes full sentences
(e. g. , UMLS and WN18RR), the model performed similarly to the state-of-the-art in cases where the training dataset had only ontological relationships, such as the /music/artist/origin relation present in the FB15k-237 dataset. One major limitation of the proposed method is the long training and inference time, which we plan to alleviate by adopting Siamese-style textual encoders (Wang et al. , 2021; Li et al. , 2022) in future work