Tuck-KGC: based on tensor decomposition for diabetes knowledge graph completion model integrating Chinese and Western medicine

Knowledge embedding

To effectively incorporate entity type information and knowledge graph triples into the vector space, we employ the TransE embedding method. The ultimate embedding of an entity is generated through the interaction of its associated entity category information auxiliary representation. For head entity embedding, we multiply the vector of head entity representation with the auxiliary entity type vector to obtain the final representation of the head entity. Similarly, we adopt the same approach to get the final representation of the tail entity. The equations are as follows: (1)${\displaystyle ∥e_s\circ e_{\tau }+w_r-e_o\circ e_{\tau }∥.}$

Where e_s, e_o ∈ ℝ^d_e represent the embedding vectors of the head entity and tail entity, e_τ represent the entity type embedding vectors, w_r ∈ ℝ^d_r represent relation embedding vectors.

Link prediction

Given a knowledge graph $G=\left(E,R,T\right)$, where E and R represent the sets of entities and relationships, respectively, and T represents the set of triples in the format of (e_s, w_r, e_o) ⊂ E × R × E; Tuck-KGC is an approach to solving the knowledge graph completion problem through the link prediction task, which aims to construct new triples using the existing set of triples T and entity types E_τ. Tuck-KGC addresses the issue of knowledge graph incompleteness by engaging in the link prediction task, with the objective of constructing a new triple $\left(e_s^{\prime },w_r^{\prime },e_o^{\prime }\right)$, where $e_s^{\prime },e_o^{\prime }\in E,w_r^{\prime }\in R$. On this basis, a suitable scoring function φ is used to determine the reasonableness of the triad and its plausibility and to predict the new relationship, where the scoring function is defined as $\phi \left(e_s,w_r,e_o\right):E\times R\times E⟶R$. The new triple $\left(e_s^{\prime },w_r^{\prime },e_o^{\prime }\right)$ is generated by replacing one of the items that belongs to the existing set of triples (e_s, w_r, e_o), where the item belongs to the same category as the replacement item. The freshly created triples are subsequently assessed through the scoring function φ, thereby converting the link prediction task into a sorting challenge. In this process, the scoring function φ evaluates each of the newly generated triples, followed by their arrangement in descending order of score.

Tuck-KGC implementation

To ensure the triad’s bidirectional prediction, we utilize the same 𝔈 embedding matrix to portray the head and tail mapping modelling A, C matrices, respectively, i.e., 𝔈 = A = C ∈ ℝ^n_e×d_e. At the same time, we map the relational embedding matrix R to the B matrix in the tensor decomposition modelling, i.e., ℜ = B ∈ ℝ^n_r×d_r. Upon this foundation, we incorporate an entity type matrix denoted as Γ ∈ ℝ^n_τ×d_τ which we use to refine the entity matrix operation. This integration of entity type information enriches the semantic representation of entities in the link prediction task. Additionally, technical abbreviations will be explained upon initial use.

In this context, n_e and n_r represent the respective quantities of entities and relations, while d_e, d_r, d_τ signify the dimensions of the embedding vectors for entities, relations, and entity types. The scoring function for the Tuck-KGC model, which is based on the tensor decomposition of fused-type embeddings, is formulated as follows: (2)${\displaystyle \phi \left(e_s,w_r,e_o,e_{\tau }\right)=\left(\left(\mathfrak{W}{\times }_1{e}_s{\times }_1{e}_{s_{\tau }}\right){\times }_2{w}_r\right){\times }_3{e}_o{\times }_3{e}_{o_{\tau }}}$

where 𝔚 ∈ ℝ^{d_e×d_r×d_e} represents the core tensor, e_s, e_o ∈ ℝ^d_e are the rows of 𝔈 representing the embedding vectors for the head entity and tail entity, e_sτore_oτ the rows of Γ representing the head entity type or tail entity type embedding vectors, w_r ∈ ℝ^d_r the rows of R representing relation embedding vectors, ×_n denotes the tensor product along the n-th mode. Equation (2) describes the methodology by which Tuck-KGC prioritizes all triples in the dataset, integrating the entities and their corresponding relationships with their specific entity types. The triple with the highest ranking is subsequently designated as the prediction outcome.

