Locality-Aware Graph-Rewiring in GNNs
Federico Barbero, Ameya Velingker, Amin Saberi, Michael Bronstein, Francesco Di Giovanni
TL;DR
This work tackles over-squashing in message-passing GNNs by proposing LASER, a locality-aware sequential graph-rewiring framework that evolves connectivity through a constrained series of graph snapshots. LASER uses a locality measure based on shortest-walk distance and a connectivity score derived from walks up to length $k$ to selectively add a sparse set of edges, controlled by a density parameter $\rho$ and enabling a continuum between standard MPNNs and multi-hop GNNs. The authors provide theoretical justification linking LASER to multi-relational and temporal GNNs and demonstrate empirically that LASER reduces over-squashing while preserving locality and sparsity, achieving state-of-the-art or competitive results on LRBG and TUDatasets with favorable scalability. The approach also clarifies connections between graph rewiring and relational GNN frameworks, suggesting practical significance for large-scale, long-range graph tasks.
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.
