Probabilistic Knowledge Graph Embeddings

1 INTRODUCTION . In 2012 , Google announced that it improved the quality of its search engines significantly by utilizing knowledge graphs ( Eder , 2012 ) . A knowledge graph is a dataset of facts represented in terms of triplets ( head , relation , tail ) , where head and tail represent entities in the world , and a relation is a property that describes the relationship between the two entities . As an example , when the head entity is ’ Paris ’ , and the relation is ’ is in ’ , then the tail entity could be ’ France ’ or ’ Europe ’ . While the number of possible relations among entities is enormous , the amount of edges in empirical knowledge graphs is often rather small . It is therefore desirable to complete the missing edges in a knowledge graph algorithmically based on patterns detected in the observed part of he graph ( Nickel et al. , 2016a ) . Such link prediction has become an important subfield of artificial intelligence ( Bordes et al. , 2013 ; Wang et al. , 2014 ; Lin et al. , 2015 ; Nickel et al. , 2016c ; Trouillon et al. , 2016 ; Wang & Li , 2016 ; Ji et al. , 2016 ; Shen et al. , 2016 ; Xiao et al. , 2017 ; Shi & Weninger , 2017 ; Lacroix et al. , 2018 ) . In statistical relational learning , link prediction is done by inferring/learning explicit rules about relationships and then applying these inferred rules in order to reason about unobserved facts ( Friedman et al. , 1999 ; Kersting et al. , 2011 ; Niu et al. , 2012 ; Pujara et al. , 2015 ) . A complementary , highly scalable approach relies on embedding models ( Kadlec et al. , 2017 ; Nguyen , 2017 ) . These models represent entities and relationships in terms of low-dimensional semantic vectors and implicitly encode the relationships in the knowledge graphs in terms of three-way mathematical operations among these vector triplets ( Yang et al. , 2015 ; Trouillon et al. , 2016 ) . In the early works , different embeddings were used for the two tasks of head and tail prediction . Lacroix et al . ( 2018 ) proposed to map these two tasks to just the tail prediction task by augmenting the data by reciprocal facts with inversed relations . In the present work , we also use this data augmentation technique to pave the way for a probabilistic interpretation of knowledge graph embedding models . The probabilistic formulation allows us to keep track of parameter uncertainty by applying Bayesian methods . We argue that the parameter uncertainty in knowledge graph embedding models can be significant , calling for a Bayesian approach . As the number of parameters grows linearly with the number of entities , and every entity is only connected to a small number of other entities by observed links , this effectively leads to a small data problem , even if the knowledge graph as a whole may be huge . For example , Figure 1 shows that for a popular large-scale knowledge base , about 80 % of entities have less than 80 observed facts . Similar observation have been made for recommender systems , where most users interact only with a small subset of the items ( Gopalan et al. , 2014 ; 2015 ; Liang et al. , 2018 ) . As in recommender systems , we show that a Bayesian approach can be highly beneficial in knowledge graph embeddings . This paper presents such a Bayesian formulation and presents an interpretation of popular embedding models as generative models of relational facts . The benefits of a Bayesian formulation of knowledge graph embeddings are two-fold . First , this formulation allows us to infer posterior distributions of embedding vectors , giving us uncertainty estimates . An arguably more important aspect is that of hyperparameter tuning , as the state-of-the-art embedding models employ a large number of hyperparameters ( regularizers ) , one for every entity ( Lacroix et al. , 2018 ) . Kadlec et al . ( 2017 ) showed that careful hyperparameter tuning can dramatically increase predictive performance . While there is no classical way to scalably tune hyperparameters on a per-entity level without some heuristic , our Bayesian approach allows for gradient-based hyperparameter optimization on the marginal likelihood of the data ( Bishop , 2006 ) . This approach , termed variational expectation maximization ( variational EM ( Bernardo et al. , 2003 ) ) , still works if variational approximations are used ( Jordan et al. , 1999 ; Blei et al. , 2017 ; Zhang et al. , 2017 ) . We propose a simple , principled , and efficient way to tune hyperparameters using variational EM . This approach lower-bounds the marginal likelihood of the data and allows this bound to be optimized for hyperparameters without running into degenerate solutions ( Bernardo et al. , 2003 ; Mandt et al. , 2017 ) . This is in contrast to point estimation , which results in degeneracies when jointly optimizing over model parameters and hyperparameters . As a result , our approach learns a macroscopic number of hyperparameters , resulting in a new state-of-the-art performance on link prediction . Our contribution in this work can be summarized as follows : • We present a probabilistic interpretation of existing knowledge graph embedding models such as ComplEx and DistMult . We reformulate these models as generative models for relational facts , which paves the way for Bayesian inference in these models . • We apply stochastic variational inference to scalably estimate an approximate posterior for each entity and relation embedding in the knowledge graph . This approach not only allows us to estimate uncertainty , but more importantly allows for gradient-based hyperparameter optimization by stochastic gradient descent on the optimized variational bound ( variational EM ) . Implementing our methodology in existing models requires only a few lines of additional code , and it does not increase the computational complexity compared to a non-Bayesian training approach . • We show experimentally that our method reaches new state-of-the-art results in link prediction . We present results on several large data sets using two popular models which we trained using variational EM . We observe that our improvements are most notable for entites with few training points , as we show by introducing a new , balanced evaluation metric . The paper is structured as follows . In Section 2 , we review the related work about knowledge graph embeddings and variational inference . In Section 3 , we present the generative model for relational facts and reinterpret existing loss functions as the maximum a posterior estimator of the generative model . In Section 4 , we propose our approximate Bayesian expectation maximization algorithm for finding the hyperparameters . Finally , in Section 5 we test the hyperparameter learning algorithm for the link prediction task on two different embedding algorithms using standard performance metrics . 2 RELATED WORK . Related work to this paper can be grouped into link prediction algorithms and variational inference . Link Prediction . Knowledge graphs and link prediction gained a lot of attention ; for a review see ( Nickel et al. , 2016b ) . Most link prediction algorithms use latent features of entities or relationships , also termed embeddings . TransE proposed by Bordes et al . ( 2013 ) is one example , where each entity and relation are embedded in a K-dimensional space . The embeddings are found such that for a given fact ( h , r , t ) , the sum of embeddings of h plus r be close to t , and far from other entities . Many models are based on tensor factorization , where an adjacency tensor is built from the triplet facts , and entities and relations are embedded to a K-dimensional space . Different methods differ in the structure of the embedded vectors . Hitchcock ( 1927 ) assume that an entity ’ s embedding depends on whether the entity appears as the head or tail of a fact . DistMult simplified the model by assuming that tails and heads have the same embeddings ( Yang et al. , 2015 ; Toutanova & Chen , 2015 ) . ComplEx extended DistMult by embedding the entities and relations to a complex space to allow for nonsymmetric relations . Lacroix et al . ( 2018 ) introduced reciprocal relations and trained with the full log-loss instead of binary loss with negative sampling , which improved the results . Link prediction algorithms are susceptible to overfitting , and the importance of hyperparameter choices has been pointed out by several authors , e.g. , ( Kadlec et al. , 2017 ; Lacroix et al. , 2018 ) . Different reguralizers have been proposed , the most standard one being the 2-norm ( Yang et al. , 2015 ) . A heuristic that often produces the best results ( e.g. , in ( Lacroix et al. , 2018 ) ) is to choose the strength of the regularizer proportional to the frequency of the entities ( see e.g. , ( Srebro & Salakhutdinov , 2010 ) ) . In contrast , our approach learns entity-dependend hyperparameters without relying on heuristics . Also , Lacroix et al . ( 2018 ) showed that the 2-norm is not a valid tensor norm , and that the 3-norm is more appropriate for triplet data . Variational Inference . Variational inference ( VI ) is a powerful technique to approximate a Bayesian posterior over latent variables given observations ( Jordan et al. , 1999 ; Blei et al. , 2017 ; Zhang et al. , 2017 ) . Besides approximating the posterior , VI also estimates the marginal likelihood of the data . This allows for gradient-based hyperparameter tuning , a technique called variational EM ( Bernardo et al. , 2003 ) , which is the main benefit of the Bayesian approach used in this paper . Our paper builds on recent probabilistic extensions of embedding models to Bayesian models , such as word ( Barkan , 2017 ) or paragraph ( Ji et al. , 2017 ) embeddings . In these works , the words are embedded into a K-dimensional space . It has been shown that using a probabilistic approach leads to better performance on smaller data , and allows these models to be combined with powerful priors , such as for time series modeling ( Bamler & Mandt , 2017 ; Jähnichen et al. , 2018 ) . Yet , the underlying probabilistic models in these papers are very different from the ones considered in our work . 3 KNOWLEDGE GRAPH EMBEDDINGS AS GENERATIVE MODELS OF FACTS . In this section , we describe a generative model for relational facts . Then , we show that existing tensor factorization models can be interpreted as a maximum a posteriori ( MAP ) estimation of the parameters of the generative model ( Bishop , 2006 ) . In Section 4 , we move to our main contribution . We use our generative model to derive an approximate Bayesian inference algorithm to optimize over hyperparameters ( regularizers ) in a principled way . Let S = { ( hi , ri , ti ) } i=1 , ... , n be the set of n triplet facts , with heads hi ∈ [ Ne ] , tails ti ∈ [ Ne ] , and relations ri ∈ [ Nr ] . Here , Ne and Nr are the number of distinct entities and relations that we model . From now on , we use the data augmentation technique used in ( Lacroix et al. , 2018 ) ( see Remark 2 in the appendix ) . It reduces the task of performing both head and tail predictions to that of performing only tail predictions by introducing reciprocal relations . Thus , given a head h and relation r , the link prediction boils down to finding the tail t in the fact ( h , r , t ) . For simplicity , we present the generative process at the example of the DistMult model ( Yang et al. , 2015 ) . Generalization to other models , such as ComplEx ( Trouillon et al. , 2016 ) or CP ( Hitchcock , 1927 ) is straight-forward . In DistMult , each entity and relation is embedded in the latent vector space RK . Let E ∈ RNe×K ( similarly R ∈ RNr×K ) be the embedding matrix for entities ( relations ) , where Ee ∈ RK ( Rr ∈ RK ) is the embedding vector for entity e ∈ [ Ne ] ( relation r ∈ [ Nr ] ) . The Score . Embedding models such as DistMult rely on a score Xh , r , t which plays a role both in the probabilistic and classical formulation of the model . This score assigns a scalar value to each combination of the embeddings Eh , Rr , and Et . In DistMult , the score is Xh , r , t = ∑K i=1 EhiRriEti . In the conventional approach , the embedding vectors are trained in such a way that observed triplets have high scores . As we show below , the score also enters the probabilistic extension of the model . Generative Process . We can now present our perspective of knowledge graph embeddings as generative models . Figure 2 shows a graphical model . The generative process is as follows : • For each entity e ∈ [ Ne ] , draw an embedding Ee ∈ RK from a prior p ( Ee|λe ) , e.g. , a normal distribution . Similarly , for each relation r ∈ [ Nr ] , draw an embedding Rr ∈ RK from a prior p ( Rr|λr ) . Here , λe and λr are hyperparameters . • Repeat for each triplet index i ∈ { 1 , . . . , n } : – Draw a head hi and a relation ri from a discrete joint distribution P ( hi , ri ) . – Draw a tail ti ∼ Multinomial ( softmaxt ( Xhi , ri , t ) ) . Above , the softmax normalizes the scores over the tail index , softmaxt ( Xh , r , t ) = exp Xh , r , t∑ t′ exp Xh , r , t′ . We see that the resulting generative model produces fact triplets conditioned on semantic embeddings of entities and relations . This concludes the generative process perspective . MAP Estimation Recovers Original Model . The cost function used for learning E and R in the existing works can be seen as maximum a posteriori ( MAP ) estimation of the generative model above , where we optimize the log joint distribution over the embeddings E and R : E ? , R ? = arg max E , R log p ( E , R , S|λ ) = arg max E , R { log p ( S|E , R ) + log p ( E , R|λ ) } . ( 1 ) In the DistMult model , the loss function with a Gaussian prior with precision λ becomes L ( E , R , λ ) = − log p ( S|E , R ) − log p ( E , R|λ ) , ( 2 ) where the negative log likelihood is just the cross entropy , − log p ( S|E , R ) = ∑ ( h , r , t ) ∈S ( −Xh , r , t + log ∑ t′ exp Xh , r , t′ ) + c , ( 3 ) and the negative log prior leads to the regularizer − log p ( E , R|λ ) = 1 2 ∑ e∈ [ Ne ] ( λe‖Ee‖2 −K log ( λe ) ) + 1 2 ∑ r∈ [ Nr ] ( λr‖Rr‖2 −K log ( λr ) ) + c. ( 4 ) Up to the terms K log ( λe/r ) in Eq . 4 , which do not affect the minimization over E and R , Eqs . 2-4 recover precisely the loss function used in ( Lacroix et al. , 2018 ) .

Summary: In this paper, authors propose a probabilistic extension of classic Neural Link Prediction models, such as DistMult and ComplEx. The underlying assumption is that the entity embeddings and the relation embeddings are sampled from a prior Multivariate Normal distribution, whose (hyper-)parameters can be estimated via maximum likelihood. In this paper, authors use Variational Inference (VI) for approximating the posterior distribution over the embeddings, and use Stochastic VI for maximising the Evidence Lower BOund (ELBO) while scaling to large datasets. In Sect. 3, authors introduce the generative process, and show how MAP estimation of the embedding matrices can recover the original models. In Sect. 4, authors start from the intractable marginal likelihood over the data (Eq. 5) for deriving the corresponding ELBO (Eq. 6), which is defined over: