UniKG: A Benchmark and Universal Embedding for Large-Scale Knowledge Graphs

Heterogeneous Graphs (HGs), also known as Heterogeneous Information Networks (HINs), consist of multiple types of nodes and links. Compared with homogeneous graphs, the heterogeneity in multi-attribute nodes and graph topology makes HGs carry richer semantics, and be more suitable to characterize a variety of complex real-world systems, such as academic networks and social networks. For this reason, methods that focus on representation learning on HGs have drawn increasing attention in recent years, and have facilitated numerous application tasks in diverse domains, including recommendation systems (Gao et al. 2022; Wu et al. 2022), malware detection systems (Zhang et al. 2022a), and healthcare systems (Cao et al. 2022).

Existing HG-related works involve the design of effective learning methods and also the construction of HG datasets, because an encyclopedic HG dataset is also crucial to promote HG learning. Regarding HG dataset construction, existing research (Sharma et al. 2022; Chen, Razniewski, and Weikum 2023; Kubeša and Straka 2023; Whitehouse et al. 2023; Ahrabian et al. 2023; Stranisci et al. 2022; Luo et al. 2023) mostly focuses on combining specific entities and relationships within knowledge graphs (KGs) to facilitate downstream tasks within the corresponding domain. Existing datasets generally consist of thousands of nodes and dozens of relationships applied in a specific domain such as finance, recommendation systems and social networks. For the learning methods on HGs, depending on the difference of methods in capturing semantic structure, the existing works can be broadly classified into aggregation-based methods (Schlichtkrull et al. 2018; Zhang et al. 2019; Xu et al. 2019) and meta-path-based methods (Wang et al. 2020; Huang et al. 2016; Sun et al. 2011; Chang et al. 2022). The aggregation-based methods primarily concentrate on iteratively aggregating neighbor information to update node representations. In contrast, meta-path-based methods focus on combining information from distinct meta-paths to learn node representations.

Although notable success has been achieved by existing works on HGs, there are still several issues to be tackled. The most important one is the lack of a large-scale HG dataset. Existing HG datasets are almost small-scale, with just thousands of nodes and dozens of relation types limited in specific knowledge domains. This greatly limits the ability to facilitate representation learning and knowledge extraction, therefore there is an urgent need for a large-scale HG dataset with encyclopedic knowledge. Furthermore, when regarding effective learning on large-scale HG datasets, another issue is the inapplicability of existing HG learning methods. On one hand, existing aggregation-based HG learning methods cost too much memory overhead when performing multi-hop aggregation in very large-scale receptive fields. Hence, they may suffer from out-of-memory problems during model training. On the other hand, for the meta-path-based methods, existing works are applied to the case where only one type of association exists between two types of nodes. However, when the association types between two nodes significantly increase, the number of metapaths may suffer explosive growth (Please see the example in Figure 1). This causes great difficulty for topology perception if only a limited number of metapaths can be sampled.

In this paper, we construct a large-scale Heterogeneous Graph benchmark dataset named UniKG from Wikidata. We release UniKG to facilitate the representation learning task of heterogeneous graphs. Overall, UniKG contains 77.31 million multi-attribute entities labels by 2000 classes, 564 million directed edges annotated by 2082 diverse association types, which significantly surpasses the scale of existing homogeneous graph datasets. To perform effective learning on the large-scale UniKG, two key measures are taken. We propose a semantic alignment strategy for multi-attribute entities, which projects the feature description of multi-attribute nodes into a common embedding space. Then we design a novel plug-and-play anisotropy propagation module (APM) as base to extend methods of large-scale homogeneous graphs to heterogeneous graphs. These two strategies enable efficient information propagation among a tremendous number of multi-attribute entities and meantimes adaptively mine multi-attribute association through the multi-hop aggregation in large-scale HGs. We set up a node classification task on our UniKG dataset that evaluate multiple baseline methods which are constructed by embedding our APM into large-scale homogenous graph learning methods. In addition, we attempt to transform knowledge of UniKG to a recommendation system to verify the capability to facilitate downstream task of diverse domains.

Our contributions can be summarized in the following three aspects:

1) We construct a large-scale Heterogeneous Graph benchmark dataset named UniKG from Wikidata which contains over 77 million nodes and 564 million edges associated with 2,082 types. To the best of our knowledge, UniKG significantly surpasses the scale of existing HG datasets.

2) We propose a plug-and-play Anisotropy Propagation Module (APM), which can adaptively mine association relationships of diverse types and extends methods of large-scale homogeneous graphs to heterogeneous graphs.

3) We transform knowledge from UniKG to recommender systems, which effectively promotes performance of recommendation task.

In this section, we introduce the classical HG datasets, as well as representative GNN methods, such as those for solving KGs, large-scale graphs and HGs tasks. Finally, we briefly discuss the limitations of existing datasets and methods and point out highlights the advantages of our proposed datasets and methods.

In this section, we introduce our proposed dataset UniKG, including construction details and data statistics. Figure 2 is a construction example. Table 1 and Table 2 show the statistics and three different scales of UniKG, respectively.

In this section, we detailed describe the novel Heterogeneous Graph Decoupling (HGD) framework, and the anisotropic propagation process of the plug-and-play Anisotropic Propagation Module (APM) on large-scale heterogeneous graphs as shown in Figure 3.

Through introducing $d_{e}$-dimensional edge embeddings, the adjacency matrix $\mathbf{A}\in\mathbb{R}^{n\times n}$ can be represented as a relation-aware adjacency matrix $\mathbf{A}_{r}\in\mathbb{R}^{n\times n\times d_{e}}$. Therefore, the Heterogeneous Graph Decoupling (HGD) framework can be represented as:

