ShallowBKGC: a BERT-enhanced shallow neural network model for knowledge graph completion

Translation-based models

Translation-based models consider the relation between a head entity and a tail entity as a translation operation in the vector space and calculate the distance between the head entity vector and the tail entity vector to measure the plausibility of a triple. Bordes et al. (2013) present the initial translation-based model TransE, which learns low-dimensional and dense vectors for every entity and relation, so that relations correspond to translation vectors operating on vectors of entities. Wang et al. (2014) present TransH, which alleviates the complex relation problem in TransE by associating each relation with a relation-specific hyperplane. Lin et al. (2015) present a path-based TransE, named PTransE, which extends TransE by relation paths. Nguyen et al. (2016) present STransE that combines SE (Bordes et al., 2011) and TransE for KGC. Sun et al. (2019) present RotatE, which defines each relation as a rotation from head entity to tail entity in the complex space. Le, Huynh & Le (2021) present RotatH that combines RotatE and TransH for KGC.

Models of this category have the advantages of simplicity, intuitiveness, and high computational efficiency. However, research has shown that they have limitations in expressive power and are not suitable for non-Euclidean spaces.

Tensor decomposition-based models

Tensor decomposition-based models use triangular norm to measure the plausibility of triplets. Yang et al. (2015) present DistMult, which considers triplets as tensor decomposition and constrains all relation embeddings to be diagonal matrices. ComplEx (Trouillon et al., 2016) extends DistMult to the complex space to better model asymmetric and inverse relations. Balažević, Allen & Hospedales (2019) present TuckER, which performs KGC based on tucker decomposition of binary tensors of known triplets. Inspired by the tucker decomposition of order-4 tensors, Shao et al. (2022) present a tensor decomposition model for temporal KGC. Zhang et al. (2024) extend tensor decomposition methods to temporal KGC.

Models of this category are proficient in capturing complex relations between entities and relations in KGs. However, as the scale of KGs grows, the computational complexity of these models may escalate rapidly.

Neural network-based models

Various neural networks have been widely explored for KGC and achieved promising performance. Dettmers et al. (2018) present a multi-layer convolutional model ConvE, which explores convolutional neural network for KGC, and uses 2D convolution over embeddings to predict missing triplets in a KG. Shang et al. (2019) present an end-to-end graph structure-aware convolutional networks model SACN that combines graph convolutional network (GCN) and ConvE for KGC. Dai Quoc Nguyen, Nguyen & Phung (2018) present ConvKB, which utilizes convolutional neural network to capture the global relationships among dimensional entries of entity and relation embeddings. CapsE (Nguyen et al., 2019) combines ConvKB with capsule network for both KGC and search personalization tasks. Schlichtkrull et al. (2018) present relational graph convolutional networks and apply them to KGC. Vashishth et al. (2020) present CompGCN, which leverages a variety of composition operations from knowledge graph embedding techniques to jointly embed both entities and relation in a graph. SHALLOM (Demir, Moussallem & Ngomo, 2021) and the prior version of the model proposed in this article ASLEEP (Jia, 2022) apply shallow neural network for KGC and achieve good performance while maintaining high efficiency.

Models of this category have significant advantages in semantic feature learning, and our proposed model belongs to this category. The main difference between them and our model is that most of them usually rely on more deeper and complex neural networks while our model employs shallow neural network, which is not computationally demanding and friendly to real-time applications. Although there are several models that use shallow neural networks for KGC, these models only use structural information and ignore the rich information contained in text. Under the premise of keeping the model as simple as possible, we consider both text and structural information for KGC.

Pre-trained language/large language-based models

Pre-trained language models and large language models have received widespread attention in many natural language processing tasks, including KGC. Yao, Mao & Luo (2019a) explore the pre-trained language model BERT for KGC. StAR (Wang et al., 2021) extends KG-BERT by taking into account structural information for KGC. Yao et al. (2023) present KG-LLM, which investigates large language models, including ChatGLM (Du et al., 2022) and LLaMA (Touvron et al., 2023) for KGC. Yang, Fang & Zhou (2023) present a constrained-prompt KGC based on large language model. Zhang et al. (2023) present KoPA, which integrates pre-trained KG structural features with large language model for KGC.

Models of this category achieve great success in KGC. However, these models usually require diverse fine-tuning strategies, and mostly cost much time in training and inference. It should be pointed out that the model we proposed also uses the pre-trained model BERT. The difference from the existing model is that in order to keep the model as simple as possible, we use BERT in a feature extraction manner, that is, the parameters in BERT are not involved in training.

