BLISS: Bandit Layer Importance Sampling Strategy for Efficient Training of Graph Neural Networks

TL;DR

BLISS presents a bandit-based layer-wise importance sampling strategy to train GNNs efficiently on large graphs. By modeling neighbor selection as a dynamic EXP3-based optimization, BLISS adaptively allocates sampling across edges and layers to minimize estimator variance while maintaining accuracy comparable to full-batch training. The approach is compatible with both GCNs and GATs, and adaptions for attentive mechanisms are proposed, including skip-connection ideas to preserve attention semantics. Empirical results on six datasets show BLISS achieving or surpassing full-batch performance with scalable training times, particularly excelling on smaller or highly heterogeneous graphs and when employing attention mechanisms.

Abstract

Graph Neural Networks (GNNs) are powerful tools for learning from graph-structured data, but their application to large graphs is hindered by computational costs. The need to process every neighbor for each node creates memory and computational bottlenecks. To address this, we introduce BLISS, a Bandit Layer Importance Sampling Strategy. It uses multi-armed bandits to dynamically select the most informative nodes at each layer, balancing exploration and exploitation to ensure comprehensive graph coverage. Unlike existing static sampling methods, BLISS adapts to evolving node importance, leading to more informed node selection and improved performance. It demonstrates versatility by integrating with both Graph Convolutional Networks (GCNs) and Graph Attention Networks (GATs), adapting its selection policy to their specific aggregation mechanisms. Experiments show that BLISS maintains or exceeds the accuracy of full-batch training.

Figures (3)

• Figure 1: The figures illustrate Node-wise vs. Layer-wise sampling. Left: Node-wise sampling selects nodes per target node, often causing redundancy (e.g., $v_4$ sampled for both $u_1$ and $u_2$), higher sampling rates (e.g., $v_4$, $v_5$), and missing edges (e.g., $u_2$–$v_4$). Right: Layer-wise sampling considers all nodes in the previous layer, preserving structure and connectivity while sampling fewer nodes.
• Figure 2: Validation Accuracy across six datasets (Citeseer, Cora, Pubmed, Flickr, Yelp, and Reddit) for BLISS and PLADIES samplers using Graph Attention Networks (GAT) and GraphSAGE (SAGE) architectures. The figure highlights the performance trends during training averaged over 5 runs. The shaded regions represent the standard deviation across runs.
• Figure 3: Validation Loss for the same datasets and models as in \ref{['fig:val_f1']}. The figure illustrates the loss trends during training, averaged over 5 runs. The shaded regions represent the standard deviation across runs.