Simple and Effective Relation-based Embedding Propagation for Knowledge Representation Learning

1 Introduction

Knowledge graphs (KGs) are potentially valuable for many applications, such as question answering [Saxena et al., 2020; Bian et al., 2021] and recommender systems [Joseph and Jiang, 2019]. Although KGs contain rich facts in the form of triplets (head entity, relation, tail entity), they are still far from complete, necessitating a demand for KG completion. A promising approach is to embed KGs into latent space and predict missing facts from existing ones. Tripletbased methods [Bordes et al., 2013; Yang et al., 2015; Sun et al., 2019; Tang et al., 2020] often regard relations as operations over entity embedding space. They can quickly scale to large KGs because of their elegant simplicity and good interpretability. However, triplet-based methods overlook global neighbor information. Recently, considerable literature has grown around graph neural networks (GNNs) to model such graph contexts. Existing studies for multi-relation KGs [Schlichtkrull et al., 2018; Vashishth et al., 2020; Cai et al., 2019; Wang et al., 2021] usually train relational GNNs from scratch and optimize triplet-based objects. Specifically, these context-based methods update the central entity with the aggregated embeddings of connected entities and relations. Thus, the time complexity is proportional to the number of triplets and the average degree of entities. The high complexity hinders their applications on real-world tasks at a large scale. How to efficiently and effectively utilize graph context in large KGs remains a challenge.

To this end, we propose Relation-based Embedding Propagation (REP) to combine triplet-based methods' simplicity and graph context strength. On the one hand, we derive inspiration from simplified GNN methods such as SGC [Wu et al., 2019] and APPNP [Klicpera et al., 2019] and remove unnecessary non-linearity. Instead of parameterizing the aggregation of neighbors with relation-specific transformation, we use triplet-based methods to capture triplet information firstly and then adapt the pre-trained embeddings with graph context. We do not use layer-specific transform matrices during embedding propagation to reduce time and space complexities. Our method incorporates graph structures after parameter training and has no activation function or backward propagation during embedding propagation, making it more computationally efficient.

On the other hand, simple embedding propagation that ignores relations, such as SGC and APPNP, is not suitable for KGs and will cause performance degeneration. Therefore, we propose a novel way of combining neighbor relations and entities. We make opposite assumptions about the head context and the tail context. For example, we regard incoming relations as additions and outgoing relations as subtractions. Then context embeddings are aggregated from neighbor entity-relation pairs. We take an update scalar to balance pre-trained embeddings and relational context embeddings. Besides, we provide theoretical and empirical evidence that our proposed method of incorporating relations can help REP obtain an accuracy gain.

We show that REP-OTE takes 5%-83% time to achieve comparable performance to the state-of-the-art GC-OTE [Tang et al., 2020] on small datasets through a series of experiments on the link prediction task. Our REP also works well on the medium and the large KGs. In particular, it brings about a 10% relative improvement on the medium OGBL- WikiKG2 [Hu et al., 2020], despite taking less than 20% time for computation. The key advantages of our REP are threefold:

• By designing the non-parametric embedding propagation without backward propagation, REP becomes computation-efficient and straightforward in utilizing graph context.
• In order to avoid information distortion during embedding propagation in KGs, we incorporate valuable relations using different triplet assumptions so that REP can maintain or mostly improve prediction quality.
• The proposed REP has significant scalability and takes less time to achieve comparable performance on large KGs, promising practical applications.

2 RelatedWork

2.1 Knowledge Representation Learning

Knowledge representation learning has been used to embed entities and relations in KGs into latent space and then infer missing facts based on existing ones. The literature falls into two major categories: (1) triplet-based methods that model triplet plausibility by score functions; (2) context-based methods that model graph context beyond triplets.

Triplet-based methods often make assumptions about entities and relations in triplets. Some regard relations as operations between connected entities in the latent space, such as addition [Bordes et al., 2013], rotation [Sun et al., 2019], and transformations in higher dimensional space [Zhang et al., 2019; Tang et al., 2020]. Others follow semantic matching and design score functions based on similarities [Yang et al., 2015; Trouillon et al., 2016]. In general, these methods are of elegant simplicity and good interpretability.

Context-based methods usually inherit the message passing framework from GNNs [Gilmer et al., 2017]. As heterogeneous relations always play an essential role in KGs, modeling such information is a promising direction. GC-OTE [Tang et al., 2020] expects the central embeddings close to the aggregation of neighbors with relation-specific orthogonal transformation. R-GCN [Schlichtkrull et al., 2018] extends GCN by introducing relation-specific weight matrices. CompGCN [Vashishth et al., 2020] follows the idea of modeling relations as composition operations. However, it makes the same assumption about entity-relation pairs in both head and tail graph contexts, which requires external parameters to denote the direction of relations.