Problem formulation

A KG is a type of multi-relational directed graph that typically consists of a collection of triplets in the form of $(h,r,t)$. It can be formally defined as $\mathcal{G}=(\mathcal{E},\mathcal{R},\mathcal{T})$, where $\mathcal{T}$ represents the set of all triplets, $\mathcal{E}$ and $\mathcal{R}$ represent the sets of all entities and relations respectively.

The objective of the KGC task is to predict missing relations in $\mathcal{G}$ based on the known triplets $\mathcal{T}$. In other words, the aim of KGC is to develop a model that accepts a query consisting of a head entity and a tail entity, $(h_i,?,t_i)$, and ranks all candidate relations $r_c\in \mathcal{R}$ to resolve the query (Lovelace & Rose, 2022). An effective KGC model should enable correct candidates to have higher rankings than incorrect candidates.

Model overview

Our proposed model ShallowBKGC takes as input an entity pair, and outputs the probability that each relation exists between the two entities. As illustrated in Fig. 1, our model consists of three key steps: (1) entity feature extraction, (2) entity-pair representation, and (3) multi-label relation modeling. The detailed calculation process of each step is as follows.

Entity feature extraction

Given an entity pair $(h,t)$, our model extracts the features of the head and tail entities by taking into account both text and structural information.

For text information, we apply the pre-trained language model BERT (Kenton & Toutanova, 2019), which has achieved great success in multiple natural language processing tasks, to extract text feature of the given entity. Figure 2 illustrates the framework of BERT for entity text feature extraction. Formally, given the text information of head and tail entities, i.e., head entity name $h\mathrm{\_}text=\{w_1^h,w_2^h,...,w_N^h\}$ and tail entity name $t\mathrm{\_}text=\{w_1^t,w_2^t,...,w_N^t\}$ (since many entities lack descriptive information and introducing additional information will increase computational complexity, we only use the name information that each entity has as text information), we first add a special classification token (CLS) and a separate token (SEP) at the beginning and end of the entity name respectively to obtain the marked entity name. Then through the tokenizer we obtain the representation of the marked entity name. Finally, we put the representation into BERT to get the text features of the head and tail entities as follows:

(1) $$\mathbf{C}=\mathrm{B}\mathrm{E}\mathrm{R}\mathrm{T}(\mathbf{C}\mathbf{L}\mathbf{S},{\mathbf{w}}_1^{\mathrm{h}},{\mathbf{w}}_2^{\mathrm{h}},...,{\mathbf{w}}_{\mathrm{N}}^{\mathrm{h}},\mathbf{S}\mathbf{E}\mathbf{P});{\mathbf{h}}_{\mathrm{t}}=\mathbf{C}$$

(2) $$\mathbf{C}=\mathrm{B}\mathrm{E}\mathrm{R}\mathrm{T}(\mathbf{C}\mathbf{L}\mathbf{S},{\mathbf{w}}_1^{\mathrm{t}},{\mathbf{w}}_2^{\mathrm{t}},...,{\mathbf{w}}_{\mathrm{N}}^{\mathrm{t}},\mathbf{S}\mathbf{E}\mathbf{P});{\mathbf{t}}_{\mathrm{t}}=\mathbf{C}$$where $\mathbf{C}\in \mathbb{R}{ }^H$ is the hidden vector of the special token (CLS), which contains the features of the entire input text. Therefore, we use it as the text feature of the given entity.

For structural information, our model receives the IDs of the given head and tail entities, and extracts the structure features of them through embedding layer as follows:

(3) $${\mathbf{h}}_s=\mathrm{E}\mathrm{m}\mathrm{b}\mathrm{e}\mathrm{d}\mathrm{d}\mathrm{i}\mathrm{n}\mathrm{g}\mathrm{\_}\mathrm{l}\mathrm{a}\mathrm{y}\mathrm{e}\mathrm{r}({\mathbf{h}}_{\mathrm{I}\mathrm{D}})$$

(4) $${\mathbf{t}}_s=\mathrm{E}\mathrm{m}\mathrm{b}\mathrm{e}\mathrm{d}\mathrm{d}\mathrm{i}\mathrm{n}\mathrm{g}\mathrm{\_}\mathrm{l}\mathrm{a}\mathrm{y}\mathrm{e}\mathrm{r}({\mathbf{t}}_{\mathrm{I}\mathrm{D}})$$where ${\mathbf{h}}_s\in \mathbb{R}{ }^d$, ${\mathbf{t}}_s\in \mathbb{R}{ }^d$ are embeddings of structural information corresponding to head and tail entities, respectively.

