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 generalpurpose 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 multitask 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.

heading: 1 Introduction
text: Knowledge graph embeddings (KGE) provide representations of the entities and relations in a knowledge graph (KG). Although a large number of KGE models have been proposed, e.g. Ge et al. (2023); Xiao et al. (2022); Bai et al. (2022), 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 (Wang et al., 2022a; Ji et al., 2021; Wang et al., 2017; Nickel et al., 2015; Bordes et al., 2013, 2011) and, indeed, pre-trained KGE models have been used to inject structured knowledge into recommender systems (El-Kishky et al., 2022; Wang et al., 2018), question answering systems (Ilyas et al., 2022) and other types of downstream applications (Ji et al., 2021).
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 (Mikolov et al., 2013) and extrinsically on their usability in downstream applications (Devlin et al., 2019; Radford et al., 2018), and where different pre-training settings are known to improve performance (Raffel et al., 2020; Liu et al., 2019).
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.
136
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 (Yu et al., 2021). 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 (Bordes et al., 2013) can be superior to KGE models with stronger performance on link prediction, such as ComplEx (Trouillon et al., 2016) or RotatE (Sun et al., 2019). 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 resources1 to promote future work in this direction.
1Available at https://github.com/uma-pi1/kge-pretraining.

heading: 2 Preliminaries and Related Work
text: 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 Nickel et al. (2015); Wang et al. (2017); Ji et al. (2021).
Link prediction. A knowledge graph G ⊆ E × R × E is a collection of (subject, predicate, object)-triples over a set E of entities and a set R of relations. Triples represent known facts such as (Austin, capitalOf, Texas). In the KGE literature, the link prediction task is defined as predicting the subject or object to questions of the form (?, capitalOf, Texas) and (Austin, capitalOf, ?), resp.
KGE models. KGE models represent each entity and each relation of a KG with a lowdimensional embedding. Each model has a scoring function s : E × R × E → R that maps each possible triple to a real-valued score. Intuitively, high scores indicate plausible triples, low scores implausible triples. For example, TransE (Bordes et al., 2013) is a translation-based model with s(i, k, j) = −∥ei + rk − ej∥, where ei ∈ Rd and rk ∈ Rd are entity and relation embeddings, resp. Scoring functions can be more involved, e.g., based on transformers (Chen et al., 2021a).
Standard training. KGE models are commonly trained on the link prediction task. For each training triple (s, p, o), models are trained such that score s(s, p, o) is high (a known positive) but score s(s, p, o′) is low for (pseudo-)negative triples (s, p, o′), where o′ ̸= o ∈ E ; similarly for subjects s′ ∈ E with negative triple (s′, p, o). Different training objectives exist, all of which follow this approach, but otherwise differ in other hyperparameters; for details, see Ali et al. (2021).
Standard evaluation. The most common evaluation protocol is entity ranking (ER), and it is also based on link prediction. Given test triple (s, p, o), models answer the link prediction queries (s, p, ?) and (?, p, o) by ranking all possible answers to each query by their scores, after filtering other known answers. Metrics such as mean reciprocal rank (MRR) and Hits@K are then computed based on the rank of the answers s and o, resp. As an evaluation method, entity ranking has been questioned in prior work (Zhou et al., 2022; Tiwari et al., 2021; Safavi and Koutra, 2020; Wang et al., 2019). In this work, we focus mostly on other evaluation tasks.
Other training approaches. Nickel et al. (2011) and Li et al. (2021) trained on the reconstruction
task, which aims at reconstructing the training set using cost functions such as ∑ t∈Gtrain∥I[(t)]− s(s, p, o)∥22, for training set Gtrain and where I[·] is a 0/1 indicator. We do not consider such methods due to excessive training costs. Chen et al. (2021b) augmented the link prediction task with relation prediction during training (but not evaluation). We extend this work by including additional pre-training tasks and by focusing on graph-structure prediction and downstream task performance instead.
Other evaluation approaches. Some works evaluate KGE models using triple classification (Socher et al., 2013; Lin et al., 2015; Wang et al., 2022b). We do not consider this task because performance estimates are typically overly optimistic and misleading unless hard negatives are used (Safavi and Koutra, 2020); such hard negatives are generally not available. Chang et al. (2020) evaluated KGE models on the relation prediction task, which we also consider as one evaluation task in this work. There is also work on probing KGE models (Meilicke et al., 2018; Allen et al., 2021; Rim et al., 2021), which focus on link prediction performance across different types of relations, e.g. symmetric. In contrast, we focus on studying whether models provide useful representations, i.e. we focus on embedding quality, not just on link prediction performance. In addition, pre-trained KGE models have been used as components in language models (He et al., 2020; Zhang et al., 2019), visual models (Baier et al., 2017), recommender systems (El-Kishky et al., 2022; Wang et al., 2018), or question answering systems (Ilyas et al., 2022). Similarly, some studies have evaluated pre-trained KGE models for entity classification or regression tasks (Pezeshkpour et al., 2018; Jain et al., 2021), as we do. We extend this line of work with a larger set of downstream tasks, and by being the first (to our knowledge) to study the impact of different pretraining methods on downstream task performance.

