CPa-WAC: Constellation Partitioning-based Scalable Weighted Aggregation Composition for Knowledge Graph Embedding

1 Introduction

Knowledge graph (KG) has gained immense popularity in recent years. A KG is a diverse multigraph consisting of more than one type of directed relation [Zamini et al., 2022] between node entities. Each KG contains a collection of facts organized as a graph to represent a group of linked entities and their semantic descriptions [Zamini et al., 2022]. First proposed in 2012 by Google, these semantic networks are widely used today, from drug interaction prediction in chemistry to fraud detection in financial transactions to entity alignment [Surisetty et al., 2022; Chaurasiya et al., 2022]. To date, many KGs have been created, such as Freebase [Bollacker et al., 2008] and Wordnet [Miller, 1995], which illustrate real-life relations between entities. Knowledge graphs are almost always incomplete [Ye et al., 2022]. Therefore, to complete such graphs, it is essential to embed them and use these embeddings to predict relations between entities. Knowledge graph embedding (KGE) is considered a foundation for several prediction tasks using KG [Ye et al., 2022].

Many methods utilize neural networks to compute embeddings of nodes and relations and use them for inference. Similarly, some methods employ deep learning to learn recursive logical rules [Meilicke et al., 2019; Cheng et al., 2022] for reasoning over KG. Methods such as graph neural networks (GNN) [Wu et al., 2020; Surisetty et al., 2023] are neural network-based frameworks that have recently been integrated into the KG domain and are primarily used to learn embeddings of nodes and relations. Relational graph convolutional networks (RGCN) [Schlichtkrull et al., 2018] was the first GCN-based method introduced in the field. Since then, more sophisticated approaches have been developed. Graph convolutional and attention networks (GAT) [Veliˇckovi´c et al., 2018] have also been quite successful in achieving high accuracy of prediction, superseding convolutional networks [Dettmers et al., 2018]. However, the information-rich embeddings from these algorithms come at a significant cost of computational and memory usage.

Several state-of-the-art algorithms such as Comp-GCN [Vashishth et al., 2019], RAGAT [Liu et al., 2021], and SEGNN [Li et al., 2022a] use high amounts of memory (GPU) and require immense amounts of training time. With a large number of nodes and edges, the time complexity to process these graphs can be considered quadratic [Wu et al., 2020] due to millions of trainable parameters. As a result, several of these algorithms can only be trained with small batch sizes. Partitioning KG into subgraphs and processing these subgraphs can be depicted as a solution to this problem [Puja et al., 2013]. Partitioning reduces the number of computations in both forward and backpropagation. At the same time, it enables the training process to be carried out parallelly. While a few methods employ partitioning and parallel training, such as the ones presented in [Kochsiek and Gemulla, 2021; Sheikh et al., 2022; Bai, 2023], it is evident that dedicated partitioning approaches are required that can partition KGs without the utilization of node or edge features.

This research presents a novel composition-based KGE method using graph neural networks that utilizes the power of KG partitioning to make the training process fast and efficient. We call the proposed algorithm constellation partitioning-based weighted aggregation composition (CPa- WAC). We harness the power of modularity maximization [Newman, 2006] through Louvain clustering to generate partitions from KG. By preserving the original graph topology, the proposed method increases link prediction performance compared to other partitioning-based state-of-the-art methods while decreasing the train time. Furthermore, we compare the Louvain algorithm to the Leiden algorithm [Traag et al., 2019] and empirically verify the effects of utilizing both algorithms on KG partitioning.

The main contributions of this paper are as follows: i) A novel KG partitioning algorithm (CPa) that utilizes fast Louvain clustering [Blondel et al., 2008] to partition the KG into several topological clusters. During partitioning, we aim to minimize the number of lost links between clusters; ii) An improved compositional-GCN algorithm, WAC. The composition operation is further coupled with the multiplication operation that is presented in [Kazemi and Poole, 2018] and a 1D convolutional network that takes advantage of feature, entity, and relation-specific weights to learn effective embeddings; iii) A global decoder framework that can use node and relationship embeddings from different clusters to achieve a global-level inference, and lastly (iv) exhaustive comparison on four distinct datasets, namely, WN18, WN18RR [Miller, 1995], FB15K, and FB15K-237 [Bollacker et al., 2008] and several state-of-the-art methods in the field, along with a detailed analysis of each comparison is presented in the paper.