Throughout the training phase, we map triples and entity types into a vector space and input the resulting vectors into the model. We then score each triple and calculate the loss value through the model’s scoring function. Parameters are updated using backpropagation. The model undergoes repeated training iterations until the loss value is minimized, and the best results are achieved. The loss function for an entity-relationship pair in comparison to all other entities is defined as: (3)${\displaystyle \mathrm{Loss}=-\frac{1}{n_e}\sum _{i=1}^{n_e}\left(y^{\left(i\right)}log\left(p^{\left(i\right)}\right)+\left(1-y^i\right)log\left(1-p^{\left(i\right)}\right)\right).}$

Where n_e represents the dimension of the entity embedding vector space, $p^{\left(i\right)}$ signifies the vector of predicted probabilities, and $y^{\left(i\right)}$ is the binary label vector. To combat overfitting, this research incorporates dropout, complemented by the Adam optimizer and label smoothing techniques during the training process. The detailed architecture of the model can be observed in Fig. 1.

Experiments

Datasets

The Dia dataset employed in this investigation comprehensively includes data from both Traditional Chinese medicine (TCM) and Western medicine diabetes data. The Western medicine dataset was sourced from the DiaKG (Chang et al., 2021) dataset available from Ali Tianchi Ruijin Hospital and from the diabetes-related section of SDKG-11 established by Zhu et al. (2022). The DiaKG dataset includes 22,050 medical entities and 6,890 annotated entity relationships, incorporating insights from 41 diabetes experts. The dataset covers various types, including medication usage, clinical cases, diagnosis and treatment methods, and clinical research. The SDKG-11 dataset, created by Zhu et al. (2022), comprises triads extracted from publications indexed in PubMed from the year 2020, which boast an impact factor of 2.0 or higher. We meticulously isolated the diabetes-related subset, which primarily consists of fundamental research and theoretical foundations of diabetes.

The Traditional Chinese Medicine diabetes dataset, used in this study is derived from the TCM guidelines for diabetes and its related complications (Diabetes Branch of the Chinese Association of Traditional Chinese Medicine, 2011a; Diabetes Branch of the Chinese Association of Traditional Chinese Medicine, 2011b). The TCM diabetes data subset also incorporates Li & Zhang (2023) perspective of syndrome differentiation based on syndrome elements, which is constructed manually following a main-predicate-object structure (Fang et al., 2017; Yang, 2022). Firstly, medical terminology entities, their attributes, and relationships between these entities are extracted from his monographs by human labor. Finally, based on TCM theories, the perspectives of Zhu (2006), and entity relationships, triples are constructed, which are ultimately merged into the TCM diabetes data subset. The knowledge graph that incorporates Zhu (2006) thoughts on syndrome differentiation and treatment can better support the construction of intelligent-aided diagnosis models based on traditional Chinese medicine theory.

Due to the presence of biomedical entities synonyms and multiple layers of semantic duplication, duplicate triples can be found in the Diabetes dataset for both Chinese and Western medicine. Initially, the MTransE (Chen et al., 2016) model was used to align entities in the raw data. Based on Gong et al. (2021) method of integrating Chinese and Western medicine knowledge to construct a diabetes knowledge graph, we also incorporated entity type information and developed a concept definition module, entity classification module, an attribute division module, and relationship matching module.