heading: 3 Graph Structure Prediction
text: 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.

heading: 3.1 Graph-Structure Tasks
text: An example and summary of the graph-structure tasks that we use in our study is given in Table 1.
We describe the queries for each task as a triple such as (s, ?, ∗), where s or o denote input entities, p denotes an input relation, ? denotes the prediction target, and ∗ acts as a wildcard. Using this notation, we consider the following tasks and queries:
• Link prediction (LP): Given a relation and a subject, predict the object (denoted (s, p, ?)). Likewise, given a relation and an object, predict the subject (denoted (?, p, o)).
• Relation prediction (REL, Chang et al. (2020); Chen et al. (2021b)): Given two entities s and p, predict the relation between them (denoted (s, ?, o)).
• Domain prediction (DOM): Given a relation, predict its domain (denoted (?, p, ∗)) or its range (denoted (∗, p, ?)).
• Entity neighborhood prediction (NBE): Given a subject entity, predict related objects (denoted (s, ∗, ?)). Likewise, given an object, predict related subjects (denoted (?, ∗, o)).
• Relation neighborhood prediction (NBR): Given a entity, predict the relations where it occurs as subject (denoted (s, ?, ∗)) and where it occurs as object (denoted (∗, ?, o)).
Note that we use the wildcard to denote existential quantification. For example, given a ground-truth KG G and domain prediction query (?, p, ∗), an entity s ∈ E is a correct answer if there exists an entity o ∈ E such that (s, p, o) ∈ G. We illustrate these new tasks in Figure 2 in Appendix A.
We chose this particular set of tasks because they are simple, they capture basic information about the graph structure beyond link prediction, and they only have one prediction target (an entity or a relation). The latter property allows efficient pre-training and evaluation, as discussed below. For this reason, we exclude tasks such as entity-pair prediction (Wang et al., 2019) (denoted (?, p, ?)) or reconstruction (Nickel et al., 2011) (denoted (?, ?, ?)). In our experimental study in Sec. 4, we found that the exclusion of some of the above pre-training tasks (e.g., LP) often improves downstream task performance. The optimal choice of tasks depends on dataset, KGE model, and downstream task, however. We leave the exploration of task selection as well as on exploring additional pre-training tasks to future work.

