Relation Prediction as an Auxiliary Training Objective for Improving Multi-Relational Graph Representations

tion ( KBC ) . In this paper , we propose a new self-supervised training objective for multi-relational graph representation learning , via simply incorporating relation prediction into the commonly used 1vsAll objective . The new training objective contains not only terms for predicting the subject and object of a given triple , but also a term for predicting the relation type . We analyse how this new objective impacts multi-relational learning in KBC : experiments on a variety of datasets and models show that relation prediction can significantly improve entity ranking , the most widely used evaluation task for KBC , yielding a 6.1 % increase in MRR and 9.9 % increase in Hits @ 1 on FB15k-237 as well as a 3.1 % increase in MRR and 3.4 % in Hits @ 1 on Aristo-v4 . Moreover , we observe that the proposed objective is especially effective on highly multi-relational datasets , i.e . datasets with a large number of predicates , and generates better representations when larger embedding sizes are used . 1 . Introduction . Aiming at completing missing entries , Knowledge Base Completion ( KBC ) , also known as Link Prediction , plays a crucial role in constructing large-scale knowledge graphs [ Nickel et al. , 2016 , Ji et al. , 2020 , Li et al. , 2020 ] . Over the past years , most of the research on KBC has been focusing on Knowledge Graph Embedding models , which learn representations for all entities and relations in a Knowledge Graph , and use them for scoring whether an edge exists or not [ Nickel et al. , 2016 ] . Numerous models and architectural innovations have been proposed in the literature , including but not limited to translation-based models [ Bordes et al. , 2013 ] , latent factorisation models [ Nickel et al. , 2011 , Trouillon et al. , 2016 , Balazevic et al. , 2019 ] , and neural network-based models [ Dettmers et al. , 2018 , Schlichtkrull et al. , 2018 , Xu et al. , 2020 ] . Other more recent research has been making complementary efforts on analysing the evaluation procedures for these KBC models . For instance , Sun et al . [ 2020 ] call for standardisation of evaluation protocols ; Kadlec et al . [ 2017 ] , Ruffinelli et al . [ 2020 ] and Jain et al . [ 2020 ] highlight the importance of training strategies and show that careful hyper-parameter tuning can produce more accurate results than adopting more elaborate model architectures ; Lacroix et al . [ 2018 ] suggests that a simple model can produce state-of-the-art results when its training objective is properly selected . Taking inspiration from these findings , this paper explores relation prediction : a simple auxiliary training objective that significantly improves multi-relational graph representation learning across several KBC models . Aside from training models to predict the subject and object entities for triples in a Knowledge Graph , we also train them to predict relation types , leading to a self-supervised training objective . Intuitively , this approach is akin to using a masked language model-like training objective [ Devlin et al. , 2019 ] instead of the commonly used auto-regressive training objective for KBC . In our experiments , we find that the new auxiliary training objective significantly improves downstream link prediction accuracy . Empirical evaluations on various models and datasets support the effectiveness of our new training objective : the largest improvements were observed on ComplEx-N3 [ Trouillon et al. , 2016 ] and CP-N3 [ Lacroix et al. , 2018 ] with embedding sizes between 2K and 4K , providing up to 9.9 % boost in Hits @ 1 and 6.1 % boost in MRR on FB15k-237 with negligible computational overhead . We further experiment on datasets with varying numbers of predicates and find that relation prediction helps more when the dataset is highly multi-relational , i.e . contains a larger number of predicates . Moreover , our qualitative analysis demonstrates improved prediction of some MANY-TOMANY [ Bordes et al. , 2013 ] predicates and more diversified relation representations . 2 . Background and Related Work . A Knowledge Graph G ⊆ E ×R×E contains a set of subject-predicate-object 〈s , p , o〉 triples , where each triple represents a relationship of type p ∈ R between the subject s ∈ E and the object o ∈ E of the triple . Here , E andR denote the set of all entities and relation types , respectively . Knowledge Graph Embedding Models A Knowledge Graph Embedding ( KGE ) model , also referred to as neural link predictor , is a differentiable model where entities in E and relation types in R are represented in a continuous embedding space , and the likelihood of a link between two entities is a function of their representations . More formally , KGE models are defined by a parametric scoring function φθ : E ×R× E 7→ R , with parameters θ that , given a triple 〈s , p , o〉 , produces the likelihood that entities s and o are related by the relationship p. Scoring Functions KGE models can be characterised by their scoring function φθ . For example , in TransE [ Bordes et al. , 2013 ] , the score of a triple 〈s , p , o〉 is given by φθ ( s , p , o ) = −‖s + p− o‖2 , where s , p , o ∈ Rk denote the embedding representations of s , p , and o , respectively . In DistMult [ Yang et al. , 2015 ] , the scoring function is defined as φθ ( s , p , o ) = 〈s , p , o〉 = ∑k i=1 sipioi , where 〈 · , · , · 〉 denotes the tri-linear dot product . Canonical Tensor Decomposition [ CP , Hitchcock , 1927 ] is similar to DistMult , with the difference that each entity x has two representations , xs ∈ Rk and xo ∈ Rk , depending on whether it is being used as a subject or object : φθ ( s , p , o ) = 〈ss , p , oo〉 . In RESCAL [ Nickel et al. , 2011 ] , the scoring function is a bilinear model given by φθ ( s , p , o ) = s > Po , where s , o ∈ Rk is the embedding representation of s and p , and P ∈ Rk×k is the representation of p. Note that DistMult is equivalent to RESCAL if P is constrained to be diagonal . Another variation of this model is ComplEx [ Trouillon et al. , 2016 ] , where the embedding representations of s , p , and o are complex vectors – i.e . s , p , o ∈ Ck – and the scoring function is given by φθ ( s , p , o ) = < ( 〈s , p , o〉 ) , where < ( x ) represents the real part of x , and x denotes the complex conjugate of x . In TuckER [ Balazevic et al. , 2019 ] , the scoring function is RELATION PREDICTION AS AN AUXILIARY TRAINING OBJECTIVE FOR KBC defined as φθ ( s , p , o ) = W ×1 s ×2 p ×3 o , where W ∈ Rks×kp×ko is a three-way tensor of parameters , and s ∈ Rks , p ∈ Rkp , and o ∈ Rko are the embedding representations of s , p , and o . In this work , we mainly focus on DistMult , CP , ComplEx , and TuckER , due to their effectiveness on several link prediction benchmarks [ Ruffinelli et al. , 2020 , Jain et al. , 2020 ] . Training Objectives Another dimension for characterising KGE models is their training objective . Early tensor factorisation models such as RESCAL and CP were trained to minimise the reconstruction error of the whole adjacency tensor [ Nickel et al. , 2011 ] . To scale to larger Knowledge Graphs , subsequent approaches such as Bordes et al . [ 2013 ] and Yang et al . [ 2015 ] simplified the training objective by using negative sampling : for each training triple , a corruption process generates a batch of negative examples by corrupting the subject and object of the triple , and the model is trained by increasing the score of the training triple while decreasing the score of its corruptions . This approach was later extended by Dettmers et al . [ 2018 ] where , given a subject s and a predicate p , the task of predicting the correct objects is cast as a |E|-dimensional multi-label classification task , where each label corresponds to a distinct object and multiple labels can be assigned to the ( s , p ) pair . This approach is referred to as KvsAll by Ruffinelli et al . [ 2020 ] . Another extension was proposed by Lacroix et al . [ 2018 ] where , given a subject s and an object p , the task of predicting the correct object o in the training triple is cast as a |E|-dimensional multi-class classification task , where each class corresponds to a distinct object and only one class can be assigned to the ( s , p ) pair . This is referred to as 1vsAll by Ruffinelli et al . [ 2020 ] . Note that , for factorisation-based models like DistMult , ComplEx , and CP , KvsAll and 1vsAll objectives can be computed efficiently on GPUs [ Lacroix et al. , 2018 , Jain et al. , 2020 ] . For example for DistMult , the score of all triples with subject s and predicate p can be computed via E ( s p ) , where denotes the element-wise product , and E ∈ R|E|×k is the entity embedding matrix . In this work , we follow Lacroix et al . [ 2018 ] and adopt the 1vsAll loss , so as to be able to compare with their results , and since Ruffinelli et al . [ 2020 ] showed that they produce similar results in terms of downstream link prediction accuracy . Recent work on standardised evaluation protocols for KBC models [ Sun et al. , 2020 ] and their systematic evaluation [ Kadlec et al. , 2017 , Mohamed et al. , 2019 , Jain et al. , 2020 , Ruffinelli et al. , 2020 ] shows that latent factorisation based models such as RESCAL , ComplEx , and CP are very competitive when their hyper-parameters are tuned properly [ Kadlec et al. , 2017 , Ruffinelli et al. , 2020 ] . In this work , we show that using relation prediction as an auxiliary training task can further improve their downstream link prediction accuracy . 3 . Relation Prediction as An Auxiliary Training Objective . In what is referred to as the 1vsAll setting [ Ruffinelli et al. , 2020 ] , KBC models are trained using a self-supervised training objective by maximising the conditional likelihood of the subject s ( resp . object o ) of training triples , given the predicate and the object o ( resp . subject s ) . More formally , KBC models are trained by maximising the following objective : arg max θ∈Θ ∑ 〈s , p , o〉∈G [ logPθ ( s | p , o ) + logPθ ( o | s , p ) ] with logPθ ( o | s , p ) = φθ ( s , p , o ) − log ∑ o′∈E exp [ φθ ( s , p , o ′ ) ] logPθ ( s | p , o ) = φθ ( s , p , o ) − log ∑ s′∈E exp [ φθ ( s ′ , p , o ) ] , ( 1 ) where θ ∈ Θ are the model parameters , including entity and relation embeddings , and φθ is a scoring function parameterised by θ . Such an objective limits predicting positions in the training objective to either the first ( s ) or the last ( o ) element of the triple . In this work , we propose relation prediction as an auxiliary task for training KBC models . The new training objective not only contains terms for predicting the subject and the object of the triple – logP ( s | p , o ) and logP ( o | s , p ) in Equation ( 1 ) – but also an objective logP ( p | s , o ) for predicting the relation type p : arg max θ∈Θ ∑ 〈s , p , o〉∈G [ logPθ ( s | p , o ) + logPθ ( o | s , p ) + λ logPθ ( p | s , o ) ] with logPθ ( p | s , o ) = φθ ( s , p , o ) − log ∑ p′∈R exp [ φθ ( s , p ′ , o ) ] , ( 2 ) where λ ∈ R+ is a user-specified hyper-parameter that determines the contribution of the relation prediction objective ; we assume λ = 1 unless specified otherwise . This new training objective adds very little overhead to the training process , and can be easily added to existing KBC implementations ; PyTorch examples are included in Appendix A . Compared to conventional approaches , relation prediction can help the model learn to further distinguish among different predicates .

Summary: The paper proposes to extend the objective function for KG-embedding learning with a term that predicts the relation (predicate) given two entities. This extension is well motivated and the experiments show that it leads to small but consistent improvements. While the technical contribution of the paper is rather small, the experimental setup is sound and there are many interesting results and analyses. Especially the clear statement and answer of the research questions is nice.