And then, we integrate entity text feature and entity structure feature through average operation, and get the entity feature as,

(5) $$\mathbf{h}=\mathrm{a}\mathrm{v}\mathrm{e}({\mathbf{h}}_{\mathrm{t}},{\mathbf{h}}_{\mathrm{s}})$$

(6) $$\mathbf{t}=\mathrm{a}\mathrm{v}\mathrm{e}({\mathbf{t}}_{\mathrm{t}},{\mathbf{t}}_{\mathrm{s}})$$

For the sake of computational convenience, and considering the consistency of tensor shapes, we intercept the first $d$ columns of text features when fusing text and structure features of entities.

Entity-pair representation

After getting the features of entities, our model integrates head entity feature and tail entity feature into entity-pair representation through the average operation and a non-linear transformation as follows:

(7) $$\mathbf{E}=\mathrm{R}\mathrm{e}\mathrm{L}\mathrm{U}(\mathbf{U}\cdot \mathrm{a}\mathrm{v}\mathrm{e}(\mathbf{h},\mathbf{t}))$$where $\mathbf{U}\in \mathbb{R}{ }^{k\times k}$ is the transformation matrix, and $\mathrm{a}\mathrm{v}\mathrm{e}()$ denotes the average function, which aims to obtain the comprehensive features among each dimensional of head and tail entity features.

Multi-label relation modeling

Since there may exist multiple relations between an entity pair, we model KGC as a multi-label learning problem. Based on the obtained entity-pair representation in the previous subsection, our model calculates the confidence scores for each relation as follows:

(8) $$\mathbf{S}=\mathbf{V}\cdot \mathbf{E}$$where $\mathbf{V}\in \mathbb{E}{ }^{|R|\times k}$ is the collection of weight vectors for each relation. Afterwards, the sigmoid function is applied on each element of the score vector $\mathbf{S}$ to compute the probability of each relation to exist:

(9) $$p_i=\frac{1}{1+e^{-s_i}},i=\{1,2,...,|R|\}$$where |R| denotes the number of relations.

Model training

We define the loss function using cross-entropy as follows:

(10) $$L=-\sum _i^{|R|}{y}_i\mathrm{l}\mathrm{o}\mathrm{g}({\mathrm{p}}_{\mathrm{i}}))+(1-{\mathrm{y}}_{\mathrm{i}})\mathrm{l}\mathrm{o}\mathrm{g}(1-{\mathrm{p}}_{\mathrm{i}})$$where $y_i\in \{0,1\}$ is the true value for relation $i$, $p_i$ is the predicted probability value for relation $i$. The loss function is optimized with Adam (Kingma & Ba, 2015), and dropout (Srivastava et al., 2014) is employed for regularization.

Datasets

We evaluate our model ShallowBKGC on three benchmark datasets: WN18RR (Dettmers et al., 2018), FB15k-237 (Toutanova & Chen, 2015) and YAGO3-10 (Mahdisoltani, Biega & Suchanek, 2015). WN18RR and FB15k-237 are derived from the lexical KG WordNet (Miller, 1995) and the real-world KG Freebase (Bollacker et al., 2008), respectively. YAGO3-10 is a dataset containing general facts from Wikipedia. The experimental datasets statistics are shown in Table 1.

Evaluation metrics

We use mean rank (MR), mean reciprocal rank (MRR) and Hits@N as evaluation metrics, in which MR is the average rank of all test triplets, MRR is the average of the reciprocal ranks, and Hits@N is the percentage of test triplets that are ranked within top N. They are formally defined as follows:

(11) $$\mathrm{M}\mathrm{R}=\frac{1}{|Triple{t}_{test}|}\sum _{i=1}^{|Triple{t}_{test}|}ran{k}_{triplet(i)}$$

(12) $$\mathrm{M}\mathrm{R}\mathrm{R}=\frac{1}{|Triple{t}_{test}|}\sum _{i=1}^{|Triple{t}_{test}|}\frac{1}{ran{k}_{triplet(i)}}$$

(13) $$\text{Hits}@\text{N}=\frac{|triplet(i)\in Triple{t}_{test}:ran{k}_{triplet(i)}\le N|}{\left|Triple{t}_{test}\right|}$$where $|Triple{t}_{test}|$ is the number of test triplets, $triplet(i)$ is the $i$-th triplet.

