Making Large Language Models Perform Better in Knowledge Graph Completion

• Making Large Language Models Perform Better in Knowledge Graph Completion

Large language model (LLM) based knowledge graph completion (KGC) aims to predict the missing triples in the KGs with LLMs and enrich the KGs to become better web infrastructure, which can benefit a lot of web-based automatic services. However, research about LLM-based KGC is limited and lacks effective utilization of LLM's inference capabilities, which ignores the important structural information in KGs and prevents LLMs from acquiring accurate factual knowledge. In this paper, we discuss how to incorporate the helpful KG structural information into the LLMs, aiming to achieve structrual-aware reasoning in the LLMs. We first transfer the existing LLM paradigms to structural-aware settings and further propose a knowledge prefix adapter (KoPA) to fulfill this stated goal. KoPA employs structural embedding pre-training to capture the structural information of entities and relations in the KG. Then KoPA informs the LLMs of the knowledge prefix adapter which projects the structural embeddings into the textual space and obtains virtual knowledge tokens as a prefix of the input prompt. We conduct comprehensive experiments on these structural-aware LLM-based KGC methods and provide an in-depth analysis comparing how the introduction of structural information would be better for LLM's knowledge reasoning ability.

🔔 News

• 2024-02 We preprint our Survey Knowledge Graphs Meet Multi-Modal Learning: A Comprehensive Survey [Repo].

🌈 Model Architecture

🔬 Dependencies

Our code is developed based on alpaca-lora. Please build the Python following the instruction in Alpaca-lora.

Some core python library config:

• Python 3.9.16

• torch 2.0.0

• transformers 4.28.0

• peft 0.3.0

• If you find the performance of fine-tuned model is close to the zero-shot results. You can check whether your LoRA checkpoint is empty. A version mismatch of the peft / transformers library may cause such an issue. See more details in this issue from the perf reposity huggingface/peft#286.

🌲 Data Preparation

Due to the size of the data, you need to download and unzip the data file data.zip from this link and put them in the data/.

📕 Training & Test

• Note: The current dataset is just a little demonstration to help you run the full pipeline. We will release the full datasets in the future.

• run KoPA tuning

# For CoDeX-S dataset
export WANDB_DISABLED=true
wandb offline
CUDA_VISIBLE_DEVICES=0 nohup python finetune_kopa.py \
    --base_model 'YOUR LLM PATH' \
    --data_path 'data/CoDeX-S-train.json' \
    --output_dir 'YOUR SAVE PATH' \
    --num_epochs 3 \
    --lora_r 64 \
    --learning_rate 3e-4 \
    --batch_size 12 \
    --micro_batch_size 12 \
    --num_prefix 1 \
    --kge_model 'data/CoDeX-S-rotate.pth' \
    --lora_target_modules='[q_proj,k_proj,v_proj,o_proj]' > log.txt &


# For FB15K-237N dataset
export WANDB_DISABLED=true
wandb offline
CUDA_VISIBLE_DEVICES=0 nohup python finetune_kopa.py \
    --base_model 'YOUR LLM PATH' \
    --data_path 'data/FB15K-237N-train.json' \
    --output_dir 'YOUR SAVE PATH' \
    --num_epochs 3 \
    --lora_r 64 \
    --learning_rate 3e-4 \
    --batch_size 12 \
    --micro_batch_size 12 \
    --num_prefix 1 \
    --kge_model 'data/FB15K-237N-rotate.pth' \
    --lora_target_modules='[q_proj,k_proj,v_proj,o_proj]' > log.txt &

You may need to fill the LLM path and save path before running. The hyper-parameters can be tuned by yourself.

• run inference

CUDA_VISIBLE_DEVICES=0 python inference_kopa.py

🤝 Cite:

Please condiser citing this paper if you use the code from our work. Thanks a lot :)

@article{DBLP:journals/corr/abs-2310-06671,
  author       = {Yichi Zhang and
                  Zhuo Chen and
                  Wen Zhang and
                  Huajun Chen},
  title        = {Making Large Language Models Perform Better in Knowledge Graph Completion},
  journal      = {CoRR},
  volume       = {abs/2310.06671},
  year         = {2023}
}
}