PyKEEN 1.0: A Python Library for Training and Evaluating Knowledge Graph Embeddings

TL;DR

PyKEEN 1.0 is a comprehensive Python library that facilitates the training, evaluation, and hyper-parameter optimization of knowledge graph embedding models, supporting diverse model configurations and efficient resource utilization.

Contribution

It re-implements and extends the original PyKEEN library with new features like inverse relation modeling and advanced hyper-parameter optimization, enhancing flexibility and performance.

Findings

Supports a wide range of KGEMs and training approaches

Includes automatic memory optimization for hardware efficiency

Provides extensive hyper-parameter tuning capabilities

Abstract

Recently, knowledge graph embeddings (KGEs) received significant attention, and several software libraries have been developed for training and evaluating KGEs. While each of them addresses specific needs, we re-designed and re-implemented PyKEEN, one of the first KGE libraries, in a community effort. PyKEEN 1.0 enables users to compose knowledge graph embedding models (KGEMs) based on a wide range of interaction models, training approaches, loss functions, and permits the explicit modeling of inverse relations. Besides, an automatic memory optimization has been realized in order to exploit the provided hardware optimally, and through the integration of Optuna extensive hyper-parameter optimization (HPO) functionalities are provided.

Tables

Table 1. Table 1: An overview of the functionalities of PyKEEN 1.0 and similar libraries. ES refers to early stopping, TA to training approach, Inv. Rels. to the explicit modeling of inverse relations, AMO to automatic memory optimization, MGS to multi-GPU support, and DTR to distributed training.

Library	Models	HPO	ES	Evaluation Metrics	Set TA	Inv. Rels.	Set Loss Fct.	AMO	MGS	DTR
AmpliGraph (Costabello et al., 2019)	6	GS	✓	MR, MRR, Hits@k	-	✓	✓	-	-	-
DGL-KE (Zheng et al., 2020)	6	-	-	MR, MRR Hits@k	-	-	✓	-	✓	✓
GraphVite (Zhu et al., 2019)	6	-	-	MR, MRR, Hits@k, AUC-ROC	-	-	-	-	✓	-
LibKGE (Ruffinelli et al., 2020)	10	GS, RS, TPE	✓	MR, MRR, Hits@k	✓	✓	✓	-	-	-
OpenKE (Han et al., 2018)	11	-	-	MR, MRR, Hits@k	-	-	✓	-	-	-
PyTorch-BigGraph (Lerer et al., 2019)	4	-	-	MR, MRR, Hits@k, AUC-ROC	-	-	✓	-	✓	✓
Pykg2vec (Yu et al., 2019)	18	TPE	-	MR, Hits@k	-	-	-	-	-	-
PyKEEN (Ali et al., 2019b)	10	RS	-	MR, Hits@k	-	-	-	-	-	-
PyKEEN 1.0	23	GS, RS, TPE	\textpdfrender TextRenderingMode=FillStroke, LineWidth=.5pt, ✓	MR, MRR AMR, Hits@k, AUC-PR, AUC-ROC	\textpdfrender TextRenderingMode=FillStroke, LineWidth=.5pt, ✓	\textpdfrender TextRenderingMode=FillStroke, LineWidth=.5pt, ✓	\textpdfrender TextRenderingMode=FillStroke, LineWidth=.5pt, ✓	\textpdfrender TextRenderingMode=FillStroke, LineWidth=.5pt, ✓	-	-

Full text

Keywords: Knowledge Graphs, Knowledge Graph Embeddings, Relational Learning

1 Introduction

Knowledge graphs (KGs) encode knowledge as a set of triples $\mathcal{K}\subseteq\mathcal{E}\times\mathcal{R}\times\mathcal{E}$ where $\mathcal{E}$ denotes the set of entities and $\mathcal{R}$ the set of relations. Knowledge graph embedding models (KGEMs) learn representations for entities and relations of KGs in vector spaces while preserving the graph’s structure. The learned embeddings can support machine learning tasks such as entity clustering, link prediction, entity disambiguation as well as downstream tasks such as question answering and recommendation (Nickel et al., 2015; Wang et al., 2017).

Most publications of KGEMs are accompanied by reference implementations, but they often lack the finesse required for general usability. Existing software packages that provide implementations for different KGEMs usually lack entire composability: model architectures (or interaction models), training approaches, loss functions, and the usage of explicit inverse relation cannot arbitrarily be combined. The full composability of KGEMs is fundamental for assessing the performance of KGEMs because it allows assessing single components individually on the model’s performance instead of attributing a performance increase solely to the model architecture, which is misleading (Ruffinelli et al., 2020; Ali et al., 2020). Besides, often only limited functionalities are provided, e.g., a small number of KGEMs are supported, or functionalities such as HPO are missing. For instance, in PyKEEN (Ali et al., 2019a, b) one of the first software packages for KGEMs, models can only be trained under the stochastic local closed-world assumption, the evaluation procedure was too slow for larger KGs, and it was designed to be mainly used through a command-line interface rather than programmatically in order to facilitate its usage for non-experts. This motivated the development of a reusable software package comprising several KGEMs and related methodologies that is entirely configurable.

Here, we present PyKEEN (Python KnowlEdge EmbeddiNgs) 1.0, a community effort in which PyKEEN has been re-designed and re-implemented from scratch to overcome the mentioned limitations, make models entirely configurable, and to extend it with more interaction models and other components.

2 System Description

In PyKEEN 1.0, a KGEM is considered as a composition of four components that can flexibly be combined: an interaction model (or model architecture), a loss function, a training approach, and the usage of inverse relations. PyKEEN 1.0 currently supports 23 interaction models, seven loss functions, four regularizers, two training approaches, HPO, six evaluation metrics, and 13 built-in benchmarking datasets. It can readily import additional datasets that have been pre-stratified into train/test/evaluation and generate appropriate splits for unstratified datasets. Additionally, we implemented an automatic memory optimization that ensures that the available memory is best utilized.