heading: 3.2 Multi-Task Ranking
text: To intrinsically evaluate whether KGE models preserve properties that are known to exist in a KG, we use the set of graph-structure prediction tasks described above to generalize the entity ranking (ER) protocol for link prediction (see Sec. 2) to a multi-task ranking (MTR) protocol. Intuitively, for each of the nine queries (LP/REL/DOM/NBE/NBR for both subject and object targets), we construct a query from each test triple, obtain a ranking of the prediction targets that do not already occur in the training/validation/test data (filtered setting), and use metrics such as MRR or Hits@K. The final metric is the micro-average over all nine queries.
We now describe how to obtain task-specific rankings. First, for a REL query of form (s, ?, o), we proceed as in Chang et al. (2020) and rank all p′ ∈ R such that (s, p′, o) /∈ Gtrain in descending order of their scores s(s, p′, o). For the other tasks, which involve wildcards, it is not immediately clear how to perform prediction using a KGE model. We first discuss scoring and ranking, then filtering of data for evaluation. Consider for example the NBR query (s, ?, ∗), where our goal is to rank relations. The perhaps simplest approach to obtain a relation ranking is to first rank all triples of form (s, p′, o′), where p′ ∈ R and o′ ∈ E , and then rank relations by their first appearance (e.g., the relation of the highest-scoring triple is ranked at the top). Generally, we make use of an extended score function that accepts wildcards (described in Algorithm 2 in Appendix A). The approach just described corresponds to using s(s, p′, ∗) = maxo′∈E s(s, p′, o′), i.e, the score of a relation p′ is the score of its most plausible triple. Although other aggregation functions are feasible,
we only consider max-aggregation because it does not make any additional assumptions on the scoring function. To filter training/validation/test data during model evaluation (as done in the literature), we remove all relations p′ such that (s, p′, o′) ∈ Gtrain for some o′ ∈ E ; i.e., we remove all prediction targets that are already implied by the filtering splits. We proceed similarly for all other tasks involving wildcards. Note that the number of score computations needed to predict entity targets for queries without wildcards is O(|E|), whereas the one for queries with wildcards is O(|E||R|). Similarly, predicting target relations costs O(|R||E|) and O(|R|) with and without wildcards, respectively. We discuss in the next section how to reduce the additional cost of using wildcards to O(|E|) or O(|R|).

heading: 3.3 Multi-Task Training
text: We generalize standard KGE model training to all of the graph-structure prediction tasks, called multi-task training (MTT). Our goal is to be able to train KGE models on multiple tasks simultaneously, while keeping training and prediction cost low. We do this by constructing a task-specific cost function for each individual task first; the final cost function is then given as a weighted linear combination of the task-specific costs (and additional regularization terms), where the weights are hyperparameters (costs increase only linearly in the number of tasks, see Table 12). We formalize the MTT training objective in Appendix A.
The task-specific cost functions for link prediction and relation prediction are obtained as in standard training (Sec. 2): for each positive triple (s, p, o) ∈ G, we construct a set of negatives according to the query (i.e., by perturbing the position of the prediction target) and then apply the loss func-
tion (e.g., cross entropy). For our proposed tasks involving wildcards, we proceed differently. Instead of performing some form of (costly) score aggregation during training, we “convert” these tasks with wildcards into tasks without wildcards. To do so, we make use of three virtual wildcard objects—one for subjects (anyS), one for relations (anyR), and one for objects (anyO)—and learn embeddings for these objects. During training, we conceptually replace wildcards by their corresponding wildcard entity and proceed as before. For example, for training triple (s, p, o) and NBR query (s, ?, ∗), we consider the virtual triple (s, p, anyO) along with query (s, ?, anyO). By doing so, we convert the NBR task into a REL task. We also use these wildcard embeddings during inference in the same way; e.g., we set s(s, p′, ∗) = s(s, p′, anyO). Instead of performing score aggregation, the model thus directly learns extended scores at the same cost (per task) as standard link prediction, i.e. O(|E|) for target entities, and O(|R|) for target relations.

