TFLEX: Temporal Feature-Logic Embedding Framework for Complex Reasoning over Temporal Knowledge Graph

</div>

Code for "TFLEX: Temporal Feature-Logic Embedding Framework for Complex Reasoning over Temporal Knowledge Graph" accepted to NeurIPS 2023.

[OpenReview] [arXiv] [Dataset: Google Drive]

Multi-hop logical reasoning over knowledge graph (KG) plays a fundamental role in many artificial intelligence tasks. Recent complex query embedding (CQE) methods for reasoning focus on static KGs, while temporal knowledge graphs (TKGs) have not been fully explored. Reasoning over TKGs has two challenges: 1. The query should answer entities or timestamps; 2. The operators should consider both set logic on entity set and temporal logic on timestamp set. To bridge this gap, we define the multi-hop logical reasoning problem on TKGs. With generated three datasets, we propose the first temporal CQE named Temporal Feature-Logic Embedding framework (TFLEX) to answer the temporal complex queries. We utilize vector logic to compute the logic part of Temporal Feature-Logic embeddings, thus naturally modeling all First-Order Logic (FOL) operations on entity set. In addition, our framework extends vector logic on timestamp set to cope with three extra temporal operators (After, Before and Between). Experiments on numerous query patterns demonstrate the effectiveness of our method.

Below is a typical multi-hop temporal complex query and its computation graph: "During François Hollande was the president of France, which countries did Xi Jinping visit but Barack Obama did not visit?". In the computation graph, there are entity set (blue circle), timestamp set (green triangle), time set projection (green arrow), entity set projection (blue arrow) and logical operators (red rectangle).

🔔 News

May. 5, 2024: Datasets are also held in 🤗 HuggingFace: ICEWS14, ICEWS05_15, GDELT
May. 1, 2024: ICEWS14 dataset is converted to json list for academic exploring.
Oct. 15, 2023: Accepted to NeurIPS 2023! We have released the datasets of TFLEX in Google Drive.

🔬 1. Install

Python (>= 3.7)
PyTorch (>= 1.8.0)
numpy (>= 1.19.2)

pip install -r requirements.txt
cd assistance
pip install -e .
cd ..

🚀 2. Get Started

❗NOTE: Download the datasets in Google Drive (~5G) and place in data folder.

./data
  - ICEWS14
    - cache
      - cache_xxx.pkl
      - cache_xxx.pkl
    - train
    - test
    - valid
  - ICEWS05-15
    - cache
      - cache_xxx.pkl
      - cache_xxx.pkl
    - train
    - test
    - valid
  - GDELT
    - cache
      - cache_xxx.pkl
      - cache_xxx.pkl
    - train
    - test
    - valid

Then run the command to train TFLEX on ICEWS14:

$ python train_TCQE_TFLEX.py --name="TFLEX_dim800_gamma15" --hidden_dim=800 --test_batch_size=32 --every_test_step=10000 --dataset="ICEWS14" --data_home="./data"

$ python train_TCQE_TFLEX.py --help
Usage: train_TCQE_TFLEX.py [OPTIONS]

Options:
  --data_home TEXT                The folder path to dataset.
  --dataset TEXT                  Which dataset to use: ICEWS14, ICEWS05_15,
                                  GDELT.
  --name TEXT                     Name of the experiment.
  --start_step INTEGER            start step.
  --max_steps INTEGER             Number of steps.
  --every_test_step INTEGER       test every k steps
  --every_valid_step INTEGER      validation every k steps.
  --batch_size INTEGER            Batch size.
  --test_batch_size INTEGER       Test batch size. Scoring to all is memory
                                  consuming. We need small test batch size.
  --negative_sample_size INTEGER  negative entities sampled per query
  --train_device TEXT             choice: cuda:0, cuda:1, cpu.
  --test_device TEXT              choice: cuda:0, cuda:1, cpu.
  --resume BOOLEAN                Resume from output directory.
  --resume_by_score FLOAT         Resume by score from output directory.
                                  Resume best if it is 0. Default: 0
  --lr FLOAT                      Learning rate.
  --cpu_num INTEGER               used to speed up torch.dataloader
  --hidden_dim INTEGER            embedding dimension
  --input_dropout FLOAT           Input layer dropout.
  --gamma FLOAT                   margin in the loss
  --center_reg FLOAT              center_reg for ConE, center_reg balances the
                                  in_cone dist and out_cone dist
  --train_tasks TEXT              the tasks for training
  --train_all BOOLEAN             if training all, it will use all tasks in
                                  data.train_queries_answers
  --eval_tasks TEXT               the tasks for evaluation
  --eval_all BOOLEAN              if evaluating all, it will use all tasks in
                                  data.test_queries_answers
  --help                          Show this message and exit.

<details> <summary>👈 🔎 Full commands for reproducing all results in the paper</summary>

