Locality-Based Mini Batching for Graph Neural Networks

Johannes Klicpera, Chendi Qian, Stephan Günnemann

Keywords: GNN, graph neural network, graphs, scalability, batching, local clustering

Abstract: Training graph neural networks on large graphs is challenging since there is no clear way of how to extract mini batches from connected data. To solve this, previous methods have primarily relied on sampling. While this often leads to good convergence, it introduces significant overhead and requires expensive random data accesses. In this work we propose locality-based mini batching (LBMB), which circumvents sampling by using fixed mini batches based on node locality. LBMB first partitions the training/validation nodes into batches, and then selects the most important auxiliary nodes for each batch using local clustering. Thanks to precomputed batches and consecutive memory accesses, LBMB accelerates training by up to 20x per epoch compared to previous methods, and thus provides significantly better convergence per runtime. Moreover, it accelerates inference by up to 100x, at little to no cost of accuracy.

One-sentence Summary: Locality-based mini batching enables large-scale training and inference for graph neural networks by precomputing fixed mini-batches based on node distances, graph partitioning, and local clustering.