Probabilistic Graph Rewiring via Virtual Nodes

TL;DR

This paper introduces IPR-MPNNs, a novel graph neural network method that uses virtual nodes for probabilistic rewiring, improving message passing, scalability, and performance over existing models like graph transformers.

Contribution

We propose IPR-MPNNs, which incorporate virtual nodes for implicit probabilistic graph rewiring, enhancing expressiveness and efficiency of message passing in GNNs.

Findings

IPR-MPNNs outperform traditional MPNNs in mitigating under-reaching and over-squashing.

The approach achieves state-of-the-art results on multiple graph datasets.

IPR-MPNNs are faster and more scalable than graph transformers.

Abstract

Message-passing graph neural networks (MPNNs) have emerged as a powerful paradigm for graph-based machine learning. Despite their effectiveness, MPNNs face challenges such as under-reaching and over-squashing, where limited receptive fields and structural bottlenecks hinder information flow in the graph. While graph transformers hold promise in addressing these issues, their scalability is limited due to quadratic complexity regarding the number of nodes, rendering them impractical for larger graphs. Here, we propose implicitly rewired message-passing neural networks (IPR-MPNNs), a novel approach that integrates implicit probabilistic graph rewiring into MPNNs. By introducing a small number of virtual nodes, i.e., adding additional nodes to a given graph and connecting them to existing nodes, in a differentiable, end-to-end manner, IPR-MPNNs enable long-distance message propagation, circumventing quadratic complexity. Theoretically, we demonstrate that IPR-MPNNs surpass the expressiveness of traditional MPNNs. Empirically, we validate our approach by showcasing its ability to mitigate under-reaching and over-squashing effects, achieving state-of-the-art performance across multiple graph datasets. Notably, IPR-MPNNs outperform graph transformers while maintaining significantly faster computational efficiency.