A knowledge graph embedding model based attention mechanism for enhanced node information integration

Single triplet representation layer

We use the single triplet representation layer to represent the initialization vector, there are two main reasons for this. First, our model is based on the graph attention network to improve the knowledge representation learning task. Although the graph attention network can constantly update and integrate the information from other nodes into its own knowledge representation, it does not pay enough attention to the interactive information contained within the triplet itself, and focuses more on the information representation of the graph structure. Secondly, GAT is sensitive to parameter initialization, therefore, we believe that we can achieve improved results by initially training the mutual semantic information between triples to acquire the initialization vector representation of entities and relationships, and then aggregating the domain information through the graph attention mechanism.

TransE is a translation-based model, which proposes that the head entity and tail entity of a triple (e_h, r, e_t) can be mapped into vector space, and the relationship between them can be expressed as e_h + r ≈ e_t. We select the TransE model method to fit information into a single triplet and pretrain entities and relationships to achieve better convergence and enhance the stability of the model. The TransE model considers that the fact expressed by a correct triplet satisfies e_h + r ≈ e_t, d(e_h + r, e_t) should reach a minimum value as far as possible. On the contrary, for a corrupted triplet, d(e_h + r, e_t) should be as large as possible, where the function d represents the distance between e_t and e_h + r mapped to the vector space, the loss function is as follows: (1)${\displaystyle L=\sum _{\left(e_h,r,e_t\right)\in S}\sum _{\left(e_h^{\prime },r,e_t^{\prime }\right)}{\left[\gamma +d\left(e_h+r,e_t\right)-d\left(e_h^{\prime }+r,e_t^{\prime }\right)\right]}_{+}}$

where ${\left[x\right]}_{+}$ denotes the positive part of x, and γ is a margin hyperparameter. In our experiment, the distance formula adopts L2 regularization constraints. The construction method of negative example triplets $\left(e_h^{\prime }+r,e_t^{\prime }\right)$ is to replace one of the head or tail entities of the correct triplet with another entity. The set $S_{\left(e_h,r,e_t\right)}^{\prime }$ of the negative example triplet is as follows: (2)${\displaystyle S_{\left(e_h,r,e_t\right)}^{\prime }=\left\{\left(e_h^{\prime },r,e_t\right)|e_h^{\prime }\in E\right\}\cup \left\{\left(e_h,r,e_t^{\prime }\right)|e_t^{\prime }\in E\right\}.}$

By minimizing the loss function, minimizing the distance of positive examples and maximizing the distance of negative examples, to achieve the goal of training the knowledge representation of entities and relations.

Key information aggregation layer

We use $h=\left\{{\overrightarrow{h}}_1,{\overrightarrow{h}}_2,\dots ,{\overrightarrow{h}}_N\right\},{\overrightarrow{h}}_i\in {\mathbb{R}}^m$ to represent the vector representation of entities obtained through the TransE model, where ${\overrightarrow{h}}_i$ represents the feature vector representation of entities, the m is the dimension of vector ${\overrightarrow{h}}_i$ and the N represents the number of nodes. Before aggregating the information from neighborhoods, we need to perform linear transformations in order to enable input features to obtain more information and enhance their expression ability. The entity vector h is linearly transformed by the feature matrix W ∈ ℝ^m′×m to obtain a new node feature vector $h\prime =\left\{{\overrightarrow{h\prime }}_1,{\overrightarrow{h\prime }}_2,\dots ,{\overrightarrow{h\prime }}_N\right\},{\overrightarrow{h\prime }}_i\in {\mathbb{R}}^{m\prime }$, where ${\overrightarrow{h\prime }}_i$ represents the processed node vector representation and m′ is the dimension of ${\overrightarrow{h\prime }}_i$.

We use the self-attention calculation method without considering the graph structure to determine the weight value of the interaction between each pair of nodes. The calculation method is to use the concatenation operation to concatenate ${\overrightarrow{h\prime }}_i$ and ${\overrightarrow{h\prime }}_j$ into 2m-dimensional vectors, then calculates dot product with transposition of vector $\overrightarrow{a}$, to construct a mutual self-attention matrix X ∈ ℝ^N×N, which $\overrightarrow{a}$ is a learning weight vector parameterized. (3)${\displaystyle X={\overrightarrow{a}}^T\left(W{\overrightarrow{h}}_i^{\prime }\left|\right|W{\overrightarrow{h}}_j^{\prime }\right).}$

Then activated by LeakyReLU function to obtain α_ij, that is the attention weight value of node j to node i. The self-attention is the interaction weight between all pairs of nodes in the global without considering the graph structure, this attention weight represents the importance of every two nodes. (4)${\displaystyle {\alpha }_{ij}=\mathrm{LeakyReLU}\left(X\right).}$

The LeakyReLU function is a neural network unit activation function that can solve both the vanishing gradients and dying ReLU problems. The mathematical expression for the LeakyReLU function is as follows: (5)${\displaystyle y=max\left(0,x\right)+leak\times min\left(0,x\right).}$

Where the parameter leak is a small constant that refers to the degree or rate of leakage, which we have chosen as 0.01 in our experiments.

The training based on triplets and the training based on graph structure largely depend on the connectivity between nodes, knowledge graphs usually contain many rare nodes with low connectivity or even no links, which are difficult to train. Generally, the knowledge graph is a complex network, and the number of neighborhoods for each entity varies enormously. Some entities have hundreds of neighborhoods, while some entities have only a few or even none at all. Entities with more neighborhoods can embed richer graph structure information, while entities with fewer neighborhoods cannot rely on graph structure information to enrich their own semantics. Therefore, we supplement nodes set with potential relationships for nodes with fewer than k neighborhoods, even if there are no connected edges between nodes, the weight value of self-attention can be used to approximately represent the impact of the whole graph on rare nodes, which strengthens the participation of rare nodes and the connection ability with the whole graph, and improves the representation ability of the model for rare nodes.