Additionally, to evaluate the model efficiency, we measure the running time of the training phase. Record the average time of three epochs of the model on the dataset, in seconds. Our experimental platform is ModelArts, and the specific configuration selected is pytorch1.8-cuda10.2-cudnn7-ubuntu18.04, and a P100 GPU (16G).

Baseline methods

We compare our model against the following state-of-the-art KGC models, including translation-based models TransE and RotatE, tensor decomposition-based models DistMult and ComplEx, neural network-based models ConvE, SHALLOM and ASLEEP, and pre-trained language/large language-based models KG-BERT, KG-ChatGLM-6B, KG-LLaMA-7B and KG-LLaMA-13B. Below we briefly introduce these models.

• TransE (Bordes et al., 2013) is the initial translation-based model that views relations as translations from head entities to tail entities on the low-dimensional space.

• DistMult (Yang et al., 2015) is a typical tensor decomposition-based model that restricts n-by-n matrices representing relations to diagonal matrices.

• ComplEx (Trouillon et al., 2016) extends DistMult to the complex space.

• RotatE (Sun et al., 2019) is an efficient ranslation-based model that represents entities as complex vectors and relations as rotations.

• ConvE (Dettmers et al., 2018) is a deep neural network-based model that applies convolutional neural network for KGC.

• SHALLOM (Demir, Moussallem & Ngomo, 2021) is a shallow neural network-based model for KGC.

• ASLEEP (Jia, 2022) improves the way SHALLOM obtains entity pair representation, and is the prior version of our proposed model.

• KG-BERT (Yao, Mao & Luo, 2019b) is a pre-trained language-based model that firstly employs BERT to KGC.

• KG-ChatGLM-6B, KG-LLaMA-7B, KG-LLaMA-13B (Yao et al., 2023) are large language-based models that perform instruction tuning with ChatGLM (Du et al., 2022) and LLaMA (Touvron et al., 2023) for KGC.

Hyperparameter optimization

We select the hyperparameters of ShallowBKGC by grid search based on Hits@1 of the relation prediction task on the validation set of each dataset. We manually specify the hyperparameter ranges: embedding size among $\{50,100,150\}$, epochs among $\{50,100,150\}$, batch size among $\{512,1000\}$, dropout rate among $\{0.25,0.5\}$, and $L_2$-normalizer among $\{0.1,0.01,0.001\}$. Table 2 shows parameters values in the experiments.

Experimental results

In this section, we compare the performance of our model ShallowBKGC with that of the baseline methods on the widely used relation prediction task. The task of relation prediction is to complete a triplet $(h,r,t)$ with $r$ missing, i.e., to predict the missing $r$ given $(h,t)$. From the relation prediction results shown in Tables 3–5, we summarize our key observations in the following section.

(1) The shallow neural network-based models, i.e., our model ShallowBKGC and the baselines ASLEEP and SHALLOM, outperform the translation-based model TransE, the complex vector-based model ComplEx and RotatE, the deep neural network-based model ConvE, the pre-trained language-based model KG-BERT, and even the large language-based model KG-ChatGLM-6B, demonstrating the effectiveness of shallow neural network for KGC. For example, compared to TransE, DistMult, ComplEx and RotatE, our model ShallowBKGC achieves 45.8%, 44.4%, 45.2%, and 50.7% relative improvements in MR on WN18RR, respectively. Compared to ConvE, our model ShallowBKGC achieves 76.5% and 38.9% absolute improvements in Hits@1 on WN18RR and FB15k-237, respectively. Compared to KG-BERT and KG-ChatGLM-6B, our model ShallowBKGC achieves 0.2% and 11.7% absolute improvements in Hits@1 on YAGO3-10, respectively. It is worth mentioning that the result of KG-ChatGLM-6B is lower than that of KG-BERT. This suggests that it is not the case that the more layers a model has or the newer the technology is, the better the results will be.

(2) Comparing our model ShallowBKGC with ASLEEP and SHALLOM, we can see that the MRR, Hits@1 and Hits@3 values of ShallowBKGC on the three datasets are better than ASLEEP and SHALLOM. This indicates that it is beneficial to take both text and structural information into account for KGC. Because the main difference between our model ShallowBKGC and the baselines ASLEEP and SHALLOM is that our model ShallowBKGC combines text and structural information for KGC, while ASLEEP and SHALLOM merely rely on structural information.