2 RelatedWork

2.1 Knowledge Graph Embedding

Several methods have been used to train a KG and obtain embeddings. Models such as the variants of Trans: TransE [Bordes et al., 2013], TransH [Wang et al., 2014], TransR [Lin et al., 2015], TransD [Ji et al., 2015], and TransG [Wang et al., 2017] were pivotal for KGE towards link prediction, node classification, and reasoning tasks. However, more complex semantic methods aimed to capture semantic relationships, such as Conv2D [Dettmers et al., 2018], RESCAL [Bordes et al., 2014], ComplEX [Zhou et al., 2017], TuckER [Balazevic et al., 2019], HAKE [Zhang et al., 2020], and SimplE [Kazemi and Poole, 2018]. While these methods are good at capturing semantic relationships, they require a high dimensionality in the embedding space [Chen et al., 2020] to embed entities and relationships for more accurate predictions.

Recent architectures such as GCN [Schlichtkrull et al., 2018] and graph attention networks (GAT) [Veliˇckovi´c et al., 2018] have been utilized in studies such as [Vashishth et al., 2019; Liu et al., 2021; Li et al., 2022a; Li et al., 2022b]. If pruned properly, these architectures produce better results than their non-GCN-based counterparts. Nevertheless, GCN and GAT-integrated models have high trainable parameters and take enormous time to train on KG datasets. Even though there are libraries available, such as Pytorch-Biggraph [Lerer et al., 2019] and DGL-KE [Zheng et al., 2020], these libraries do not address the scalability of GNN-based knowledge graph embedding algorithms.

2.2 Reducing Training Time of Knowledge Graphs

One way to reduce training time lies in the task of KG augmentation. KG augmentation can be used to learn embeddings in a lower dimensional vector space. As demonstrated by [Wang et al., 2022a], KG augmentation, when used in sequence with a Trans or semantic series embedding model, can reduce the processing time and increase the overall performance of these models. Methods such as GreenKGC [Wang et al., 2022b] and DGL-KE [Zheng et al., 2020] have addressed the reduction of embedding dimensionality and the training time using feature pruning, partitioning, parallel training, and multi-GPU training.

Partitioning is critical to scalable GNN-based KG embedding generation. Partitioning can be used to split a KG into multiple subgraphs without destroying the original topology of the graphs. This statement is backed by the analysis of Jain et al. in their work presented in [Jain et al., 2021], which showcased that the semantic representations of embeddings are not foolproof over the entire graph. Instead, semantic features are only contained locally, implying that partitioning will not distort the overall graph structure. Ontologybased KG partitioning has been employed by Bai in his work in [Bai, 2023], which breaks the KG into multi-partitions. The method shows high performance on partitioned datasets, compared to the performance on the full dataset. Similarly, METIS [Karypis and Kumar, 1998] and k-means clustering have been utilized in [Wang et al., 2022b] and [Zheng et al., 2020], respectively, followed by Trans series algorithms to partition KG into disjoint clusters to reduce the training time. Apart from traditional k-means and METIS partitioning, some novel KG partitioning methods are presented in the works [Sheikh et al., 2022] and [Priyadarshi and Kochut, 2021]. The work in [Sheikh et al., 2022] presents an edge-cut partitioning method that can aggregate information from only essential nodes from connected subgraphs after partitioning. Similarly, the work presented in [Priyadarshi and Kochut, 2021] has proposed a workload-aware knowledge graph partitioning that analyses training triples and the connected KG to form partitions.

Despite several studies on KG partitioning-based embedding models, properly partitioning a KG with the least cross partition edges remains challenging. Furthermore, a framework is required to merge these individual embeddings into one complete graph structure. Considering these methods in the field, a dedicated KG partitioning architecture is required that ensures the least cross-cluster relations. Furthermore, developing a lightweight GCN-based architecture is essential, which can be coupled with partitioning to learn effective embeddings from KG faster. In this paper, we address this problem by proposing a novel KG partitioning method and a composition-based aggregation method to achieve faster training while keeping the prediction accuracy high.

3 Materials andMethods