heading: 4 Experimental Study
text: 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 (Raffel et al., 2020; Liu et al., 2019). 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 Ruffinelli et al. (2020). We implemented everything in LibKGE (Broscheit et al., 2020), used four benchmark datasets commonly used in recent work (Ge et al., 2023; Xiao et al., 2022; Zhu et al., 2022), 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 Table 1 (LP, REL, DOM, NBE, NBR), and evaluated models on each of these tasks using filtered MRR, and aggregated these metrics into
multi-task ranking MRR (MTR). We selected standard (STD) models with LP and MTT models with both the LP and MTR task, all on validation.
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 (El-Kishky et al., 2022; Ilyas et al., 2022). We chose four popular models: TransE (Bordes et al., 2013), DistMult (Yang et al., 2015), RotatE (Sun et al., 2019) and ComplEx (Trouillon et al., 2016). These are not the latest KGE models, but they are common baselines in recent work (Ge et al., 2023; Gui et al., 2022; Chao et al., 2021), and are common choices for pre-trained models (Zhu et al., 2022; El-Kishky et al., 2022; Ilyas et al., 2022). They can also reach SOTA performance with reasonable embedding sizes (Ruffinelli et al., 2020), allowing us to scale our study, and with larger embedding size (Lacroix et al., 2018) and additional training objectives (Chen et al., 2021b), ComplEx outperforms more involved models, e.g. the transformer-based HittER model (Chen et al., 2021a). 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 (Zhu et al., 2021).
Downstream Tasks Setup. To extrinsically evaluate our pre-trained models, we collected/created data for 35 downstream tasks on FB15K-237, YAGO3-10 and WIKI5M (examples in Table 2). For downstream models, we used scikit-learn (Pedregosa et al., 2011) models that use only entity embeddings from pre-trained KGE models as input features. We used multilayer perceptrons (MLP), logistic regression, KNN, and random forests for classification, and linear regression and MLP for regression, and treated the choice of downstream model as a hyperparameter. For entity classification, we report weighted F1 (as Jain et al. (2021)) aggregated across all classification tasks (denoted EC). For regression, we chose relative squared error (RSE) (defined in Sec. B.2), as it allows meaningful averaging across different regression tasks (denoted REG, lower values are better). We report mean and standard deviation over 3 training runs. As baseline, we included KE-GCN (Yu et al., 2021), a state-of-the-art GNN for entity classification. In contrast to KGEs, this model is directly trained on the downstream task (i.e., no pre-training) and uses the KG for inference. Tuning, training, evaluation was done as with KGEs and downstream models.

heading: 4.1 Results on Graph-Structure Prediction
text: In Table 3, we report test MRR of graph-structure prediction tasks using standard (STD) and MTT training. We report both training approaches with LP for model selection, as we found this to often produce better downstream performance with MTT (such “cross-over” selection was not useful for STD training, see Table 11). We report MTT with MTR for model selection, and results for WNRR and WIKI5M in Appendix C (Tables 13 and 14, resp.)
Every model is able to capture more information about the KG when trained on multiple tasks simultaneously. For a given model, the improvement can be large, often by a factor of 2x and up to 10x depending on model, task and dataset (or even larger when MTT is used with MTR for model selection, see Table 13). This suggests that, unless trained for it, KGE models often fail to capture graph structure beyond what is necessary to perform link prediction. MTT models had slightly lower performance on LP, but the decrease was often small and outweighed by significantly improved performance on other tasks. Moreover, the best models for LP with STD training are often far outperformed on other tasks by other STD models with lower LP performance, suggesting that good LP performance is not indicative of general KG representation. For example, the best LP performance on FB15K-237 is ComplEx STD, but RotatE STD outperforms it considerably on REL and TransE STD on DOM. Similar observations also hold for the best models on MTR, but the compromise on other tasks is significanly smaller. In general, MTT improved significantly on STD for graph structure prediction and can thus be used so models simultaneously learn more properties in a KG.
Note that our performance on LP, even with MTT, is comparable and sometimes better than recently published works that use comparable embedding sizes (Yang et al., 2022; Dong et al., 2022), or even larger embedding sizes (Gui et al., 2022).
Discussion. From a training perspective, these results are not surprising, as STD training only focuses on the LP task. However, these results do challenge studies that describe KGEs as generally capturing semantic properties of a KG (Ge et al., 2023; Xiao et al., 2022; Gui et al., 2022; Nickel et al., 2015; Bordes et al., 2013), which were likely inspired by work on capturing properties of words despite not directly training for it (Mikolov et al., 2013; Bordes et al., 2013). In addition, some of the new tasks are similar enough to link prediction that the results are indeed unexpected. For example, a good link prediction model may be able to answer (Austin, capital of, ?) and (?, capital of, Texas), yet it may not be able to predict that capital of is a relation connected to Austin and/or Texas (NBR). Similar arguments can be made for other tasks. Generally, if the goal is purely link prediction, STD training is more suitable. But we show empirically in Sec. 4.3 that the choice of training objective has an impact on the learned representations and that including the LP task during pre-training is often detrimental for downstream performance.