2.2 Simplified Graph Neural Networks

Recent progress in simplifying GNNs for homogeneous graphs has led to a surge of improvements in scalability. SGC [Wu et al., 2019] empirically observes that local averaging is more critical and removes non-linearity. APPNP [Klicpera et al., 2019] utilizes a propagation schema derived from personalized PageRank and achieves linear computational complexity in the number of edges. Reminiscent of our method, C&S [Huang et al., 2021] uses graph structure as a post-processing mechanism but propagates labels to correct predictions.

3 Relation-based Embedding Propagation

Our approach starts with simple triplet-based methods, which learn embeddings for entities and relations through minimizing objectives under triplet assumptions. After, REP takes the pre-trained embeddings as input and improves them by incorporating neighbor information in graph context. As illustrated in Figure 1, REP consists of (1) relational graph context functions efficiently and effectively aggregate relationspecific neighbor information; (2) the entity adaptation computes new embeddings based on pre-trained and context embeddings.

Notations. We consider knowledge graph as a collection of triplets T = {(h, r, t) | h, t ∈ E, r ∈ R} with E as the entity set and R as the relation set. Let AH be the adjacent head list of the knowledge graph and AT be the adjacent tail list. Specifically, the head adjacent list AHi of the entity i is composed of entity-relation pairs (e, r) with (ei, r, e) ∈ T . In contrast, the tail adjacent list ATi contains pairs (e, r) with (e, r, ei) ∈ T . Furthermore, corresponding bold letters denote embeddings.

3.1 Relation-based Context Function

A key motivation behind our relational context functions is that heterogeneous relations in KGs play a crucial role in understanding entities' meanings. Triplet-based methods usually take relations as operations between entities such as additions, multiplications, rotations, and orthogonal transformations. Thus, ignoring relations during propagation can bring information distortion. We design relational context functions according to their assumptions to keep consistent with pre-trained embeddings. As shown in Figure 1, there are two kinds of context.

Head Graph Context includes incoming head-relation pairs for a central entity e, where e acts as a tail entity. As tripletbased methods often apply relation operations to head entities, we can directly use their assumptions for the headrelation pairs. We choose four typical triplet-based methods, which regard relations as additions, multiplications, rotations, and orthogonal transformations, respectively, and define corresponding head context functions gh(h, r) in Table 1.

Tail Graph Context consists of all outgoing tail-relation pairs connected to the central entity. KGs are directed graphs, where a valid (h, r, t) does not mean (t, r, h) is also a real fact. Therefore, we can not use the same assumptions as of the head graph context. Propagating information of tails and relations to heads utilize de facto inverse edges (t, r-1 , h) in KGs. These inverse edges are usually modeled through augmenting relations with inverse ones [Dettmers et al., 2018] or using external parameters to represent relations' directions [Schlichtkrull et al., 2018], which requires external computing resources. To make our method concise and informative, we make opposite assumptions about the tail-relation pairs. Specifically, embeddings of tails and relations are combined by the inverse operations of those for the head graph context, namely subtraction, multiplication, inverse rotation, and inverse orthogonal transformation. In this way, the tail context functions gt(t, r) are defined in Table 1.

Here we introduce how the context functions are derived in detail. For the methods that minimize the distance between t and the combination of h and r, we set their score functions to the desired 0. Then we get ∥h + r - t∥ = 0 for additionbased TransE, ∥h⊙r-t∥ = 0 for rotation-based RotatE, and ∑Li=1 ∥diag(exp(sr,(i)))ϕ(Mr,(i))h(i) - t∥ = 0 for orthogonal transformation-based OTE. After, we can deduce h and t from these equations. In particular, the relation matrices in OTE are orthogonal so that we can get their inverse by simple transpose. For DistMult that maximizes generalized dot product, we expect two elements as close as possible: either t = h ⊙ r or h = t ⊙ r.

Context Aggregation aims to combine neighbor information. As this work aims to improve the scalability of graph context-based methods, we use averaging for aggregation, which helps to keep embeddings on the same scale. Then the head context embedding ˜h(k) i of ei is defined as

and the tail context embedding ˜t(k) i is

3.2 Entity Adaptation

The propagation schema derived from personalized PageRank [Klicpera et al., 2019] can be formulated as Eq. (3).

where E(0) is the input feature matrix, and A is the symmetrically normalized adjacency matrix with self-loops. Here we denote embeddings after k updates as E(k) .

However, storing both E(k) and E(0) is very demanding on memory for large KGs. We do not store the input feature matrix E(0) , i.e., pre-trained embeddings in our method, but use them for embedding initialization. Besides, we separate selfloops from A and use different weights for entity and context embeddings to retain inherent information in pre-trained embeddings. Especially, Eq. (3) can not utilize relations in KGs, so we design relational context functions in previous Section 3.1.

In summary, we use an update scalar α ∈ [0, 1) to balance the trade-off between the pre-learned triplet information and the graph context information. Then the final embedding of entity ei is computed according to Eq. (4).

where e(0) is the pre-trained entity embeddings. In particular, instead of updating parameters with gradient descent algorithms, we directly use the computation results for link prediction. As the number of relations is relatively small, we fix the pre-trained embeddings of heterogeneous relations.

3.3 Theoretical Analysis

In this section, we provide theoretical analysis from the parameter update perspective. Generally, triplet-based methods start with their score function fr(h, t) and learn embeddings through minimizing the margin-based ranking criterion:

where γ denotes the parameter margin, [x]+ takes the positive part of x. A negative sample (h′ , r, t′ ) is constructed by uniformly replacing h or t in (h, r, t) with other entities in E. Parameters of embeddings are updated by gradient descent methods. The typical stochastic gradient descent algorithm can be formulated as Eq. (6).

where β denotes the learning rate. The motivation behind the margin-based ranking criterion is to make valid triplets have relatively higher scores by enlarging the distance to negative samples. Therefore, when embeddings are fully optimized by this criterion, the distances between valid triplets and negative samples reach a local optimum.

Proof. Here we prove that our REP further improves the objective L by maximizing scores of valid triplets. Without loss of generality, we take TransE as an example. It interprets relations as additions in latent space, i.e., t is expected to be the nearest neighbor of h + r. Accordingly, the score function is defined as fr(h, t) = -∥h + r - t∥. Then the objective of maximizing scores of valid triplets is formulated as Eq. (7).

Assuming that parameters are updated by stochastic gradient descent method, we first calculate the partial derivative for hi as shown in Eq. (8).

In this case, Eq. (6) can be reformulated as Eq. (9) by replacing the partial derivative with Eq. (8).

Let α = 1 - 2β, this update function is equivalent to Eq. (4) when entities act as heads, where the head context function is t - r. For the tail entity ti, we have

where the addition operation is consistent with the tail context function in Table 1. Above all, REP acts the same as using stochastic gradient descent to maximize valid triplets' scores, which further optimizes the triplet-based objective L of pretrained embeddings.

4 Experiments

4.1 Datasets

We use four datasets at diverse scales, as reported in Table 2. Large scale. WikiKG90M-LSC [Hu et al., 2021] of KDD Cup 2021 1 is extracted fromWikidata. Aside from triplets, it provides text features learned by RoBERTa [Liu et al., 2019], which are combined by concatenation when training tripletbased methods. It also provides 1000 randomly sampled entity candidates for each tail in the validation/test set.

Medium scale. OGBL-WikiKG2 [Hu et al., 2020] is a link prediction dataset from Open Graph Benchmark (OGB) 2 . It is constructed from the relational statements in Wikidata after filtering out rare entities. Moreover, 1000 randomly sampled candidates for tail entities are provided likewise.

Small scale. FB15k-237 and WN18RR [Dettmers et al., 2018] are the most widely used datasets for link prediction. FB15k-237 mostly describes people-related facts, while WN18RR consists of hyponym and hypernym relations.

1 https://ogb.stanford.edu/kddcup2021/

2 https://ogb.stanford.edu/

4.2 EvaluationMetrics

Link prediction aims to predict missing facts based on existing triplets, namely to predict valid heads for (t, r) or valid tails for (h, r). Specifically, we first corrupt test triplets and construct candidate triplets using entity candidates. As small datasets do not provide candidates, we use all entities e that (h, r, e) /∈ T and (e, r, t) /∈ T as candidates for a test triplet (h, r, t). Then we compute the plausibility of the test triplet and its candidate triplets and sort them by descending order. Evaluation metrics are based on the rank of test triplets, including Mean Reciprocal Rank (MRR) and Hits@K which denote the proportion of ranks less than K.

4.3 Implementation Details

Experiments were conducted on Intel Xeon Gold 6271C CPUs and Tesla V100 SXM2 GPUs. We use the public code to reproduce triplet-based methods, while the results of context-based methods are from original papers. The code and details of REP are released in Graph4KG 3 .

4.4 Results and Discussion

Simplicity. For model parameters, REP requires no parameters except for embeddings of entities and relations, while context-based methods inherit from GNNs and require multiple layer-specific weight matrices besides embeddings. For the computation complexity, we conduct speed experiments to compare REP-OTE and GC-OTE. GC-OTE also requires no external parameters. In theory, REP-OTE and GC-OTE have the same time complexity O(|T |D), where |T | is the number of triplets and D denotes the average degree. However, as REP has no backward propagation, its constant is far smaller than GC-OTE. The speedup numbers in Table 3 empirically prove this. We count the time consumed to traverse all training data. To make a fair comparison, the time of REPOTE consists of the time REP takes and the time used by OTE in an epoch during pre-training. We can note that REP brings more acceleration when the data scale is larger.

