Greener GRASS: Enhancing GNNs with Encoding, Rewiring, and Attention

TL;DR

GRASS tackles robust graph learning by uniting three complementary components: random walk-based graph encoding, random rewiring via superimposed random regular graphs, and a graph-tailored additive attention. The RRWP family encodings, including the efficient D-RRWP variant, provide rich structural signals to nodes and edges, while the rewiring increases long-range connectivity and mitigates information bottlenecks; the attention mechanism leverages edge representations to drive selective, directionally aware aggregation. Through extensive benchmarking and ablations, GRASS demonstrates state-of-the-art or competitive performance across standard GNN datasets and long-range graph benchmarks, with notable gains on ZINC (e.g., MAE improvement of $20.3 ext{%}$). The work also discusses its scalability, stochastic outputs, and reproducibility, highlighting practical trade-offs and providing a public codebase for replication and future research. Overall, GRASS offers a cohesive framework that enhances GNNs by synergistically encoding structure, enriching topology, and refining attention to graph-structured data.

Abstract

Graph Neural Networks (GNNs) have become important tools for machine learning on graph-structured data. In this paper, we explore the synergistic combination of graph encoding, graph rewiring, and graph attention, by introducing Graph Attention with Stochastic Structures (GRASS), a novel GNN architecture. GRASS utilizes relative random walk probabilities (RRWP) encoding and a novel decomposed variant (D-RRWP) to efficiently capture structural information. It rewires the input graph by superimposing a random regular graph to enhance long-range information propagation. It also employs a novel additive attention mechanism tailored for graph-structured data. Our empirical evaluations demonstrate that GRASS achieves state-of-the-art performance on multiple benchmark datasets, including a 20.3% reduction in mean absolute error on the ZINC dataset.

Figures (5)

• Figure 1: The structure of GRASS. Prior to training, GRASS precomputes (D-)RRWP encodings. At each training iteration, it rewires the input graph and adds distinct embeddings to added edges.
• Figure 2: Visualization of the proposed random rewiring technique. Solid lines denote existing edges of the input graph, and dashed lines denote added edges. (a) An example of the input graph $G$ that has poor connectivity. (b, c) Two among all possible instances of the randomly rewired graph $H$ with $r=2$. They have better connectivity than the input graph.
• Figure 3: Simplified visualization of the GRASS attention mechanism. (a) The edge aggregator extracting node relations to update edge representations. (b) The attentive node aggregator weighted by edge representations. For simplicity, attention from a node to itself, residual connections, and activation functions are omitted here. Figure \ref{['fig:attention']} provides a more detailed visualization.
• Figure 4: Visualized ablation study results for the number of added edges per node, with random regular and non-regular graphs, on ZINC. Error bars represent one standard error of the mean. The setup is identical to that described in Table \ref{['table:rr_ablation_zinc']}.
• Figure 5: The structure of an attention layer of GRASS. Node aggregation is attentive, with attention weights derived from edge representations. Edge aggregation is done through an MLP. For simplicity, biases are not shown here.