Friendly Neighbors: Contextualized Sequence-to-Sequence Link Prediction (2023)

Abstract: We propose KGT5-context, a simple sequence-to-sequence model for link prediction (LP) in knowledge graphs (KG). Our work expands on KGT5, a recent LP model that exploits textual features of the KG, has small model size, and is scalable. To reach good predictive performance, however, KGT5 relies on an ensemble with a knowledge graph embedding model, which itself is excessively large and costly to use. In this short paper, we show empirically that adding contextual information — i.e., information about the direct neighborhood of the query entity — alleviates the need for a separate KGE model to obtain good performance. The resulting KGT5-context model is simple, reduces model size significantly, and obtains state-of-the-art performance in our experimental study.

## introduction a knowledge graph (kg) is a collection of facts describing relations between real-world entities. facts are represented in the form of subjectrelation-object @xcite ) triples such as @xcite . in this paper, we study the link prediction (lp) problem, which is to infer missing links in the kg. we focus on kgs in which the entities and relations have textual features, such as mention names or descriptions. @xcite made a case for large language models (lm) for this task. they proposed the kgt5 model, which posed the link prediction problem as a sequence-to-sequence (seq2seq) task. the main advantages of this approach are that (i) it allows for small model sizes, and (ii) it decouples inference cost from the graph size (and, in particular, the number of entities). they found that kgt5's performance was particularly strong when predicting the object of new relations for a query entity (e.g., the birthplace of a person), but fell short of alternative approaches when predicting additional objects for a known relation (e.g., additional awards won by someone). to avoid this problem, @xcite used an ensemble of kgt5 with a large knowledge graph embedding (kge) model (complex @xcite ). this ensemble did reach good performance but destroyed both advantages (i) and (ii) of using a lm. in fact, kge models learn a lowdimensional representation of each entity and each relation in the graph @xcite @xcite . consequently, model size and lp cost are linear in the number of entities in the graph, which can be expensive to use for large-scale kgs. for example, the currently bestperforming model @xcite for the large-scale wikikg90mv2 benchmark @xcite consists of an ensemble of 85 kge models; each taking up more than 86 gb of space for parameters. though kge model sizes can be reduced by using compositional embeddings based on text mentions @xcite @xcite , inference cost remains high for large graphs. we propose and study kgt5-context, which expands on kgt5 by providing contextual information about the query entity-i.e., information about the direct neighborhood of the query entity-to facilitate link prediction. our work is motivated by the kge model hitter @xcite , which follows a similar approach; we use the seq2seq model kgt5 instead of a transformer-based kge model. kgt5-context is very simple: the only change to kgt5 is that we add a verbalization of the neighborhood of the query entity to the description of a given lp task; see fig. 1 for an example. kgt5-context retains advantages (i) and (ii) of kgt5. we performed an experimental study using the wikidata5m @xcite and figure 1 : overview of kgt5-context (at bottom) and comparison to kgt5 (on top); real example from wiki-data5m, best viewed in color. kgt5-context differs from kgt5 in that it appends the neighboring relations and entities of yambaó (a drama movie) to the verbalized query. both models then apply t5, sample predictions from the decoder, map the samples to entities, and rank by sample logit scores. wikikg90mv2 @xcite benchmarks. we found that-without further hyperparameter tuning-kgt5-context reached or exceeded stateof-the-art performance on both benchmarks using a significantly smaller model size than alternative approaches. the simple kgt5-context model thus provides a suitable baseline for further research. ## expanding kgt5 with context given a query (s, r, ?) and a kg, lp is the task to predict new answer entities, i.e., the ? slot of the query. an example is given in fig. 1 . kgt5 @xcite treats link prediction as a seq2seq task. it exploits available textual information for entities and relations, such as mention names (for both entities and relations) or descriptions. kgt5's architecture is based on the encoderdecoder transformer model t5 @xcite . it uses canonical mentions to verbalize the lp query to a text sequence of form "predict tail: <subject @xmath0. to predict answers, kgt5 samples (exact) candidate mentions from the decoder; the cost of sampling is independent of the number of entities in the kg. to train kgt5, @xcite use standard training techniques for llms: kgt5 is trained on facts in the kg and asked to generate the true answer using teacher forcing and a crossentropy loss. kgt5-context (ours) proceeds in the same way as kgt5 but extends the verbalization of the query. in particular, we append a textual sequence of the one-hop neighborhood of the query entity s to the verbalized query of kgt5. as a result, the query entity is contextualized, an approach that has been applied successfully to kge models before @xcite . kgt5-context simplifies the predic-tion problem because additional information that is readily available in the kg is provided along with the query. in the example of fig. 1 , the contextual information states that yambaó is a mexican movie. this information is helpful; e.g., it already rules out the top two predictions of kgt5, which incorrectly suggest that yambaó is a piece of music. for a more detailed analysis, see sec. ## experimental study we conducted an experimental study to investigate (i) to what extent integrating context in terms of the entity neighborhood into kgt5 improves link prediction performance, (ii) whether the use of context can mitigate the necessity for an ensemble of the text-based kgt5 model with a kge model, and (iii) for what kind of queries context is helpful. we found that: ## conclusion we proposed and studied kgt5-context, a sequence-to-sequence model for link prediction in knowledge graphs. kgt5-context extends the kgt5 model of @xcite by using contextual information of the query entity for prediction. kgt5-context is simple, small, and scalable, and it obtained or exceeded state-of-the-art performance in our experimental study. it thus provides a suitable baseline for further research in this area. a natural direction, for example, is to explore approaches that integrate contextual information in a less naive way than kgt5-context does.