Better Together: Enhancing Generative Knowledge Graph Completion with Language Models and Neighborhood Information (2023)

Abstract: Real-world Knowledge Graphs (KGs) often suffer from incompleteness, which limits their potential performance. Knowledge Graph Completion (KGC) techniques aim to address this issue. However, traditional KGC methods are computationally intensive and impractical for large-scale KGs, necessitating the learning of dense node embeddings and computing pairwise distances. Generative transformer-based language models (e.g., T5 and recent KGT5) offer a promising solution as they can predict the tail nodes directly. In this study, we propose to include node neighborhoods as additional information to improve KGC methods based on language models. We examine the effects of this imputation and show that, on both inductive and transductive Wikidata subsets, our method outperforms KGT5 and conventional KGC approaches. We also provide an extensive analysis of the impact of neighborhood on model prediction and show its importance. Furthermore, we point the way to significantly improve KGC through more effective neighborhood selection.

## introduction knowledge graphs (kg) represent structured information in the form of a directed multi-relational graph with entities as the graph's nodes and relations between entities as the graph's edges. entities in such graphs could be both real-world objects or abstract concepts. kgs are represented as a set of triplets with representation entities and a relation between them, i.e. (head entity, relation, tail entity) or (h, r, t) for short. due to kgs' recent rapid development, they are now widely used in a variety of applications, including data mining, information retrieval, question answering, and recommendation systems. however, most real-world knowledge graphs suffer from incompleteness. knowledge graph completion (kgc) algorithms attempt to fill in the gaps in the kg in order to solve this problem. kg link prediction commonly involves learning knowledge graph embeddings (kge), which capture structural and semantic information about kg nodes, relations, and their neighborhoods in a way that a simple distance-based function would be able to distinguish real triplets from false ones. however, building a unique embedding for each node and relation leads to a linear increase in model size and inference time as the kg size grows. another drawback is that kg updating requires full re-training of kges, rather than simply learning new embeddings for updated entities or relations. application of language models (lms) to kgc tasks is a promising direction @xcite . state-of-the-art performance in most natural language processing (nlp) tasks has recently been demonstrated by pre-trained language models like bert @xcite and more recent large or foundational language models @xcite @xcite . language models-based approaches, such as kg-bert @xcite , bertrl @xcite , kgt5 @xcite and genkgc @xcite consider a triplet in the kg as a text sequence. while distance-based methods incorporate more structural information in dense embeddings, language models tend to capture more semantic information contained in both relations and entities of kg, enabling such models to embed previously unseen entities and operate on new facts based on semantic information contained in triplets textual representation. thus, kgt5 positions kg link prediction as sequence-to-sequence tasks. kgt5 model trained to generate target node using textually verbalized triplets achieves competitive performance and offers a scalable and simple framework for kgc and qa tasks. current methods based on language models (kg-bert, bertrl, kgt5, and genkgc) rely only on information that can be extracted from lms parameters. however, since kg itself is available, we build a method that relies on both internal language model knowledge and external information from kg. for the link prediction task (h, r, ?), we propose to extract nodes and relations adjacent to h from the kg, i.e. the neighborhood. node h with its neighborhood better together represents h itself: neighborhood may contain clues such as palo alto and united states of america for the triplet (steve jobs, place of birth, ?) on figure 1 . though some efforts have been made to incorporate neighborhood into the kgc algorithms @xcite , our proposed pipeline inherits the benefits of generative llms compared to kge approaches: (1) scalability and size of the model, (2) applicability to both transductive and inductive kgs due to an ability to generalize unseen entities, (3) no need to rank all possible candidate triplets due to direct generation of tail entity. also, a concurrent work by the kgt5 authors, kgt5context @xcite , proposed a similar idea of integrating node neighborhood in the context of the generative lm model, supporting the main concerns and results of our study. our contributions are as follows: (1) we propose a generative graph completion method which effectively integrates language models with graph neighborhood information. (2) we show that incorporating graph information into multiple generative language models achieve strong performance improvement of link prediction in both transductive and inductive settings. moreover, our approach achieves state-of-the-art results on the inductive ilpc dataset. (3) we analyze and evaluate the impact of adding neighborhood information on the performance of the models. ## methods one can represent kg more formally as @xmath0, e, r), where e is a set of entities, r is a set of relations and t is a set of triplets @xmath1. therefore, the formulation of the kgc task is the following: for each triplet t @xmath2, r, t), where @xmath3, @xmath4, @xmath5, answer both (h, r, ?) and (?, r, t) queries on the graph g/t i . the kgc tasks can have different variations depending on the kg type: inductive and transductive. in a transductive setting the kgc task is formulated as follows: given a train kg g @xmath6, e, r) train the model to predict on the inference kg g inf @xmath7, e, r), therefore in the train and inference kg there is a same set of entities and relations. in inductive kgc the model is trained on the kg g @xmath8, r, e train ) and predicts the inference kg g inf @xmath9, r, e inf erence ), so the trained model observed only on the entities from train subgraph of kg. we followed the kgt5 @xcite approach to the kgc task using an encoder-decoder transformer based on t5 @xcite . the link prediction task was posed as follows: for each triplet (entity 1 , relation, entity 2 ) from kg we had to predict the tail in the queries (entity 1 , relation, ?) and (entity 2 , relation -1 , ?). to generate input for the model, we converted such queries into a text form. apart from the query verbalization only as in kgt5, we formed a neighborhood for the node and relations for each triplet in kg. thus, verbalized queries and their neighborhoods in the form of text were fed to the input of the model for training (see figure 1 ) and later to predict incomplete triplets. we form a neighborhood of nodes that have an edge with the triplet head node (1-hop neighbors). the number of adjacent nodes may be too large to fit into the context size of selected language models. therefore, we sort them based on relation semantic similarity. we take only those verbalized nodes and relations that fit into the max sequence length. we provide more details on neighborhood formation in appendix b and verbalization in appendix c. to evaluate our approach to link prediction tasks we used publically available transductive and inductive kgs. specifically, we used the wiki-data5m @xcite dataset, which is one of the largest publicly available transductive kg up-to-date and large and small versions of ilpc dataset @xcite , the largest fully inductive benchmark. ilpc inference graph size is challenging for modern gnn methods with inductive capabilities @xcite . apart from being challenging to solve by common approaches, the ilpc kgs are much more sparse than the transductive wikidata5m. the graph size and sparseness of the data make the task hard for existing methods and result in modest performance metrics. still, there were no existing solutions that beat the baseline provided in the ilpc paper. both used datasets are the subsets of the real-world kg -wikidata. datasets' statistics are provided in the we compared models with and without neighborhoods to evaluate our hypothesis that neighborhoods help to better together represent graph node. in addition, unlike the kgt5, we experimented with different model configurations, both fine-tuning an existing pre-trained model and training from scratch as common knowledge obtained during language models pre-training might be useful on real-world knowledge graphs like wikidata. we used the t5-small with 60m parameters as our base model with both pretrained and randomly initialized weights to evaluate the effect of pretraining of the language model on the kg task. we added additional "[sep]" tokens to the tokenizer with 30k tokens vocabulary to separate query and neighborhood. the model was trained using adamw optimizer @xcite , 1e-5 learning rate, 10% dropout, and batch size equals 320. the number of training steps was equal to 5m for the wikidata5m dataset and 2m for all others. for multi-gpu training, we used the horovod parallelization framework @xcite . in cases where a batch size equal to 320 did not fit gpu, we used gradient accumulation steps. we also benchmarked ope-nai gpt-3.5-turbo foot_0 zero-shot link prediction and evaluate the benefits of a neighborhood on llm (see appendix g). the code for all the described experiments is available at the link foot_1 . as kgt5, our model produces a probability distribution over all possible tokens in the vocabulary at each step of decoding. such distribution is penalized for not matching the actual distribution, i.e., for the difference between the real target token at position i and the generated token at step i using ce loss. such an approach is principally different from distance-based methods as it does not require generating negative samples to make them further from the target prediction for each training query. in this way, training iteration becomes independent of the kg size. to form predictions, we first verbalize the query (head, relation, ?) and pass it as an input to the generative model to produce an output sequence, which will be a verbalization of the target entity id. for inference, we employed two strategies. the first one was sampling k most probable sequences from the decoder. we follow the kgt5 @xcite setting and use a sample size of 50. we apply this approach to wikidata5m datasets as it is transductive. the second option for inference was to search for the closest entities in kg to the generated one. we apply the second approach to ilpc inductive datasets as language models have to generate previously unseen entities. more details on inference procedure are in appendix d. ## results the hypothesis behind our approach was that neighborhood information incorporated in a context fed to the language model should enhance link prediction of generative lm. tables 1 and 2 demonstrate the performance of our technique on wikidata5m and ilpc datasets under various conditions, including training a t5 model with and without neighbors, using a pre-trained model or training a model from scratch, and zero-shot with gpt-3.5-turbo. these results support the hypothesis that the quality of link prediction increases with the use of neighbors in the context of the generative model in both transductive and inductive settings. moreover, our approach demonstrated sota performance on the ilpc dataset, while on wikidata we achieved the best result among generative link prediction models, which is still comparable to scores of sota graph and lm methods with many more parameters. relatively small absolute values of scores on the ilpc dataset can be explained by its fully inductive setup and larger sparseness compared to wikidata5m. still, there are no publicly available solutions that exceed the baseline score provided by ilpc authors. our solution overperforms the existing baseline, is elegantly simple, and more interpretable than gnn-based solutions. experiments on pre-trained models revealed that pre-training speeds up convergence but degrades also, we demonstrated that neighbors' information enhances performance across different generative models. this includes open-source models such as t5, as well as proprietary models like openai gpts that are distributed as a model-as-aservice. although both types of models have their own advantages and limitations, it should be noted that openai gpts are black box models that cannot be modified or interpreted further. usage of openai gpts is restricted in the case of more specialized or proprietary kgs that can store secure or non-public information. in the context of our study, openai's gpt-3.5-turbo has a much larger size and cost, compared to smaller and transparent t5, still our approach improves the performance of both models. to get an understanding of how proposed models learn to use neighboring triplets we compared predictions of models trained with and without graph information in the input. one possible case when neighbors help in link prediction is when a target entity is somehow mentioned in the model input. because of neighborhood filtering, the target triplet never appears in the model input, however in some cases target entity can occur with different relations or as a substring of another entity in a neighborhood or even in the task itself ( the analysis shows that adding relevant graph context information greatly enhances link prediction. hence one way of improving performance is increasing neighborhood size by adding k-hop neighbors. due to the quadratic complexity of attention and the limited input size of pretrained transformers, processing exponentially growing neighborhoods would be computationally expensive. to deal with large input sizes one can use efficient t5 architectures with input sizes up to 32k @xcite and 64k @xcite tokens or recurrent approaches @xcite that can sequentially process effectively unlimited neighborhood sizes. another potential way of subgraph selection is to add relevant but more distant knowledge graph information, such as random walks or more sophisticated methods @xcite . ## conclusions and future work our study found that incorporating kg neighborhood information into a generative model's context can enhance the quality of the link prediction task. we trained the t5-small model on transductive wikidata5m and inductive ilpc datasets with various configurations and showed that coupling the node and its neighborhood information together better represents the node. our approach significantly outperformed baselines on the largest realworld kgcs, with a comparable or fewer number of parameters. we also demonstrated that ope-nai's gpt-3.5-turbo performance on the kgc task also benefits from neighborhood information without prior fine-tuning. in addition to the quality improvement, we analyzed the impact of neighborhood imputation, highlighting the significance of selecting relevant triplets or extending the model's context in future works. furthermore, future studies could explore the impact of model size, the effect of neighborhood imputation on the datasets from other domains, and neighborhood selection strategies.