Initially, we input Chinese and Western medicine entities into the entity concept layer and established the disease layer, clinical performance layer, and treatment layer. Subsequently, we used the entity classification module to input various generic entities into the entity concept layer, classifying them based on diabetes symptoms and the corresponding treatments and protective measures. The classification process is to classify various common entities related to traditional Chinese and Western medicine according to their relevance to diabetes symptoms and their treatment and prevention methods (Wei, 2021). Subsequently, we input the types of diabetes symptoms along with the corresponding TCM treatments and prescriptions as documented in the TCM treatment process. We then provide the clinical manifestations of diabetes, its complications, and the corresponding treatments and protective methods in Western medicine. Additionally, we list the distinct entities associated with TCM and Western medicine separately. The attribute classification module categorizes the TCM and Western medicine entities in the treatment layer according to their attributes. The relationship matching module, was employed to identify the potential relationships between Chinese and Western medicine entities, thereby providing a theoretical foundation for the interconnection between these entities and the generic entities. Finally, the construction module connects the entities based on the entity’s inherent relationship in the entity concept layer and the types determined by the entity classification module, resulting in the formation of a ternary-based knowledge graph. This process yielded a total of 10,294 medical entities and 12,863 pairs of entity relationships, as shown in Table 2. For a detailed breakdown of the entity concept layer, please refer to Appendix S1 (Tables SA1–SA3).

The construction of the Dia dataset, integrating diagnostic and therapeutic knowledge of diabetes from both Traditional Chinese medicine and Western medicine, will be divided at random into training, validation, and test subsets in this research with proportion of 0.6, 0.2, and 0.2, respectively.

Evaluation indicators

In this paper, we reference the evaluation metrics introduced by Balažević, Allen & Hospedales (2019) to demonstrate the efficacy of our model in the link prediction task.

We utilize standard evaluation metrics commonly employed in the link prediction literature, namely, mean reciprocal rank (MRR) and Hits@k (where k = 1, 3, 10). The mean reciprocal rank is computed as the average of the reciprocals of the mean ranks assigned to the correct triple among all candidate triples: (4)${\displaystyle \mathrm{MRR}=\frac{1}{\left|T_{\mathrm{test}}\right|}\sum _{\left(h,r,t\right)\in T_{\mathrm{test}}}\frac{1}{2}\left(\frac{1}{{\mathrm{rank}}_{r,t}\left(h\right)}+\frac{1}{{\mathrm{rank}}_{h,r}\left(t\right)}\right)}$

where rank_r,t(h) denotes head entity ordering, and rank_h,r(t) denotes tail entity ordering. The value of MRR ∈ (0, 1) ranges from MRR and the larger the calculated value, the better the performance of the model link prediction. Hits@k quantifies the percentage of instances where a true triple is positioned among the top k candidate triples: (5)${\displaystyle \mathrm{Hits}@k=\frac{1}{\left|T_{\mathrm{test}}\right|}\sum _{\left(h,r,t\right)\in T_{\mathrm{test}}}\frac{1}{2}\left(\left|\left\{\left(h,r,t\right)|{\mathrm{rank}}_{r,t}\left(h\right)\le k\right\}\right|+\left|\left\{\left(h,r,t\right)|{\mathrm{rank}}_{r,t}\left(t\right)\le k\right\}\right|\right).}$

Hits@k focuses on the overall ranking, with larger values indicating that the model performs better in the link prediction task.

Baseline and parameter

We compare Tuck-KGC with other link prediction models, including TransE (Bordes et al., 2013), DistMult (Yang et al., 2014), ComplEx (Trouillon et al., 2016), SimplE (Kazemi & Poole, 2018), relational graph convolutional neural networks (R-GCN) (Schlichtkrull et al., 2018), ConvE (Dettmers et al., 2018), InteractE (Vashishth et al., 2020), and RotatE (Sun et al., 2019), where entities and relations are embedded into vector space through linear operations; DistMult (Yang et al., 2014), is a general neural network framework for multi-relational representation learning; ComplEx (Trouillon et al., 2016) extends DistMult to the complex domain; SimplE (Kazemi & Poole, 2018) which encodes each entity with two separate embedding vectors entity actual connection in the knowledge graph; R-GCN (Schlichtkrull et al., 2018) is an extension of graph convolutional networks for relational data; ConvE (Dettmers et al., 2018), which utilizes 2D convolution to learn deep features of entities and relations; InteractE (Vashishth et al., 2020) replaces the simple feature reshaping used in ConvE with check reshaping and circular convolution; RotatE (Sun et al., 2019) charactering each relation as a rotation from a source entity to a target entity in a complex vector space.