The methodology is divided into three stages: Partition generation (3.2), Weighted Aggregation Composition Convolution (3.3), and the partitioning-based decoder framework for link prediction, (3.4). The proposed CPa-WAC algorithm is summarized in Figure 1. The code is made available for the research community 1 , and summarized in Algorithm 1.

3.1 ProblemDefinition and Notations

A knowledge graph is a semantic network, represented as KG = {E,R, T}. Here E, R, and T represent entities, relations, and triples, respectively. The triples in a KG can be regarded as T = {ϱ, µ, ξ}, where ϱ, µ, and ξ symbolizes a head, relation, and tail, respectively. The relations are directed and represent a link between two entities. This study aims to empirically verify that partitioning can be used to speed up KGE without destroying the structure of the KG and jumbling the inference logic.

3.2 Constellation Partitioning: CPa

As the first step towards speed-up, the proposed method employs Louvain clustering (LC) [Blondel et al., 2008] as its base. The essence of LC lies in community detection/subgraph generation, which partitions a given homogeneous graph based on the density of its edges. This method exclusively suits graphs without node and edge attributes, making it a perfect fit for KG partitioning. The proposed constellation partitioning (CPa) algorithm utilizes the LC algorithm on KG to partition it into several topological structures. The process starts by creating a symmetric adjacency matrix (A). If a link exists between two nodes in a graph, then a 1 is added at positions in A representing the connection. This process is repeated for every type of relation for all nodes. This weighted adjacency matrix can be fed to the LC algorithm to obtain the initial number of clusters using modularity score maximization. The modularity score Mc for each cluster c is given as:

Here, αin is the summation of the total weights of all links contained only within cluster c. Similarly, αall represents the summation of the weights of links that the nodes of cluster c have with each other and also with nodes from other clusters. The variable ω is the total weight of all links in the graph. Once the initial number of clusters are formed using the LC algorithm, using a hard threshold of δ, clusters having only a few entities are all merged into one cluster. This groups those few outlier entities in KG that have very few links and can distort the embedding process. Here, threshold δ is a KGspecific hyperparameter.

Despite creating one merged cluster with a threshold δ, some clusters with a small number of entities may remain. This is due to the LC algorithm's limitation in partitioning heterogeneous directed graphs. Therefore, merging these small clusters with bigger and denser clusters is essential in avoiding loss of structural information. To do so, we utilize three levels of hierarchical merging using Φ=γ ×β ×δ. Here, δ is the initial threshold from the previous step, γ is the incremental step, and β is also a hyperparameter to control the number of clusters created. This finally gives us the C number of clusters. The logic is that if the number of entities in a cluster is below a threshold Φ, it will be merged with a cluster

1 https://github.com/ganzagun/CPa-WAC

Algorithm 1 Louvain Constellation Partitioning

• 1: Initialize: A, A' as square matrices of zeroes with each dimension equal to the number of nodes, N
• 2: Initialize: γ, β, δ, σ 3: For i = 1 : length(T) 4: Add 1 to positions in A(T(i,0),T(i,2)) 5: Add 1 to positions in A(T(i,2),T(i,0)) 6: A′(T(i,0),T(i,2))=1 7: A′(T(i,2),T(i,0))=1 8: End 9: Apply LC to A 10: Maximize modularity using Equation 1 11: Cluster outliers with threshold δ 12: For j = 1 : 3 (for three levels of hierarchical merging) 13: Obtain k number of clusters and store in vector L to map each entity to a cluster 14: Create a list U with all entities present in each cluster 15: For m= 1 : N 16: B = A′(m) × L 17: The label at position A′(m,m) = l 18: For p = 1: k 19: x = Number of entities in B with label p 20: Store x in vector F at position F(p) 21: End 22: Find argmax(F) whose label is not l 23: Store p in vector M at position Mm. This is the cluster that is adjacent to node m 24: End 25: For q = 1 : k 26: Identify positions in vector L with label q 27: Take the corresponding positions from vector M and store them in vector D 28: Based on Majority voting of nodes in vector D select the nearest cluster, y 29: Φ= γ × β × δ 30: If length(U(y)) < σ and length(U(q)) < Φ 31: Merge the two clusters by applying labels from the bigger cluster y, to the smaller cluster q 32: End 33: End 34: γ= 2 * γ (incremental step) 35: End

