RADE: Random Add-Drop Edge as a Regularizer
Abstract
Graph Neural Networks (GNNs) suffer from overfitting and over-squashing of long-range information. Stochastic graph augmentations (e.g., edge deletion) regularize training against overfitting but can introduce train-inference misalignment and do not improve over-squashing. In contrast, rewiring methods improve connectivity to mitigate over-squashing, but are not designed to regularize training. We propose Random Add-Drop Edge (RADE), a stochastic graph augmentation method that jointly drops and adds edges to address both overfitting and over-squashing simultaneously. RADE is provably designed to align training and inference so that random augmentations regularize training without distribution shift, while supporting long-range communication at inference. We further propose and study a mini-batch gradient-norm balancing algorithm that adapts deletion and addition rates during training, rendering RADE hyperparameter-free in practice. Experiments on node- and graph-classification benchmarks show that RADE is a strong regularizer and mitigates over-squashing. Ablations support the roles of train-inference alignment, adaptive rate selection, and the complementary effects of random edge deletion and edge addition.
Lay Summary
Many datasets can be represented as networks: people connected in social networks, atoms connected in molecules, or papers connected by citations. Modern machine learning models can learn from these networks by passing information along their connections, but they face two common problems. First, they can overfit by relying too much on the training data. Second, useful information from far-away parts of the network may be difficult to carry through many intermediate connections. We propose RADE, a training method that randomly removes some existing connections and adds some new ones while the model learns. Removing connections helps prevent overfitting, while adding connections can create shortcuts that make long-range information easier to use. A key challenge is that training on randomly changed networks can make the model behave differently at test time, so RADE includes corrections that keep training and testing aligned. It also automatically adjusts how many connections to remove or add during training, reducing the need for manual tuning. Across several network learning tasks, RADE improves performance and is especially helpful when long-range information matters.