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

TL;DR

This paper proposes a prompt-based KG foundation model via in-context learning, namely KG-ICL, to achieve a universal reasoning ability and indicates that the proposed KG-ICL outperforms baselines on most datasets, showcasing its outstanding generalization and universal reasoning capabilities.

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: https://github.com/nju-websoft/KG-ICL.

Figures (6)

• Figure 1: Overview of the in-context KG reasoning foundation model. (A) Given the query and KG, we extract prompt graphs as context for the query relation "player in league". The entities and relations in the prompt graphs are mapped to the unified tokens. (B) We employ a message passing neural network to encode the prompt graph and readout the relation representations as the prompts. (C) Then we use the prompts to initialize the representations of entities and relations in the KG. After KG encoding, we score the candidate entities according to their embeddings in the last layer.
• Figure 2: MRR results on various KGs.
• Figure 3: MRR with different numbers of examples.
• Figure 4: Average MRR results of relation subgroups. Relations in the inductive and transductive dataset groups are divided into 6 subgroups based on the number of training samples, and the results represent the average scores for the relations within their respective subgroups. The percentage on the right side of each data bar indicates the proportion of relations in that subgroup to the total number of relations in their respective groups.
• Figure 5: Case study on prompt graphs. The left side shows some relation paths extracted from two prompt graphs of NELL-995 and FB15k-237. The right side depicts a heatmap where cosine similarities between relations in two prompt graphs are pairwise computed.