Aggregation Buffer: Revisiting DropEdge with a New Parameter Block

TL;DR

The paper challenges DropEdge's effectiveness by showing that intrinsic aggregation biases in GNNs limit robustness gains, then introduces Aggregation Buffer AGG_B as a modular post-processing block to enhance edge-robustness across architectures. By modeling a discrepancy bound and designing two essential conditions (edge-awareness and stability), AGG_B refines aggregation outputs through a degree-normalized linear transform and is trained in a separate stage to preserve backbone knowledge. The approach yields consistent improvements on 12 node classification benchmarks and generalizes to multiple GNN families, addressing degree bias and structural disparity as a unified mechanism. The work provides practical, reusable tooling (AGG_B) and code for improving robustness without disrupting deployed models, with potential extensions toward end-to-end training.

Abstract

We revisit DropEdge, a data augmentation technique for GNNs which randomly removes edges to expose diverse graph structures during training. While being a promising approach to effectively reduce overfitting on specific connections in the graph, we observe that its potential performance gain in supervised learning tasks is significantly limited. To understand why, we provide a theoretical analysis showing that the limited performance of DropEdge comes from the fundamental limitation that exists in many GNN architectures. Based on this analysis, we propose Aggregation Buffer, a parameter block specifically designed to improve the robustness of GNNs by addressing the limitation of DropEdge. Our method is compatible with any GNN model, and shows consistent performance improvements on multiple datasets. Moreover, our method effectively addresses well-known problems such as degree bias or structural disparity as a unifying solution. Code and datasets are available at https://github.com/dooho00/agg-buffer.

Figures (4)

• Figure 1: DropEdge generates various reduced rooted subgraphs for center nodes (*) by randomly removing edges.
• Figure 2: Accuracy and loss terms on test data during the training of a GCN at PubMed. We illustrate the average of 10 independent runs, with shaded regions representing the minimum and maximum values. While DropEdge decreases the robustness term compared to standard GNN, it leads to increasing the bias term, eventually resulting in similar test accuracy to standard GNNs.
• Figure 3: Illustration of $\text{AGG}_{B}$ and its training scheme. After the integration into a pre-trained GNN, $\text{AGG}_{B}$ is trained using ${\mathcal{L}}_{\text{RC}}$ with DropEdge, while the pre-trained parameters remain frozen.
• Figure 4: Changes in two different approximations of the robustness loss, $\mathbb{E}_{P}[\log Q({\bm{y}}_i|\mathcal{G}_i) - \log Q({\bm{y}}_i|\tilde{\mathcal{G}}_i)]$, during training of the base GCN (top row) and AGG$_B$ (bottom row). Each curve represents the average over 10 independent runs, with shaded areas indicating the minimum and maximum values. Blue represents the robustness term in our proposed robustness-controlled loss, where $P$ is approximated by $Q$. Orange represents the label-based approximation, where P is approximated using ground-truth labels. Both approximations exhibit similar trends: robustness loss gradually emerges during GCN training and is further optimized during AGG$_B$ training.

Theorems & Definitions (13)

• Definition 3.1: Rooted Subgraph
• Definition 3.2: Reduced Rooted Subgraph
• Definition 3.3: Discrepancy bound
• Lemma 3.4
• Lemma 3.5
• Theorem 3.6
• Lemma 3.7
• Theorem 3.8
• Theorem 3.9
• Theorem 4.1