$$\begin{cases}\mathbf{Z}=\delta([\mathbf{X}\mathbf{\Theta}_{0}|\mathbf{C}\mathbf{\Theta}_{1}|...|\mathbf{C}^{K}\mathbf{\Theta}_{K}])\\ \mathbf{Y}=\sigma(\mathbf{Z}\mathbf{W})\end{cases}$$

(2)

where, for each $k\in[1,K]$, $\mathbf{C}^{k}\in\mathbb{R}^{n\times d}$ denotes the anisotropic propagation matrix of the $k$-th hop. We formulate the process of anisotropic propagation as a function $\hat{f}(\cdot)$, then $\mathbf{C}^{k+1}=\hat{f}(\mathbf{A}_{r},\mathbf{C}^{k})$ and $\mathbf{C}^{0}=\mathbf{X}$ is the original feature matrix.

To integrate the anisotropic edge features into the message matrix, we first upgrade dimension $\mathbf{C}^{k}$ from $\mathbb{R}^{n\times n}$ to $\mathbb{R}^{n\times n\times d}$, and the matrix is denoted as:

$$\mathbf{H}^{k}=\mathbf{C}^{k}\otimes\mathbf{I}_{n}\vspace{-0.03cm}$$

(3)

where $\mathbf{I}_{n}\in\mathbb{R}^{n\times n}$ is the identity matrix. The $i$-th row of $\mathbf{H}^{k}$ represents a message matrix $\mathbf{M}\in\mathbb{R}^{n\times d}$ from node $i$ to others. Based on the $\mathbf{H}^{k}$, through the tensor product of the anisotropic message matrix and its coefficient matrix, we can obtain the anisotropic propagation matrix at the $k+1$-th hop by the attention mechanism, which can be expressed as:

$$\mathbf{C}^{k+1}=\mathbf{B}^{k}\times_{2}\hat{\mathbf{Q}}$$

(4)

where $\times_{2}$ represents the tensor product on the second dimension, $\mathbf{C}^{k+1}\in\mathbb{R}^{n\times d}$ denotes the $(k+1)$-th hop anisotropic propagation matrix, $\mathbf{B}^{k}$ is the message propagation matrix, and $\hat{\mathbf{Q}}$ represents the edge-aware coefficient matrix. Specifically, we perform element-wise Hadamard multiplication $\odot$ between $\mathbf{A}_{r}$ and $\mathbf{H}^{k}$ to derive the message propagation matrix $\mathbf{B}^{k}\in\mathbb{R}^{n\times n\times d}$:

$$\mathbf{B}^{k}=\mathbf{A}_{r}\odot\mathbf{H}^{k}$$

(5)

where, $\mathbf{B}^{k}$ describes messages between each node-pairs. The $l_{2}$-norm matrix of $\mathbf{B}^{k}$ on the third dimension can be represented as $\mathbf{Q}\in\mathbb{R}^{n\times n}$:

$$\vspace{-0.03cm}\mathbf{Q}=||\mathbf{B}^{k}||_{2}\vspace{-0.03cm}$$

(6)

then, calculate the edge-aware coefficient matrix $\hat{\mathbf{Q}}\in\mathbb{R}^{n\times n}$ for the messages between each node-pair by row-wise softmax:

$$\vspace{-0.03cm}\hat{\mathbf{Q}}=\text{softmax}(\mathbf{Q})\vspace{-0.03cm}$$

(7)

Finally, perform tensor multiplication along the second dimension between the message propagation matrix $\mathbf{B}^{k}$ and the edge-aware coefficient matrix $\hat{\mathbf{Q}}$, which means that messages of $k$-th hop are anisotropically propagate to the corresponding nodes of $k+1$-th hop. By iterating this procedure $K$ times, we obtain a sequence of anisotropic propagation matrices $[\mathbf{C}^{1},...,\mathbf{C}^{K}]$, which serves as the input to the post-classifier of the decoupling-based graph neural network.

This method supports fast multi-attribute aggregation and can extends decoupled methods of the large-scale homogeneous graphs to the large-scale heterogeneous graphs in a edge-aware manner.

In real-world, nodes tend to be annotated by multi-label. We use multi-label loss functions, such as the Binary Cross-Entropy (BCE) loss:

$$\mathcal{L}=-\frac{1}{N}\sum_{i=1}^{N}\left[\mathbf{y}_{i}\cdot\log(\hat{\mathbf{y}}_{i})+(1-\mathbf{y}_{i})\cdot\log(1-\hat{\mathbf{y}}_{i})\right]$$

(8)

where $N$ is the number of nodes, $\mathbf{y}_{i}$ denotes the ground truth vector of node $i$, and $\hat{\mathbf{y}}_{i}$ represents the predicted label vector.

In this section, we introduce two experiments to evaluate the proposed dataset and method, along with the ability to facilitate other downstream tasks. The detailed descriptions of recommendation task are reported in Appendices.

We construct a large-scale heterogeneous graph benchmark dataset UniKG, along with associated representation learning methods: the Heterogeneous Graph Decoupling framework and a plug-and-play Anisotropy Propagation Module. Our proposed construction methods model knowledge graphs as large-scale heterogeneous graphs, and the scale of the constructed UniKG significantly surpasses the scale of existing heterogeneous graph datasets. The integration of the decoupling-based methods of homogeneous graphs with Heterogeneous Graph Decoupling framework is easy to implement and applicable to large-scale heterogeneous graphs with competitive performance. The real-world knowledge of UniKG can facilitate diverse downstream tasks.