[Paper Review] Bayesian Graph Neural Networks with Adaptive Connection Sampling

Arman Hasanzadeh, Ehsan Hajiramezanali+5 more|arXiv (Cornell University)|Jun 7, 2020

Advanced Graph Neural Networks41 references17 citations

TL;DR

This paper proposes Graph DropConnect (GDC), a unified framework for adaptive connection sampling in graph neural networks (GNNs) that jointly learns sampling rates during training. By modeling connection drop rates via a hierarchical beta-Bernoulli construction and using continuous relaxation with ARM gradients, GDC enables Bayesian-like uncertainty quantification and significantly reduces over-smoothing and over-fitting, outperforming fixed-rate methods like DropEdge and DropOut in semi-supervised node classification on benchmark datasets.

ABSTRACT

We propose a unified framework for adaptive connection sampling in graph neural networks (GNNs) that generalizes existing stochastic regularization methods for training GNNs. The proposed framework not only alleviates over-smoothing and over-fitting tendencies of deep GNNs, but also enables learning with uncertainty in graph analytic tasks with GNNs. Instead of using fixed sampling rates or hand-tuning them as model hyperparameters in existing stochastic regularization methods, our adaptive connection sampling can be trained jointly with GNN model parameters in both global and local fashions. GNN training with adaptive connection sampling is shown to be mathematically equivalent to an efficient approximation of training Bayesian GNNs. Experimental results with ablation studies on benchmark datasets validate that adaptively learning the sampling rate given graph training data is the key to boost the performance of GNNs in semi-supervised node classification, less prone to over-smoothing and over-fitting with more robust prediction.

Motivation & Objective

• To address over-smoothing and over-fitting in deep GNNs, which limit their performance and depth.
• To enable uncertainty quantification (UQ) in GNN predictions, a capability missing in standard GNNs.
• To unify existing stochastic regularization techniques—like DropEdge and DropOut—under a single adaptive framework.
• To develop a differentiable, learnable connection sampling mechanism that improves GNN robustness and generalization.
• To establish a Bayesian interpretation of GDC, linking it to Monte Carlo uncertainty estimation and Bayesian GNNs.

Proposed method

• Proposes Graph DropConnect (GDC), a stochastic regularization technique that adaptively samples edges in GNNs using learnable drop rates per message-passing channel.
• Models the connection sampling rate as a hierarchical beta-Bernoulli process to enable Bayesian learning of edge dropout probabilities.
• Employs continuous relaxation and the Augment-REINFORCE-Merge (ARM) gradient estimator to enable end-to-end backpropagation through discrete sampling operations.
• Derives a Bayesian approximation of GNN training via GDC, where Monte Carlo inference over sampled graphs yields predictive uncertainty estimates.
• Applies GDC to GCN and GraphSAGE architectures, enabling channel-wise, independent sampling masks for improved representation diversity.
• Introduces block-wise masking to reduce memory usage by sharing sampling masks across feature channels.

Experimental results

Research questions

• RQ1 Can adaptive, learnable connection sampling improve GNN generalization and reduce over-smoothing compared to fixed-rate stochastic regularization?
• RQ2 Can GDC provide meaningful uncertainty quantification in GNN predictions, comparable to Bayesian deep learning in CNNs?
• RQ3 How does GDC compare to existing methods like DropEdge and DropOut in terms of performance and robustness on semi-supervised node classification?
• RQ4 Does the Bayesian interpretation of GDC hold, such that Monte Carlo sampling over graphs approximates posterior predictive distributions?
• RQ5 What is the trade-off between model performance, memory usage, and computational cost when varying the number of feature blocks in GDC?

Key findings

• GCN-BBGDC with adaptive GDC achieves 83.3% accuracy on Cora with 32 blocks, outperforming standard GCN-DO and GCN-DE, which show significant performance drops beyond 4 layers.
• The TV (total variation) of hidden layer outputs in GCN-BBGDC increases during training, indicating reduced over-smoothing, whereas GCN-DO shows decreasing TV, signaling over-smoothing.
• Performance of GCN-BBGDC degrades only slightly with increasing depth (up to 16 layers), while GCN-DO drops to 64.5% accuracy at 16 layers, demonstrating superior robustness.
• Ablation studies show that learning only the DropEdge rate (without DropOut) substantially improves performance, highlighting the importance of adaptive edge sampling.
• Increasing the number of blocks in GDC improves accuracy (from 82.2% with 2 blocks to 83.3% with 32 blocks), confirming enhanced flexibility and representation diversity.
• The predictive uncertainty derived via Monte Carlo sampling in GDC is more reliable than that from standard DropOut, as validated by improved UQ performance.