MEIM: Multi-partition Embedding Interaction Beyond Block Term Format for Efficient and Expressive Link Prediction

1 Introduction

Knowledge graphs are used to represent relational information between entities. There are large real world knowledge graphs such as YAGO [Mahdisoltani et al., 2015] containing millions of entities. These knowledge graphs and their representations can be used in artificial intelligent applications such as semantic queries and question answering [Tran and Takasu, 2019b] [Tran, 2020].

Knowledge graph embedding aims to predict the missing relations between entities in knowledge graphs. They usually represents a triple (h, t, r) as embeddings and use a score function to compute its matching score S(h, t, r). The score function defines the interaction mechanism such as bilinear map between the embeddings and the interaction pattern specifying how embedding entries interact with each other.

On large real world knowledge graphs, a good tradeoff between efficiency and expressiveness is crucial. Tensordecomposition-based models usually provide good trade-off by designing special interaction mechanisms with sparse and expressive interaction patterns. They have been subsumed by the recent multi-partition embedding interaction (MEI) model [Tran and Takasu, 2020], that divides the embedding vector into multiple partitions for sparsity, automatically learns the local interaction patterns on each partition for expressiveness, then combines the local scores to get the full interaction score. The trade-off between efficiency and expressiveness can be systematically controlled by changing the partition size and learning the interaction patterns through the core tensor of the block term tensor format.

However, by looking from two perspectives beyond the scope of block term format, that is, ensemble boosting effects and max-rank relational mapping, we identify two drawbacks of MEI and related tensor-decomposition-based models. First, using one core tensor for all partitions leads to similar local interactions, that may harm the ensemble effects of the full interaction in MEI. Second, the relational mapping matrices in MEI are generated from relatively small relation embeddings, that may make the model prone to degenerate.

In this paper, we propose the novel Multi-partition Embedding Interaction iMproved beyond block term format (MEIM) model with two techniques, namely independent core tensor to improve ensemble effects and max-rank mapping by soft orthogonality to improve model expressiveness. This introduces new aspects to tensor-decomposition-based models exploiting ensemble boosting effects and max-rank relational mapping, in addition to multi-partition embedding by MEI. MEIM improves expressiveness while still being highly efficient, helping it to outperform strong baselines and achieve new state-of-the-art results on difficult link prediction benchmarks using fairly small embedding sizes.

In general, our contributions include the following.

• We propose MEIM, a novel tensor-decomposition-based model with independent core tensor and max-rank mapping by soft orthogonality to improve expressiveness.
• We extensively experiment to show that the proposed model is highly efficient and expressive, and achieves state-of-the-art results.
• We analyze the proposed model to clarify its characteristics and demonstrate the critical advantages of its soft orthogonality compared to previous models.

2 Background

2.1 Notations and Definitions

A knowledge graph is a collection of triples D, with each triple denoted as a tuple (h, t, r), where h and t are head and tail entities in the entity set E and r belongs to the relation set R. A knowledge graph can also be represented by a thirdorder binary data tensor G ∈ {0, 1}|E|×|E|×|R| , where each entry ghtr = 1 ⇔ (h, t, r) exists in D.

The contextual link prediction task, also called link prediction, aims to predict the connection between two entities given a relation as the context.

2.2 Tensor-Decomposition-based Knowledge Graph Embedding Models

There are several knowledge graph embedding models that adapt tensor formats to represent the knowledge graph data tensor and adapt tensor decomposition methods to solve the link prediction task. This approach has led to some of the best models in terms of efficiency and expressiveness.

RESCAL [Nickel et al., 2011] is an early model that adapts Tucker2 decomposition to compute the embedding interaction score by a bilinear map

where h, t ∈ RD , and Mr ∈ RD×D are the embeddings of h, t, and r, respectively, with D being the embedding size. RESCAL is expressive, but the mapping matrix Mr grows quadratically with embedding size, making the model expensive and prone to overfitting.

