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

AI summary

Overview

Research area: Graph Neural Networks (GNNs), specifically efficient sampling methods for scalable training on large graphs, with multi-armed bandits drawn from online learning.

Technical level: Intermediate. Readers need basic familiarity with GNN message passing and Monte Carlo estimation, but the core idea is intuitive without deep math.

Scope: The paper proposes BLISS, a layer-wise neighbor sampling strategy that uses adversarial bandits (EXP3) to adaptively select the most informative nodes at each layer, evaluated against the PLADIES sampler on six benchmark graphs with GraphSAGE and GATv2 backbones.

What This Paper Is About

Training GNNs on large graphs is expensive because every node must aggregate information from all its neighbors, and the receptive field grows exponentially with depth (neighborhood explosion). Existing approaches sample neighbors to cut this cost, but most rely on static importance scores computed once and never updated, which can systematically under-sample nodes that become important as embeddings evolve. BLISS reframes layer-wise neighbor selection as a bandit problem, so the sampling distribution continuously adapts to which neighbors actually reduce estimator variance during training.

Key Contributions

1. Bandit formulation of layer-wise sampling. The authors model each edge as a bandit "arm" and define the reward as the neighbor's contribution to reducing the variance of the node representation estimator. This extends prior node-wise bandit sampling (BS-GNN) to the layer-wise regime, where sampling decisions are made jointly across all nodes in a layer, reducing redundancy and preserving connectivity.

2. The BLISS algorithm. A top-down procedure that computes layer-wise node sampling probabilities from per-node edge sampling distributions, selects k nodes per layer, runs a GNN forward pass, then updates edge weights via EXP3 and recomputes sampling probabilities. The exploration rate η controls the balance between exploiting high-reward neighbors and exploring under-sampled regions.

3. Adaptation to attentive architectures. Because attention coefficients cannot be computed exactly from a sampled subset, BLISS introduces an adjusted feedback attention term that uses the bandit sampling probabilities as a surrogate for normalization over the full neighborhood. The authors also extend PLADIES to GATs by forcing seed nodes to always be selected (infinite sampling probability), creating skip connections that guarantee every node retains at least one neighbor for attention to operate on.

4. Theoretical and empirical comparison. Complexity, memory, and variance analyses for full-batch, GraphSAGE, LADIES, and BLISS, plus experiments across Cora, Citeseer, Pubmed, Flickr, Yelp, and Reddit demonstrating that BLISS maintains or exceeds full-batch accuracy while running with layer-wise sampling overhead.

Main Findings

• BLISS generally matches or beats the state-of-the-art layer-wise baseline. Across six datasets and two architectures, BLISS achieves higher test F1 than PLADIES in most configurations, with GAT showing the most consistent gains.

• Largest gains appear on attentive models and smaller or heterogeneous graphs. With GAT on Citeseer, BLISS reaches 70.6% test F1 versus PLADIES at 68.3%; on Pubmed, 73.1% versus 71.8%. On Yelp with GraphSAGE, BLISS scores 52.9% versus 50.2%, suggesting the adaptive policy captures nuanced class relationships in heterogeneous graphs.

• Gains are marginal or absent on dense, uniform graphs. On Flickr and Reddit the two samplers are essentially tied (e.g., Reddit SAGE: 96.2% for both; Reddit GAT: 94.9% versus 95.0%). On Citeseer and Flickr with GraphSAGE, PLADIES is marginally higher, so the claim of consistent superiority holds only for the GAT experiments and the aggregate picture.

• Cross-architecture comparison favors BLISS. Comparing BLISS on GAT against the original PLADIES on GraphSAGE (the architecture it was designed for) shows large gaps: Citeseer 70.6% versus 60.1%, Pubmed 73.1% versus 55.7%.

• Theoretical variance reduction. The variance bound for BLISS includes a (1−η)⁻¹ factor arising from the bandit exploration term, and the analysis argues that dynamically prioritizing informative neighbors lowers estimator variance relative to static importance sampling.

• Runtime overhead is modest. Per-iteration training time for BLISS is slightly higher than PLADIES across all datasets (e.g., Citeseer GAT 0.065s versus 0.055s; Reddit GAT 0.207s versus 0.156s). The authors note the implementation uses naive Python loops and is unoptimized.

• One exception to the general trend. GAT on the Yelp dataset overfitted for both samplers, so those results were left out of the evaluation to keep conditions uniform.

Methodology in Plain English

The authors start from the observation that a GNN layer's output is an expected value over neighbors. If you sample a subset of neighbors and correct each sample by the inverse of its sampling probability, you get an unbiased estimate of that expectation, and the quality of the estimate depends on how well the sampling distribution matches the true importance of each neighbor.