heading: 4.2 Results on Downstream Tasks
text: Table 3 also reports downstream performance using models pre-trained with STD and MTT. The EC column reports mean weighted F1 across all classification datasets, and REG column reports mean RSE across all regression datasets. We report on individual downstream tasks in Appendix C.
The best overall downstream task performance across all KGE models was achieved by MTT in all cases, and often combined with LP for model selection. While the margin compared to STD was sometimes small (e.g., EC on YAGO3-10) and sometimes large (e.g., REG on FB1K-237), training only for link prediction (STD) resulted in worse average downstream performance compared to MTT more often than not (especially when considering MTT with MTR selection, see Table 13). Nevertheless, for a given KGE model, STD training did perform better at times. In addition, we found that the best models for both LP and MTR are often not the best models in downstream applications. Perhaps more importantly, the best downstream performance often comes from models with weaker LP
performance (e.g. RotatE on EC in FB15K-237) or weaker MTR performance (e.g. ComplEx on REG in FB15K-237). This is more clearly visible in Table 15 in Appendix C. This is problematic, as it suggests that MTR and, perhaps more importantly, LP are often inadequate tasks to guide the choice of the more suitable KGE models for downstream applications. Ultimately, we conclude that the choice of pre-training objective clearly has an impact on downstream performance, but it is unclear how to make this choice in order to maximize downstream performance.
Downstream Baseline Performance. Compared to KE-GCN, KGE models clearly outperform KE-GCN almost every time (except in REG on YAGO3-10) These results suggest that the information captured by KGE models during pre-training is useful for simple downstream models to be competitive with, and even outperform, more involved downstream models that train directly on the task.

heading: 4.3 Impact of Pre-Training Task Selection
text: Table 4 summarizes our results about the impact that pre-training task selection has on downstream tasks. To keep computational costs feasible, we
focused on FB15K-237. We explored performance using MTT without either the LP, REL, DOM, NBE, or NBR pre-training task, and without LP+REL or without DOM+NBR. We report models and sets of tasks relevant for our discussion. For details, see Table 16.
Impact on Graph-Structure Tasks. We found that for graph-structure predictions, excluding a task generally led to lower performance on that task, as expected. It may also, however, lead to a boost in performance on other tasks. For example, RotatE performs best on DOM when the standard LP task is excluded from the training objective.
Impact on Downstream Tasks. For downstream performance, the choice of pre-training tasks has a significant impact, but good choices differ between KGE models and the type of downstream task. For example, compared to full MTT training, using a subset of tasks led to improvements almost every time. Surprisingly, excluding the LP task during pre-training improved downstream performance half of the time compared to STD and full MTT training, suggesting that pretraining with LP can often be detrimental to downstream performance.

heading: 4.4 Data Efficiency Tests
text: To see whether KGE models that capture more information during pre-training are more beneficial as downstream data becomes scarce, we tested models in a few-shot scenario. For classification, we sampled n positive and n negative examples per class, where n ∈ {3, 5, 10}. Figure 1 shows the results for the YAGO3-10 classification tasks (higher is better). We found that as less data becomes available, the average performance of STD models becomes considerably lower compared to pre-trained MTT models, except TransE, where performance difference is not as significant. We observed the same pattern in FB15K-237 (see Fig. 3).
The few-shot scenario applied to regression tasks produced unsatisfactory models almost every time. We thus constructed a different scenario with scarce training data. We randomly sampled n% of the training set, where again n ∈ {3, 5, 10} (see Figures 5 and 7 in Appendix C, lower is better). For most models, the trend observed with a complete training set is mostly maintained, suggesting that pre-trained MTT models are not always more beneficial with less training data. Still, at no point do models pre-trained with STD become a better choice. Overall, although not every time, we observed the clear trend that MTT models are more data efficient than STD models, especially for the classification tasks in our tests.

heading: 5 Conclusion
text: 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.