The most simple model is DistMult [Yang et al., 2015], that uses a sparse diagonal mapping matrix Mr = diag(r), where r ∈ RD is the relation embedding vector. It is an adaptation of CP decomposition [Kolda and Bader, 2009], with the score function written as a trilinear product

Recent models aim to be more expressive than DistMult but still efficient by designing new special interaction mechanisms between the embeddings. It has been shown that these models are equivalent to bilinear model with sparse blockdiagonal mapping matrix Mr ,

where the matrix block Mr,k have a special interaction pattern resulted from the specific interaction mechanism [Tran and Takasu, 2020]. For example, ComplEx [Trouillon et al., 2016] uses complex-valued trilinear product, resulting in the 2-dimensional rotation-scaling pattern on complex plane,

Sparsity and the interaction pattern are crucial concepts in previous tensor-decomposition-based models. They are generalized by the recent multi-partition embedding interaction (MEI) model [Tran and Takasu, 2020], that divides the embedding vector into multiple partitions, learns the local interaction patterns on each partition, then sums the local scores to get the full interaction score. The trade-off between efficiency and expressiveness can be systematically controlled by changing the partition size and learning the interaction patterns through the core tensor of the block term tensor format.

3 Multi-partition Embedding Interaction iMproved Beyond Block Term Format

In this section, we propose the MEIMmodel that introduces new aspects to tensor-decomposition-based models exploiting ensemble boosting effects and max-rank relational mapping, in addition to multi-partition embedding by MEI, as illustrated in Figure 1.

3.1 The Score Function

Inspired by the generality of MEI, we adopt the multipartition embedding representation, with Tucker format [Tucker, 1966] for local interaction and block term format [De Lathauwer, 2008] for full interaction. Following convention of MEI, triple (h, t, r) provides the embedding vectors h, t ∈ RDe , and r ∈ RDr that are treated as multi-partition embedding matrices H, T ∈ RK×Ce , and R ∈ RK×Cr , respectively. The score function of MEIM is written as

where hk:, tk:, and rk: are the embedding partitions k; Wk ∈ RCe×Ce×Cr is the core tensor at partition k; ¯×n denotes the n-mode tensor product with a vector; and MW,r,k ∈ RCe×Ce is the bilinear mapping matrix.

Note that Eq. 5 shows the sum of local interaction scores Sk(h, t, r; θ). Eq. 6 shows the block term format. Eq. 8 shows the bilinear format with block-diagonal mapping matrix, in which each block MW,r,k is generated by Wk ¯×3rk:. Eq. 8 can be seen as a dynamic linear neural network [Tran and Takasu, 2020], where the hidden layer MW,r,k is generated by another hyper neural network [Ha et al., 2016] with the core tensor Wk as its weights and rk: as its input.

The main novelties of MEIM are in better parameterization of the core tensor and mapping matrix, by looking from two perspectives beyond the scope of block term format, that is, ensemble boosting effects and max-rank relational mapping.

3.2 Core Tensor for Ensemble Boosting

Technically, the score function in Eq. 5 can be seen as an ensemble system of K local interactions by summing their scores. This works in a similar manner to gradient boosting [Mason et al., 1999] because the full interaction score is computed at training time and gradients are back-propagated to optimize all local interactions together, enabling them to implicitly minimize the residual error of each other.

The MEI model briefly mentioned this perspective but did not actively exploit the ensemble boosting effects. MEI used only one core tensor for all partitions by enforcing the following shared core tensor constraint:

where W ∈ RC×C×C is the learned shared core tensor, although block term format may have different core tensors. MEI learns one interaction pattern for all partitions, similar to that of other tensor-decomposition-based models such as ComplEx. However, this may be harmful to the performance of an ensemble system because of the similar local interaction patterns instead of combining different patterns. To exploit the ensemble boosting effects, it is crucial to promote independence and difference between the local interactions.

Independent Core Tensor Parameterization

There are several potential ways to resolving the problem of independence and difference between the local interactions. First, we may enforce the predicted scores to be different:

