Improving Knowledge Graph Representation Learning by Structure Contextual Pre-training

Similar to pre-trained language model, SCoP model in this paper includes two steps: pre-training and fine-tuning. The SCoP model is based on structure contextual triples to learn the structural information of knowledge graph.

From the perspective of model structure, SCoP consists two modules that integrate and extract the deep structure information of KG. In this paper, bold characters denote vector representations, for example, $\mathbf{h}$ is the vector representation of entity $h$.

Given a KG $\mathcal{G}=\{\mathcal{E},\mathcal{R},\mathcal{T}\}$, $\mathcal{E}$, $\mathcal{R}$ and $\mathcal{T}$ are the set of entities, relations and triples respectively. And $\mathcal{T}=\{(h,r,t)|h\in\mathcal{E},r\in\mathcal{R},t\in\mathcal{E}\}$.

At the pre-training step, all parameters of the model are first trained on a general task which is triple classification in this paper, as the existence of a triple is the most widely available information in a knowledge graph.

At the fine-tuning step, we generate model variants based on the specific downstream task and the parameters in SCoP are fine-tuned with its dataset, as shown in Figure 2. This process is lightweight and fast to better results.

To prove the effective of pre-training knowledge graph, we conduct experiments on the whole knowledge graph WordNet, named WN-all. We also construct two downstream datasets for fine-tuning. Detailed statistics of datasets are shown in Table 2.

We pre-train SCoP on WN-all via triple classification task, aiming to judge whether the triple is correct with label $l$. The loss is defined as $\mathcal{L}=CrossEntropy(log(softmax(\mathbf{a})),l)$.

We fix the length of input contextual triple sequence with 256 and the length of each contextual triple of $h$, $r$, $t$ is 84. We set the self-attention with 6 layers, 3 heads and the hidden dimension of representation vectors of 192 in Transformer in A-Mod, and learning rate of $2e^{-5}$, training batch size of 32, learning rate warm-up over the first 1,000 steps and a dropout probability of 0.1 on all layers.

Entity type prediction concentrates to complete a pair (entity, entity type) when its type is missing, which aims to verify the capability of our model for inferring missing entity type instances. Under this definition, we extract entity and its type from triples with relation $\_hypernym$ in WN-all as the dataset for this task, shown in Table 2.

To fine-tune on entity type prediction task and its dataset, we input the structure contextual triples sequence of entity type pair $\tau$ and use the output vector $\mathbf{a_{\tau}}\in\mathbb{R}^{d}$ of A-Mod to compute the score function $s_{\tau}=f(\rm{ent},\rm{type})={\rm softmax}(\mathbf{a}_{\tau}\mathbf{W})$ and entity type classification loss: $\mathcal{L}_{\tau}=CrossEntropy(log(s_{\tau}),l_{\tau})$, where $\mathbf{W}\in\mathbb{R}^{d\times 2}$ is the classification weight and $s_{\tau}\in\mathbb{R}^{2}$ is a 2-dimensional real vector. Label $l_{\tau}\in\{0,1\}$ is $1$ if entity type pair $\tau$ is positive while $0$ if negative. The Figure 2 (a) shows the structure of fine-tuning SCoP for entity type prediction task.

We adopt the Mean Reciprocal Rank(MRR) and percentage of triples ranked in top $k$ (Hit@$k$) where $k\in\{1,3,10\}$ as evaluation metrics, which are commonly adopted in previous KG works. See Section 4 for details of baseline models.

Experimental results are presented in the left part of Table 1. Compared with the results of fine-tuning stage (even rows with ”-ft” suffix in table), SCoP-ft outperforms other models on all evaluation metrics, especially obtaining 3.8% prediction accuracy improvement on Hit@1 over the prior state of the art. Besides, we observe that SCoP-ft model has better performance than SCoP-pr while on the contrary, the results of other models after fine-tuning decreased. It proves that the structure contextual triples, and the structural information that they represent in essence, can effectively improve the effect of entity type prediction.

Entity alignment task aims to find two entities that refer to the same thing in the real world. Similarly, we generate the dataset for entity alignment consisting of two entities with the same meaning from triples with relation $\_similar\_to$, such as (nascent, emergent) and (sleep, slumber). Similar to the fine-tuning progress of entity type prediction, we compute the score function $s_{\omega}$ for alignment entity pair $\omega$ and its classification loss $\mathcal{L}_{\omega}$ in the same way. Figure 2 (b) shows its model structure and the right part of Table 1 shows the results.

The experimental results show that SCoP model still has the greatest advantage in Hit@1 metric, which is 93.15% and surpasses the second result by 3.3%. On the entity type prediction task, the results from pre-training step of SCoP model are not always optimal. For example, RotatE-pr has a better result with Hit@10 metric on entity type prediction task. However, those results of SCoP model outperform all other models at pre-training step on entity alignment task, let alone the overwhelming results after fine-tuning.

We observed that the data of Hit@1, Hit@3, and Hit@10 in SCoP model are very close, and the interval between them is less than 1% while those of other models have clear gaps. Therefore, we further investigated the distribution of prediction score and the triple classification task, whose aim is to judge whether a given triple ($h,r,t$) is correct or not, as shown in Figure 2 (c).

From Figure 3, it is not difficult to see that the prediction scores of the positive samples and negative samples are relatively evenly distributed in the whole score space predicted by the baseline model like TransE, and most of the positive examples are in the upper part of the space while the negative in the lower part. The score distribution of ComplE and RotatE resembles TransE, so these similar figures are not repeatedly shown. However, The prediction scores of positive and negative samples in SCoP model are close to their upper and lower boundaries respectively and only a small part of the predicted values are distributed in the middle region. Figure 3 shows that structural information can bring significant gains to triple classification tasks, and it is easier to set the margin $\gamma$ to distinguish between positive and negative samples. The experiment shows that, different margins $\gamma$ among $\{20,40,60,80\}$ percent of maximum prediction score, hardly affect the prediction accuracy of SCOP with 92.6% numerically, while that of TransE varies greatly, from 13.6% to 91.1%.