KGLM : Integrating Knowledge Graph Structure in Language Models for Link Prediction (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 fine-tuned 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 pre-training 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.

## introduction knowledge graph (kg) is defined as a directed, multi-relational graph where entities (nodes) are connected with one or more relations (edges) @xcite . 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 @xcite , both manually curated kgs like dbpedia @xcite , wordnet @xcite , kids @xcite , and card @xcite , and automatically curated ones like free-base @xcite , knowledge vault 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. @xcite , and nell @xcite exist. however, these kgs often suffer from incompleteness. for example, 71% of the people in freebase have no known place of birth @xcite . 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) @xcite 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 @xcite . 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, daugh-terof, bill gates). other popular types of models are latent feature models such as transe @xcite , transh @xcite , and ro-tate @xcite where entities and relations are converted into a latent space using embeddings. generate pre-training data input knowledge graph pre-training corpus of hop 1 e2 e3 e4 e0 e1 r2 r4 r 3 r 1 r1 r0 e0 r1 e1 triple 1 e0 r1 e1 triple 2 triple 12 e3 r4 e4 pre-train language model token embeddings position embeddings entity/relation type embeddings pre-trained language model e tok e4 e [s] e tok r4 e tok e3 e 0 e 4 e 2 e 1 e 3 e e e e e r4 - foot_0 e [/s] figure 2 : proposed pre-training approach of the kglm. first, both the forward and inverse triples are extracted from the knowledge graph to serve as the pre-training corpus. we then continue pre-training the language model, roberta in our case, using the masked language model training objective, with an additional entity/relation-type embedding layer. the entity/relation-type embedding scheme shown here corresponds to the kglm ger , the most fine-grained version where both the entity and relation types are considered unique. note that the inverse relation denoted by -1 is different from its forward counterpart. for demonstration purposes, we assume all entities and relations to have a single token. 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 @xcite and roberta @xcite have shown state-of-the-art performance in all of the natural language processing (nlp) tasks. as a natural extension, models like kg-bert @xcite and bertrl @xcite 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 @xcite models are generalizable to unseen entities or relations @xcite , 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 @xcite 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 @xcite 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 2 background 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 @xmath0, r), where {h, t} ∈ e is the set of all entities and @xmath1. 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 @xcite as ## proposed approach in this work, we propose to continue pre-training, instead of pre-training from scratch, the language model roberta large @xcite ) 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 t e different entity types depending on the schema of the knowledge graph, (e.g., t @xmath2, company, and location in figure 1 ). note that many knowledge graphs do not specify the entity types, in which case t @xmath3. for the set of relations r, there exist t @xmath4, where n r is the number of unique relations in the knowledge graph and the multiplier of 2 comes from forward and inverse directions (e.g., t @xmath5. to be specific, we convert a triple (h, r, t) to a sequence of tokens w (h,r,t) |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 wn18rr fb15k-237 umls method hits @1 hits @3 hits @10 mr mrr hits @1 hits @3 hits @10 mr mrr hits@10 mr model type: not based on language models transe .043 .441 .532 2300 .243 .198 .376 .441 323 .279 .989 1.84 transh .053 .463 .540 2126 .279 .306 .450 .613 219 .320 --distmult .412 .470 .504 7000 .444 .199 .301 .446 512 .281 .846 5.52 complex .409 .469 .530 7882 .449 .194 .297 .450 546 .278 .967 2.59 conve .390 .430 .480 5277 .46 .239 .350 .491 246 .316 .990 1.51 rotate .428 .492 .571 3340 .476 .241 .375 .533 177 .338 --gaat .424 .525 .604 1270 .467 .512 .572 .650 187 .547 --lineare .453 .509 .578 1644 .495 .264 .391 .545 155 .357 --quatde .438 .509 .586 1977 .489 .268 .400 .563 90 .365 --model type: based on language models kg-bert .041 .302 .524 97 .216 --.420 153 -.990 1.47 star .243 .491 .709 51 .401 .205 .322 .482 117 .296 .991 1.49 kglm base .305 .518 .730 47.97 .445 -------kglm gr .330 .538 .741 40.18 .467 .200 .314 .468 125.9 .289 .995 1.19 unlike the segment embeddings in the kg-bert and star that were used to mark the input tokens with either the entity (s e ) or relation (s r ), the ertype embedding now replaces its functionality. finally, we pre-train the model using the masked language model (mlm) training objective @xcite . 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 ## conclusion 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, kglm ger , on domainspecific knowledge graphs like kids @xcite with entity types in their schema.