Multi-Task Learning for Knowledge Graph Completion with Pre-trained Language Models (2020)

Abstract: As research on utilizing human knowledge in natural language processing has attracted considerable attention in recent years, knowledge graph (KG) completion has come into the spotlight. Recently, a new knowledge graph completion method using a pre-trained language model, such as KG-BERT, is presented and showed high performance. However, its scores in ranking metrics such as Hits@k are still behind state-of-the-art models. We claim that there are two main reasons: 1) failure in sufficiently learning relational information in knowledge graphs, and 2) difficulty in picking out the correct answer from lexically similar candidates. In this paper, we propose an effective multi-task learning method to overcome the limitations of previous works. By combining relation prediction and relevance ranking tasks with our target link prediction, the proposed model can learn more relational properties in KGs and properly perform even when lexical similarity occurs. Experimental results show that we not only largely improve the ranking performances compared to KG-BERT but also achieve the state-of-the-art performances in Mean Rank and Hits@10 on the WN18RR dataset.

## introduction a knowledge graph (kg) is a graph-structured knowledge base, where real-world knowledge is represented in the form of triple (h, r, t): (head entity, relation, tail entity) which means h and t have a relationship r. entities and the relation in a triple are denoted as nodes and an edge of the graph, respectively. in recent years, natural language processing (nlp) has benefited from utilizing kgs in various applications such as language modeling @xcite ), question answering @xcite , and machine reading @xcite . since there has been an increasing demand for high-quality knowledge, the reliability of kg has also become important. therefore, knowledge graph completion (a.k.a. link prediction), which identifies whether the triple in kg is valid or not, has been actively investigated. several studies on the knowledge graph completion have been conducted @xcite @xcite . they presented methods to model the connectivity patterns between entities in kg, and score functions to define the validity of the triple. however, these methods only consider graph structure and relational information depending on existing kg. thus, they cannot predict well on triples that contain less frequent entities. recently, addressing the sparseness problem of previous models, @xcite proposed a method called kg-bert for knowledge graph completion, using entity descriptions and pre-trained language models. even though kg-bert significantly improved mean ranks using preliminary linguistic information from bert @xcite , the results in other ranking metrics such as mrr and hit@k are still behind the state-of-the-art models. we claim that there are two major reasons for this problem. first, kg-bert misses lots of relation information in kgs. while previous state-of-the-art methods aimed to model relational properties in graphs, kg-bert only uses binary cross entropy loss to predict valid or invalid triples for the link prediction task. next, kg-bert has difficulty in picking out the answer entity between lexically similar candidates. for example, given head entity and relation as (take a breather, derivationally related for, ) and the correct tail entity as "breathing time", kg-bert predicts "snorkel breather" and "breath" as top scores because of the lexical similarity by "breath". this problem leads to lower performance in mrr and hits@k. in this paper, we propose an effective multi-task learning method to overcome these problems. we devise a multi-task framework by adding two tasks (relation prediction and relevance ranking) to link prediction, our target task. in the relation prediction, the model is trained to predict the relationship between given two entities, which helps the model learn more relational properties. in the relevance ranking, the model is trained by the margin ranking loss to make a gap between the valid triple and lexically similar candidates. we evaluate the proposed method on two popular datasets wn18rr and fb15k-237, and experimental results show that our method could improve ranking performance by a large margin compared to kg-bert. notably, our method achieves state-of-the-art performances in mean rank and hits@10 on the wn18rr dataset. ## proposed method in this section, we propose a multi-task learning for knowledge graph completion. as shown in figure 1 , we follow a multi-task learning framework in mt-dnn @xcite , and use the pre-trained bert model as a shared layer. we combine three tasks: link prediction, relation prediction, and relevance ranking. each task has a classification layer @xmath0. following @xcite , every input sequence has a [cls] token at the head of sentence, and [sep] token is used as a separator. ## experiments datasets we evaluated the proposed multi-task learning method on two benchmark datasets wn18rr @xcite and fb15k-237 @xcite . each dataset consists of a set of triples in the form of (h, r, t). wn18rr is a subset of wordnet, which is a lexical database of english. thus, entities in wn18rr are words or short phrases, and there exists 11 relations between two words, such as hypernym and similar to. fb15k-237 is a subset of freebase @xcite , a largescale graph database including general human knowledge. fb15k-237 has more general entities, such as lincoln and monaco, and relations are longer and more complex than wn18rr. we used the same entity descriptions with @xcite : synset definitions from wordnet for wn18rr and descriptions from @xcite for fb15k-237. baselines we mainly compare our method with kg-bert @xcite , and also provide a comparison with several outstanding models: transe @xcite , distmult @xcite , complex @xcite , conve @xcite , and rotate @xcite . ## related work a common approach for the knowledge graph completion is learning vector embeddings of the entities and the relationships in kg @xcite @xcite @xcite . the most widely used method is transe @xcite , which models the relationships as translations in low-dimensional vector space. @xcite and @xcite proposed the embedding models using a convolutional neural network. recent research has shown that the relation in complex vector space can infer the connectivity patterns: symmetry/antisymmetry, inversion, and composition @xcite . on the one hand, @xcite proposed kg-bert that uses pre-trained language models (plm) with entity descriptions. it can capture the contextualized meaning of entities and significantly improve mean ranks with rich linguistic information from plm. multi-task learning has gained popularity over a decade in natural language processing @xcite @xcite of various tasks. it aims to regularize deep learning models from overfitting by sharing parameters of different tasks while jointly training them. with the advent of powerful plms such as bert @xcite and xlnet @xcite , a multi-task learning scheme is applied by sharing pre-trained parameters of these models when training different tasks simultaneously. ## conclusion and future work we propose an effective multi-task learning method for knowledge graph completion by combining relation prediction and relevance ranking tasks with link prediction. experimental results demonstrate that our method outperforms previous strong baselines, and we largely improve mrr and hits@k compared to the previous kg-bert model. in the future, we plan to investigate how to combine pre-trained language models and graph embedding methods to fully utilize the prior linguistic information of pre-trained models and graph structural information.