Table of Contents
Fetching ...

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.

Locality-Aware Graph-Rewiring in GNNs

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 to selectively add a sparse set of edges, controlled by a density parameter 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.
Paper Structure (45 sections, 4 theorems, 26 equations, 4 figures, 9 tables, 2 algorithms)

This paper contains 45 sections, 4 theorems, 26 equations, 4 figures, 9 tables, 2 algorithms.

Key Result

Proposition 5.1

Let $v,u\in\mathsf{V}$ with $d_\mathsf{G}(v,u) = r$, and assume that there exists a single path of length $r$ connecting $v$ and $u$. Assume that LASER adds an edge between $v$ and some node $j$ belonging to the path of length $r$ connecting $v$ to $u$, with $d_\mathsf{G}(v,j) = \ell < r$. Then for

Figures (4)

  • Figure 1: Difference between spectral (left), spatial (middle), and LASER (right) rewirings in green with respect to the blue node of reference. Spectral rewirings are sparse and connect distant nodes. Spatial rewirings are able to retain local inductive biases at the cost of sparsity. LASER remains both local and sparse by optimizing over the edges to be added.
  • Figure 2: Total effective resistance and graph distance, as measured by the Frobenius norm, when varying the number of snapshots from $2$ to $5$ and $\rho$ from $0.1$ to $1$.
  • Figure 3: Performance on the Peptides tasks when varying the number of snapshots from $2$ to $5$ and $\rho$ from $0.1$ to $1$. The baseline indicates a standard GCN model.
  • Figure 4: Average Precision (AP) on Peptides-func obtained by varying the number of snapshots for LASER and FOSR, respecting the 500k parameter budget. For LASER , we fix $\rho = 0.1$. For FOSR, each snapshot contains $10$ iterations.

Theorems & Definitions (6)

  • Proposition 5.1
  • Proposition 5.2
  • Proposition A.1
  • proof
  • Proposition A.2
  • proof