PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Huu Tan Mai, Cuong Xuan Chu, Heiko Paulheim, Daria Stepanova

arXiv:2609.34499v1cs.LGcs.AI

TL;DR

GNNs offer strong KG representations but their message-passing computation remains costly, particularly for large graphs. This paper extends RSPMM to support more expressive composition operators. The implementations deliver lossless, task- and model-independent scalability gains and can help GNNs rival state-of-the-art methods with fewer resources.

  • Problem

    GNN message passing is computationally expensive, and existing scalability approaches do not address its cost or readily support expressive message functions.

    Source: §1 · p. 2

  • Method

    The paper extends RSPMM to compute a wider range of composition operators in O(|ℰ|d) space, including block-diagonal multiplication, Givens rotation, and circular correlation.

    Source: §1 · p. 2, §2 · p. 3

  • Results

    Across models and tasks, the RSPMM extension yields up to ×10.7 reductions in memory footprint while preserving model equivalence, and the implementations can rival state-of-the-art methods with substantially fewer resources.

    Source: §3 · p. 4, §4 · p. 6

  • Takeaways & Limitations

    The lossless scalability gains can let practitioners afford better architectures within the same computational budget and tackle larger-scale KG problems.

    Source: §4 · p. 6

    For aggregation with min or max, dedicated CUDA kernels would be needed, and these are outside the work's scope.

    Source: §2 · p. 3

Abstract

from arXiv · show

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.

1. Introduction

KG GNNs can be more predictive than shallow KG embeddings because they learn from neighborhoods, but message passing makes them computationally expensive. The introduction identifies edge-wise message materialization as a scalability bottleneck and motivates extending RSPMM beyond pointwise operators.

1. Introduction

  • GNNs use neighborhood message passing for higher predictive accuracy than shallow KG embeddings, which remain the de facto approach for web-scale graphs but have accuracy limitations.The introduction describes shallow embeddings as computationally simple and supported by dedicated software.Source: §1 · p. 1
  • The KG adjacency tensor assigns zero weight to every non-existent entity-relation-entity triple.This is the notation used to specify the message-passing problem.Source: §1 · p. 1
  • Traditional gather-scatter message passing can materialize an O(|𝒯|d)-sized tensor, making memory use a major obstacle to scalable GNN learning.The figure contrasts this approach with fused RSPMM implementations that aim for lower space complexity when |ℰ| ≪ |𝒯|.Source: §1 · p. 2, Figure 1 · p. 2
  • Prior vocabulary reduction and subgraph sampling do not address costly message passing, while existing RSPMM for + and ⊙ is not readily available for expressive messages.The existing approach computes in O(|ℰ|d) space, which is beneficial when |ℰ| ≪ |𝒯|.Source: §1 · p. 2
  • The proposed RSPMM extension supports more expressive composition operators, including 2 × 2 block-diagonal matrix multiplication, Givens rotation, and circular correlation.Many KG GNNs use composition-based messages and can benefit from the generic optimization.Source: §1 · p. 2

2. Approach

The approach extends RSPMM to more composition operators while retaining O(|ℰ|d) space complexity. It focuses on sum aggregation and motivates the expanded operator choices by their relational expressiveness.

2. Approach

  • The approach focuses on sum aggregation, which is popular in GNNs and has desirable linearity and distributivity properties.The paper assumes an even embedding dimension.Source: §2 · p. 2
  • The extension supports a wider variety of composition operators in O(|ℰ|d) space complexity.The operator catalogue is inspired by shallow KG embeddings.Source: §2 · p. 3, Table 1 · p. 3
  • 2 × 2 block-diagonal multiplication reduces to two RSPMM⊙ calls, while complex multiplication and Givens rotation/reflection use this formulation and circular correlation uses RSPMM⊙C.These constructions retain O(|ℰ|d) space complexity.Source: §2 · p. 3
  • Although min and max aggregation are theoretically possible, dedicated CUDA kernels for them are outside this work’s scope.The paper’s implementations focus on sum aggregation.Source: §2 · pp. 2–3
  • The choice of composition operator matters because different KG embeddings can express different relational patterns.The paper notes that TransE models inverses and compositions but not symmetric relations well, whereas RotatE can model all three.Source: §2 · pp. 3–4

3. Experiments

Experiments evaluate existing GNN implementations and prototypical models on KG tasks, including settings with frozen language-model entity features. The RSPMM extension reduces runtime and memory for existing models, while stress tests report large-dataset outcomes.

Setup

  • The experiments cover entity classification and KGC, comparing existing GNN message-passing costs, prototypical models, and GNNs against language-model embedding methods.Experiments used an NVIDIA A100 GPU and AMD Epyc 7543 CPU; runs exceeding 24 hours were cut off.Source: §3 · p. 4

Prototypical Models

  • The prototypical models modify CompGCN and AdaProp, and the text-feature experiments use frozen all-mpnet-base-v2 entity-description embeddings.MODEL-A changes relation updates and performs transformations before message passing; MODEL-B uses query-dependent initializations and a shared entity scorer.Source: §3 · p. 4

Results

  • Up to ×10.7 lower memory footprint and ×32.2 lower runtime were achieved across existing GNN models and tasks without sacrificing performance.The extension preserves theoretical equivalence; on dblp with 4.5M entities and 22M edges, R-GCN used up to 33.8 GiB during full hyperparameter search and reached 0.72 final accuracy, while the version without RSPMM ran out of memory.Source: §3 · p. 4, Figure 2 · p. 5
  • For KGC on WN18RR and FB15k-237, RSPMM implementations are substantially more resource-friendly than counterparts in their categories and perform competitively.R-MPNNs with RSPMM and initial entity features perform comparably or better than methods involving LM training, with substantially less memory and runtime.Source: §3 · p. 4, Table 3 · p. 5

Stress Tests

  • On AM+, the RSPMM-based implementation used 12.6 GiB over 11 hours, versus 21.5 GiB over 14 hours without RSPMM; both achieved .92 final accuracy.On DBLP, the implementation achieved .72 final accuracy with a maximum footprint of 33.8 GiB over 3 days; base R-GCN ran out of memory.Source: §3 · p. 5

4. Conclusion

The conclusion presents an RSPMM extension for diverse composition operators, readily applicable to existing KG GNNs. It reports significant lossless scalability gains from implementation changes and says these gains can support better architectures within fixed budgets and SOTA comparisons with substantially fewer resources.

  • The RSPMM extension supports diverse composition operators and applies readily to many existing KG GNNs.Source: §4 · p. 6
  • Switching implementations gives relational GNNs significant, lossless scalability gains, independent of task or model.Source: §4 · p. 6
  • At the same computational budget, these gains allow better architectures that can rival SOTA methods with substantially fewer resources.Source: §4 · p. 6
  • The authors view RSPMM as a tool to speed iteration and tackle larger-scale KG problems.Source: §4 · p. 6

Declaration on Generative AI

The authors state that they did not employ any Generative AI tools.

Related papers

Loading 2609.34499v1…