Effectiveness. Results of medium and large datasets are reported in Table 4. We have observed that when there are more than 80K entities and 1.6M triplets, context-based methods like CompGCN will be out of memory on a 32GB GPU to learn embeddings with dimension 200. Thus, we only report their results on small-scale datasets. As we can see, REPaugmented methods perform best on all metrics. In particular, REP-TransE significantly surpasses TransE and is competitive with OTE, the state-of-the-art triplet-based method with higher complexity. It averagely achieves a 10% relative improvement on OGBL-WikiKG2 compared with four tripletbased baselines. This result may suggest that graph context is of great use for knowledge graphs at a medium scale. The improvement of WikiKG90M-LSC is not as significant as that on OGBL-WikiKG2. We reason that the text features used during pre-training provide external information to entities with few triplets, so graph context is possible to be superfluous. Results on small datasets are reported in Table 5. In this case, REP-OTE outperforms all triplet-based methods and has comparable prediction quality to the state-of-the-art context-based methods. Above all, REP can bring stable improvements to triplet-based methods and achieve competitive performance as context-based methods with less cost.

3 https://github.com/PaddlePaddle/PGL/tree/main/apps/Graph4KG

Practical Effect. We assume in previous experiments that triplet-based methods have fully trained pre-trained embeddings. In this section, we explore how REP performs on not converged embeddings. Specifically, we apply REP to the embeddings trained by triplet-based methods after 0.25N, 0.5N, 0.75N, and N steps separately if the embeddings converge at step N. The MRR results of these embeddings are reported in Figure 2. Experiments on the four triplet-based methods show that REP brings a significant performance boost to not converged embeddings. Particularly, tripletbased methods with REP can achieve comparable results as the converged ones with just half of the training steps. In other words, REP allows obtaining high-quality embeddings within less time.

How to select α? The update scalar α determines how much neighbor information contributes to the central entities. To investigate the impact of α, we set α as values ranging from 0.95 to 0.99 with step 0.01. Then we evaluate embeddings after once to 15 times entity adaptations separately and draw heat maps using their MRRs in Figure 3. As illustrated in Section 3.3, α = 1 - 2β, where β denotes learning rate. Figure 3 also demonstrates this point. When α increases, REP needs more iterations to get the best MRR.

Are multi-hop neighbors always helpful? For REP, the central entity receives neighbor information k-hops away after k times entity adaption. We investigate how multi-hop neighbors influence embedding quality by changing the number of hops from 1 to 20. We set α = 0.98 to control variables. Results of TransE and RotatE on datasets at different scales are plotted in Figure 4. As we can see, MRRs of OGBL-WikiKG2 grow almost linearly as the number of hops increases, maybe because of its triplet sparsity. Moreover, the curves of both methods on FB125k-237 and WikiKG90MLSC are nearly concave. It shows that enlarging the receptive field contributes to embedding quality, but neighbors too far away would hurt performance when triplets are dense, or there is external information.

4.5 Influence of Relations

We have theoretically analyzed the necessity of relations in Section 3.3. In this section, we provide some empirical evidence. Embedding Propagation (EP) methods like APPNP have shown strong performance on homogeneous graphs with single relation. However, when applied to KGs, they can not utilize types of relations and regard relations as identity matrices. In order to prove the significance of relations, we conduct ablation studies on the FB15k-237 dataset. Figure 5 plots the evaluation results of four triplet-based methods. Base denotes MRRs of fully trained triplet-based methods. EP propagates entity embeddings to neighbor entities directly and ignores relations during this process. We also apply EP to the pretrained embeddings as our REP method. As shown in Figure 5, embedding propagation without relations (EP) causes significant performance deterioration. It proves that the relationbased context functions designed in Section 3.1 play a critical role during embedding propagation.

5 Conclusion

We proposed the novel method REP to utilize graph context in KGs during the post-training process. The key idea is incorporating relational graph structure information into pretrained triplet-based embeddings. For simplicity, we use local averaging to conduct non-parametric embedding propagation without backward propagation. For effectiveness, we design graph context functions for head-relation neighbors and tailrelation neighbors, respectively. As a result, REP can reduce information distortion during embedding propagation in KGs. Furthermore, such design brings REP excellent scalability, which has practical significance in utilizing graph context for large KGs in the real world. Experimental results also show that by enriching pre-trained triplet-based embeddings with graph context, REP improves or maintains prediction quality with less time cost.

Acknowledgments

This work has been supported by National Key Research and Development Program of China (2018AAA0101900).