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

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.

Paper

Similar papers

Peer review

Reviewer JyNi6/10 · confidence 4/52024-07-08

Summary

This paper introduces KG-ICL, a model that facilitates generalized reasoning over knowledge graphs via in-context learning. KG-ICL first extracts example facts relevant to the query from the knowledge graph to generate prompt graphs. These prompt graphs are then encoded using a unified tokenizer and message passing neural network to produce relation representations. Next, the generated embeddings of relation entities are then integrated with knowledge graph to generate the score of candidate entities. Extensive experiments on 43 different knowledge graphs under transduction and inductive validate the effectiveness of KG-ICL.

Strengths

1. The paper is well-structured. 2. Extensive experimental results demonstrate that KG-ICL can consistently outperform baseline models. 3. Codes are provided for reproducibility.

Weaknesses

1. The motivation for the ICL setting over KGs is not clearly motivation. It appears to be essentially the same as the inductive learning setting. Furthermore, the results shown in Figure 3 indicate that the number of examples has no significant effect on performance, which contradicts the hypothesis of ICL. 2. In the section 4.1, the authors propose to sample $M$ example and extract the prompt graph within $k$ hop paths. However, it's unclear whether the number of $M$ and $k$ will affect the model performance. Unfortunately, the authors provide little explanation for this. 3. The paper introduces the unified tokenizer as a component of the KG-ICL model, yet the motivation for its use is not clear. Additionally, the ablation studies (Table 2) show that excluding the unified tokenizer or the token representation results in only a marginal performance decline. This raises questions about the necessity and practical advantage of using the unified tokenizer for the reasoning task.

Questions

1. How KG-ICL generate the prompt graph? Does it directly sample the example from the knowledge graph and then connect them? 2. It is ambiguous what is the difference between the w/o unified tokenizer and w/o token represent. 3. In this paper, the task is assuming that a model is pre-trained using a set of source KGs. Is this a standard task setting? How realistic is this setting for solving some real problem?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

The authors have discussed the limitations of their work in the paper.

Reviewer 4T5d7/10 · confidence 5/52024-07-09

Summary

The paper studies inductive KG reasoning with unseen entities and relations at inference time and introduces KG-ICL, an in-context learning model for KG completion. Following Ultra [1], KG-ICL employs a two-stage approach: (1) obtaining relational representations based on the given graph; (2) performing entity prediction on the given graph using those relational representations. The main difference with ULTRA is Stage 1, that is, instead of creating a graph of relations and learning tokens for meta-relations, KG-ICL mines several (around 5) examples of k-hop subgraphs around each relation type and applies a different labeling trick learning a different vocabulary of transferable tokens. Nodes are labeled following the distance encoding of GraIL [2] (tokens represent all combinations of pairwise distances up to hop $k$), relations are labeled simply by a binary indicator whether a given relation is the query relation or not (two tokens). Applying a GNN over each subgraph, their final representations are mean pooled to get a singe tensor of relational representations. In Stage 2, those representations are used in a standard inductive pipeline for initializing edge type vectors, putting the query relation vector on the starting head node, and running another GNN to get the final predictions. Experimentally, KG-ICL demonstrates promising results and outperforms Ultra on inductive datasets while being marginally better on larger transductive datasets.

Strengths

**S1.** Foundation models for inductive KG reasoning on any unseen graph is a timely and important topic. New non-trivial approaches in this field are rare and this paper does a good job presenting and explaining the details clearly. **S2.** A different approach for obtaining relational representations through few-shot subgraph examples of each relation (resembling the “in-context learning” scenario) and learning a transferable vocabulary - here it is an extended GraIL-based labeling with distance encoding for nodes and binary same/not-same indicator for relations (although Table 2 shows that learning tokens does not bring much benefits and labeling strategy is more important). **S3**. Informative ablations, experiments on the Ultra benchmark involving 57 datasets

Weaknesses