with a higher number of entities with the most links. We call this the nearest linked neighbor (NLN). The three incremental thresholds applied to the clusters ensure steady growth of entities in highly dense clusters. To avoid entity explosion, a second threshold of σ is used to cap the maximum number of entities in a cluster. If the number of entities in a cluster > σ, then a smaller cluster will not be merged with it even if this cluster is its NLN. Hence, for clusters containing a higher number of entities, smaller clusters that are designated as NLN will be merged iteratively with them. Thus, a relatively similar number of entities is obtained for each cluster while the overall modularity of the graph remains high.

3.3 Weighted Aggregation Composition (WAC) Convolution

In this section, we present an improved version of compositional message-passing [Vashishth et al., 2019], namely Weighted Aggregation Composition (WAC) Convolution. We harness the power of graph attention layers [Liu et al., 2021] to obtain the embeddings of entities and relations by passing through two different composition functions for message aggregation. These functions are denoted by,

where eu is the embedding of the entity number u and r is the relation embeddings. The variables, wu, wv, and wb, denote the learnable weight vectors needed to transform the entities and relationships. Here, hu and hv represent the composition operations needed to transform a particular entity, once updating itself using its entity weight and self-loop (Eq. 2) and a second time by the weights and messages from its neighbors (Eq. 3). The term g represents the current state of embeddings, v represents the neighbor number of the uth entity, b represents the relation type number out of z number of relations that connects the uth entity to its neighbors, and f represents the total number of neighbors that the uth entity is connected to, and, '×' is element-wise feature multiplication.

Once hv and hu are calculated, an activation function is used for each level of message aggregation. The Gaussian error linear unit (GELU) activation function ψ is chosen to update the message-passing function. Next, the summation of neighborhood messages (hv) is passed through an attention layer after normalizing the messages with the degree matrix G. The attention layer, Γ, adds specific attention to all features for each entity. Finally, we take a weighted sum of hv and hu to produce the embeddings eu. The entire process can be summarized using the following equation,

where g + 1 represents the next stage of the entity embeddings, and ζ and θ are scalar weights and are separate for each level of aggregation. The relation embeddings are updated using,

where vb represents a separate learnable weight vector from wb and can be updated using backpropagation to emphasize more dominant relations in the KG.

Finally, a 1D convolutional neural network (1D-CNN) decodes the embedding. It uses a multiplication operation of entities and relation embeddings, similar to the work in [Kazemi and Poole, 2018]. This module consists of several filters and a dense layer, which is trained iteratively to produce meaningful features from obtained embeddings. Batch normalization is utilized twice, once after the convolutional process and the second after passing through the dense layer. The binary cross-entropy (BCE) loss function trains models on each cluster of entities and relationships separately.

3.4 Global Decoder (GD) Framework & Inference

Once the embeddings for entities and relations are obtained, the method moves to its last stage where a separate framework is created and trained for inference. This module is named the global decoder (GD) framework entire training data is fed to this network a second time to train the weights of the MLP

This framework concatenates the feature vectors for all nodes and relations. Let an entity embedding ecu for uth entity from a cluster c be of dimension 1×s. Considering C generated clusters, embedding ecu is upscaled to a dimension 1×Cs. In upscaled ecu , a portion of values is the same embedding, while the remaining 1×(C-1)s values are zeros. These features are then fed into an MLP after multiplying the embeddings with trainable weight matrices, We and Wr. Here, We and Wr are the weight matrix that projects the upscaled entity and relation embeddings to lower dimensions, respectively. These weights are then trained along with the weights for the MLP in an end-to-end fashion utilizing a multiclass BCE loss. The final model and the embeddings of nodes and relations are then used for inference. For testing, the nodes and relations of the test set are fed to the model to predict the tail. Similarly, the reverse is done for predicting the head.

4 Experimental Setting

4.1 Implementation and Database Details

All experiments have been conducted on an I7-13700, 2.1 GHz system with 32 GB RAM and NVIDIA RTX A2000 12 GB GPU. The AdamW optimizer is utilized to train the weights of the proposed architecture for a total of 400 epochs for all partition-based experimentation. Consequently, for comparing the WAC convolution to other architectures, we show results for a model trained on 400 and 1500 epochs (as some papers in literature use 1500 epochs as their setting). Furthermore, all state-of-the-art models have been implemented using the same hyperparameter settings as the proposed architecture. This setting includes batch number, regularization rate, weight decay, learning rate, and the same embedding dimensions for each model.

