Augmenting Knowledge Graphs for Better Link Prediction

We start by obtaining the set of values $v_{i}$ which are associated with entities $e\in{E}$ of a given attribute $a$, namely $(e,a,v_{i})\in G$. As dates come in different granularities (e.g., year, month, or day), we transform all dates with a granularity of year or finer (e.g., month) to year values. We sort these values in an ascending order, obtaining $V=\{v_{1},v_{2},...,v_{m}\}$, such that $v_{i}\leq v_{i+1}$. We next describe how we create bins based on the values in $V$, resulting in a dictionary $D:V\rightarrow\mathcal{E}$ per attribute, which maps each value in $V$ to one or multiple bin entities in $\mathcal{E}$. The discretization of KGA consists of two components: 1) definition of bin intervals; and 2) specification of bin levels.

Our fixed (equal width) binning strategy divides the entire space of values into bins of equal width. The length of each bin $i$ is $k=\frac{(z^{+}-z^{-})}{b}$, and the $i$-th bin contains the values $z_{i}=[z^{-}+ik,z^{-}+(i+1)k)$.

The quantile (equal frequency) binning strategy splits [$z^{-},z^{+}$] based on the distributional density, by ensuring that each bin $i$ covers the values for the same number of entities. This allows denser areas of the distribution to be represented with more bins. For a attribute $a$ defined for $N_{a}$ entities, each bin will cover approximately $n=\frac{N_{a}}{b}$ entities.

The single strategy creates only one set of disjoint bins, as described previously.

The overlapping strategy creates bins with intervals that overlap. Starting from a single auxiliary set of $2b$ bins, the overlapping bins are created by merging two neighboring bins at a time. The $i$-th merged bin consists of merging bin $b_{i}$ and $b_{i+1}$, the $(i+1)$-th merges $b_{i+1}$ and $b_{i+2}$, and so on. This process results in $2b-1$ overlapping bins. Overlapping bins are illustrated in Figure 1, middle. Here, the auxiliary bins (not shown) would be $[[1,1826],[1826,1882],[1882,1912],…]$. The derived overlapping bins merge two neighboring bins at a time - merging the first two bins produces $[1,1882]$, merging the second and the third produces $[1826,1912]$, etc.

The hierarchy strategy creates multiple binning levels, signifying different granularity levels. A level $l$ has $b^{l}$ bins. The $0$-th level is the coarsest, and it contains a single bin with the entire set of values $[z^{-},z^{+})$. The levels grow further in terms of detail: considering $b=2$, the first level has 2 bins, the second one has 4, and the third one has 8. An example for a quantile-based hierarchy binning is shown in Figure 1 bottom. The entire set of birth years at level $0$ ($[1,2021]$) is split into two bins at level $1$: $[1,1935)$ and $[1935,2021]$. Level $2$ splits the bins at level $1$ into even smaller, more precise bins: $[1,1882)$, $[1182,1935)$, $[1935,1966)$ and $[1966,2021]$.

The created bins provide a mapping from a value $v_{i}\in V$ to a set of bins $b_{i}\in B$. In order to augment the original graph $G$, we perform two operations: chaining of bins and linking bins to corresponding entities.

1. Bin chaining defines connections between two bins $b_{i}$ and $b_{j}$, where $b_{i},b_{j}\in B$. The links between bins depend on the bin level setting. Both single bins and overlapping bins are chained by connecting each bin $b_{i}$ to its predecessor $b_{i-1}$ and successor $b_{i+1}$, see Figure 1 (top and middle). The bins in the hierarchy augmentation mode are connected both horizontally and vertically: 1) horizontally each bin $b_{i}$ connects to its predecessor $b_{i-1}$ and successor $b_{i+1}$; 2) vertically each bin is connected to its coarser level and finer level correspondents (see Figure 1 bottom).