BLISS works from the last layer back to the first. At each layer, it maintains a weight for every edge, initialized uniformly. These weights are converted into per-node sampling probabilities, which are then aggregated into a layer-wide probability distribution over all candidate nodes in the previous layer. A fixed number of nodes is drawn from that distribution, and the GNN performs its normal forward pass using only the sampled nodes.

After the forward pass, BLISS computes a reward for each sampled edge based on how much that neighbor contributed to the node's representation, weighted by the attention or aggregation coefficient and normalized by the sampling probability. These rewards feed into the EXP3 bandit update, which multiplicatively increases weights on high-reward edges and mixes in a small amount of uniform exploration to prevent any neighbor from being starved. The updated weights become the sampling distribution for the next training step, so the policy tracks which neighbors matter as embeddings change.

For attention-based models, the authors handle the fact that the normalization term in softmax attention depends on all neighbors, not just the sampled ones. They approximate it by summing the sampling probabilities of the selected neighbors, and they weight the unnormalized attention scores accordingly. For the PLADIES extension, they force seed nodes to be selected with probability one so that every node in the sampled subgraph has a valid attention context.

Evaluation uses 3-layer GNNs with hidden dimension 256, ADAM at learning rate 0.002, and bandit hyperparameters η = 0.4 and δ = η/10⁶. Each configuration is run five times with different seeds, reporting mean and standard deviation of test F1.

Why This Matters

Impact on research. Sampling is one of the central bottlenecks for scaling GNNs to production-sized graphs, and the layer-wise category has largely relied on static importance heuristics since LADIES. BLISS shows that treating the sampling distribution as something to be learned online, rather than computed once, yields measurable accuracy improvements without abandoning the linear-time complexity that makes layer-wise sampling attractive. It also demonstrates that bandit methods transfer cleanly from node-wise to layer-wise sampling, opening a bridge between two previously separate lines of work.

Real-world applications:

• Recommendation systems, where user-item interaction graphs can have billions of edges and neighbor importance shifts as user behavior evolves, making adaptive sampling more valuable than fixed importance scores.
• Drug discovery and molecular property prediction, where graphs are smaller but heterogeneous and class relationships are nuanced, matching the regime where BLISS shows its largest gains (Yelp, Pubmed).
• Social network and fraud detection, where dense uniform subgraphs coexist with sparse informative structures often found in anomalous communities.
• GNN-augmented language models, an application area the authors explicitly name as a future direction, where retrieval graphs over documents or entities are large and dynamic.

Industry relevance. Any deployment where full-batch GNN training is infeasible but accuracy cannot be sacrificed stands to benefit. The modest per-iteration overhead reported (roughly 15–30% in an unoptimized implementation) is a reasonable trade for the accuracy gains on attention-based models, and the method integrates into existing sampling pipelines without architectural changes.

Future Directions

• Advanced bandit algorithms. The paper uses EXP3, a relatively simple adversarial bandit method. Contextual or combinatorial bandits (CMAB) could exploit structure across edges and layers to converge faster or produce tighter sampling policies, which the authors flag as an explicit next step.

• Optimized implementation. The current prototype uses naive Python loops, and the reported runtimes include that overhead. A vectorized or C++ implementation could narrow or close the gap with PLADIES and make the method practical at much larger scale.

• Extension to GNN-augmented LLMs and vision tasks. The authors name these domains as targets, but the paper provides no experiments outside standard node classification benchmarks. Whether the bandit reward signal transfers to retrieval graphs or vision GNNs remains open.

• Explaining the architecture-dependence of gains. BLISS's advantage is largest on GAT and on heterogeneous graphs, and near zero on Reddit and Flickr. The paper attributes this to variance reduction and adaptive importance, but a more precise characterization of when dynamic sampling pays off, and when the added complexity is not worth it, would help practitioners decide whether to adopt it.

Target Audience

Researchers and practitioners working on scalable GNN training, particularly those already using or evaluating layer-wise samplers such as LADIES, PLADIES, or LABOR. It is also relevant to readers interested in applying bandit methods to systems problems where the reward signal is noisy and non-stationary. The paper assumes comfort with GNN message passing, importance sampling, and basic probability, so it is best suited to graduate students and engineers with some machine learning background rather than complete newcomers. Readers focused on subgraph sampling (GraphSAINT, Cluster-GCN) or node-wise sampling (GraphSAGE) will find the comparison setup less directly applicable, since the authors deliberately restrict their baseline to the layer-wise category.

Authors’ 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.