As summarized in Table 1, we use the most widely used Wordnet [Miller, 1995] and Freebase [Bollacker et al., 2008] for our experiments. Two versions of each dataset are used, namely WN18 and WN18RR for Wordnet and FB15K and FB15K-237 for Freebase. These four datasets capture different denseness and scales of the graph. Particularly, FB15K and FB15K-237 are large-scale KG, making them vital for research on scalable KG.

4.2 EvaluationMetrics

We evaluate the proposed method using popular KGE evaluation metrics of mean reciprocal rank (MRR) and Hits@K. Furthermore, partitioning is evaluated using a novel proposed evaluation metric, that is, tail and head in the same partition (THP). The metrics are defined as follows:

where Q is the number of reciprocal ranks, d is the element number, and rank is the position of the highest-ranked answer (1, 2, 3, ..., X) for X answers returned in a query. Here, o+ is the number of positives occurring in the top-X positions, and Sall is the total number of triples in the dataset. For C number of obtained partitions, THP can be defined as the number of triples Sn with the head and tail within the same partition compared to the total number of triples, Sall.

5 Results and Analysis

This section talks about results and comparisons from proposed CPa-WAC and other state-of-the-art methods in the field of KG partitioning and optimization.

5.1 Partition Analysis

Using THP as an evaluation metric for partitioning, we evaluate the proposed partitioning method CPa over four KG datasets in Table 2. Over varying cluster counts, we observe that over 70% of the training triples have their head and tail within the same partition. This even true at as high as eight separate clusters. For the test and validation sets, except for dataset FB15k-237, all datasets show a high THP over 70%. Despite THP being inversely proportional to the number of clusters, the proposed CPa method can retain much of the original structure of the KG even after partitioning. This is particularly evident for dense KG such as FB15k and WN18.

5.2 EmbeddingMethod Analysis

As the next step, we analyze the effectiveness of the proposed WAC Convolution method on link prediction and compare it to state-of-the-art methods that employ GNN in their architectures. The comparison is done in terms of the MRR and average time per epoch of training (T/epoch) in seconds. For this experiment, we consider the whole KG for each dataset with predefined training, valid, and test sets.

From Table 3, we observe that the proposed method shows slightly better performance in terms of MRR. Alongside marginally better performance, the reduction in time for each epoch is considerably high, giving WAC an edge over the rest in terms of faster training of embeddings. The proposed WAC not only improves the MRR for denser graphs such as FB15k and WN18 but also speeds up the training immensely for these graphs. For sparser graphs such as FB15k-237 and WN18RR, while the MRR does not show much improvement, the performance speedup is twice compared to SEGNN [Li et al., 2022a] and approximately 1.7 times compared to RAGAT [Liu et al., 2021].

5.3 Partition Performance Analysis

We analyze the acceleration rate in comparison to the performance drop by applying the proposed method on several KG. Figure 2 depicts the effects of CPa coupled with LC and WAC on MRR and Hits@K for C number of clusters. In this experiment, we train a separate architecture to infer links using our global decoder framework. In this setting, the obtained features from partitions are used to retrain the dense layers as described in Section 3.4.

From the experiments, we can conclude that disjoint subgraphs created by the proposed CPa-WAC method retains the essence of the original structure. With a global decoder, CPaWAC outperforms, with the observed MRR and Hit@K

higher than the whole KG (for 4 and 5 different partitions). From Figure 2, we see that despite a slight reduction in MRR

and Hits@K with increasing partitions, the dip is small up to 6 partitions. Note that setting C=1 translates to processing the whole KG together. In most cases, the acceleration increases with the increase in the number of partitions. We further employ parallelization of partitions on a single GPU by invoking subprocesses. As the number of entities and relationships is far less for each partition compared to the original KG, a single GPU can process these subgraphs separately. The system can be used to train models on eight separate partitions simultaneously with the same memory cost compared to processing the entire graph using WAC. Figure 2(d) represents the speedup obtained during the training process. This speedup is calculated as the ratio of time to processes C number of partitions parallelly against the processing time of the entire KG together using WAC. While the acceleration remains between 2 to 4 times the average training time for most datasets, FB15k shows the most acceleration going up to 4.5 times for 7 and 8 clusters. However, using a multi-GPU system can further accelerate the process. Figure 2(c) shows that for denser KG (FB15K and WN18), the acceleration is better than their sparser counterparts (FB15K-237 and WN18RR). This implies that density also affects speed-up.