**W1.** The main problem of KG-ICL is a different pre-training dataset mixture (inductive FB V1, inductive NELL V1, transductive Codex Small) which makes it hard to directly compare the performance numbers against Ultra (that used only transductive FB15k237, WN18RR, and Codex Medium). For a fair comparison, KG-ICL should have been either pre-trained on the same datasets as Ultra, or Ultra should have been re-trained on the new mixture. The new pre-training mixture consists of smaller graphs and is better suited for inference on smaller inductive datasets (where KG-ICL gains are the largest). While it is possible to pull the argument that LLMs and FMs are trained on different datasets (often undisclosed to the public) and evaluated on the same benchmarks, I believe the Ultra benchmarking datasets for KG reasoning are open and transparent enough for conducting fair evaluations. **W2** Since KG-ICL focuses on the zero-shot inference performance (line 239), one missing experiment is to measure the performance as a function of training graphs in the pre-training mixture - for instance, Ultra provides several checkpoints with the growing inductive inference performance with more datasets added to the training. References: [1] Galkin et al. Towards Foundation Models for Knowledge Graph Reasoning. ICLR 2024. [2] Teru et al. Inductive relation prediction by subgraph reasoning. ICML 2020. [3] Huang et al. A Theory of Link Prediction via Relational Weisfeiler-Leman on Knowledge Graphs. NeurIPS 2023.

Questions

**Q1.** Why were 4 inductive (HM) and 10 fully-inductive (ISDEA) datasets omitted from the main results table? Those are all viable datasets and I don’t see a major reason to present the average over 43 datasets instead of 57. **Q2.** Line 168: the formula for the total number of tokens $\frac{(k+1)(k+2)}{2} - 2$ seems to be incorrect? Setting $k=2$, the formula gives 4 tokens but we only have $(0,1), (1,0), (1,1)$ distances. Similarly, for $k=4$, there are 15 combinations of distances (without (0,0)) but the formula gives 13 options **C1**. Line 190: citations on conditional message passing GNNs miss [3] that theoretically formalized C-MPNNs.

Rating

7

Confidence

5

Soundness

3

Presentation

4

Contribution

3

Limitations

Addressed

Reviewer xJpJ7/10 · confidence 3/52024-07-12

Summary

This paper aims to build a foundation model for knowledge graphs, to have a universal reasoning ability across diverse knowledge graphs including the unseen entities and relations. Specifically, given the query, the proposed approach first extracts its relevant prompt graphs and then map their entities and relations to the predefined tokens (versatile for any entities and relations). After that, it performs two message passing, in order to encode the prompt graphs and perform knowledge graph reasoning with them. The authors extensively validate the proposed approach on 43 knowledge graphs, showcasing its effectiveness over the strong knowledge graph foundational model baseline.

Strengths

* Building a foundation model for knowledge graphs is a very important task. * The proposed (in-context learning style) approach to extract prompt graphs (for the given query) and use them to perform knowledge graph reasoning is reasonable. * The proposed approach outperforms the previous knowledge graph foundation model by large margins. * This paper is well-written.

Weaknesses

* The technical contribution of the proposed approach over the prior knowledge graph foundation work [6] looks marginal, which is also not clearly discussed in the paper. In my view, the major improvement (and the novelty) of the proposed approach is to extract multiple prompt graphs (relevant to the query) and use them for knowledge graph reasoning, instead of using only one target subgraph. * The performance comparison between high and low resources relations (or entities) is worthwhile to present (in the main paper).

Questions

Please see my weaknesses above.

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

The authors discuss the limitations and potential negative societal impact of their work.

Reviewer 4T5d2024-08-08

Response

Thank you for the clarifications and new experimental results, I am pretty satisfied with the response and increasing the score to 7.

Reviewer JyNi2024-08-09

Thanks for the response

Thanks for the response. My concerns are well addressed. I will increase my score to 6.

Authorsrebuttal2024-08-09

Grateful Thanks to Reviewer JyNi

Dear Reviewer JyNi, We sincerely appreciate your recognition of our efforts to address your concerns. Thank you for taking the time to provide insightful suggestions that will help us further refine our paper. Best regards, Authors

Reviewer xJpJ2024-08-12

Thank you for your response, which addresses all of my concerns. This is a good paper and I increase the rating (from 6) to 7.

Authorsrebuttal2024-08-13

Appreciation to Reviewer xJpJ

Dear Reviewer xJpJ, We sincerely appreciate your support and for revising the rating. Thank you once again for your invaluable time and effort during the review and discussion periods. Best regards, The Authors

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC