Scalable GNN-based Knowledge Graph Representation Learning with Efficient Message Passing

2026-09-28

AI Summary

This work addresses the high computational complexity bottleneck of graph neural networks on large-scale knowledge graphs caused by message-passing mechanisms. We propose a general optimization framework based on an extension of relational sparse matrix multiplication. By reformulating the message-passing process as block-diagonal matrix operations and introducing universal mathematical transformations, our approach transcends the limitations of traditional subgraph sampling, enabling lossless and efficient recomputation of richer expressive functions. The proposed method requires no task-specific overhead and achieves near state-of-the-art performance on tasks such as link prediction while substantially reducing runtime and memory costs. Consequently, this framework significantly enhances the scalability of knowledge graph representation learning.

Abstract

Graph neural networks (GNNs) excel at representation learning on Knowledge Graphs (KGs), achieving stateof-the-art performance on tasks like link prediction or entity classification. However, their high computational complexity, inherent to their user-defined message passing (MP) algorithm, still prohibits their widespread adoption, especially for large KGs. Current efforts to mitigate the scalability bottlenecks of GNNs on KGs, such as subgraph sampling, are often task- and model-specific, and do not reliably guarantee lossless (if applicable) runtime/space reductions. To address this, we extend Relational Sparse Matrix Multiplication (RSPMM), originally designed to losslessly lower the space complexity of composition-based MP with pointwise composition functions, to support more expressive functions (e.g., 2x2 block-diagonal matrix multiplication, Givens rotation, circular correlation). Our method delivers significant task-independent reductions in runtime and space for current GNNs on KGs and facilitates efficient re-implementations of GNNs that maintain near state-of-the-art performance on challenging KG tasks, for a fraction of computational costs.

Problem

• Knowledge Graph
• Graph Neural Networks
• Scalability
• Message Passing
• Representation Learning

Innovation

• Knowledge Graph
• Graph Neural Networks
• Message Passing
• Relational Sparse Matrix Multiplication
• Scalability