Flock (PyTorch)

Flock: A Knowledge Graph Foundation Model via Learning on Random Walks
Jinwoo Kim*, Xingyue Huang*, Krzysztof Olejniczak, Kyungbin Min, Michael Bronstein, Seunghoon Hong, İsmail İlkan Ceylan (* equal contribution)
arXiv 2025

This codebase contains training and evaluation scripts for Flock, including pretraining scripts and checkpoints on FB15k-237, WN18RR, and CoDEx Medium, zero-shot and finetuning evaluation over 54 knowledge graphs, and our new synthetic dataset Petals.

The codebase has been tested with NVIDIA H100 GPUs.

Setup

We recommend using the official PyTorch Docker image with CUDA support.

docker pull pytorch/pytorch:2.8.0-cuda12.6-cudnn9-devel
docker run -it --gpus all --ipc host --name flock -v /home:/home pytorch/pytorch:2.8.0-cuda12.6-cudnn9-devel bash

Assuming the codebase is located at ~/flock inside Docker container, install the required packages:

cd ~/flock
pip3 install torch_geometric easydict pybind11
bash install_walker.sh

Pretraining and Evaluation

Entity prediction pretraining

cd ~/flock
python3 src_entity/pretrain.py --config src_entity/config/pretrain_3g.yaml --gpus [0]

Entity prediction evaluation (using the checkpoint at checkpoints/flock_entity.pth)

cd ~/flock
mkdir -p ~/flock/output-entity

# zero-shot
CKPT=~/flock/checkpoints/flock_entity.pth bash scripts/entity_zeroshot.sh

# finetuning
CKPT=~/flock/checkpoints/flock_entity.pth bash scripts/entity_finetune.sh

Relation prediction pretraining

cd ~/flock
python3 src_relation/pretrain.py --config src_relation/config/pretrain_3g.yaml --gpus [0]

Relation prediction evaluation (using the checkpoint at checkpoints/flock_relation.pth)

cd ~/flock
mkdir -p ~/flock/output-relation

# zero-shot
CKPT=~/flock/checkpoints/flock_relation.pth bash scripts/relation_zeroshot.sh

# finetuning
CKPT=~/flock/checkpoints/flock_relation.pth bash scripts/relation_finetune.sh

Synthetic Dataset Generation

Execute the following commands to generate the Petals synthetic dataset at data/petals/:

python3 generate_petals.py

References

Our implementation is based on the code from the following repositories:

• RWNN and graph-walker for random walks and their anonymizations
• MOTIF and ULTRA for entity prediction experiments
• TRIX for relation prediction experiments

Citation

If you find our work useful, please consider citing it:

@article{kim2025flock,
  author    = {Jinwoo Kim and Xingyue Huang and Krzysztof Olejniczak and Kyungbin Min and Michael Bronstein and Seunghoon Hong and İsmail İlkan Ceylan},
  title     = {Flock: A Knowledge Graph Foundation Model via Learning on Random Walks},
  journal   = {arXiv},
  volume    = {abs/2510.01510},
  year      = {2025},
  url       = {https://arxiv.org/abs/2510.01510}
}