However, because this forces the predicted scores to be not the same, the local interactions cannot be correct at the same time. Thus, this is an overstrong constraint that may cause more harm than good. Another way is to explicitly enforce different parameters of the local interaction functions, such as the core tensors:

However, such explicit difference constraints do not guarantee different interactions because of symmetric swapping or scaling reparameterizations of the core tensors and the embeddings. For example, scaling the core tensor by x and the corresponding embedding by 1/x do not change the score.

Therefore, in this paper, we do not use difference constraints, but instead independence constraints, denoted as:

For simplicity, we let the model implicitly learn independent core tensors from data automatically by removing the constraint in Eq. 9. Using independent core tensors enables the model to learn independent and possibly different local interactions to improve the ensemble boosting effects. More advanced methods based on further studying the ensemble boosting effects are left for future work.

Note that MEI is easy to implement in common deep learning frameworks by using only matrix product. However, independent core tensor makes the multi-partition tensor product in MEIM more difficult to implement efficiently. To resolve this problem, we stack together K independent core tensors as a fourth-order tensor:

where W ∈ RK×C×C×C is used for the Einstein sum notation technique to create an efficient implementation.

3.3 Max-Rank Relational Mapping

MEI is a contextual link prediction model with two components, the multi-partition relation-based contextual mapping and the simple dot product matching. The MEI model used a linear mapping where MW,r,k can be any Ce ×Ce square matrix. However, when some columns of MW,r,k are dependent, it becomes a singular matrix with rank smaller than Ce. Such singular matrix would map the head embedding to a subspace with lower dimension, thus reduce the embedding space effective size. It may make the score function degenerate to always produces score 0 for some entities, thus, drastically reduces the expressiveness of the model. This may become a critical issue in MEI because MW,r,k is generated from a relatively small relation embedding partition rk: as shown in Eq. 8. To resolve this problem, the mapping matricesMW,r,k need to have max rank, that may not always be full rank Ce, but should be as large as possible.

Max-Rank Mapping Matrix by Soft Orthogonality

There are several different types of full rank matrices. A particularly interesting case is the orthogonal matrices, that act as the rotational transformations when the determinant is 1 and an additional reflection across the origin when the determinant is -1. These matrices have linearly independent column vectors with unit Frobenius norms, written as the constraints:

where I ∈ RCe×Ce is the identity matrix. In addition, note that the mapping matrices are generated by the core tensors and the relation embedding vectors. To reduce the search space and make optimization easier, we also constrain each relation partition rk: to have unit norm:

Because of the wayMW,r,k are generated, it may be difficult to enforce the mapping matrices to be strictly orthogonal. Even if it is possible, it may not always be preferable because the main objective is link prediction. Geometrically speaking, the orthogonal matrices form the Stiefel manifold [Sok´oł and Park, 2020], we need to make sure that the mapping matrices stay as close to this manifold as possible without hurting the link prediction accuracy. Using Lagrangian relaxation, the hard constraints in Eq. 14 and Eq. 15 can be converted to the following soft orthogonality loss term:

where λortho, λunitnorm, and p are hyperparameters. To penalize large differences from unit norm, p = 3 is used. When λortho is small, the model tends to learn linear mappings. When it is large, the model tends to learn strictly orthogonal mappings.

The soft orthogonality loss enables the model to balance between orthogonality and link prediction. Minimizing this loss effectively influences the column vectors to be orthogonal and have unit norms, pushes the matrix MW,r,k close to the Stiefel manifold and maximizes its rank in balance with the link prediction objective.

3.4 Connections to Previous Models

Compared to the MEI model, MEIM inherits its key benefit, that is, the systematic trade-off between efficiency and expressiveness by controlling the partition size and learning the interaction patterns. About computational cost, the number of parameters in MEIM is similar to that of MEI, except that the independent core tensors use KC 2 e Cr instead of C 2 e Cr parameters, but this becomes less relevant when the number of entities and relations are large, and MEIM is as fast as MEI by using Einstein sum notation technique for implementation of independent core tensor product.

