Source-linked AI summary
Bayesian Graph Neural Networks with Adaptive Connection Sampling
Arman Hasanzadeh, Ehsan Hajiramezanali, Shahin Boluki, Mingyuan Zhou, Nick Duffield, Krishna Narayanan, Xiaoning Qian
TL;DR
GNNs struggle with deep-model over-smoothing, over-fitting, and uncertainty quantification. The paper introduces adaptive connection sampling, Graph DropConnect, whose Bayesian approximation and benchmark results support more robust semi-supervised node classification.
Problem
Deep GNNs are limited by over-smoothing and over-fitting, while current GNN implementations do not provide uncertainty quantification.
Method
Graph DropConnect adaptively samples connections and learns sampling rates jointly with GNN parameters, with a Bayesian approximation interpretation.
Results
BBGDC outperforms state-of-the-art stochastic regularization techniques in accuracy across all benchmark datasets.
Takeaways & Limitations
Adaptive sampling can improve semi-supervised node classification while alleviating over-smoothing and over-fitting and enabling predictive uncertainty estimation.
Takeaways & Limitations
Learning Bernoulli parameters remains challenging, and sequential ARM optimization is left for future study.
Abstract
from arXiv · showhide
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.
1. Introduction
GNNs achieve strong graph-learning results but are limited by over-smoothing, over-fitting, and absent uncertainty quantification. GDC addresses these issues through adaptive connection sampling and a Bayesian approximation interpretation.
- GNNs face two central limitations: deep models suffer over-smoothing and over-fitting, while standard implementations lack output uncertainty quantification.
- Existing remedies include DropOut, DropEdge, and node sampling, but graph topology makes independent feature dropout insufficient for controlling over-smoothing.Combining DropEdge with DropOut has shown potential to alleviate these problems.
- GDC unifies stochastic GNN regularization by adaptively sampling connections, with existing DropOut, DropEdge, and node sampling as special cases.It regularizes neighborhood aggregation separately at each channel.
- Learning GDC sampling rates jointly with GNN parameters improves task-specific regularization, including substantial gains from learning only the DropEdge rate without DropOut.
- GDC provides a Bayesian approximation in which Monte Carlo GNN outputs estimate predictive posterior uncertainty, extending the interpretation to neighborhood-sampling models such as GraphSAGE.
2. Preliminaries
Bayesian neural networks provide uncertainty estimates but are computationally difficult, motivating approximate methods such as dropout. GNNs additionally use graph diffusion and stochastic sampling, which can improve efficiency while creating over-smoothing and over-fitting concerns.
- Bayesian neural networks: Bayesian neural networks model uncertainty by placing priors over parameters, but Bayesian inference remains computationally intractable and approximate methods are costlier than non-Bayesian models.
- Dropout: Test-time dropout enables uncertainty quantification by interpreting network outputs as Monte Carlo samples from a predictive distribution.
- GNN foundations: GCN graph convolution performs Laplacian smoothing by mixing node and neighbor features, improving classification but potentially causing over-smoothing.
- GNN foundations: Over-parameterized GCNs can overfit limited training data, fitting training examples well while generalizing poorly to test data.
- Stochastic regularization & reduction: Deep GNN performance is constrained by increasing model complexity and empirically observed over-fitting and over-smoothing on large graphs.
- Stochastic regularization & reduction: DropOut, DropEdge, random-walk neighborhood sampling, and node sampling address deep-GNN limitations through stochastic regularization or computational reduction.
3. Graph DropConnect
Graph DropConnect (GDC) is an adaptive connection-sampling framework that unifies existing stochastic regularization methods and provides a Bayesian approximation for GNNs.
- GDC applies independent random masks to edges for each channel, regularizing neighborhood aggregation at every GNN layer.Its sparse mask has Bernoulli-distributed nonzero elements, with layer-specific sampling rates π_l.
- GDC generalizes DropOut, DropEdge, and node sampling by allowing flexible masks over channels, edges, and nodes.The framework adds a free parameter for adjusting the binary mask across these connection types.
- Connection sampling can be transformed into edge-specific random weights, so GDC learns different message-passing weights for each graph edge.The edge set includes input-graph edges and self-loops, and each edge is paired with a corresponding weight parameter.
- GDC defines an approximating variational distribution for the Bayesian GNN posterior, with KL divergence serving as a training regularizer.The factorized formulation includes weight decay and an additional regularization term for the sampled connections.
- When the sampling rate is optimized, the Bernoulli entropy term matters and can push the drop rate toward 0.5, which may be undesirable in some cases.This pressure is absent when π_l is fixed and the entropy term can be omitted.
- Adaptive GDC learns sampling rates jointly with GNN parameters from graph data rather than fixing them as hyperparameters.Optimizing binary masks is challenging because standard reparameterization is not directly applicable and score-function estimators can have high variance.
4. Variational Beta-Bernoulli GDC
This section develops a hierarchical beta-Bernoulli formulation of GDC that learns connection drop rates jointly with GNN parameters. It derives variational optimization strategies while noting computational and estimator trade-offs.
- Adaptive formulation: GDC can learn drop rates jointly with GNN parameters instead of fixing them as hyperparameters.The adaptive mechanism uses graph data and is framed as a hierarchical Bayesian construct.
- Variational formulation: Connection masks can enforce sparsity in the sampled adjacency matrices, a property identified as necessary for regularizing deep GNNs.The masked adjacency is formed from A and the random mask Z^(l).
- Inference: Hierarchical beta-Bernoulli GDC makes Gibbs-sampling inference computationally demanding on large graph datasets, motivating efficient variational inference.The computational burden is stated specifically for large datasets.
- Variational formulation: The variational distribution factorizes over layerwise drop rates and conditionally independent edge masks, using Kumaraswamy and Bernoulli components.The formulation defines q(Z^(l), π_l) = q(Z^(l) | π_l)q(π_l), with edge independence given π_l.
- Optimization: Concrete relaxation enables efficient optimization of binary masks but introduces bias, whereas ARM directly optimizes Bernoulli parameters with unbiased, low-variance gradients.ARM requires two forward passes, creating a computational trade-off.
5. Connection to Random Walk Sampling
This section connects GDC with random-walk neighborhood sampling in GNNs. It shows that random-walk-sampled GNN outputs admit a Bayesian approximation interpretation when evaluated with Monte Carlo sampling.
- Random-walk formulation: Random-walk sampling reduces input-graph size and model complexity in GNNs such as GraphSAGE.The sampled neighborhoods form connected subgraphs.
- Variational structure: In the random-walk setting, GDC masks depend on previous-layer masks because each sampled neighborhood is a connected subgraph.This dependence yields an autoregressive variational decomposition.
- Bayesian interpretation: GraphSAGE and other random-walk-sampled GNNs are Bayesian GNN approximations when their outputs are calculated using Monte Carlo sampling.With fixed Bernoulli parameters, weight gradients can be estimated by Monte Carlo integration.
6. Sampling Complexity
This section analyzes GDC’s sampling cost and describes strategies to reduce computation. It also identifies the benchmark table as reporting semi-supervised node-classification accuracy.
- Sampling complexity: Variational inference can require |E| × f_l × f_(l+1) random samples per GNN layer, or |E| × f_l in constrained GDC.Large graphs or many filters can therefore increase space complexity and computation time.
- Complexity reduction: Sampling one mask for a feature block reduces the number of draws compared with sampling separately for every feature.The paper proposes blockwise sampling to address high memory and computation costs.
- Implementation trade-offs: Multiplying masks with a pre-computed normalized adjacency matrix reduces computation time with negligible experimental performance effect.Asymmetric masks would instead increase the number of samples by a factor of two.
7. Numerical Results
Across citation-graph benchmarks, BBGDC improves node-classification accuracy over stochastic regularization baselines, while adaptive connection sampling supports uncertainty estimation and deeper GNNs with less over-smoothing.
- 7.1. Numerical Results: BBGDC outperforms state-of-the-art stochastic regularization techniques in accuracy across all benchmark datasets.The comparison uses semi-supervised node classification on Cora, Citeseer, and Cora-ML.
- 7.1.2. Discussion: Learning drop rates improves BBDE over DropEdge and DropOut- plus-DropEdge baselines, while GDC’s greater flexibility yields further improvement.This ablation separates the contribution of learnable sampling rates from the flexibility of adaptive connection sampling.
- 7.2. Uncertainty Quantification: GCN-BBGDC consistently outperforms GCN-DO in PAvPU across uncertainty thresholds from 0.5 to 1 of maximum predictive uncertainty.The evaluation uses 20 Monte Carlo samples and predictive entropy on Cora.
- 7.3. Over-smoothing and Over-fitting: GCN-BBGDC maintains increasing hidden-layer total variation during training, unlike GCN-DO’s decreasing second-layer variation associated with poor performance.Lower total variation indicates that neighboring-node representations are closer and may reflect over-smoothing.
- 7.3. Over-smoothing and Over-fitting: BBGDC accuracy improves through four hidden layers and then declines more gradually than competing methods as depth increases to 16 layers.For GCN-DO, accuracy drops to 69.50% at 8 layers and 64.5% at 16 layers.
- 7.4. Effect of Number of Blocks: Accuracy increases with the number of feature blocks in 4-layer GCN-BBGDC, creating a performance–memory and computational-complexity trade-off.The evaluated configurations use 2, 16, and 32 blocks on Cora.
8. Conclusion
The paper presents GDC as a unified adaptive connection-sampling framework that generalizes stochastic GNN regularization, mitigates over-smoothing and over-fitting, and supports uncertainty estimation through a Bayesian approximation.
- GDC jointly trains connection-sampling parameters with GNN parameters, rather than using fixed sampling rates, and approximates Bayesian GNN training.
- Experiments show that GDC improves semi-supervised node-classification performance while alleviating over-smoothing and over-fitting.
- GDC produces better uncertainty quality than DropOut in GNNs.
A. Ablation Study: Global versus Local
The study extends learnable GDC from a shared layer-level sampling distribution to edge- and layer-specific distributions, but local learning encounters posterior collapse.
- A. Ablation Study: Global versus Local: Local learnable GDC assigns a distinct connection-sampling distribution to each edge at each layer, generalizing one distribution shared across all edges in a layer.Its variational beta-Bernoulli formulation and KL term follow the global learnable GDC derivation.
- A. Ablation Study: Global versus Local: Training local learnable GDC on citation datasets reduces accuracy while driving KL divergence to zero for every tested prior, indicating posterior collapse.The passage attributes this issue to over-parametrization in the local learnable GDC.
- A. Ablation Study: Global versus Local: The supplied dataset-statistics reference identifies Table 4 but does not provide its numerical contents.No dataset values can be recovered from the supplied passage.
B. Datasets and Implementation Details
The experiments use PyTorch implementations, a single RTX 2080 GPU node, and three citation-network benchmark datasets with standard preprocessing and fixed data splits.
- B. Datasets and Implementation Details: All models are implemented in PyTorch and simulations run on a single NVIDIA GeForce RTX 2080 GPU node.The evaluated methods include GCN-BBDE, GCN-BBGDC, and baseline models.
- B. Datasets and Implementation Details: The evaluation uses three standard citation-network benchmark datasets with preprocessing and splits following prior work.For Cora and Cora-ML, the split uses 140 training, 500 validation, and 1,000 test nodes.
C. GDC versus Other Stochastic Regularization Techniques
GDC samples individual connections in GCN layers, contrasting with feature dropout, edge dropout, and node sampling, which remove features, node edges, or whole nodes.
- C. GDC versus Other Stochastic Regularization Techniques: The GCN layer’s baseline connections are sparse because they follow the input graph topology, with node-specific connections highlighted across layers.The schematic uses four nodes and two input and output features.
- C. GDC versus Other Stochastic Regularization Techniques: GDC independently drops connections across layers, distinguishing it from methods that remove features, edges between nodes, or nodes.Kept connections are shown in gray and dropped connections with dashed lines in the proposed schematic.
- C. GDC versus Other Stochastic Regularization Techniques: DropOut drops features at each layer, so the corresponding feature connections are removed while other features remain active.The schematic represents dropped features as faded circles and dropped connections as dashed lines.
- C. GDC versus Other Stochastic Regularization Techniques: DropEdge drops edges between nodes, removing all connections between their corresponding channels with a symmetric mask.The example drops the edges between nodes 1 and 2 and between nodes 1 and 4.
- C. GDC versus Other Stochastic Regularization Techniques: FastGCN drops nodes, thereby removing every connection incident to each sampled-out node.The figure depicts dropped nodes as faded nodes and their connections as dashed lines.