Rewiring with Positional Encodings for Graph Neural Networks

Published: 08 Nov 2023, Last Modified: 08 Nov 2023Accepted by TMLREveryoneRevisionsBibTeX
Abstract: Several recent works use positional encodings to extend the receptive fields of graph neural network (GNN) layers equipped with attention mechanisms. These techniques, however, extend receptive fields to the complete graph, at substantial computational cost and risking a change in the inductive biases of conventional GNNs, or require complex architecture adjustments. As a conservative alternative, we use positional encodings to expand receptive fields to r-hop neighborhoods. More specifically, our method augments the input graph with additional nodes/edges and uses positional encodings as node and/or edge features. We thus modify graphs before inputting them to a downstream GNN model, instead of modifying the model itself. This makes our method model-agnostic, i.e., compatible with any of the existing GNN architectures. We also provide examples of positional encodings that are lossless with a one-to-one map between the original and the modified graphs. We demonstrate that extending receptive fields via positional encodings and a virtual fully- connected node significantly improves GNN performance and alleviates over-squashing using small r. We obtain improvements on a variety of models and datasets and reach competitive performance using traditional GNNs or graph Transformers.
Submission Length: Regular submission (no more than 12 pages of main content)
Changes Since Last Submission: We have revised the submission to address feedback received from the reviewers and action editor.
Assigned Action Editor: ~Yujia_Li1
License: Creative Commons Attribution 4.0 International (CC BY 4.0)
Submission Number: 1429