heading: 6 Limitations
text: In our study, we explored the use of different selfsupervised tasks for training KGE models. However, as a first step, we tested models using only a limited set of simple pre-training tasks. Aside from the link prediction task that is almost exclusively used in the literature, we also included the relation prediction task (as already done by Chen et al. (2021b)), as well a new set of tasks that we proposed (see Table 1). However, other pretraining tasks are possible and should be explored, e.g. self-supervised tasks such as predicting the nhop neighborhood of an entity, or even objectives that resemble downstream tasks, such as predicting the size of a neighborhood. It is also possible to combine such objectives with supervised training objectives during training, as already done in previous work (Aribandi et al., 2022).
Another limitation of our work is the small variety in types of downstream tasks. While we focused on entity-level classification and regression tasks, the impact of different pre-training approaches on more involved downstream applications should be explored. Some examples would be testing the use of pre-trained KGEs in recommender systems as in El-Kishky et al. (2022), or question answering systems as in Ilyas et al. (2022).
Finally, while we take the first steps into exploring alternatives for pre-training KGE models, our work does not find a concrete solution to the problem, which may indeed by challenging, as models need to encode hundreds or thousands of different, and often uncorrelated, relation types between entities. We observed the impact that different pretraining tasks have both on capturing properties of a graph, as well as in downstream application performance. In particular, we found that training with more tasks is beneficial for capturing more properties of a KG, and often for improving downstream performance. However, we have no concrete suggestions on how to pre-train KGE models more generally. Different pre-training tasks should be explored in the context of different types of downstream tasks, so that we may better understand the relation between pre-training KGEs and their quality as KG representations in downstream applications. As part of our work, we provide all of our code as well as our collection of downstream task data, to create opportunities for future research into this unexplored question.

heading: A Multi-Task Training and Evaluation
text: We illustrate how training objectives are constructed using more than one training task, i.e. query. To this end, we define both the standard training objective (STD) based on link prediction and our proposed multi-task objective (MTT) as follows. Let To = {(t, l)} be the set of relevant positive and negative examples t and corresponding label l induced by the link prediction query (s, p, ?) in a given training set. Let Ts be the analogous set of examples for query (?, p, o). For some loss function L, the STD training approach optimizes the following objective function (we omit the penalty term for brevity):
f(θ) = argmin θ   1 |Ts| ∑
(t,l)∈Ts L(s(t), l) +
1 |To| ∑
(t,l)∈To L(s(t), l)
  (1)
where s is a KGE score function parameterized by model parameters θ. We generalize this objective to define the following multi-task training (MTT) objective:
f(θ) = argmin θ
1
N
∑
Ti∈T
∑
(t,l)∈Ti λiL(s(t), l) (2)
where T = {T1, T2, . . .} is a superset of training examples for queries Ti, N is the sum of the cardinalities of each Ti and λi a hyperparameter that controls the impact of query i in the training objective. Chen et al. (2021b) have already followed this training approach by adding the relation prediction task, i.e. (i, ?, j) to Eq. 1. They set λs = λo = 1
and tune λr. Note that Equations 1 and 2 do not describe the exact training objective with some loss functions, e.g. some losses require a positive and corresponding set of negatives to compute a loss value. However, the MTT objective can be reformulated for every loss function commonly used to train KGE models. We provide such a general description of the MTT approach in Algorithm 1. As with loss functions, the MTT approach is agnostic to the choice of model and training task.
W.r.t. to evaluation, Algorithm 2 describes the extended score function described in Section 3.2.

