Beyond Link Prediction: On Pre-Training Knowledge Graph Embeddings (2024)

Abstract: Knowledge graph embeddings (KGEs) provide low-dimensional representations of the entities and relations in a knowledge graph (KG) in order to reason about the KG and to inject structured knowledge into various downstream applications. Most prior work, however, focuses almost exclusively on training and evaluating KGE models for the task of link prediction. In this work, we explore KGE models as general-purpose representations of KGs and study their suitability (i) for more generally capturing properties of the KG and (ii) for downstream tasks such as entity classification and regression. For (i), we designed a new set of graph-structure prediction tasks to assess whether models capture different structures in the graph. For (ii), we investigate whether models provide useful features for a variety of downstream tasks. We found that strong link prediction performance was neither an indication that models generally capture patterns in the graph, nor that they were more useful in downstream tasks. As a result, we included our proposed graph-structure prediction tasks as additional training objectives and found that models trained with this multi-task approach generally, but not always, performed better at both graph-structure prediction and downstream tasks. However, the most suitable choice of pre-training tasks varies across KGE models and types of downstream tasks, suggesting opportunities for more research into the relation between pre-training KGE models and their usability on downstream applications.

## introduction knowledge graph embeddings (kge) provide representations of the entities and relations in a knowledge graph @xcite . although a large number of kge models have been proposed, e.g. @xcite ; @xcite ; @xcite , most prior work focuses on the task of link prediction, i.e., answering questions such as (austin, capitalof, ?) by reasoning over an incomplete kg. in addition to link prediction, it is often argued that kges can provide representations that capture semantic properties of the entities @xcite @xcite @xcite and, indeed, pre-trained kge models have been used to inject structured knowledge into recommender systems @xcite ), question answering systems @xcite and other types of downstream applications @xcite . despite their use as kg representations in downstream applications, the question of whether pretrained kge models are generally useful representations of kgs-i.e. representations that are useful beyond the link prediction task-remains largely unexplored. specifically, it is not well-understood how different pre-training settings affect these representations. this stands in contrast with representation learning of natural language, where representations are intrinsically tested for known linguistic properties @xcite and extrinsically on their usability in downstream applications @xcite , and where different pre-training settings are known to improve performance @xcite . in this work, we study the suitability of kge models as general-purpose kg representations. first, we intrinsically assess whether kge models capture known properties of the graph, by evaluating their performance on basic graph-structure prediction tasks. we focus on new tasks that are similar to link prediction, but that test different forms of structural knowledge, such as predicting the relation of a triple (e.g., the relationship between austin and texas), the domain and range of a relation (e.g., whether austin is a capital), and the entity and relation neighborhood of an entity (e.g., which entities are related to austin). we found that commonly trained kge models often performed poorly on such tasks, challenging the intuition that kge models preserve the structure of a kg. second, we extrinsically evaluate whether kge models are useful pre-trained representations for node-level downstream tasks such as entity classification (e.g., the profession of a person) or regression (e.g., the rating of a movie). we conduct an empirical study using 35 downstream tasks on three different kgs. we found that kge models often perform decent on these tasks, almost always exceeding the performance of graph neural networks that train directly on the downstream task, such as ke-gcn @xcite . however, the kge models with best downstream task performance were often not the best-performing models for link prediction. for example, the basic transe model @xcite can be superior to kge models with stronger performance on link prediction, such as complex @xcite or rotate @xcite . this suggests that good link prediction performance is not necessarily indicative of good downstream task performance. both of these findings suggest that the focus on link prediction tasks is too narrow for pre-training kge models, i.e., to provide generally useful representations of a kg. we thus included the graphstructure prediction tasks discussed above as additional training objectives. the resulting multitask kge models had significantly better overall performance for graph-structure prediction tasks, suggesting that the learned representations capture more information about the graph, at the cost of a small drop in link prediction performance. perhaps more importantly, when using pretrained kges in downstream tasks, we found that multi-task training often (but not always) improved downstream performance, especially as data becomes scarce. in fact, excluding the link prediction task during pre-training resulted in better downstream performance more often than not. however, capturing more information about the graph did not directly translate to better downstream performance, as the best performing models in downstream applications were often those that were not pre-trained using all possible tasks. in general, the best choice of pre-training tasks depends on the dataset, kge model, and type of downstream task, suggesting opportunities for more research to better understand how to pre-train kge models so they provide generally useful kg representations. we provide all of our resources 1 to promote future work in this direction. 1 available at https://github.com/uma-pi1/kge-pretraining . ## preliminaries and related work we briefly describe kge models, training and evaluation methods for link prediction, as well as prior work on other tasks. for a more comprehensive discussion, please see surveys from @xcite ; @xcite . ## graph structure prediction in this section, we describe the new graph-structure tasks used in our study. specifically, how we use them to test whether kge models preserve known properties in a kg, and how we adapted kges to efficiently train on these tasks. ## experimental study to our knowledge, no prior work has studied the impact that different training objectives have on kg embedding quality, despite this being common practice, e.g. in language models @xcite . we conducted a large experimental study with the following goals: (i) to assess whether kge models capture various properties of a kg by intrinsically evaluating their performance on new graph-structure prediction tasks, (ii) to determine whether (and by how much) kges improve their performance on these tasks when simultaneously trained for them, and (iii) to assess the impact that different pre-training approaches have on downstream tasks by extrinsically evaluating pre-trained kge models. we briefly describe our experimental setup here, for details, see sec. b.1. pre-training setup. for training and evaluating kges, we closely follow @xcite . we implemented everything in libkge @xcite , used four benchmark datasets commonly used in recent work @xcite @xcite , all models were trained under the same conditions (as much as possible) and tuned with a large hyperparameter space using random search. for mtt training, we used all tasks in choice of kge models. we focused on models that provide entity representations, so we may test their quality in downstream tasks, as done in the industry @xcite . we chose four popular models: transe @xcite , distmult @xcite , ro-tate @xcite and complex @xcite . these are not the latest kge models, but they are common baselines in recent work @xcite @xcite , and are common choices for pre-trained models @xcite @xcite . they can also reach sota performance with reasonable embedding sizes @xcite , allowing us to scale our study, and with larger embedding size @xcite and additional training objectives (chen et al., 2021b), complex outperforms more involved models, e.g. the transformer-based hitter model @xcite . some recent models achieve better performance on link prediction, but focus exclusively on that task and do not directly provide entity representations for downstream tasks, e.g. hitter (chen et al., 2021a) and nbfnet @xcite . ## conclusion to explore kge models as general-purpose representions of kgs, we designed a new set of graphstructure prediction tasks for intrinsic evaluation. we found that standard kge models are not good at predicting simple structures in the graph, challenging the intuition that these models generally capture properties in a kg. in addition, we extrinsically evaluated pre-trained kge models on several entity-level downstream tasks. we found that link prediction was not indicative of good downstream performance, and that multi-task pre-training was generally better for downstream tasks, often when excluding link prediction during pre-training. however, the best choice of pre-training tasks depends on both kge model and downstream task, suggesting more research is needed into pre-training kges to obtain generally-useful kg representations.