The Tuck-KGC model is implemented in Pytorch unifying the Tuck-KGC and baseline models on NVIDIA Gefiorce RTX3090Ti GPU. We establish the embedding dimension for both entities and relations as d_e = d_r = 400, the embedding dimension of entity types to d_τ = 400, the learning rate selected from the set {0.01,0.005,0.003,0.001,0.0005,0.0003}, and the learning decay rate to be chosen from {1.0,0.998,0.995,0.99}. Experimental outcomes reveal optimal performance at a learning rate and learning rate decay of (0.0003,1.0). Furthermore, our experiments demonstrate the dropout value of (0.3,0.4,0.4,0.5) can effectively mitigates overfitting. The baseline models mentioned are implemented using OpenKE (Han et al., 2018) toolkit. Table 3 shows the detailed parameter list of the Tuck-KGC model.

Influence of parameter

To test the performance of the Tuck-KGC model, we conducted a comparative analysis focusing on the entity embedding dimension d_e, the total number of training rounds N, and the sensitivity to the sparsity of the knowledge graph on the Dia dataset. We established the dimensions for entity embedding as $d_e,d_{\tau }\in \left\{100,200,300,400,500,600,700\right\}$ and the number of training rounds as $N\in \left\{100,200,300,400,500,600\right\}$. The experimental parameters were harmonized with those previously described in the baseline and parameter configurations, except for the study-specific variables. The results of the experiment are visually depicted in Fig. 2.

Figure 2A shows the variation in the Hits@10 metric for our Tuck-KGC model as the entity vector dimensionality is altered. It is observed that with the increase in the entity embedding dimension, each model exhibits varying degrees of enhancement. The optimal performance is achieved at an entity embedding dimension of 400. Beyond this threshold, the model’s efficacy tends to decline and stabilize. Our findings indicate that an entity embedding dimension of 400 yields the most favorable results, highlighting the critical role of this parameter in the model’s performance. An insufficiently low dimension may result in underfitting, impeding the model’s ability to effectively capture entity information, while an excessively high dimension can lead to overfitting, thus compromising the model’s performance.

Figure 2B shows the variation trend of the Hits@1indicator for our Tuck-KGC model as the training iterations increase. It is evident from the illustration that with the progression of training iterations, each model exhibits varying degrees of enhancement, culminating at the optimal performance when the iteration count reaches 300. Subsequently, surpassing this threshold leads to a gradual decline and stabilization of the model’s outcomes. The findings suggest that the model achieves peak efficacy at 300 training iterations, highlighting the critical role of iteration count in influencing model performance. An insufficient number of iterations may result in underfitting, impeding the model’s ability to effectively assimilate entity information, while an excessive number can lead to overfitting, detrimental to model performance.

To assess the sensitivity of Tuck-KGC to the sparsity of the knowledge graph, we randomly ignore triples from the Dia dataset and evaluate the effect on the entire test set. Figure 3 shows the test results of Tuck-KGC and the basic model on Hits@3. As depicted in Fig. 3, as the ratio of omitted triples increases, the performance of all models experiences a decline to varying extents. However, our method outperforms other baseline models throughout this decline, indicating a greater robustness to sparsity in our model compared to the baseline.

Visualization of clustering entity type representations

We cluster the type embeddings using Kmeans and further implement dimensionality reduction using t-distributed scholastic neighbor embedding (t-SNE) for 2D visualization. As shown in Fig. 4A, different types of entities are clustered into separate categories in TuckER, while some clusters are close to each other because these entities share many common types. Figure 4B shows the clustering of entity embeddings of Tuck-KGC. It can be clearly observed that entity clustering with fused entity types can better distinguish entities, indicating that entity type embeddings can reflect the characteristics of entities. In other words, when an entity represents different meanings in different triples, it is difficult to distinguish. However, when given different types, the same entity has different entity type embeddings when encoded, so it is distinguishable. For example, we introduced a specific case in the “Introduction”, the hypoglycemia entity type in the triple hypoglycemia, ADE_Disease, renal failure is a symptom of the disease, but hypoglycemia entity type in the triple fasting, Reason_Disease, hypoglycemia is the disease. These visualizations explain the effectiveness of using learned entity type embeddings in our knowledge graph.