LibKGE is a PyTorch-based library for efficient training, evaluation, and hyperparameter optimization of knowledge graph embeddings (KGE). It is highly configurable, easy to use, and extensible. Other KGE frameworks are listed below.

The key goal of LibKGE is to foster reproducible research into (as well as meaningful comparisons between) KGE models and training methods. As we argue in our ICLR 2020 paper (see video), the choice of training strategy and hyperparameters are very influential on model performance, often more so than the model class itself. LibKGE aims to provide clean implementations of training, hyperparameter optimization, and evaluation strategies that can be used with any model. Every potential knob or heuristic implemented in the framework is exposed explicitly via well-documented configuration files (e.g., see here and here). LibKGE also provides the most common KGE models and new ones can be easily added (contributions welcome!).

For link prediction tasks, rule-based systems such as AnyBURL are a competitive alternative to KGE.

# retrieve and install project in development mode
git clone https://github.com/uma-pi1/kge.git
cd kge
pip install -e .

# download and preprocess datasets
cd data
sh download_all.sh
cd ..

# train an example model on toy dataset (you can omit '--job.device cpu' when you have a gpu)
kge start examples/toy-complex-train.yaml --job.device cpu

1. Features
2. Results and pretrained models
3. Using LibKGE
4. Currently supported KGE models
5. Extending LibKGE
6. Known issues
7. Changelog
8. Other KGE frameworks
9. How to cite

• Training
  • Training types: negative sampling, 1vsAll, KvsAll
  • Losses: binary cross entropy (BCE), Kullback-Leibler divergence (KL), margin ranking (MR), squared error (SE)
  • All optimizers and learning rate schedulers of PyTorch supported and can be chosen individually for different parameters (e.g., different for entity and for relation embeddings)
  • Learning rate warmup
  • Early stopping
  • Checkpointing
  • Stop (e.g., via Ctrl-C) and resume at any time
  • Automatic memory management to support large batch sizes (see config key train.subbatch_auto_tune)
• Hyperparameter tuning
  • Grid search, manual search, quasi-random search (using Ax), Bayesian optimization (using Ax)
  • Highly parallelizable (multiple CPUs/GPUs on single machine)
  • Stop and resume at any time
• Evaluation
  • Entity ranking metrics: Mean Reciprocal Rank (MRR), HITS@k with/without filtering
  • Drill-down by: relation type, relation frequency, head or tail
• Extensive logging and tracing
  • Detailed progress information about training, hyper-parameter tuning, and evaluation is recorded in machine readable formats
  • Quick export of all/selected parts of the traced data into CSV or YAML files to facilitate analysis
• KGE models
  • All models can be used with or without reciprocal relations
  • RESCAL (code, config)
  • TransE (code, config)
  • TransH (code, config)
  • DistMult (code, config)
  • ComplEx (code, config)
  • ConvE (code, config)
  • RelationalTucker3/TuckER (code, config)
  • CP (code, config)
  • SimplE (code, config)
  • RotatE (code, config)
  • Transformer ("No context" model) (code, config)
• Embedders
  • Lookup embedder (code, config)
  • Projection embedder (code, config)

We list some example results (filtered MRR and HITS@k on test data) obtained with LibKGE below. These results are obtained by running automatic hyperparameter search as described here.

These results are not necessarily the best results that can be achieved using LibKGE, but they are comparable in that a common experimental setup (and equal amount of work) has been used for hyperparameter optimization for each model. Since we use filtered MRR for model selection, our results may not be indicative of the achievable model performance for other validation metrics (such as HITS@10, which has been used for model selection elsewhere).

We report performance numbers on the entire test set, including the triples that contain entities not seen during training. This is not done consistently throughout existing KGE implementations: some frameworks remove unseen entities from the test set, which leads to a perceived increase in performance (e.g., roughly add +3pp to our WN18RR MRR numbers for this method of evaluation).

We also provide pretrained models for these results. Each pretrained model is given in the form of a LibKGE checkpoint, which contains the model as well as additional information (such as the configuration being used). See the documentation below on how to use checkpoints.

LibKGE supports large datasets such as Wikidata5M (4.8M entities). The result given below was found by automatic hyperparameter search similar to the one used for the smaller datasets above, but with some values fixed (training with shared negative sampling, embedding dimension: 128, batch size: 1024, optimizer: Adagrad, regularization: weighted). We ran 30 pseudo-random configurations for 20 epochs, and then reran the configuration that performed best on validation data for 200 epochs.

The result given below was found by the same automatic hyperparameter search used for Wikidata5M. We reran the configuration that performed best on validation data for 400 epochs.

CoDEx is a Wikidata-based KG completion benchmark. The results here have been obtained using the automatic hyperparameter search used for the Freebase and WordNet datasets, but with fewer epochs and Ax trials for CoDEx-M and CoDEx-L. See the CoDEx paper (EMNLP 2020) for details.