Locality-Based Mini Batching for Graph Neural Networks

1 Introduction . Modern neural networks commonly use stochastic mini-batch training to leverage large datasets and accelerate convergence . This strategy becomes highly non-trivial for connected data , since creating mini batches requires selecting a meaningful subset from the dataset , despite its connectedness . Graph neural networks ( GNNs ) typically rely on sampling a set of nodes from the graph to resolve this issue . However , graph sampling requires non-contiguous memory accesses , which significantly slows down training for large datasets . This severely limits their applicability to real-world graphs , which often consist of millions or billions of nodes . The main question behind the connected mini-batching problem is : How do we choose the best nodes for constructing the next mini batch ? To answer this question , we introduce the concept of primary and auxiliary nodes . Primary nodes are those for which we compute a prediction in this batch , typically a set of training nodes . Auxiliary nodes only help with computing the primary nodes ’ outputs . This distinction allows us to provide a meaningful neighborhood for every node ’ s prediction , while ignoring irrelevant parts of the graph . Note that primary nodes in one batch can be auxiliary nodes in another batch . This distinction splits the main question into two parts : 1 . How do we choose the primary nodes for a mini batch ? 2 . How do we choose the auxiliary nodes for a given set of primary nodes ? Having split the problem like this , we see that most previous works actually only focus on the second question and just choose a uniformly random subset of nodes as primary nodes ( Hamilton et al. , 2017 ; Zou et al. , 2019 ) . However , the nodes in a graph are not independent , and a better approach can lead to large improvements in runtime . In this work , we propose one simple answer for both questions : Locality . Choosing a set of locally connected nodes is advantageous both from a computational and a predictive perspective . It allows us to share computation between multiple nodes , keeps memory accesses local , and reduces the mini batch ’ s memory footprint . The predictions of most GNNs already leverage the fact that nearby nodes are more important . Previous works have even shown that incorporating locality can improve GNN accuracy ( Klicpera et al. , 2019b ; Huang et al. , 2021 ) . Batching nearby primary nodes together can create synergies since one primary node can leverage another one ’ s auxiliary nodes . More specifically , we propose to use either graph partitioning or node distances to select mini batches of primary nodes , and local clustering to select their auxiliary nodes . We then use the subgraph induced by these nodes as a mini batch . See Fig . 1 for an overview of this process . Importantly , these mini batches can be computed a priori , and loaded from a cache to ensure efficient memory access . During training , we counteract the effect of correlated mini batches with learning rate and batch scheduling . Overall , our method achieves an up to 20x improvement in time per epoch , with similar final accuracy . This faster time per epoch more than makes up for any slow-down in convergence per step . Its speed advantage grows even further for the common setting of low label ratios , since our method avoids computation on irrelevant parts of the graph . As opposed to most previous works , our method can successfully be used for both training and inference . It accelerates inference by up to 100x compared to previous methods that achieve similar accuracy . In summary , our core contributions are : • Locality-based mini batching ( LBMB ) : A general mini-batching method that works for a variety of GNNs and datasets . It substantially accelerates both training and inference without sacrificing accuracy , especially for small label ratios . • We examine the impact of fixed , correlated mini batches on gradient estimation , and propose training methods to mitigate these effects . • An extensive and fair experimental evaluation of scalable training methods , covering four datasets , two GNNs , and both training and inference . 2 Background . Graph neural networks . We consider a graph G = ( V , E ) with node set V and ( possibly directed ) edge set E . N = |V| denotes the number of nodes , E = |E| the number of edges , and A ∈ RN×N the adjacency matrix . GNNs use one embedding per node hu ∈ RH and edge e ( uv ) ∈ RHe , and update them in each layer via message passing between neighboring nodes . Most GNNs can be expressed via the following equations : h ( l+1 ) u = fnode ( h ( l ) u , Agg v∈Nu [ fmsg ( h ( l ) u , h ( l ) v , e ( l ) ( uv ) ) ] ) , ( 1 ) e ( l+1 ) ( uv ) = fedge ( h ( l+1 ) u , h ( l+1 ) v , e ( l ) ( uv ) ) . ( 2 ) The node and edge update functions fnode and fedge , and the message function fmsg can be implemented using e.g . linear layers , multi-layer perceptrons ( MLPs ) , and skip connections . The node ’ s neighborhood Nu is usually defined directly by the graph G ( Kipf & Welling , 2017 ) , but can be generalized to consider larger or even global neighborhoods ( Klicpera et al. , 2019b ; Alon & Yahav , 2021 ) , or feature similarity ( Deng et al. , 2020 ) . The most common aggregation function Agg is summation , but mean , min , standard deviation , and other alternatives have also been explored ( Corso et al. , 2020 ; Geisler et al. , 2020 ) . Edge embeddings e ( uv ) are not present in many GNNs , but some variants rely on them exclusively ( Chen et al. , 2019 ) . GNNs commonly use node features X ∈ RN×F for the input embeddings h ( 0 ) u , but they can also be augmented with positional encodings ( You et al. , 2019 ; Dwivedi et al. , 2020 ) or node IDs ( Vignac et al. , 2020 ) . See App . A for related work in scalable GNNs . 3 Locality-based mini batching . An efficient GNN mini-batching method needs to consider both computational aspects and convergence per gradient step . Previous works have largely been focused on sampling methods that improve GNN convergence , and treated computational issues mostly as an afterthought . By focusing more on computational aspects like memory access times , we can accelerate training by multiple orders of magnitude , and therefore more than make up for any disadvantage in convergence per step . If we skip details like caching the computational perspective is rather simple : Save as many operations and memory accesses as possible while maintaining a good GNN approximation . When the auxiliary nodes of different primary nodes in a batch are shared , we only have to compute their embeddings once and save computation and memory accesses . We therefore aim to group together primary nodes in a way that maximizes the number of shared auxiliary nodes . In locality-based mini batching ( LBMB ) , we do this in two steps : 1 . Obtain the k most important auxiliary nodes for each primary node . Batch together primary nodes so that the union of their auxiliary nodes is smallest , yielding the primary node partition PI . 2 . Select the most important auxiliary nodes SII for each subset of primary nodes SI ∈ PI . This process yields a single , fixed set of batches . We thus only need to perform it once during preprocessing . We then cache each mini batch in consecutive blocks of memory , thereby circumventing expensive random data accesses . This significantly accelerates training , allows efficient distributed training , and enables even expensive node selection procedures . In contrast , most previous methods select both primary and auxiliary nodes randomly in each epoch ( Ying et al. , 2018 ; Zeng et al. , 2020 ) , which incurs significant overhead . Our experiments show that our more efficient memory accesses clearly outweigh the slightly better gradient estimates gained from re-sampling in each epoch ( see Sec . 5 ) . We will next describe the details of LBMB ’ s two main steps . 3.1 Primary node partitioning . Optimal partitioning . In this step we are interested in finding the partition PI with the highest number of shared auxiliary nodes . Naïvely , we could find PI by comparing the auxiliary node overlaps achieved by every possible partition . This is clearly intractable since the number of partitions increases exponentially with the number of primary nodes . An obvious way of accelerating this is a greedy approach , in which we iteratively put those nodes into a batch that has the largest overlap . Unfortunately , this would still require computing the overlap for every pair of primary nodes . This requires a quadratic runtime O ( N2 ) , which is intractable for large datasets . We must therefore rely on heuristics to obtain a scalable , well-performing primary node partitioning algorithms . Distance-based partitioning . We propose two methods that leverage graph locality as a heuristic for effective node partitioning . The first is based on node distances . In this approach we first compute the pairwise node distances between nodes that are close in the graph . A common node distance measure in undirected graphs are random walks with restart , or personalized PageRank ( PPR ) ( Page et al. , 1998 ) . The PPR matrix is given by Πppr = α ( IN − ( 1− α ) D−1A ) −1 , ( 3 ) with the teleport probability α ∈ ( 0 , 1 ] and the diagonal degree matrix Dij = ∑ k Aikδij . The proximity between nodes u and v is then given by Πppruv . Calculating this inverse is obviously infeasible . However , we can approximate Πppr with a sparse matrix Π̃ppr in time O ( 1εα ) per column , with error εdeg ( u ) ( Andersen et al. , 2006 ) . Next , we greedily construct the partition PI from Π̃ppr . To do so , we start by putting every node u into a separate batch { u } . We then sort all elements in Π̃ppr by magnitude , independent of their row or column . We scan over these values in descending order , considering the value ’ s indices ( u , v ) and merging the batches containing the two nodes . Finally , we randomly merge any small leftover batches . We stay within a memory constraint by only merging batches that stay below a maximum batch size . Note that the resulting partition is unbalanced . This method achieves well-overlapping batches and can efficiently add incrementally incoming primary nodes , e.g . in a streaming setting . Our experiments show that this method achieves a good compromise between well-overlapping batches and good gradients for training ( see Sec . 5 ) . Graph partitioning . For our second method , we note that partitioning primary nodes into overlapping mini batches is closely connected to partitioning graphs . We can thus leverage the extensive amount of research on this topic by using the METIS graph partitioning algorithm ( Karypis & Kumar , 1998 ) to find a partition of primary nodes PI . Note that this approach completely skips step 1 in the above LBMB process . We found that graph partitioning yields roughly a two times higher overlap of auxiliary nodes than distance-based partitioning , thus leading to significantly more efficient batches . However , it also introduces a strong bias in batch selection that we found to be detrimental for training ( see Sec . 5 ) . Note that LBMB with graph partitioning is closely related to Cluster-GCN ( Chiang et al. , 2019 ) . However , in contrast to Cluster-GCN our method ignores irrelevant parts of the graph and can obtain overlapping mini batches . This significantly accelerates training on small training sets and improves the accuracy of primary nodes close to the partition boundary .

Summary: This paper proposes locality-based mini batching (LBMB), a method for extracting batches for GNN training. LBMB aims to reduce the computation cost of GNN training by treating primary (output) nodes and auxiliary nodes separately.  LBMB can be used for both training and inference, and speeds up both tasks substantially while maintaining comparable accuracy.