Exploring the impact of adaptive rewiring in Graph Neural Networks

TL;DR

This work tackles the computational burden of Graph Neural Networks by investigating sparsity as a regularizer and introducing adaptive rewiring to dynamically reconfigure connectivity during training. By combining Erdős-Rényi initialization ($\epsilon$) with fixed-rate ($\zeta_f$) and adaptive rewiring ($\zeta_a$) and an early-stopping mechanism, the authors demonstrate efficiency gains and competitive accuracy on MUTAG, PROTEINS, and a real-world N-1 power-grid contingency task. Key findings show that moderate sparsity improves generalization, excessive sparsity degrades performance, and adaptive rewiring often yields smoother training and higher final accuracy, with GINE particularly robust in high-sparsity N-1 scenarios. The results highlight the practical potential of sparse-to-sparse GNNs for real-time infrastructure analytics, reducing memory and compute while preserving predictive performance, and suggest directions for optimizing adaptive sparsity in broader graph-learning problems.

Abstract

This paper explores sparsification methods as a form of regularization in Graph Neural Networks (GNNs) to address high memory usage and computational costs in large-scale graph applications. Using techniques from Network Science and Machine Learning, including Erdős-Rényi for model sparsification, we enhance the efficiency of GNNs for real-world applications. We demonstrate our approach on N-1 contingency assessment in electrical grids, a critical task for ensuring grid reliability. We apply our methods to three datasets of varying sizes, exploring Graph Convolutional Networks (GCN) and Graph Isomorphism Networks (GIN) with different degrees of sparsification and rewiring. Comparison across sparsification levels shows the potential of combining insights from both research fields to improve GNN performance and scalability. Our experiments highlight the importance of tuning sparsity parameters: while sparsity can improve generalization, excessive sparsity may hinder learning of complex patterns. Our adaptive rewiring approach, particularly when combined with early stopping, proves promising by allowing the model to adapt its connectivity structure during training. This research contributes to understanding how sparsity can be effectively leveraged in GNNs for critical applications like power grid reliability analysis.

Figures (6)

• Figure 1: Information flow in a single step of message passing (MP) in a Graph Neural Network (GNN). The Message stage involves green source nodes sending messages to lilac target nodes. Aggregate stage combines these messages from the lilac nodes. The Update stage integrates the aggregated messages into the source node's representation (which transitions to red during this phase), propagating information.
• Figure 2: This figure illustrates a Graph Neural Network (GNN) framework incorporating sparsity for graph classification. The process starts with an input graph which is provided to a GNN. The weight matrices can be initialized using the Erdős-Rényi method (option A) to introduce sparsity (i.e., $\epsilon$). If no sparsity is applied, the model proceeds with weight matrices. Then during training, the GNN parameters are optimized via gradient based learning, while the sparse connectivity serves as a starting point for subsequent rewirings. Sparsity is handled in two ways: a Fixed (option B) approach maintains fixed sparsity (i.e., $\zeta_f$), while an Adaptive (option C) approach modifies the sparsity pattern using an adjustment rule (i.e., $\zeta_a$) throughout training. If neither fixed-rate nor adaptive rewiring is selected, the model follows the Base training mode, as shown in the upper branch of the diagram. The output of the GNN is a graph classification.
• Figure 3: The diagram where sparsity (1) and rewiring (2) is applied to the GNN. Here, the Erdos-Renyi mask with sparsity level ($\epsilon$) is applied directly to the model’s weight matrices ($MLP$), reducing the number of trainable connections before training begins. Rewiring is introduced during the message-passing (MP) process. Rewiring ($\zeta$) updates the model connectivity during training by modifying the weights $W_{edges}$/$W_{nodes}$ in GCN or $MLP_3$/$MLP_4$ in GINE. After the MP blocks, the OUTPUT layer aggregates the learned node and edge embeddings to produce a final graph-level representation via a linear layer (GCN) or $MLP_{out}$ (GINE).
• Figure 4: Training accuracy of different regularization techniques for the GIN model on the PROTEINS dataset, showing the performance of different sparsity and rewiring configurations: sparsity ($\epsilon=0.5$), fixed-rate rewiring ($\zeta_f=0.3$), adaptive-rate rewiring ($\zeta_a=0.3$), the baseline ($\epsilon=0$, $\zeta=0$), and with dropout. The results are averaged over 50 runs.
• Figure 5: Heatmaps illustrating the combined impact of model sparsity ($\epsilon$), and the average $\zeta$ value between fixed-rate, and adaptive-rate rewiring on GNN performance for the MUTAG (top) and PROTEINS (bottom) datasets. Left: GCNE model. Right: GINE model. The results are averaged over 50 runs, and the colorbars indicate model test accuracy.