On the other hand, MEIM introduces new parameterization techniques to address some major drawbacks of MEI. With independent core tensor, MEIM can learn independent local interaction patterns to improve ensemble effects. With soft orthogonality, MEIM balances the orthogonality and link prediction objectives to learn max-rank mapping that improves expressiveness of the model.

Note that, coincidentally, some knowledge graph embedding models also try to use orthogonal mapping, not to maximize the mapping matrix rank but to model some relational patterns, such as symmetry, inversion, and composition [Sun et al., 2019]. There are many different ways to design and use orthogonal mapping in knowledge graph embedding. For example, RotatE [Sun et al., 2019] uses rotation in the complex plane, QuatE [Zhang et al., 2019] uses rotation in the quaternion space, GC-OTE [Tang et al., 2020] uses rotation by the Gram Schmidt process, and RotH [Chami et al., 2020] uses rotation by the Givens matrix. However, there are some important differences to our method. First, all of them use rigid orthogonality where the mapping matrix is strictly orthogonal. This may be suboptimal because orthogonal mapping may not always be useful for link prediction, and even when it is useful, strictly orthogonal mapping may reduce the potential result. In contrast, we approach the problem from the perspective of maximizing the mapping matrix rank, thus, we use soft orthogonality as a means instead of rigid orthogonality. Soft orthogonality can avoid the above problems because we can choose how to balance the orthogonality and the link prediction objectives for each dataset, as we will show this is crucial for knowledge graph embedding in experiments. Second, MEIM has a totally different model architecture compared to previous orthogonal models, because the mapping matrix in MEIM is not learned directly but generated. Finally, MEIM is a semantic matching model, with the matching score computed by inner product, whereas other orthogonal models are translation-based model, in which the score is computed by distance.

4 Loss Function and Learning

The model is learned via the link prediction task that can be modeled as a multi-class classification problem. There are two directions, classifying correct tail entity t among all entities given a (h, r) pair, and classifying correct head entity h among all entities given a (t, r) pair. Following recent practice, we use the 1-vs-all and k-vs-all sampling methods with the softmax cross-entropy loss function [Dettmers et al., 2018] [Ruffinelli et al., 2020].

First, using 1-vs-all and k-vs-all sampling, we define the groundtruth categorical distributions over all entities E given (h, r) and given (t, r), denoted ˆ ph ˙ tr and ˆ p˙htr , respectively. Second, we compute the corresponding predicted distributions ph ˙ tr and p˙htr using the softmax function on the matching scores. The link prediction loss is the cross-entropy summed over training data D:

The final loss function is the sum of link prediction loss and soft orthogonality loss:

with λortho tuned to balance the two loss terms. For regularization, we only used dropout [Srivastava et al., 2014] and batch normalization [Ioffe and Szegedy, 2015]. All parameters in the model are learned in an end-to-end fashion by minimizing the loss using mini-batch stochastic gradient descent with Adam optimizer [Kingma and Ba, 2014].

We did not use some recent complemental techniques for training, regularization, and additional losses that can be used with existing methods including MEIM. Adding them to MEIM to further improve the result is left for future work.

5 Experiments

5.1 Experimental Settings

Datasets. We use three standard benchmark datasets, as shown in Table 1. WN18RR [Dettmers et al., 2018] is a subset of WordNet containing lexical information. FB15K-237 [Toutanova and Chen, 2015] is a subset of Freebase containing general facts. In addition, YAGO3-10 [Mahdisoltani et al., 2015] is a large and very competitive dataset containing general facts fromWikipedia.

Evaluations. We evaluate on the link prediction task. For each true triple (h, t, r) in the test set, we replace h and t by every other entity to generate corrupted triples (h′ , t, r) and (h, t′ , r), respectively. The model tries to rank the true triple (h, t, r) before the corrupted triples based on the score S. We compute MRR (mean reciprocal rank) and H@k for k ∈ {1, 3, 10} (how many triples correctly ranked in the top k) [Trouillon et al., 2016]. The higher MRR and H@k are, the better the model performs. Filtered metrics are used to avoid penalization when ranking other true triples before the current target triple [Bordes et al., 2013].

