NeurIPS 2024

A Prompt-Based Knowledge Graph Foundation Model for Universal In-Context Reasoning

Yuanning Cui, Zequn Sun, Wei Hu

Main figure for A Prompt-Based Knowledge Graph Foundation Model for Universal In-Context Reasoning

Abstract

Extensive knowledge graphs (KGs) have been constructed to facilitate knowledge-driven tasks across various scenarios. However, existing work usually develops separate reasoning models for different KGs, lacking the ability to generalize and transfer knowledge across diverse KGs and reasoning settings. In this paper, we propose a prompt-based KG foundation model via in-context learning, namely KG-ICL, to achieve a universal reasoning ability. Specifically, we introduce a prompt graph centered with a query-related example fact as context to understand the query relation. To encode prompt graphs with the generalization ability to unseen entities and relations in queries, we first propose a unified tokenizer that maps entities and relations in prompt graphs to predefined tokens. Then, we propose two message passing neural networks to perform prompt encoding and KG reasoning, respectively. We conduct evaluation on 43 different KGs in both transductive and inductive settings. Results indicate that the proposed KG-ICL outperforms baselines on most datasets, showcasing its outstanding generalization and universal reasoning capabilities. The source code is accessible on GitHub:.

Benchmark Results

Tested on 2 datasets

Tables show presented results in global context (best 3 papers + GCN + MLP).

Knowledge Graphs

#ModelMRR ↑Paper
1
FedSage
0.8426 ±0.0026
2
FED-PUB
0.8348 ±0.0025
3
FedAvg
0.8252 ±0.0023
7
GCN
0.8171 ±0.0030
34
KG-ICL finetune
0.5360
This paper
58
ULTRA pre-train
0.4800
This paper
59
ULTRA finetune
0.4800
This paper
80
KG-ICL pre-train
0.4550
This paper

FB15k-237

Global leaderboard →

Default metric is MRR. The paper reports 5 values on MRR and 5 values on H@10 .

#ModelMRR ↑ defaultH@10 ↑Paper
1
sum, Msg_r^3, f(t)=0
0.8500
—
2
sum, Msg_r^3, f(t)=t
0.8440
—
3
PNA, Msg_r^1, f(t)=0
0.8010
—
69
KG-ICL pre-train
0.3290
0.4330
H@10 0.4330
This paper
73
ULTRA finetune
0.3250
0.4110
H@10 0.4110
This paper
77
KG-ICL finetune
0.3240
0.4160
H@10 0.4160
This paper
76
ULTRA pre-train
0.3240
0.3980
H@10 0.3980
This paper
81
Supervised SOTA
0.2930
0.3370
H@10 0.3370
This paper
SOTA performance (not significantly different from the best) Proposed method (this paper) Baseline / comparison (this paper) Global context (other papers)