Locality-Aware Graph-Rewiring in GNNs

TL;DR

This paper introduces a locality-aware graph-rewiring framework for GNNs that reduces over-squashing, respects graph locality, and maintains sparsity, leading to improved performance on real-world benchmarks.

Contribution

The paper proposes a novel locality-aware graph-rewiring method that balances over-squashing reduction, locality preservation, and sparsity, addressing fundamental trade-offs in existing techniques.

Findings

The proposed rewiring framework outperforms existing methods on several benchmarks.

It effectively reduces over-squashing while preserving graph locality.

The approach maintains graph sparsity, enabling scalable GNN training.

Abstract

Graph Neural Networks (GNNs) are popular models for machine learning on graphs that typically follow the message-passing paradigm, whereby the feature of a node is updated recursively upon aggregating information over its neighbors. While exchanging messages over the input graph endows GNNs with a strong inductive bias, it can also make GNNs susceptible to over-squashing, thereby preventing them from capturing long-range interactions in the given graph. To rectify this issue, graph rewiring techniques have been proposed as a means of improving information flow by altering the graph connectivity. In this work, we identify three desiderata for graph-rewiring: (i) reduce over-squashing, (ii) respect the locality of the graph, and (iii) preserve the sparsity of the graph. We highlight fundamental trade-offs that occur between spatial and spectral rewiring techniques; while the former often satisfy (i) and (ii) but not (iii), the latter generally satisfy (i) and (iii) at the expense of (ii). We propose a novel rewiring framework that satisfies all of (i)--(iii) through a locality-aware sequence of rewiring operations. We then discuss a specific instance of such rewiring framework and validate its effectiveness on several real-world benchmarks, showing that it either matches or significantly outperforms existing rewiring approaches.