Flock: A Knowledge Graph Foundation Model via Learning on Random Walks

2025-10-01

AI Summary

This work addresses zero-shot link prediction in knowledge graphs, where models must generalize to unseen entities and relations. We propose a probabilistic node-relation equivariant representation learning framework based on random walk sequences: structural-aware walk sequences are sampled and jointly encoded via sequence modeling and learnable pooling to capture the co-occurring structural semantics of nodes and relations; distribution-level equivariance constraints are introduced to overcome the expressive limitations of conventional deterministic equivariance, enabling fine-grained discrimination between structurally similar but semantically distinct relations. Evaluated on the newly constructed diagnostic benchmark Petals, our method achieves 100% zero-shot accuracy. Moreover, it attains state-of-the-art performance across 54 cross-domain knowledge graphs, significantly outperforming existing approaches.

Technology Category

• Reasoning under Uncertainty: Relational Probabilistic Models
• Machine Learning: Structured Learning
• Data Mining & Knowledge Management: Linked Open Data, Knowledge Graphs & KB Completion

Application Category

• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Semantics and Knowledge: Scalable techniques for the creation, curation, publication, maintenance, and consumption of large, Web-based, structured, reusable, knowledge graphs and ontologies
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

We study the problem of zero-shot link prediction on knowledge graphs (KGs), which requires models to generalize over novel entities and novel relations. Knowledge graph foundation models (KGFMs) address this task by enforcing equivariance over both nodes and relations, learning from structural properties of nodes and relations, which are then transferable to novel graphs with similar structural properties. However, the conventional notion of deterministic equivariance imposes inherent limits on the expressive power of KGFMs, preventing them from distinguishing structurally similar but semantically distinct relations. To overcome this limitation, we introduce probabilistic node-relation equivariance, which preserves equivariance in distribution while incorporating a principled randomization to break symmetries during inference. Building on this principle, we present Flock, a KGFM that iteratively samples random walks, encodes them into sequences via a recording protocol, embeds them with a sequence model, and aggregates representations of nodes and relations via learned pooling. Crucially, Flock respects probabilistic node-relation equivariance and is a universal approximator for isomorphism-invariant link-level functions over KGs. Empirically, Flock perfectly solves our new diagnostic dataset Petals where current KGFMs fail, and achieves state-of-the-art performances on entity- and relation prediction tasks on 54 KGs from diverse domains.

Problem

• Addressing zero-shot link prediction on knowledge graphs with novel entities
• Overcoming deterministic equivariance limits through probabilistic node-relation equivariance
• Developing universal approximator for isomorphism-invariant link-level functions

Innovation

• Probabilistic equivariance breaks deterministic symmetry limits
• Random walk sampling with sequence encoding protocol
• Learned pooling aggregates node-relation representations universally