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 introduces a new self-supervised objective for improving link prediction in knowledge graphs. The idea presented in the paper naturally extends previous self-supervised objectives (e.g., TransE, CP, DistMult, etc.) by adding a training objective to predict the relation given the head and the tail entity in a knowledge graph triplet. The idea is really simple and it is surprising that previous works haven't tried this loss in addition to predicting the head/tail entity given the relation and the other entity.