AI Research

KGPFN: Unlocking the Potential of Knowledge Graph F... | AI Research

Key Takeaways

KGPFN: Unlocking the Potential of Knowledge Graph Foundation Model via In-Context Learning Knowledge graphs (KGs) are essential for representing complex rela...
Knowledge graph (KG) foundation models aim to generalize across graphs with unseen entities and relations by learning transferable relational structure.
However, most existing methods primarily emphasize relation-level universality, while in-context learning, the other pillar of foundation models remains under-explored for KG reasoning.
We propose KGPFN, a KG foundation model using Prior-data Fitted Network that unifies transferable relational regularities with inference-time in-context learning from structured context.
KGPFN first learns relation representations via message passing on relation graphs to capture cross-graph relational invariances.

Paper AbstractExpand

Knowledge graph (KG) foundation models aim to generalize across graphs with unseen entities and relations by learning transferable relational structure. However, most existing methods primarily emphasize relation-level universality, while in-context learning, the other pillar of foundation models remains under-explored for KG reasoning. In KGs, context is inherently structured and heterogeneous: effective prediction requires conditioning on the local context around the query entities as well as the global context that summarizes how a relation behaves across many instances. We propose KGPFN, a KG foundation model using Prior-data Fitted Network that unifies transferable relational regularities with inference-time in-context learning from structured context. KGPFN first learns relation representations via message passing on relation graphs to capture cross-graph relational invariances. For query-specific reasoning, it encodes local neighborhoods using a multi-layer NBFNet as local context. To enable ICL at global scale, it constructs relation-specific global context by retrieving a large set of instances of the query relation together with their local neighborhoods, and aggregates them within a Prior-Data Fitted Network framework that combines feature-level and sample-level attention. Through multi-graph pretraining on diverse KGs, KGPFN learns when to instantiate reusable patterns and when to override them using contextual evidence. Experiments on 57 KG benchmarks demonstrate that KGPFN achieves strong adaptation to previously unseen graphs through in-context learning alone, consistently outperforming competitive fine-tuned KG foundation models. Our code is available at this https URL .

KGPFN: Unlocking the Potential of Knowledge Graph Foundation Model via In-Context Learning
Knowledge graphs (KGs) are essential for representing complex relational data, but traditional models often struggle to generalize to new, unseen graphs without being retrained from scratch. While recent foundation models have made progress in learning transferable relational patterns, they often overlook "in-context learning"—the ability of a model to adapt its predictions based on a few examples provided at inference time. KGPFN addresses this gap by introducing a framework that combines general relational knowledge with structured, context-aware reasoning, allowing the model to adapt to new graphs instantly without needing further fine-tuning.

Bridging Local and Global Context

To reason effectively, KGPFN treats context as a two-part system. First, it captures "local context," which consists of the immediate neighborhood around a query entity. This is crucial because the same relational pattern might be valid in one part of a graph but misleading in another; the local neighborhood provides the necessary evidence to decide whether a pattern should be applied or overridden. Second, it captures "global context" by retrieving a set of representative instances of the query relation from across the graph. This provides a broader view of how a relation typically behaves, helping the model understand the typical structural constraints associated with that specific relation.

The Prior-Data Fitted Network Approach

KGPFN utilizes the Prior-data Fitted Network (PFN) paradigm to process this context. Instead of relying on static, pre-learned entity embeddings, the model is trained to perform "amortized Bayesian inference." During inference, it takes the query and a set of labeled context examples—both local and global—and processes them through a series of attention mechanisms. By combining feature-level and sample-level attention, the model learns to weigh the importance of different context examples, effectively "learning" the task at hand in a single forward pass.

Performance and Generalization

The model’s architecture includes a multi-layer NBFNet to encode local neighborhoods and a message-passing framework to learn transferable relation representations. By pretraining on a diverse set of knowledge graphs, KGPFN learns to identify when to rely on reusable structural patterns and when to prioritize the specific evidence provided in the context. Experiments across 57 different benchmarks demonstrate that this approach is highly effective, consistently outperforming competitive models that require fine-tuning. This confirms that incorporating structured in-context learning is a powerful strategy for building more flexible and capable knowledge graph foundation models.

Comments (0)

No comments yet

Be the first to share your thoughts!