heading: B Experimental Settings
text: B.1 Experimental Setup: Pre-Training KGEs
Knowledge graphs. We chose four commonly used benchmark datasets for evaluating KGE models: FB15K-237 (Toutanova and Chen, 2015), WNRR (Dettmers et al., 2018), YAGO310 (Mahdisoltani et al., 2014), and WIKIDATA5M (WIKI5M) (Wang et al., 2021). Each dataset is associated with a training, a validation and a test split. FB15K-237 and WNRR are designed to be harder benchmarks for link prediction. YAGO3-10 and WIKI5M are considerably larger. Dataset statistics are summarized in Table 5.
KGE training. We used LibKGE (Broscheit et al., 2020) for STD training (LP only) as a baseline and added MTT/MTR model training/evaluation. All KGE models were trained for a maximum of 200 epochs with early stopping on validation MRR checked every 10 epochs. We used cross-entropy as loss function, as it systematically outperformed other losses in most prior studies. We used 1vsAll training with FB15K-237 and WNRR (to achieve good results) and NegSamp with YAGO3-10 and WIKI5M to scale to these larger datasets. Models were selected w.r.t. performance (MRR) on the validation data. We selected STD models with LP task and MTT models with the MTR task. For MTT training, we used all tasks in Table 1.
KGE evaluation. As with training, we evaluate KGE models with respect to each of the five graph-structure prediction tasks in Table 1 (LP, REL, DOM, NBE, NBR) using filtered MRR on test data. We also aggregate these metrics into the multi-task ranking MRR (MTR).
KGE hyperparameters. We closely follow the approach of the experimental study of Ruffinelli et al. (2020) to perform hyperparameter selection.
Algorithm 1: Multi-task Training (MTT) Require: T : set of training triples,
E : set of entities in knowledge graph K θ: model parameters, Q: set of (q, w) pairs of training queries and corresponding weights
Ensure: Updated model parameters θ 1 foreach q, w ∈ Q do 2 N ← construct set of negatives for q using T 3 Tall ← T ∪N 4 sall ← COMPUTE_SCORES(Tall) 5 lq ← w ∗ COMPUTE_LOSS(sall, Tall) // loss weighted by w 6 θ ← UPDATE_PARAMETERS(θ, lq)
We performed 30 random trials using SOBOL sampling (Bergstra and Bengio, 2012) over a large search space to tune several hyperparameters, e.g. regularization, embedding size, batch size, dropout, initialization, and task weights (each in [0.1, 10.0], log scale). To keep our study feasible, we reduced the maximum batch and embedding size for larger datasets and expensive models. The full search space can be found in Table 6.
B.2 Relative Squared Error For evaluating regression performance, we chose relative squared error (RSE), defined as follows:
RSE = ∑N
i=1(yi − ŷi)2∑N i=1(yi − ȳ)2
(3)
where N is the number of evaluation examples, yi are targets to predict, ŷi are model predictions, and ȳ = 1N ∑ n yi, i.e. the mean of targets to predict. We chose RSE because it is interpretable and allows meaningful averaging across the different regression tasks (denoted REG). An RSE value of 1 is equivalent to the performance of a model that
predicts the average of the dependent variable in the evaluation data; lower values are better.
B.3 Experimental Setup: Downstream Tasks
Downstream tasks. We collected or created data for 35 downstream tasks on FB15K-237, YAGO310 or WIKI5M (see Tables 7 and 8). This includes the datasets of Jain et al. (2021) for entity classification on FB15K-237 and YAGO3-10, which aim to predict the types of entities at different granularities. For regression, we use the datasets of Pezeshkpour et al. (2018) for YAGO3-10, which consist of temporal prediction tasks (e.g., the year an event took place), and the dataset of Huang et al. (2021) for node importance prediction. We also created several regression tasks for FB15K-237 from the multimodal data of García-Durán et al. (2018) by predicting literals associated to entities (e.g., a date, a person’s height, the rating of a movie). To create regression tasks for WIKI5M, we followed the same approach using numerical relations extracted from Wikidata (Van Veen, 2019). Datasets statistics are given in Tables 7 and 8.
KGE models. Since we are interested in pretrained KGE models, we used the KGE models trained for the experiments discussed in Sec. 4.1. Thus, no information from downstream tasks was used for KGE model training and selection; i.e. the same KGE model is used for all downstream tasks in each experiment. For model selection, we selected STD models with LP task (the standard approach), but combined MTT models with the LP task or the MTR task. Further improvements may be made by using downstream tasks during training (Aribandi et al., 2022) at the cost, perhaps, of obtaining less general representations; we leave such exploration to future work.
Downstream models. We use scikit-learn (Pedregosa et al., 2011) using only the node embeddings of the pre-trained KG model as input features. For classification, we use multilayer perceptrons (MLP), logistic regression, KNN, and random forests. For regression, we use MLP and linear regression.
Downstream training. Each model was trained using 5-fold cross validation and selected based on mean validation performance across folds (see below). We then retrained the selected model on the union of the training and validation split (if present). To tune hyperparameters, we use 10 trials of random search with SOBOL sampling for each downstream model. The search space for each downstream model is given in Table 9. Note that we treat the choice of downstream model as a hyperparameter as well.
Downstream evaluation. For entity classification, we report weighted F1, as in Jain et al. (2021), aggregated across all classification tasks (denoted EC). For regression, we chose relative squared error (RSE) because it is interpretable and allows meaningful averaging across the different regression tasks (denoted REG). An RSE value of 1 is equivalent to the performance of a model that predicts the average of the dependent variable in the evaluation data; lower values are better. For each
metric, we report the mean and standard deviation over 3 training runs of the downstream model.
Downstream baselines. We include KEGCN (Yu et al., 2021), a recent GNN with stateof-the-art results for graph alignment and entity classification. In contrast to KGEs, this model is directly trained on the downstream task (i.e., no pretraining) and needs to access the KG to perform predictions. For regression tasks, we use a linear layer after the final convolutional layer of KE-GCN, as this led to better performance in our experiments compared to using a single dimensional output in the final convolution layer as done by Huang et al. (2021). We tune hyperparameters using 30 SOBOL trials (as for KGE models); the search space is shown in Table 9. For training, evaluation, and model selection, we follow the approach for our downstream models (e.g,. 5-fold CV).