2. Bin assignment links the generated bins $B(v)$ for a given attribute value $v$ to its corresponding entity $e$. A numeric triple $(e,a,v)$ is replaced with a set of triples $(e,a,b)$, where $b\in B(v)$. The updated set of triples modifies the original graph $G$ into an augmented graph $G^{\prime}$. Figure 1 presents examples of bin assignment. The birth year (attribute P569) of Barack Obama (entity Q76) is $1961$. For single level bins (Figure 1, top), the birth year of Obama is placed into bin $[1935,1966)$. For overlapping bins (Figure 1, middle), the birth year is placed into the bins $[1935,1966)$ and $[1951,1982)$. And, when using hierarchy bin levels (Figure 1, bottom), the birth year is placed into bins $[1,2021]$, $[1935,2021]$ and $[1935,1966)$.

$G\prime$ can now be used to perform entity and numeric LP. KGA is trained by simply replacing $G$ with $G^{\prime}$ in the input to the base embedding model. Link prediction of entities with KGA is performed by selecting the entity node $e$ with a highest probability. Numeric link prediction is performed by selecting the bin $b$ with the highest score as a predicted range; and obtaining its median, $m$, as an approximate predicted value.

We use benchmarks based on three KGs to evaluate LP performance. We show data statistics in Table 7 in the appendix.

1. FB15K-237 Toutanova and Chen (2015) is a subset of Freebase Bollacker et al. (2008), which mainly covers facts about movies, actors, awards, sports, and sports teams. FB15K-237 is derived from the benchmark FB15K by removing inverse relations, because a significant number of test triples in FB15K could be inferred by simply reversing the triples in the training set Dettmers et al. (2018). We use the same split as Toutanova and Chen (2015). For numeric prediction we follow Kotnis and García-Durán (2019).

2. YAGO15K Liu et al. (2019) is a subset of YAGO Suchanek et al. (2007), which is also a general-domain KG. YAGO15K contains 40% of the data in the FB15K-237 dataset. However, YAGO15K contains more valid numeric triples than FB15K-237. For entity LP, we use the data split proposed in Lacroix et al. (2020), while for numeric prediction we follow Kotnis and García-Durán (2019).

3. DWD. We introduce DARPA Wikidata (DWD), a large subset of Wikidata Vrandečić and Krötzsch (2014) that excludes prominent domain-specific Wikidata classes such as review article (Q7318358), scholarly Article (Q13442814), and chemical compound (Q11173). DWD describes over 37M items (42% of all items in Wikidata) with approximately 166M statements. DWD is several orders of magnitude larger than any of the previous LP benchmarks, thus providing a realistic evaluation dataset with sparsity and size akin to the size of modern hyperrelational KGs. We split the DWD benchmark at a 98-1-1 ratio, for both entity and numeric link prediction. Given the size of the DWD benchmark, 1% corresponds to a large number of data points (over 1M statements).

For entity LP on FB15K-237 and YAGO15K, we preserve a checkpoint of the model every 5 epochs for DistMult, ComplEx, every 10 epochs for ConvE, TuckER, and every 50 steps for TransE, RotatE. We use the MRR on the validation set to select the best model checkpoint. We report the filtered MRR, Hits@1, and Hits@10 of the best model checkpoint. For DWD, we preserve a checkpoint for each epoch. We use the MRR on the validation set to select the best model checkpoint, and report unfiltered MRR and Hits@10 Lerer et al. (2019). Filtered MRR evaluation requires discarding all positive edges when generating corrupted triples, which is not scalable for large KGs. For DWD, we report these metrics by ranking positive edges among $C=500$ randomly sampled corrupted edges. For numeric LP, we report MAE for each KG attribute.

For entity link prediction, we evaluate KGA with the following six embedding models: TransE Bordes et al. (2013), DistMult Yang et al. (2014), ComplEx Trouillon et al. (2016), ConvE Dettmers et al. (2018), RotatE Sun et al. (2019), and TuckER Balažević et al. (2019). We run KGA with $b\in[2,4,8,16,32]$ and show the best result for each model. We compare KGA to the embedding-based LP predictions on the original graph, and to prior literal-aware methods: KBLN García-Durán and Niepert (2017), MTKGNN Tay et al. (2017), and LiteralE Kristiadi et al. (2019). We compare our numeric LP performance to the methods NAP++  Kotnis and García-Durán (2019) and MrAP Bayram et al. (2021). As NAP++ is based on TransE embeddings, we focus on our results with this embedding model, and we provide the results for the remaining five embedding models in the appendix. As mentioned above, we take the median of the most probable bin to be the predicted value by KGA.

The more computationally demanding embedding models (ConvE, RotatE, and TuckER) cannot be run on DWD. The size of DWD is prohibitive for ConvE and TuckER because they depend on 1-N sampling, where batch training requires to load the entire entity embedding into memory.⁵⁵5With a hidden size of 200, loading the entity embedding in memory requires at least 42.5M * 200 * 4 = 34 GB GPU memory, which is challenging for most GPU devices today. RotatE is even more memory-intensive, because of its hidden size of 2000, which requires an order of magnitude more memory compared to ConvE. As LiteralE and KBLN are based on these methods, they can also not run on DWD (see appendix for more implementation details). Thus, on DWD, we compare KGA’s performance on the models ComplEx, TransE, and DistMult against their base model performance.

We provide details on the parameter values, software, and computational environments in the technical appendix.

Main results LP results on FB15K-237 and YAGO15K are shown in Table 1. KGA outperforms the vanilla model and the literal-aware LP methods by a comfortable margin with all six embedding models. The best overall performance is obtained with the TuckER and RotatE models for FB15K-237, and the ConvE and RotatE models for YAGO15K, which is in line with the relative results of the original models. On FB15K237, the largest performance gain is obtained for DistMult (2.7 MRR points), and the performance gain is around 1 MRR point for the remaining models. The comparatively smaller improvement for TuckER on FB15K-237, brought by KGA and other literal-aware methods, could be due to TuckER’s model parameters being particularly tuned for this dataset. This intuition is supported by the much higher contribution of KGA for TuckER on the YAGO15K dataset. KGA brings a higher MRR gain on the YAGO15K dataset, which shows that KGA is able to learn valuable information from the additional 24% literal triples which we added to YAGO15K. On YAGO15K, KGA improves the performance of the most recent models (ConvE, RotatE, and TuckER) by over 2 MRR points. These results highlight the impact of KGA’s approach of augmenting KG embedding models with discretized literal values.

However, the improvement brought by KGA for the models TransE, DistMult, and ComplEx is between 0.3 and 0.7 MRR points, which is much lower than its impact on the smaller datasets. We note that the results in Table 1 show the best configuration for KGA, while the results in Table 2 show a single configuration (QOC with 32 bins). Therefore, we hypothesize that the performance of KGA on DWD can be improved with further tuning of the number of bins and the discretization strategies applied. For computational reasons, we investigate the impact bin sizes and discretization strategies on the FB15K dataset, and leave the analogous investigation for DWD to future work.

Main results Next, we investigate the ability of KGA to predict quantity and year values directly. We compare KGA-QOC with 32 bins to the baselines MrAP and NAP++ on the FB15K-237 and YAGO15K benchmarks in Table 9. We observe that both KGA and MrAP perform better than NAP++.⁶⁶6The results from the original NAP++ paper differ from those obtained by reproducing experiments, see Bayram et al. (2021). KGA performs better than MrAP on the majority of the attributes: 7 out of 11 in FB15K-237, and 5 out of 7 attributes in YAGO15K. Closer investigation reveals that our method is superior in decreasing the error for years and attributes with a large range of values, such as area and population, which could be attributed to the quantile-based binning strategy. MrAP outperforms our model on the latitude and longitude attributes on both datasets. Given the low MAE error of MrAP, we think that such a regression model that operates on the raw numeric values of triples can learn to reliably predict latitude and longitudes based on other structured information captured by the base embedding model, while literal information might. notimprove prediction further.

We show the numeric LP results for the five most populous quantity attributes and the five most populous year attributes in DWD in Table 6. We compare KGA and a linear regression (LR) model to a baseline which selects the median value for an attribute. Both KGA and the LR model perform better than the median baseline, yet, the LR model outperforms KGA on this dataset for 9 out of 10 attributes. These results reinforce the findings in Table 2 that KGA requires further tuning in order to be applied to DWD.