# ICEWS14
CUDA_VISIBLE_DEVICES=0 python train_TCQE_TFLEX.py --name="TFLEX_dim800_gamma15" --hidden_dim=800 --test_batch_size=32 --every_test_step=10000 --dataset="ICEWS14"
CUDA_VISIBLE_DEVICES=0 python train_TCQE_X+ConE.py --name="X+ConE_dim800_gamma15" --hidden_dim=800 --test_batch_size=32 --every_test_step=10000 --dataset="ICEWS14"
CUDA_VISIBLE_DEVICES=0 python train_TCQE_X-1F.py --name="X-1F_dim800_gamma15" --hidden_dim=800 --test_batch_size=32 --every_test_step=10000 --dataset="ICEWS14"
CUDA_VISIBLE_DEVICES=0 python train_TCQE_X_without_entity_logic.py --name="X_without_entity_logic_dim800_gamma15" --hidden_dim=800 --test_batch_size=32 --every_test_step=10000 --dataset="ICEWS14"
CUDA_VISIBLE_DEVICES=0 python train_TCQE_X_without_time_logic.py --name="X_without_time_logic_dim800_gamma15" --hidden_dim=800 --test_batch_size=32 --every_test_step=10000 --dataset="ICEWS14"
CUDA_VISIBLE_DEVICES=0 python train_TCQE_X_without_logic.py --name="X_without_logic_dim800_gamma15" --hidden_dim=800 --test_batch_size=32 --every_test_step=10000 --dataset="ICEWS14"
CUDA_VISIBLE_DEVICES=0 python train_TCQE_Query2box.py --name="Query2box_dim800_gamma15" --hidden_dim=800 --test_batch_size=32 --every_test_step=10000 --dataset="ICEWS14"
CUDA_VISIBLE_DEVICES=0 python train_TCQE_BetaE.py --name="BetaE_dim800_gamma15" --hidden_dim=800 --test_batch_size=32 --every_test_step=10000 --dataset="ICEWS14"
CUDA_VISIBLE_DEVICES=0 python train_TCQE_ConE.py --name="ConE_dim800_gamma15" --hidden_dim=800 --test_batch_size=32 --every_test_step=10000 --dataset="ICEWS14"

CUDA_VISIBLE_DEVICES=0 python train_TCQE_Query2box.py --name="Query2box_dim800_gamma15" --hidden_dim=800 --test_batch_size=32 --every_test_step=10000 --dataset="ICEWS14" --resume=True --eval_tasks="Pe,Pe2,Pe3,e2i,e3i"
CUDA_VISIBLE_DEVICES=0 python train_TCQE_BetaE.py --name="BetaE_dim800_gamma15" --hidden_dim=800 --test_batch_size=32 --every_test_step=10000 --dataset="ICEWS14" --resume=True --eval_tasks="Pe,Pe2,Pe3,e2i,e3i,e2i_N,e3i_N,Pe_e2i_Pe_NPe,e2i_PeN,e2i_NPe,e2u,Pe_e2u"
CUDA_VISIBLE_DEVICES=0 python train_TCQE_ConE.py --name="ConE_dim800_gamma15" --hidden_dim=800 --test_batch_size=16 --every_test_step=10000 --dataset="ICEWS14" --resume=True --eval_tasks="Pe,Pe2,Pe3,e2i,e3i,e2i_N,e3i_N,Pe_e2i_Pe_NPe,e2i_PeN,e2i_NPe,e2u,Pe_e2u"

# ICEWS05-15
CUDA_VISIBLE_DEVICES=0 python train_TCQE_TFLEX.py --name="TFLEX_dim800_gamma15" --hidden_dim=800 --test_batch_size=16 --every_test_step=10000 --dataset="ICEWS05_15"
CUDA_VISIBLE_DEVICES=0 python train_TCQE_X+ConE.py --name="X+ConE_dim800_gamma15" --hidden_dim=800 --test_batch_size=16 --every_test_step=10000 --dataset="ICEWS05_15"
CUDA_VISIBLE_DEVICES=0 python train_TCQE_X-1F.py --name="X-1F_dim800_gamma15" --hidden_dim=800 --test_batch_size=16 --every_test_step=10000 --dataset="ICEWS05_15"
CUDA_VISIBLE_DEVICES=0 python train_TCQE_X_without_entity_logic.py --name="X_without_entity_logic_dim800_gamma15" --hidden_dim=800 --test_batch_size=16 --every_test_step=10000 --dataset="ICEWS05_15"
CUDA_VISIBLE_DEVICES=0 python train_TCQE_X_without_time_logic.py --name="X_without_time_logic_dim800_gamma15" --hidden_dim=800 --test_batch_size=16 --every_test_step=10000 --dataset="ICEWS05_15"
CUDA_VISIBLE_DEVICES=0 python train_TCQE_X_without_logic.py --name="X_without_logic_dim800_gamma15" --hidden_dim=800 --test_batch_size=16 --every_test_step=10000 --dataset="ICEWS05_15"
CUDA_VISIBLE_DEVICES=0 python train_TCQE_Query2box.py --name="Query2box_dim800_gamma15" --hidden_dim=800 --test_batch_size=16 --every_test_step=10000 --dataset="ICEWS05_15"
CUDA_VISIBLE_DEVICES=0 python train_TCQE_BetaE.py --name="BetaE_dim800_gamma15" --hidden_dim=800 --test_batch_size=16 --every_test_step=10000 --dataset="ICEWS05_15"
CUDA_VISIBLE_DEVICES=0 python train_TCQE_ConE.py --name="ConE_dim800_gamma15" --hidden

TFLEX

Install / Use

README

TFLEX: Temporal Feature-Logic Embedding Framework for Complex Reasoning over Temporal Knowledge Graph

🔔 News

🌍 Contents

🔬 1. Install

🚀 2. Get Started