/TuckER

TuckER: Tensor Factorization for Knowledge Graph Completion

Primary LanguagePythonMIT LicenseMIT

TuckER: Tensor Factorization for Knowledge Graph Completion

This codebase contains PyTorch implementation of the paper:

TuckER: Tensor Factorization for Knowledge Graph Completion. Ivana Balažević, Carl Allen, and Timothy M. Hospedales. Empirical Methods in Natural Language Processing (EMNLP), 2019. [Paper]

TuckER: Tensor Factorization for Knowledge Graph Completion. Ivana Balažević, Carl Allen, and Timothy M. Hospedales. ICML Adaptive & Multitask Learning Workshop, 2019. [Short Paper]

Link Prediction Results

Dataset MRR Hits@10 Hits@3 Hits@1
FB15k 0.795 0.892 0.833 0.741
WN18 0.953 0.958 0.955 0.949
FB15k-237 0.358 0.544 0.394 0.266
WN18RR 0.470 0.526 0.482 0.443

Running a model

To run the model, execute the following command:

 CUDA_VISIBLE_DEVICES=0 python main.py --dataset FB15k-237 --num_iterations 500 --batch_size 128
                                       --lr 0.0005 --dr 1.0 --edim 200 --rdim 200 --input_dropout 0.3 
                                       --hidden_dropout1 0.4 --hidden_dropout2 0.5 --label_smoothing 0.1

Available datasets are:

FB15k-237
WN18RR
FB15k
WN18

To reproduce the results from the paper, use the following combinations of hyperparameters with batch_size=128:

dataset lr dr edim rdim input_d hidden_d1 hidden_d2 label_smoothing
FB15k 0.003 0.99 200 200 0.2 0.2 0.3 0.
WN18 0.005 0.995 200 30 0.2 0.1 0.2 0.1
FB15k-237 0.0005 1.0 200 200 0.3 0.4 0.5 0.1
WN18RR 0.003 1.0 200 30 0.2 0.2 0.3 0.1

Requirements

The codebase is implemented in Python 3.6.6. Required packages are:

numpy      1.15.1
pytorch    1.0.1

Citation

If you found this codebase useful, please cite:

@inproceedings{balazevic2019tucker,
title={TuckER: Tensor Factorization for Knowledge Graph Completion},
author={Bala\v{z}evi\'c, Ivana and Allen, Carl and Hospedales, Timothy M},
booktitle={Empirical Methods in Natural Language Processing},
year={2019}
}