Note that MEIM can solve other tasks by converting them to contextual link prediction task. For example, by defining the alignment relation ra, we have the alignment triples of the form (e1, e2, ra), that can be used directly in MEIM for solving entity alignment. These tasks are left for future work.

Implementations. The MEIM model is implemented as a neural network using PyTorch. Following MEI, we use embedding size K = 3, C = 100 on WN18RR and FB15K-237, and K = 5, C = 100 on YAGO3-10. They are equivalent in terms of model size to ComplEx with embedding size 190 on WN18RR, 250 on FB15K-237, and 270 on YAGO3-10, that are relatively small compared to most related work. By preliminary experiments, we use k-vs-all sampling on WN18RR and 1-vs-all sampling on other datasets. We use batch size 1024 and learning rate 3e-3, with exponential decay 0.995 on YAGO3-10 and 0.99775 on other datasets. Dropout and batch normalization are used on the input and hidden layers, that is, on h and h⊤Mr, respectively. Other hyperparameters are tuned by grid search, with grids [0, 0.75] and step size 0.01 for drop rates, {1, 1e-1, 1e-2, 1e-3, 1e-4, 1e-5, 0} for λortho, {1, 1e-1, 1e-2, 1e-3, 5e-4, 1e-4, 0} for λunitnorm. The following values were found. On WN18RR, input drop rate is 0.71, hidden drop rate is 0.67, λortho = 1e-1, λunitnorm = 5e-4. On FB15K-237, input drop rate is 0.66, hidden drop rate is 0.67, λortho = 0, λunitnorm = 0. On YAGO3-10, input drop rate is 0.1, hidden drop rate is 0.15, λortho = 1e-3, λunitnorm = 0. Hyperparameters were tuned to maximize the validation MRR.

Baselines. The main baselines are MEI and tensordecomposition-based models such as ComplEx, RotatE, QuatE that we aim to improve. MEI results are reproduced in this paper using comparable model sizes and recent training techniques that lead to better performance than previously reported. We also compare to translation-based models such as TransE, neural-network-based models such as ConvE, and recent models such as RotH. We tried to evaluate against popular and strong baselines with comparable settings and evaluation protocol for fair and informative comparisons. These exclude incomparable results that use extra data, excessively large model size, complemental techniques, or inappropriate evaluation protocol [Sun et al., 2020]. Note that even at normal model size, MEIM can get state-of-the-art results.

5.2 Parameter Efficiency

We first compare the results and model sizes of MEIM to that of popular baselines as shown in Table 2. We can see that MEIM strongly outperforms these baselines while using roughly similar number of parameters.

In addition, Table 3 shows the scalability of MEIM and other baselines at different model sizes. MEIM outperforms the baselines at every model size and achieves the best result overall. This demonstrates that MEIM achieves a better balance between efficiency and expressiveness.

5.3 Main Results

The main results for link prediction are shown in Table 4. In general, MEIM achieves good result on all three datasets, including the very large and difficult YAGO3-10 dataset. Considering the most important and robust metric MRR, MEIM outperforms all baselines by a large margin on all datasets.

The strongest baselines on FB15K-237 and YAGO3-10 include the MEI model. This demonstrates that MEI was a strong model, but MEIM successfully addresses the drawbacks of MEI and its subsumed tensor-decomposition-based models to significantly improve the result on all datasets.

Another strong baseline is the RotH model on WN18RR, but the results on FB15K-237 are quite weak. We notice that most baselines usually have good results on one dataset and poor results on other datasets. This sets MEIM apart as it can achieve good results on all three datasets.

5.4 Analyses and Discussions

Now we look into the characteristics of the model to understand why it gives good result.

Ablation Study

Table 5 shows the main ablation results of MEIM, in which each feature is removed from the full model to see how much it contributes to the final results. For comparison, we also ablate the multi-partition embedding inherited from MEI.

In general, we can see that removing the features decreases the results, which confirms their effectiveness. In particular, the independent core tensor works well on all three datasets. It strongly contributes to the full model results on FB15K237, and slightly but consistently contributes to the final results on WN18RR and YAGO3-10. The effects of soft orthogonality is more complicated. On WN18RR and YAGO3-10, soft orthogonality is crucial to achieve the best results. However, it is harmful on FB15K-237 and was not included in the full model. This mixed effect guarantees further detailed investigation as follows.

The Effects of Soft Orthogonality

Table 6 shows the results when λortho is changing from no orthogonality to strong orthogonality. We see that strong orthogonality works well on WN18RR, average orthogonality works well on YAGO3-10, but none orthogonality works on FB15K-237. These results empirically demonstrate the weakness of previous models compared to MEIM.

For example, models with rigid orthogonality such as QuatE and RotH get strong results on WN18RR but weak on FB15K-237, whereas models without orthogonality such as RESCAL, TuckER, and MEI get strong results on FB15K237 but weak on WN18RR. In contrast, MEIM with soft orthogonality can get strong results on both WN18RR and FB15K-237. More importantly, the large and difficult YAGO3-10 dataset requires average orthogonality. Thus, only MEIM can get the best results on this dataset. This shows a crucial advantage of MEIM, because its soft orthogonality can be tuned to work optimally on each dataset.

6 RelatedWork

Knowledge graph embedding methods can be categorized based on how they compute the matching score.

Tensor-decomposition-based models adapt tensor representation formats such as CP format, Tucker format, and block term format to represent the knowledge graph data tensor [Kolda and Bader, 2009]. They make up most of recent state-of-the-art models such as ComplEx [Trouillon et al., 2016], SimplE [Kazemi and Poole, 2018], TuckER [Balaˇzevi´c et al., 2019], and MEI [Tran and Takasu, 2020]. They often achieve good trade-off between efficiency and expressiveness. Especially, the recent MEI model provides a systematic solution for this trade-off. However, they have some limitations affecting the ensemble effects and making them prone to degenerate, that we address in this paper.

Neural-network-based models use a neural network to compute the matching score, such as ConvE [Dettmers et al., 2018] using convolutional neural networks, CompGCN [Vashishth et al., 2020b] using graph convolutional networks. These models are generally more expensive but not always get better results than tensor-decomposition-based models.

Translation-based models use geometrical distance to compute the score, with relation embeddings act as the translation vectors, such as TransE [Bordes et al., 2013]. These models are efficient and intuitive, but they have limitations in expressive power [Kazemi and Poole, 2018].

There are several ways to use orthogonality in knowledge graph embedding, such as RotatE [Sun et al., 2019] using complex product, Quaternion [Tran and Takasu, 2019a] and QuatE [Zhang et al., 2019] using quaternion product, GCOTE [Tang et al., 2020] using the Gram Schmidt process, and RotH [Chami et al., 2020] using the Givens matrix. These models usually compute the score using distance, differently from MEIM. Moreover, they use rigid orthogonality that is harmful on some datasets and potentially suboptimal on other datasets compared to soft orthogonality.

7 Conclusion

In this paper, we introduce new aspects to tensordecomposition-based models exploiting ensemble boosting effects and max-rank relational mapping, in addition to multipartition embedding by MEI. We propose the MEIM model with two new techniques, namely independent core tensor to improve ensemble effects and max-rank mapping by soft orthogonality to improve expressiveness. MEIM achieves stateof-the-art link prediction results, including the large and difficult YAGO3-10 dataset, using fairly small embedding sizes. Moreover, we analyze and demonstrate the limitations of previous rigid orthogonality models and show that MEIM with soft orthogonality works well on multiple datasets.

For future work, it is promising to continue research on combining deep learning techniques and tensor decomposition. It is also interesting to use MEIM to solve tasks such as entity alignment and recommendation by converting them to contextual link prediction.

Acknowledgments

This work was supported by the Cross-ministerial Strategic Innovation Promotion Program (SIP) Second Phase, "Bigdata and AI-enabled Cyberspace Technologies" by the New Energy and Industrial Technology Development Organization (NEDO).