Adaptive Bandit Cluster Selection for Graph Neural Networks

Ghadir Ayache, Thomas Hugues, Chris Xu, Julia Zhang, Diane Hu

Abstract: Graph neural networks (GNNs) have shown successful performance in graph-based applications; however, training large-scale GNNs remains challenging. Current SGD-based algorithms suffer from a high computational cost that exponentially grows with the number of GNN layers. Moreover, they require a large memory to store both the entire graph and the embedding of each node. In this work, we propose a novel, bandit-based graph sampling scheme that improves GNN training efficiency, which results in faster training times and improved accuracy. Our sampling strategy consists of partitioning the graph into densely connected clusters and prioritizing more informative clusters with respect to the learning objective. The cluster selection is carried by a bandit-based graph sampling algorithm that defines the reward based on the robustness of the induced stochastic gradient. Our goal is to select more often the clusters that guide our optimization objective towards the optimal solution more quickly by processing snapshots of the graph while maintaining a theoretical guarantee of approaching optimal learning performance. We employ our algorithm on a recommendation module at the e-commerce platform Etsy and on the public benchmark Yoochoose dataset, where we show an improvement in the rate of convergence and the recommendation accuracy over graph sampling baselines.