Context-Enhanced Entity and Relation Embedding for
Knowledge Graph Completion (Student Abstract)

Knowledge graphs(KGs) store a wealth of knowledge from real world into structured graphs, which consist of collections of triplets, and each triplet $(h,r,t)$ represents that head entity $h$ is related to tail entity $t$ through a relation type $r$. KGs have play a significant role in AI-related applications such as recommendation systems, question answering, and information retrieval. However, the coverage of knowledge graphs nowadays is still far from complete and comprehensive, researchers have proposed a number of knowledge graph completion(KGC) methods to predict the missing links/relations in the incomplete knowledge graph.

Most state-of-the-art methods for KGC are usually based on knowledge graph embeddings, which normally assign an embedding vector to each entity and relation in the continuous embedding space and train the embeddings based on the existing triplets, i.e., to make the scores for observed graph triplets higher than random ones. However, most of them fail to utilize the context/neighborhood of the entity or relation, which may contain rich and valuable information for KGC.

Recently, some researches have proved the significance of contextual information of entity and relation in KGC. For example, A2N (Bansal et al. 2019) and RGHAT (Zhang et al. 2020) propose attention-based methods which leverage the contextual information for link prediction by attending to the neighbor entities and lead to more accurate KGC. PathCon (Wang, Ren, and Leskovec 2020) considers relational context of the head/tail entity and relational paths between head and tail entity in one model, and finds that they are critical to the task of relation prediction. However, these researches just utilize entity context or relation context, which may lead to information loss.

In this paper, we aim to take full advantage of both the entity context and relation context for enhancing the KGC task. Specifically, different from the neighborhood definition in traditional KG topology, for each element in each triplet, we extract the pair composed of the other two elements as one neighbor in its context. Then we propose an efficient model, named AggrE, to alternately aggregate the information of entity context and relation context in multi-hops into entity and relation, and learn context-enhanced entity embeddings and relation embeddings. Then we use the learned embeddings to predict the missing relation $r$ given a pair of entities $(h,?,t)$ to complete knowledge graphs.

Given a knowledge graph $G=\{(h,r,t)\}\subseteq E\times R\times E$, where $E$ and $R$ are entity set and relation set respectively. Firstly, we extract contexts of entities and relations from the existing triplets, as shown in Figure 1. For an entity $e_{i}$ in $G$, we define the entity context of $e_{i}$ as $C_{e_{i}}=\{(r_{j},e_{k})|(e_{i},r_{j},e_{k})\in G\vee(e_{k},r_{j},e_{i})\in G\}$, which is actually the set of neighbor entities with their corresponding relations of $e_{i}$. For a relations $r_{i}$ in $G$, we define the relation context of $r_{i}$ as $C_{r_{i}}=\{(e_{j},e_{k})|(e_{j},r_{i},e_{k})\in G\}$, which is actually the two endpoints of $r_{i}$. Denote $\mathbf{e}^{(0)}_{i}$ and $\mathbf{r}^{(0)}_{i}$ as the randomly initialized embedding of $e_{i}$ and $r_{i}$ respectively, our intuition is to aggregate the contextual information into the representations of each entity and relation to help the prediction. We define the aggregation functions as:

where $0\leq l<L$ and $L$ is the number of aggregation layers, $\mathbf{e}^{(l+1)}_{i}$ and $\mathbf{r}^{(l+1)}_{i}$ is the $l$-th layer’s output embedding of $e_{i}$ and $r_{i}$, $\odot$ is the element-wise product, $\alpha_{i,j,k}^{(l)}$ and $\beta_{i,j,k}^{(l)}$ are layer-specific softmax normalization constants, $\alpha_{i,j,k}^{(l)}$ represent how important each entity context is for $e_{i}$: $\alpha_{i,j,k}^{(l)}=\frac{\exp(s^{(l)}_{i,j,k})}{\sum\nolimits_{(r_{j^{\prime}},e_{k^{\prime}})\in C_{e_{i}}}\exp(s^{(l)}_{i,j^{\prime},k^{\prime}})}$, and $\beta_{i,j,k}^{(l)}$ represent how important each relation context is for $r_{i}$:$\beta_{i,j,k}^{(l)}=\frac{\exp(s^{(l)}_{j,i,k})}{\sum\nolimits_{(e_{j^{\prime}},e_{k^{\prime}})\in C_{r_{i}}}\exp(s^{(l)}_{j^{\prime},i,k^{\prime}})}$. The $s^{(l)}_{i,j,k}$ represents the score for each possible triplet $(e_{i},r_{j},e_{k})$ after $l$ layers’ aggregation, and we use the same score function with DistMult(Yang et al. 2014) to calculate the triplet scores:

Where $Diag(\mathbf{r}^{(l)}_{j})$ is a diagonal matrix with $\mathbf{r}^{(l)}_{j}$ in its diagonal. After $L$ layers’ aggregation, we can obtain the final output $\mathbf{e}^{(L)}_{i}$ and $\mathbf{r}^{(L)}_{i}$ of each entity and relation, which contain neighbor information from their $L$-hops contexts. Then we conduct a softmax loss function on the final triplet scores to compute the likelihood of predicting the correct relations:

where $R$ is the set of relations. We use a mini-batch Adam optimizer to minimize $\mathcal{L}$. The difference between our aggregation model and Graph Neural Network(GNN) is that instead of using complex matrix transformation, we use element-wise products to obtain neighborhood information and add it directly to central nodes, as the embedding itself can be regarded as trainable transformation parameters. Also our model is expected to be more efficient and suitable for larger knowledge graph.

We conduct experiments on two widely used KG benchmarks: FB15K-237 and WN18RR. Noted that the trainable parameters in our model are only entity and relation embeddings, for a fair comparison, we choose 5 traditional baselines with a small amount of parameters. We use Mean Reciprocal Rank (MRR, the mean of all the reciprocals of predicted ranks), Mean Rank (MR, the mean of all the predicted ranks), and Hit@3(the proportion of correctly predicted entities ranked in the top 3 predictions) as evaluation metrics. In the experiment, we set the embedding dimensionality as 256, the learning rate as 5e-3, l2 penalty coefficient as 1e-7, batch size as 512 and a maximum of 20 epochs. We set the number of aggregation layers $L$ as 2 for WN18RR and 4 for FB15K-237 repectively.

As shown in Table 1, The results indicate that AggrE significantly outperforms all the baselines on two benchmarks, which indicates the effectiveness of AggrE. Specifically, the improvement is rather significant for WN18RR, where the links between entities are more sparse than in FB15K-237, this may because without using extra parameters other than embeddings, AggrE is less prone to overfitting. Besides, AggrE achieves substantial improvements against DistMult on all the metrics, and noted that they have the similar objective functions, it indicates that aggregating contextual information for entities and relations is valuable and can great improve the performance of prediction.

In this paper, we specify a novel definition on the context/neighborhood of entity and relation in KGs, and propose a multi-layer aggregation models to compose contextual information to embeddings for KGC. In the future, we will explore more possible aggregation functions in our model.

This research was supported by the NSFC under Grant 61836013, National Key R&D Plan of China (2016YFB0501901), Beijing Nova Program of Science and Technology (Z191100001119090).