RAwR: Role-Aware Rewiring via Approximate Equitable Partition

TL;DR

RAwR is a novel, efficient graph rewiring method that enhances long-range interaction modeling in GNNs by leveraging equitable partitions and spectral role metrics, achieving state-of-the-art results.

Contribution

Introduces RAwR, a role-aware rewiring framework using approximate equitable partitions to improve GNN performance on long-range tasks, with theoretical insights and a new spectral metric.

Findings

RAwR improves GNN performance on diverse benchmarks.

The method reduces structural bottlenecks and oversquashing.

Spectral Role Lift (SRL) effectively identifies optimal partitions.

Abstract

While Graph Neural Networks (GNNs) have demonstrated significant efficacy in node classification tasks, where predictions rely on local neighborhood information, the performance of GNNs often drops when prediction tasks depend on long-range interactions. These limitations are attributed to phenomena such as oversquashing, where structural bottlenecks restrict signal propagation across the network topology. To address this challenge, we introduce RAwR, a computationally efficient rewiring framework that augments the input graph with a quotient graph derived from equitable partitions. This approach facilitates accelerated communication between nodes that share identical structural roles, as identified by the Weisfeiler-Leman graph coloring, and thereby reduces the total effective resistance of the system. Furthermore, by employing an approximate definition of the equitable partition, RAwR enables a controllable reduction of the quotient graph, which, in its most condensed state, recovers the conventional Master Node rewiring technique. Empirical evaluations across a diverse suite of benchmarks -- including homophilic, heterophilic, and synthetic long-range datasets -- demonstrate that RAwR achieves state-of-the-art results. Our contribution is further supported by an analytical investigation using a teacher-student model of linear GNNs, which elucidates the theoretical foundations of role-based rewiring. This analysis leads to the formulation of Spectral Role Lift (SRL), a metric designed to identify the optimal approximate equitable partition for maximizing predictive performance.