JAKET: Joint Pre-training of Knowledge Graph and Language Understanding

1 INTRODUCTION . Pre-trained language models ( PLM ) leverage large-scale unlabeled corpora to conduct selfsupervised training . They have achieved remarkable performance in various NLP tasks , exemplified by BERT ( Devlin et al. , 2018 ) , RoBERTa ( Liu et al. , 2019b ) , XLNet ( Yang et al. , 2019 ) , and GPT series ( Radford et al. , 2018 ; 2019 ; Brown et al. , 2020 ) . It has been shown that PLMs can effectively characterize linguistic patterns in text and generate high-quality context-aware representations ( Liu et al. , 2019a ) . However , these models struggle to grasp world knowledge about entities and relations ( Poerner et al. , 2019 ; Talmor et al. , 2019 ) , which are very important in language understanding . Knowledge graphs ( KGs ) represent entities and relations in a structural way . They can also solve the sparsity problem in text modeling . For instance , a language model may require tens of instances of the phrase “ labrador is a kind of dog ” in its training corpus before it implicitly learns this fact . In comparison , a knowledge graph can use two entity nodes “ labrador ” , “ dog ” and a relation edge “ is a ” between these nodes to precisely represent this fact . Recently , some efforts have been made to integrate knowledge graphs into PLM . Most of them combine the token representations in PLM with representations of aligned KG entities . The entity embeddings in those methods are either pre-computed based on an external source by a separate model ( Zhang et al. , 2019 ; Peters et al. , 2019 ) , which may not be easily aligned with the language representation space , or directly learned as model parameters ( Févry et al. , 2020 ; Verga et al. , 2020 ) , which often have an over-parameterization issue due to the large number of entities . Moreover , all the previous works share a common challenge : when the pre-trained model is fine-tuned in a new domain with a previously unseen knowledge graph , it struggles to adapt to the new entities , relations and structure . Therefore , we propose JAKET , a Joint pre-trAining framework for KnowledgE graph and Text . Our framework contains a knowledge module and a language module , which mutually assist each other by providing required information to achieve more effective semantic analysis . The knowledge module leverages a graph attention network ( Veličković et al. , 2017 ) to provide structure-aware entity embeddings for language modeling . And the language module produces contextual representations as initial embeddings for KG entities and relations given their descriptive text . Thus , in both modules , content understanding is based on related knowledge and rich context . On one hand , the joint pre-training effectively projects entities/relations and text into a shared semantic latent space , which eases the semantic matching between them . On the other hand , as the knowledge module produces representations from descriptive text , it solves the over-parameterization issue since entity embeddings are no longer part of the model ’ s parameters . In order to solve the cyclic dependency between the two modules , we propose a novel two-step language module LM1 and LM2 , respectively . LM1 provides embeddings for both LM2 and KG . The entity embeddings from KG are also fed into LM2 , which produces the final representation . LM1 and LM2 can be easily established as the first several transformer layers and the rest layers of a pre-trained language model such as BERT and RoBERTa . Furthermore , we design an entity context embedding memory with periodic update which speeds up the pre-training by 15x . The pre-training tasks are all self-supervised , including entity category classification and relation type prediction for the knowledge module , and masked token prediction and masked entity prediction for the language module . A great benefit of our framework is that it can easily adapt to unseen knowledge graphs in the finetuning phase . As the initial embeddings of entities and relations come from their descriptive text , JAKET is not confined to any fixed KG . With the learned ability to integrate structural information during pre-training , the framework is extensible to novel knowledge graphs with previously unseen entities and relations , as illustrated in Figure 1 . We conduct empirical studies on several knowledge-aware natural language understanding ( NLU ) tasks , including few-shot relation classification , question answering and entity classification . The results show that JAKET achieves the best performance compared with strong baseline methods on all the tasks , including those with a previously unseen knowledge graph . 2 RELATED WORK . Pre-trained language models have been shown to be very effective in various NLP tasks , including ELMo ( Peters et al. , 2018 ) , GPT ( Radford et al. , 2018 ) , BERT ( Devlin et al. , 2018 ) , RoBERTa ( Liu et al. , 2019b ) and XLNet ( Yang et al. , 2019 ) . Built upon large-scale corpora , these pretrained models learn effective representations for various semantic structures and linguistic relationships . They are trained on self-supervised tasks like masked language modeling and next sentence prediction . Recently , a lot of efforts have been made on investigating how to integrate knowledge into PLMs ( Levine et al. , 2019 ; Soares et al. , 2019 ; Liu et al. , 2020 ; Guu et al. , 2020 ) . These approaches can be grouped into two categories : 1 . Explicitly injecting entity representation into the language model , where the representations are either pre-computed from external sources ( Zhang et al. , 2019 ; Peters et al. , 2019 ) or directly learned as model parameters ( Févry et al. , 2020 ; Verga et al. , 2020 ) . For example , ERNIE ( THU ) ( Zhang et al. , 2019 ) pre-trains the entity embeddings on a knowledge graph using TransE ( Bordes et al. , 2013 ) , while EAE ( Févry et al. , 2020 ) learns the representation from pre-training objectives with all the other model parameters . K-BERT ( Liu et al. , 2020 ) represents the entities by the embeddings of surface form tokens ( i.e . entity names ) , which contains much less semantic information compared with description text . Moreover , it only injects KG during fine-tuning phase instead of joint-pretraining KG and text . 2 . Implicitly modeling knowledge information , including entity-level masked language modeling ( Sun et al. , 2019b ; Shen et al. , 2020 ) , entity-based replacement prediction ( Xiong et al. , 2019 ) and knowledge embedding loss as regularization ( Wang et al. , 2019b ) . For example , besides tokenlevel masked language modeling , ERNIE ( Baidu ) ( Sun et al. , 2019b ) uses phrase-level and entitylevel masking to predict all the masked slots . KEPLER ( Wang et al. , 2019b ) calculates entity embeddings using a pre-trained language model based on the description text , which is similar to our work . However , they use the entity embeddings for the knowledge graph completion task instead of injecting them into the language model . Some works ( Ding et al. , 2019 ; Lv et al. , 2020 ) investigated the combination of GNN and PLM . For example , Lv et al . ( 2020 ) uses XLNet to generate initial node representation based on node context and feeds them into a GNN . However , these approaches do not integrate knowledge into language modeling , and they are designed for specific NLP tasks such as reading comprehension or commonsense reasoning . In comparison , we jointly pre-train both the knowledge graph representation and language modeling and target for general knowledge-aware NLU tasks . 3 METHOD . In this section , we introduce the JAKET framework of joint pre-training knowledge graph and language understanding . We begin by defining the mathematical notations , and then present our model architecture with the knowledge module and language module . Finally , we introduce how to pretrain our model and fine-tune it for downstream tasks . The framework is illustrated in Figure 2 . 3.1 DEFINITION . A knowledge graph is denoted by KG = ( E , R , T ) , where E = { e1 . . . eN } is the set of entities and R = { r1 . . . rP } is the set of relations . T = { ( et1i , rt2i , et3i ) |1 ≤ i ≤ T , et1i , et3i ∈ E , rt2i ∈ R } stands for the set of head-relation-tail triplets . Nv = { ( r , u ) | ( v , r , u ) ∈ T } represents the set of neighboring relations and entities of an entity v. We define V = { [ MASK ] , [ CLS ] , [ EOS ] , w1 . . . wV } as a vocabulary of tokens and the contextual text x = [ x1 , x2 , . . . , xL ] as a sequence of tokens where xi ∈ V . In the vocabulary , [ MASK ] is the special token for masked language modeling ( Devlin et al. , 2018 ) and [ CLS ] , [ EOS ] are the special tokens indicating the beginning and end of the sequence . We define F as the dimension of token embeddings , which is equal to the dimension of entity/relation embeddings from the KG . The text x has a list of entity mentions m = [ m1 , . . . , mM ] , where each mention mi = ( emi , smi , omi ) : emi is the corresponding entity and smi , omi are the start and end index of this mention in the context . In other words , [ xsmi , . . . , xomi ] is linked with entity emi 1 . We assume the span of mentions are disjoint for a given text sequence . As entities in the knowledge graph are represented by nodes without context , we use entity description text to describe the concept and meaning of entities . For each entity ei , its description text xei describes this entity . The mention of ei in xei is denoted as mei = ( ei , sei , o e i ) , similarly defined as above . For instance , the description text for the entity “ sun ” can be “ [ CLS ] The Sun is the star at the center of the Solar System [ EOS ] ” . Then the mention is mSun = ( Sun , 3 , 3 ) . If there are multiple mentions of ei in its description text , we choose the first one . If there ’ s no mention of ei in its description text , we set sei = o e i = 1 . Similarly , we define relation description text as the text that can describe each relation . 3.2 KNOWLEDGE MODULE . The goal of the knowledge module ( KM ) is to model the knowledge graph to generate knowledgebased entity representations . To compute entity node embeddings , we employ the graph attention network ( GAT ) ( Veličković et al. , 2017 ) , which uses the self-attention mechanism to specify different weights for different neighboring nodes . However , the vanilla GAT is designed for homogeneous graphs with singlerelation edges . To leverage the multi-relational information , we adopt the idea of composition operator ( Vashishth et al. , 2019 ) to compose entity embeddings and relation embeddings . In detail , in the l-th layer of LM , we update the embedding E ( l ) v of entity v as follows : E ( l ) v = LayerNorm K⊕ k=1 σ ∑ ( r , u ) ∈Nv αkv , r , uW kf ( E ( l−1 ) u , Rr ) + E ( l−1 ) v ( 1 ) αkv , r , u = exp ( LeakyReLU ( aT [ W kE ( l−1 ) v ⊕W kf ( E ( l−1 ) u , Rr ) ] ) ) ∑ ( r′ , u′ ) ∈Nv exp ( LeakyReLU ( aT [ W kE ( l−1 ) u ⊕W kf ( E ( l−1 ) u′ , Rr′ ) ] ) ) ( 2 ) where LayerNorm stands for layer normalization ( Ba et al. , 2016 ) . ⊕ means concatenation and K is the number of attention heads . W k is the model parameter and Rr is the embedding of relation r. Note that the relation embeddings are shared across different layers . The function f ( · , · ) : RF × RF → RF merges a pair of entity and relation embeddings into one representation . Here , we set f ( x , y ) = x + y inspired by TransE ( Bordes et al. , 2013 ) . More complicated functions like MLP network can also be applied . The initial entity embeddings E ( 0 ) and relation embeddings R are generated from our language module , which will be introduced in Section 3.4 . Then , the output entity embeddings from the last GAT layer are used as the final entity representations EKM . Note that the knowledge graph can be very large , making the embedding update over all the entities in Equation ( 1 ) not tractable . Thus we follow the minibatch setting ( Hamilton et al. , 2017 ) : given a set of input entities , we perform neighborhood sampling to generate their multi-hop neighbor sets and we compute representations only on the entities and relations that are necessary for the embedding update .

Summary: This paper proposed a new language modeling pretraining method that leverages the knowledge graph information. Specifically, the paper replaces the entity embedding in one hidden layer of BERT context embedding, with the corresponding graph attention embedding that is obtained from the knowledge graph. The pretraining tasks contain not only the language related tasks (like predicting masked tokens), but also the knowledge graph tasks like entity classification or relation type prediction. Experiments on few-shot learning tasks, question answering and entity classification show better performance over other pretraining counterparts.