Composable KGEMs

To ensure the composability of KGEMs, the interaction models, loss functions, and training approaches are separated from each other and implemented as independent submodules, whereas the modeling of inverse relations is handled by the interaction models. Our modules can be arbitrarily replaced because we ensured through inheritance that all interaction models, loss functions, and training approaches follow unified APIs, which are defined by pykeen.model.Model, pykeen.loss.Loss, and pykeen.training.TrainingLoop. Currently, we provide implementations of 23 interaction models, the most common loss functions used for training KGEMs including the binary-cross entropy, cross entropy, mean square error, negative-sampling self-adversarial loss, and the softplus loss, as well as the local closed-world assumption and the stochastic local closed-world assumption training approach (Nickel et al., 2015). It is known that some interaction models (e.g., ConvE) benefit from being explicitly trained with inverse relations, i.e., for each relation $r\in\mathcal{R}$ an inverse relation $r_{inv}$ is introduced, and the task of predicting the head entity of a $(r,t)$-pair becomes the task of predicting the tail entity of the corresponding inverse pair $(t,r_{inv})$. Therefore, in PyKEEN 1.0, we enable users to train each interaction model with explicit inverse relations.

Evaluation

KGEMs are usually evaluated on the task of link prediction. Given $(h,r)$ (or $(r,t)$), all possible entities $\mathcal{E}$ are considered as tail (or head) and ranked according to the KGEMs interaction model. The individual ranks are commonly aggregated to mean rank, mean reciprocal rank, and hits@k. However, these metrics have been realized differently throughout the literature based on different definitions of the rank (Berrendorf et al., 2020), leading to difficulties in reproducibility and comparability (Akrami et al., 2018). The three most common rank definitions are the average rank, optimistic rank, and pessimistic rank. In PyKEEN 1.0, we explicitly compute the aggregation metrics for all common rank definitions, average, optimistic, and pessimistic, allowing inspection of differences between them. This can help to reveal cases where the model predicts exactly equal scores for many different triples, which is usually an undesired behavior. In addition, we support the recently proposed adjusted mean rank (Berrendorf et al., 2020), which allows comparing results across differently sized datasets, as well as offering an interface to use all metrics implemented in scikit-learn (Pedregosa et al., 2011), including AUC-PR and AUC-ROC.

Hyper-Parameter Optimization

We integrated Optuna (Akiba et al., 2019) as the hyper-parameter optimization (HPO) framework to enable PyKEEN 1.0 to take advantage of its wide range of HPO functionalities (i.e., grid search, random search, tree-parzen estimator). To optimize the hyper-parameters on the validation set, we implemented early stopping. Besides, we implemented an HPO workflow that enables users to effectively find an appropriate set of hyper-parameters on the validation set, and train and evaluate the final model on the test set $n$ times to measure the robustness of the model.

Automatic Memory Optimization

Allowing high computational throughput while ensuring that the available hardware memory is not exceeded during training and evaluation requires the knowledge of the maximum possible training and evaluation batch size for the current model configuration. However, determining the training and evaluation batch sizes is a tedious process, and not feasible when a large set of heterogeneous experiments are run. Therefore, we implemented an automatic memory optimization step that computes the maximum possible training and evaluation batch sizes for the current model configuration and available hardware before the actual experiment starts. If the user-provided batch size is too large for the used hardware, the automatic memory optimization determines the maximum sub-batch size for the training.

Extensibility

Because we defined a uniform API for each interaction model, any new model can be integrated by following the API of the existing models (pykeen.models). Similarly, the remaining components, e.g., regularizers, and negative samplers follow a unified API, so that new modules can be smoothly integrated.

Community Standards

PyKEEN 1.0 relies on several community-oriented tools to ensure it is accessible, reusable, reproducible, and maintainable. It is implemented for Python 3.7+ using the PyTorch package. It comes with a suite of thorough unit tests that are automated with PyTest, Tox, run in a continuous integration setting on Travis-CI, and are tracked over time using codecov.io. Code quality is ensured with flake8 and careful application of the GitHub Flow development workflow. Documentation is quality checked by doc8, built with Sphinx, and hosted on ReadTheDocs.org.

3 Comparison to Related Software

Table 1 depicts the most popular KGE frameworks and their features. It shows that PyKEEN 1.0 compared to related software packages emphasize on both, full composability of KGEMs and extensive functionalities, i.e., a large number of supported interaction models, and extensive evaluation and HPO functionalities. Finally, PyKEEN 1.0 is the only library that performs an automatic memory optimization that ensures that the memory is not exceeded during training and evaluation. GraphVite, DGL-KE, and PyTorch-BibGraph focus on scalability, i.e., they provide support for multi-GPU/CPU or/and distributed training, but focus less on compositionality and extensibility. For instance, PyTorch-BigGraph supports only a small number of interaction models that follow specific computation blocks.

4 Availability and Maintenance

PyKEEN 1.0 is publicly available under the MIT License at https://github.com/pykeen/pykeen, and is distributed through the Python Package Index. It will be maintained by the core developer team that is supported by the Smart Data Analytics research group (University of Bonn), Fraunhofer IAIS, Enveda Therapeutics, Munich Center for Machine Learning (MCML), Siemens, and the Technical University of Denmark (section for Cognitive Systems and section for Statistics and Data Analysis). The project is funded by the German Federal Ministry of Education and Research (BMBF) under Grant No. 01IS18036A and Grant No. 01IS18050D (project MLWin) as well as the Innovation Fund Denmark with the Danish Center for Big Data Analytics driven Innovation (DABAI) which ensures the maintenance of the project in the next years.