Source code for the paper Language Models as Knowledge Embeddings

[June 2023] We recently identified a data leakage issue in our code that, during prediction, we inadvertently leaked degree information about the entities to be predicted. This unintentionally provided a shortcut for the model, which affected the experimental results to some extent. We have fixed this issue and re-conducted our experiments, and updated the paper accordingly. The revised results do not impact the majority of the paper's conclusions and contributions. The method continues to achieve state-of-the-art (SOTA) performance on the WN18RR, FB13, and WN11 datasets, compared to previous works. However, on the FB15k-237 dataset, the model's performance has declined to a certain extent and underperforms state-of-the-art structured-based methods. We sincerely apologize for this error.

Updated Results:

• PyTorch version >= 1.7.1
• NumPy version >= 1.19.5
• transformers
• tqdm
• Python version >= 3.6

Run main.py to train or test our models.

An example for training for triple classification:

python main.py --batch_size 16 --plm bert --data wn18rr --task TC

An example for training for link prediction:

python main.py --batch_size 16 --plm bert --contrastive --self_adversarial --data wn18rr --task LP 

The arguments are as following:

• --bert_lr: learning rate of the language model.
• --model_lr: learning rate of other parameters.
• --batch_size: batch size used in training.
• --weight_decay: weight dacay used in training.
• --data: name of the dataset. Choose from 'fb15k-237', 'wn18rr', 'fb13' and 'umls'.
• --plm: choice of the language model. Choose from 'bert' and 'bert_tiny'.
• --load_path: path of checkpoint to load.
• --load_epoch: load the checkpoint of a specific epoch. Use with --load_metric.
• --load_metric: use with --load_epoch.
• --link_prediction: run link prediction evaluation after loading a checkpoint.
• --triple_classification: run triple classification evaluation after loading a checkpoint.
• --self_adversarial: use self-adversarial negative sampling for efficient KE learning.
• --contrastive: use contrastive LMKE.
• --task: specify the task. Choose from 'LP' (link prediction) and 'TC' (triple classification).

The datasets are put in the folder 'data', including fb15k-237, WN18RR, FB13 and umls.