(3) The RTs of our model ShallowBKGC on three benchmark datasets significantly outperform the baselines, which shows the efficiency of our model. It should be noted that in order to minimize the impact of programming differences, we use OpenKE (Han et al., 2018) to reproduce the running times of TransE, DistMult, ComplEx and RotatE. The RTs of ConvE, SHALLOM, ASLEEP and KG-BERT are obtained from their corresponding source codes. Due to permission issues, the RTs of KG-ChatGLM-6B,KG-LLaMA-7B and KG-LLaMA-13B are missing. However, from the perspective of the number of layers and parameter scale, the RTs of these models are likely to be larger than KG-BERT. More formally, the time complexity of our model is $O(d_e)$ (where $d_e$ represents the dimension of entities), which is the same as that of the baselines, except the language-based baselines. From Wang et al. (2021), we can see that the most relevant language-based baseline KG-BERT’s time complexity is $O(L_t^2|E{|}^2|R|)$, where $L_t$ is the length of triple text, |E| and |R| are the numbers of entities and relations respectively. Additionally, the space complexity of our model is $O(L_e|E|d_{token}+|E|d_e)$, where $L_e$ is the length of entity text, $d_{token}$ is the dimension of entity text tokens.

(4) It is worth noting that the results of ConvE are significantly lower than those of the other models, probably because it relies on an improper pre-trained model for initialization, and is trained on entity prediction task (i.e., given $(h,r)$ predict $t$, or given $(r,t)$ predict $h$) but tested on relation prediction task. It has been demonstrated that the initialization, hyperparameter optimization, and training strategies have significant effects on prediction performance (Demir & Ngomo, 2021; Ruffinelli, Broscheit & Gemulla, 2020). In contrast, our model ShallowBKGC is as simple as possible, it does not require retraining the pre-trained model BERT, special hyperparameter optimization approach, or complex training strategy, thus minimizing model uncertainty.

Fine-grained performance analysis

To further verify the capacity of our model from a fine-grained perspective, we plot the percentage of each relation on WN18RR in Fig. 3, and report the Hits@N and MRR performance on each relation in Figs. 4 and 5, respectively.

From these figures, we can observe that:

(1) From the perspective of Hits@N, there are five relations, i.e., _similar_to, _member_of_domain_region, _instance_hypernym, _member_meronym, and _derivationally_related_form, Hits@1 values exceed 90%. Ten relations Hits@3 values exceed 90%. In particular, there are four relations, i.e., _similar_to, _member_of_domain_region, _verb_group and _hypernym, Hits@3 values reached 100%. Thus, the results are consistent with Table 4, which further demonstrates the effective of our model at a fine-grained level.

(2) There are eight relations with MRR values exceeding 90%. They are _similar_to,_member_of_domain_usage, _member_of_domain_region, _synset_domain_topic_of, _instance_hypernym, _member_meronym, _derivationally_related_form and _hypernym. Moreover, five of these eight relations Hits@1 values exceed 90%. This experimental result once again demonstrates the effectiveness of our model and the consistency of the results at a fine-grained level.

Ablation study

We conduct ablation studies to provide a more detailed analysis of the effectiveness of each part of our model. The models used for comparison are the following ones: (a) ShallowBKGC-Text is the model that merely rely on text information, and (b) ShallowBKGC-Structure is the model that merely rely on structural information.

Table 6 shows the relation prediction results on three datasets, i.e., FB15k-237, WN18RR, and YAGO3-10. From which we can see that:

(1) ShallowBKGC outperforms ShallowBKGC-Text and ShallowBKGC-Structure in all three datasets, indicating that considering both text and structural information is beneficial to KGC. This is also consistent with the results of the previous experiments.

(2) ShallowBKGC-Structure achieves better results than ShallowBKGC-Text, the main reason is that we set the entity text feature obtained by BERT untrainable to reduce the computational complexity, this sacrifices the performance to a certain extent. It is worth further explaining that from Tables 5 and 6, we can see that ShallowBKGC-Text outperforms KG-ChatGLM-6B, considering that the latter has far more parameters than the former, but the experimental results are very close, which also reflects the efficiency of our model.

(3) ShallowBKGC-Text has the shortest RT because it has the fewest parameters. The RTs of ShallowBKGC-structure and ShallowBKGC are close, indicating that our model do not spend much time fusing text and structural information.