Training Heterogeneous Graph Neural Networks using Bandit Sampling

Graph neural networks (GNNs) have gained significant attention across diverse areas due to their superior performance in learning graph representations. While GNNs exhibit superior performance compared to other methods, they are primarily designed for homogeneous graphs, where all nodes and edges are of the same type. Training a GNN model for large-scale graphs incurs high computation and storage costs, especially when considering the heterogeneous structural information of each node. To address the demand for efficient GNN training, various sampling methods have been proposed. In this paper, we propose a sampling method based on bandit sampling, an online learning algorithm with provable convergence under weak assumptions on the learning objective. To the best of our knowledge, this is the first bandit-based sampling method applied to heterogeneous GNNs with a theoretical guarantee. The main idea is to prioritize node types with more informative connections with respect to the learning objective. Compared with existing techniques for GNN training on heterogeneous graphs, extensive experiments using the Open Academic Graph (OAG) dataset demonstrate that our proposed method outperforms the state-of-the-art in terms of the runtime across various tasks with a speed-up of 1.5-2x, while achieving similar accuracy.