[Paper Review] Rewiring with Positional Encodings for Graph Neural Networks

Rickard Brüel‐Gabrielsson, Mikhail Yurochkin+1 more|arXiv (Cornell University)|Jan 29, 2022

Advanced Graph Neural Networks12 citations

TL;DR

This paper proposes a model-agnostic method to enhance Graph Neural Networks (GNNs) by rewiring graphs with $r$-hop neighborhood edges and incorporating positional encodings as node/edge features, significantly improving performance and alleviating over-squashing with minimal computational cost. The approach achieves competitive results using small $r$ values and a virtual CLS-node, outperforming standard GNNs and Transformers on multiple benchmarks.

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.

Motivation & Objective

• To address over-squashing and under-reaching in GNNs by expanding receptive fields beyond 1-hop neighborhoods without modifying the GNN architecture.
• To explore efficient, computationally feasible alternatives to full-graph attention mechanisms that scale quadratically with node count.
• To demonstrate that small increases in receptive field size ($r$) with appropriate positional encodings yield significant performance gains.
• To provide a systematic analysis of the compute-performance trade-offs in receptive field expansion using various positional encodings.
• To show that the method is compatible with both standard GNNs and graph Transformers, enabling broader applicability.

Proposed method

• The method augments the input graph by adding edges between all nodes within $r$-hop distance in the original graph topology.
• Positional encodings—such as shortest-path lengths, spectral features, and adjacency matrix powers—are used as node and edge features on the augmented graph.
• A virtual CLS-node (global source and sink) is added to improve long-range message passing and reduce over-squashing.
• The modified graph is then fed into any existing GNN or Transformer-based GNN model without architectural changes, making the method model-agnostic.
• The approach preserves the sparsity of the original graph for small $r$, avoiding the quadratic cost of full-graph attention.
• The method enables a controlled trade-off between computational cost and performance by tuning $r$ and choosing among different positional encoding schemes.

Experimental results

Research questions

• RQ1 Can expanding the receptive field to $r$-hop neighborhoods via graph rewiring improve GNN performance while remaining computationally efficient?
• RQ2 How do different positional encoding schemes (shortest paths, spectral, adjacency) compare in performance and efficiency when used with $r$-hop rewiring?
• RQ3 To what extent does adding a CLS-node improve performance and mitigate over-squashing in GNNs with limited receptive fields?
• RQ4 What is the optimal value of $r$ for balancing performance gains and computational cost across diverse graph datasets?
• RQ5 Does increasing $r$ reduce the reliance on homophily as an inductive bias in GNNs?

Key findings

• Using $r=2$ or $r=3$ with adjacency or spectral positional encodings and a CLS-node significantly improves GNN performance on multiple benchmarks, including Cora, PubMed, and ZINC.
• The CLS-node achieves performance comparable to $r=2$ or $r=3$ with $r=1$, demonstrating its effectiveness in enhancing long-range connectivity with minimal cost.
• Runtime scales approximately with graph density, and increasing $r$ leads to a substantial performance boost even at low densities, with optimal performance often achieved when average density exceeds 0.5.
• Spectral positional encodings perform worse than adjacency encodings on ZINC but are competitive on other datasets, indicating dataset-dependent sensitivity.
• The method reduces reliance on homophily as an inductive bias, suggesting that expanded receptive fields can compensate for low homophily in real-world graphs.
• The approach achieves competitive performance with standard GNNs and graph Transformers, demonstrating broad compatibility and effectiveness across architectures.