In the normalization processing of the neighborhood information aggregation layer, we consider the weight values of k key nodes with most import information for the central node, which k is a controllable hyperparameter. We consider two cases here: For the nodes with more than or equal to k neighborhoods in the graph, only the k neighborhoods with the highest interaction weight need to be filtered and calculated for the neighborhoods. For the nodes with fewer than k neighborhoods, we divide the set into two parts, the first part is all the neighborhoods of the node, for the remaining insufficient part, we filter the most important key nodes from α_ij to supplement the set to k neighborhoods, these two parts are combined to form a key node set of rare nodes, as shown in Fig. 2.

Our specific method for aggregating node information is as follows: (i) If the current node i has c neighborhoods (c ≥ k), our method of forming the key node set K is to select the first k nodes with high attention weight values from the neighborhoods of node i. (ii) If the current node i has c neighborhoods (c < k), our method of forming the key node set K is including all neighborhoods, and then calculate the global node’s attention weight value for the current node i. Merge the highest k − c nodes into the set K, if there are neighborhoods for the current node i, skip it and continue searching downwards.

After obtained the key node set K, we use the softmax function to normalize and calculate the attention weight of all neighborhoods, so that the coefficients can be easily compared between different nodes: (6)${\displaystyle {\beta }_{ij}=softma{x}_j\left(e_{ij}\right)=\frac{exp\left({\alpha }_{ij}\right)}{\sum _{n\in K}exp\left({\alpha }_{in}\right)}.}$

β_ij is the normalized attention coefficients of the central node and its key nodes. Finally, we aggregate the information of k key nodes, introduce the weight matrix W, and obtain the hidden representation of nodes through linear transformation: (7)${\displaystyle {\overrightarrow{h}}_i=\sigma \left(\sum _{j\in K}{\beta }_{ij}W{\overrightarrow{h}}_j\right)+W{\overrightarrow{h}}_i.}$

We adopt the method of adding the vector representation of the central node after the weight matrix W transformation to strengthen the semantics of the central node, ${\overrightarrow{h}}_i$ is the final vector which contains the information of the central node and the aggregated neighborhood information. The KATKG model filters k key nodes for aggregation instead of aggregating all neighborhoods, it can limit the aggregation of redundant information in the case of too many neighborhoods and improves the aggregation effect for rare nodes. Figure 3 shows the structure of the key information layer.

Relation update layer

After key information aggregation layer, the hidden representation of entities has been trained, but the hidden representation of the relation still remains as the vector representation before training. Therefore, we introduce a relation update layer to update the relation embedding vector. According to the translation hypothesis proposed by Bordes et al. (2013), we update the embedded representation of the relation by traversing the difference of entity vectors between different triplets of the same relation. (8)${\displaystyle r_i^{\prime }=\frac{1}{n}\sum _{n\in R_i}\left({e_t}_n-{e_h}_n\right)+r_i.}$

where R_i is a set of all triplets that include relation i, r_i is the embedding representation of the relation obtained from the single triplet representation layer, and $r_i^{\prime }$ is the embedding representation of the finally relation.

Evaluation metrics

To evaluate the effectiveness of our model, we conducted link prediction experiments on the WN18RR and FB15k-237 datasets. In these experiments, we construct corrupted triplets by replacing either the head or tail entity of each correct triplet. By minimizing the loss function, we train the vector representations of entities and relations. To assess the performance of our model, we utilize several evaluation metrics. We use the Hits@k (Hit@1, Hit@3, Hit@10), mean ranking MR and mean reciprocal rank MRR to evaluate the model. MR represents the average ranking of all correct triples, so a lower value indicates better model performance. MRR represents the sum of the reciprocal rankings of all correct triplets, so a higher value indicates better performance of the model; Hits@k indicates that the correct triplet ranking is less than the average proportion of k, as the same of MRR, a higher value indicates better performance of the model.

Implementation details

In the training process, we use the Adam optimizer and select the following hyperparameters: entity and relation embedding dimension m among 50, 100, 200, layer depth L among 1, 2,  hidden layer dimension among m_d 100, 200, 300, fixed learning rate λ among 0.01, 0.005, 0.001, and margin γ among 1, 2. In the experiment, we need to select key k neighborhoods to aggregate information, so we perform a statistical analysis of the number of neighborhoods of each entity in each set. The details are shown in Table 2.

In the WN18RR dataset, the number of nodes with 2 neighborhoods is the largest, the average number of neighborhoods of a node is 4.24, therefore, we chose k value is among 3, 5, 10 in the WN18RR dataset. In the FB15k-237 dataset, the average number of neighborhoods for nodes is 37.42, and most number of nodes have 11 edges, therefore, the k value is among 5, 10, 20 on the FB15k-237 dataset.

We select the above parameters for training, and retain the best parameters when MRR obtains the best performance on each verification set. For WN18RR, the optimal parameters are entity and relation embedding dimension 50, layer depth 2, hidden layer dimension 100, learning rate 0.005, margin γ = 2, and number of neighborhoods k = 3. For FB15k-237, the optimal parameters are entity and relation embedding dimension m = 50, layer depth L = 2, hidden layer dimension m_d = 200, learning rate λ = 0.01, margin γ = 1, and number of neighborhoods k = 10, the ranges and values of key experimental parameters are shown in Table 3 for clarity.