Influence-Based Mini-Batching for Graph Neural Networks

TL;DR

This paper introduces influence-based mini-batching (IBMB), a novel method for constructing mini-batches in graph neural networks that significantly accelerates inference and training by leveraging influence scores to optimize batch composition.

Contribution

The paper proposes IBMB, a new approach that models batch construction based on influence scores, leading to faster inference and training without sacrificing accuracy.

Findings

IBMB accelerates inference by up to 130x.

IBMB speeds up training by up to 18x per epoch.

IBMB achieves up to 17x faster convergence per runtime.

Abstract

Using graph neural networks for large graphs is challenging since there is no clear way of constructing mini-batches. To solve this, previous methods have relied on sampling or graph clustering. While these approaches often lead to good training convergence, they introduce significant overhead due to expensive random data accesses and perform poorly during inference. In this work we instead focus on model behavior during inference. We theoretically model batch construction via maximizing the influence score of nodes on the outputs. This formulation leads to optimal approximation of the output when we do not have knowledge of the trained model. We call the resulting method influence-based mini-batching (IBMB). IBMB accelerates inference by up to 130x compared to previous methods that reach similar accuracy. Remarkably, with adaptive optimization and the right training schedule IBMB can also substantially accelerate training, thanks to precomputed batches and consecutive memory accesses. This results in up to 18x faster training per epoch and up to 17x faster convergence per runtime compared to previous methods.