Source-linked AI summary
PyKEEN 1.0: A Python Library for Training and Evaluating Knowledge Graph Embeddings
Mehdi Ali, Max Berrendorf, Charles Tapley Hoyt, Laurent Vermue, Sahand Sharifzadeh, Volker Tresp, Jens Lehmann
TL;DR
Existing KGE software often lacks fully composable models and broad functionality, limiting flexible assessment of model components. PyKEEN 1.0 redesigns the library around configurable KGEM components while adding extensive evaluation, HPO, and automatic memory optimization. The resulting library supports a wide range of interaction models, training approaches, loss functions, datasets, and explicit inverse-relation modeling.
Problem
Existing KGEM implementations often lack arbitrary composability across architectures, training approaches, loss functions, and inverse relations, while some also provide limited models or omit HPO.
Method
PyKEEN 1.0 redesigns KGEM software around flexibly combinable components, unified APIs, Optuna HPO, explicit evaluation rank definitions, and automatic memory optimization.
Results
PyKEEN 1.0 supports 23 interaction models, seven loss functions, four regularizers, two training approaches, six evaluation metrics, and 13 benchmarking datasets.
Takeaways & Limitations
PyKEEN 1.0 provides a configurable and extensible framework for training and evaluating KGEMs across varied model, optimization, and dataset configurations.
Abstract
from arXiv · showhide
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.
1. Introduction
Knowledge graph embeddings represent entities and relations in vector spaces for tasks such as link prediction and recommendation, but existing implementations often lack composability and broad functionality. PyKEEN 1.0 was redesigned to address these limitations through configurable components and additional capabilities.
- Knowledge graph embeddings learn vector representations that preserve graph structure and support tasks including clustering, link prediction, disambiguation, question answering, and recommendation.
- Existing KG embedding implementations often lack the composability needed to combine model architectures, training approaches, loss functions, and inverse relations arbitrarily.This limits assessment of individual component effects on model performance.
- Existing packages also commonly provide limited model coverage or omit functionalities such as hyper-parameter optimization.
- PyKEEN 1.0 was redesigned and re-implemented from scratch as a community effort to make models entirely configurable and extend interaction models and other components.
2. System Description
PyKEEN 1.0 organizes KGEMs into flexibly combinable components and adds broad evaluation, optimization, memory-management, and extensibility features. Its implementation uses standardized interfaces and community-oriented development practices.
- Composable KGEMs: PyKEEN 1.0 flexibly combines interaction models, loss functions, training approaches, and explicit inverse-relation modeling.
- PyKEEN 1.0 supports 23 interaction models, seven loss functions, four regularizers, two training approaches, HPO, six evaluation metrics, and 13 benchmarking datasets.
- Composable KGEMs: Unified APIs separate interaction models, loss functions, and training approaches so modules can be independently replaced or extended.
- Evaluation: PyKEEN 1.0 computes mean rank, mean reciprocal rank, and hits@k under average, optimistic, and pessimistic rank definitions.This explicitly exposes differences caused by alternative rank definitions.
- Hyper-Parameter Optimization: Optuna integration provides grid, random, and tree-parzen HPO, with early stopping and repeated final evaluation to measure robustness.
- Automatic Memory Optimization: Automatic memory optimization computes maximum training and evaluation batch sizes for the current configuration and available hardware before experiments begin.It can determine a maximum sub-batch size when the user-provided training batch is too large.
- Community Standards: The library uses Python 3.7+, PyTorch, automated tests, continuous integration, code-quality checks, and documented development workflows.
3. Comparison to Related Software
Compared with related KGE frameworks, PyKEEN 1.0 emphasizes full KGEM composability alongside extensive functionality and automatic memory optimization, while some alternatives prioritize scalability.
- PyKEEN 1.0 emphasizes full KGEM composability, extensive functionalities, and automatic memory optimization compared with related KGE software packages.The library is described as the only compared package with automatic memory optimization preventing memory overflow during training and evaluation.
- GraphVite, DGL-KE, and PyTorch-BigGraph focus on multi-GPU, CPU, or distributed training, with less emphasis on compositionality and extensibility.
4. Availability and Maintenance
PyKEEN 1.0 is publicly available under the MIT License through GitHub and the Python Package Index, with maintenance supported by a developer team and partner organizations.
- PyKEEN 1.0 is publicly available under the MIT License on GitHub and through the Python Package Index.
- The project will be maintained by a core developer team supported by research groups, companies, and the Technical University of Denmark.