heading: C Additional Experimental Results
text: Model selection using downstream information. To explore whether results can improve by using downstream information to select models, Table 10 reports performance on FB15K-237 of some KGE models using both training approaches in combination with either LP for model selection (which consistently provided better results for these models with both training approaches) or by selecting directly on the metric used to evaluate the downstream task (weighted F1 for entity classification and RSE for regression). We found that model selection with the downstream task metric provides only marginal benefits for both STD and MTT and can in fact be detrimental, likely due to overfitting on validation data. This indicates that model selection without information about downstream tasks—i.e., using LP or MTR—may be preferrable to using downstream information. This is beneficial, as including downstream information during pre-training or model selection would likely make the resulting representations less general.
Overall, we found that full MTT training with LP for model selection was a suitable choice, but further improvements are possible by dataset-, modeland task-specific choices of pre-training task and validation objective, as discussed in the next section.
Further model selection approaches. For completeness, we also explored the impact of further combinations of model selection methods with both STD and MTT training. To explore whether there
would be improvements in STD models when selecting them based on performance on the MTR task, Table 11 reports downstream performance of some KGE models using STD training combined with either LP or MTR for model selection. We see that the combination of STD with MTR leads to lower downstream performance almost every time.

heading: Future Work
text: , 2022) at the cost, perhaps, of obtaining less general representations; we leave such exploration to future work. Downstream models. We use scikit-learn (Pedregosa et al. , 2011) using only the node embeddings of the pre-trained KG model as input features. For classification, we use multilayer perceptrons (MLP), logistic regression, KNN, and random forests. For regression, we use MLP and linear regression. Downstream training. Each model was trained using 5-fold cross validation and selected based on mean validation performance across folds (see below). We then retrained the selected model on the union of the training and validation split (if present). To tune hyperparameters, we use 10 trials of random search with SOBOL sampling for each downstream model. The search space for each downstream model is given in Table 9. Note that we treat the choice of downstream model as a hyperparameter as well. Downstream evaluation. For entity classification, we report weighted F1, as in Jain et al. (2021), aggregated across all classification tasks (denoted EC). For regression, we chose relative squared error (RSE) because it is interpretable and allows meaningful averaging across the different regression tasks (denoted REG). An RSE value of 1 is equivalent to the performance of a model that predicts the average of the dependent variable in the evaluation data; lower values are better. For each
metric, we report the mean and standard deviation over 3 training runs of the downstream model. Downstream baselines. We include KEGCN (Yu et al. , 2021), a recent GNN with stateof-the-art results for graph alignment and entity classification. In contrast to KGEs, this model is directly trained on the downstream task (i. e. , no pretraining) and needs to access the KG to perform predictions. For regression tasks, we use a linear layer after the final convolutional layer of KE-GCN, as this led to better performance in our experiments compared to using a single dimensional output in the final convolution layer as done by Huang et al. (2021). We tune hyperparameters using 30 SOBOL trials (as for KGE models); the search space is shown in Table 9. For training, evaluation, and model selection, we follow the approach for our downstream models (e. g,. 5-fold CV)