/MTL-KGC

code for the paper 'Multi-Task Learning for Knowledge Graph Completion with Pre-trained Language Models' (COLING20).

Primary LanguagePythonApache License 2.0Apache-2.0

MTL-KGC

This is PyTorch implementation of the Multi-Task Learning for Knowledge Graph Completion with Pre-trained Language Models.

Train

Train multitask learning with link prediction (LP), relation prediction (RP), and relevance ranking (RR). if you get AssertionError: Default process group is not initialized, please try with python -m torch.distributed.launch

python run_bert_multitask.py \
    --do_train \
    --task_list lp,rp,rr \
    --data_dir ./data/wn18rr \
    --bert_model bert-base-cased \
    --max_seq_length 128 \
    --output_dir ./output_dir \
    --num_train_epochs 5.0 \
    --learning_rate=2e-5 \
    --tb_log_dir=runs/log_dir

Train link prediction only

python run_bert_multitask.py \
    --do_train \
    --task_list lp \
    --data_dir ./data/wn18rr \
    --bert_model bert-base-cased \
    --max_seq_length 128 \
    --output_dir ./output_dir \
    --num_train_epochs 5.0 \
    --learning_rate=2e-5 \
    --tb_log_dir=runs/log_dir

Evaulation

python run_bert_multitask.py \
    --do_predict \
    --data_dir ./data/wn18rr \
    --bert_model {test_model} \
    --max_seq_length 128 \
    --output_dir ./output_dir \

Performance

WN18RR

MRR MR HITS@1 HITS@3 HITS@10
LP (Yao et al, 2019) 0.219 108 0.095 0.243 0.497
LP + RP 0.302 112 0.177 0.353 0.560
LP + RR 0.277 97 0.130 0.341 0.576
LP + RP + RR 0.331 89 0.203 0.383 0.597

FB15k-237

MRR MR HITS@1 HITS@3 HITS@10
LP (Yao et al, 2019) 0.237 145 0.144 0.260 0.427
LP + RP 0.262 138 0.169 0.289 0.447
LP + RR 0.247 143 0.154 0.272 0.434
LP + RP + RR 0.267 132 0.172 0.298 0.458