We compare the best performance achieved by the proposed method for each dataset and compare them with stateof-the-art methods of KG partitioning. These results are illustrated in Table 4. The comparative methods are selected as all of these use partitioning to achieve faster training. Overall, the proposed method achieves the best performance and second best for the FB15k-237 dataset, as observed by improved MRR and Hit@10. Furthermore, on the FB15k-237 dataset, [Sheikh et al., 2022] reports speed up using multiGPU distributed training to be four times. In comparison, the proposed method shows a similar speed-up at 3.7 times with just a single GPU with a considerable improvement in MRR and Hit@10. Hence, we conclude that the proposed method shows an overall edge in performance compared to other methods in the field that combine partitioning and KGE.

6 Ablation Study

In this section, we conduct an ablation study to analyze the contribution of the CPa component in the proposed KGE method on three aspects. (i) We compare the test MRR and Hit@K for an alternate setup: border entity addition (BEA). This method utilizes the addition of triples with a head or a tail in another partition to the training data of each partition. This will enable some entities to be present in more than one cluster. This experiment aims to study the effects of overlapping partitioning and deduce its effect on the performance of KGE; (ii) We conduct another study that utilizes separate models for inference. In this setting, we save each model separately for each partition and use the saved model (SM) for inference using a reasoning framework. Under this scenario, if the head of a test triple falls within a particular partition, then the tail will only be selected from the same partition. Thus, in this setup, the architecture can not predict any link that exists between partitions, and so the MRR for these triples is considered as 0. When inferred through separate embedding models, the drop in performance is significant; (iii) We rerun our experiments utilizing a different partitioning method, namely, Leiden partitioning (LeP) [Traag et al., 2019]. Given both are modularity maximization methods, both can be used due to quicker runtime and find application in community detection in homogenous graphs. This algorithm utilizes modularity vertex partitioning that only considers links between nodes. Other methods like METIS and spectral clustering are computationally expensive.

Table 5 shows the outcome of different iterations on the performance of the architecture on four partitions. In this table, the LC+CPa+WAC+GD signifies the results from the proposed algorithm. As seen from the results, adding border entities degrades the performance of KG when relationship types are less (while increasing the training time slightly). However, for denser KG such as FB15k that have a high number of relationship types, the addition of border entities enhances the performance by a small margin at the expense of an increase in training time. Therefore, we conclude that for sparse graphs containing a lower number of relationship types, border entity addition has overall negative effects on the performance of the KGE algorithm. Contrarily, for dense KG with a high number of relationship types, positive effects on performance can be achieved at expense of higher time.

Furthermore, using Leiden clustering to initialize the partitions yields better performance for denser graphs such as FB15k while the performance remains comparatively the same for sparser graphs such as WN18. Similarly, LeP shows a reduction in the training time compared to the LC for both datasets. This highlights that the proposed constellation partitioning algorithm (CPa) works perfectly when coupled with any homogenous graph partitioning algorithm that only considers the modularity of graphs.

7 Conclusion

In this paper, we present a combination of a novel method of partitioning KG along with an improved GCN-based embedding architecture. Results show that the proposed architecture performs well with a significant reduction in time complexity for GCN-based embedding generation architectures. Additionally, a comparison of three state-of-the-art methods with the proposed method has also been provided in this study, highlighting their embedding performance on KGE. The results have conclusively pointed toward the superior performance of the proposed method, both at the partitioning and the embedding levels. For future studies, we intend to develop a sophisticated reasoning algorithm that would replace the global merging algorithm to reduce the retraining time. From ethical standpoint, we recommend selecting train time vs. performance trade-off based on the metrics, properties of the graph, and criticality of the application.

Acknowledgments

This work was supported by Mitacs through the Mitacs Accelerate Program.