On Batch-size Selection for Stochastic Training for Graph Neural Networks

1 INTRODUCTION . Training large neural networks is often time consuming . In many real world scenarios training might take hours or even days to converge Radford et al . ( 2018 ) ; Devlin et al . ( 2018 ) . As a consequence , the identification of strategies to reduce the training time while retaining accuracy is an important research objective . The most popular training algorithms for deep learning are Stochastic Gradient Descent ( SGD ) and its variants such as RMSProp or Adam Graves ( 2013 ) ; Kingma & Ba ( 2014 ) . These algorithms work in an iterative manner , such that in each epoch , the data is first partitioned into minibatches and then weight updates are calculated using only the data in each minibatch . It has been observed that the size of the minibatches plays a crucial role in the network ’ s accuracy , generalization capability and converge time ( Keskar et al . ( 2016 ) ; He et al . ( 2019 ) ; McCandlish et al . ( 2018 ) ; Radiuk ( 2017 ) ) . For typical deep learning tasks , practitioners have observed that small batch sizes , e.g. , { 4 , 16 , . . . , 512 } , lead to a better generalization performance and training efficiency Keskar et al . ( 2016 ) . For Graph Neural Networks ( GNNs ) selecting the appropriate batch size remains more of a mystery , and to the best of our knowledge , there has been no published work that focuses on batch size selection for GNNs . The small batch size guidelines for conventional NNs do not carry over because the batches are used to approximate the graph aggregations or convolutions . The approximation error propagates and leads to a much more substantial variance in the gradients than is observed for NNs . In practice , based on released code , we see that implementations tend to either use the largest batch size that can fit into memory Li et al . ( 2020 ) or use a small batch size similar to those for non-graph settings Chen et al . ( 2018 ) ; Zou et al . ( 2019 ) . In this work , we explore the choice of batch size for graph neural networks . By means of a theoretical investigation , we develop guidelines for the choice of batch size that depend on the average degree and number of nodes of the graph . These guidelines lead to intermediate batch sizes , considerably larger than the small NN batch sizes but much smaller than the maximum size dictated by memory limits of a modern GPU . We provide empirical results that demonstrate that the batch sizes derived using our guidelines provide an excellent trade-off between training time and accuracy . Substantially smaller sizes may lead to faster convergence but reduced accuracy ; using larger sizes can achieve similar accuracy but training may take much longer to converge . 2 RELATED WORK . Graph Neural Networks ( GNNs ) have become increasingly popular in addressing graph-based tasks ( Kipf & Welling , 2016 ; Hamilton et al. , 2017 ; Defferrard et al. , 2016 ; Gilmer et al. , 2017 ; Ying et al. , 2018 ) . One major line of research aims to improve the expressiveness of GNNs via 1 ) advanced aggregation functions ( Veličković et al. , 2017 ; Monti et al. , 2017 ; Liu et al. , 2019 ; Qu et al. , 2019 ; Pei et al. , 2020 ) 2 ) deeper architecture ( Li et al. , 2019 ; 2020 ) ; and 3 ) adaptive graph structure ( Li et al. , 2018 ; Vashishth et al. , 2019 ; Zhang et al. , 2019 ) . However , training a large-scale GNN model remains challenging because of the large memory consumption , long convergence time , and heavy computation ( Chiang et al. , 2019 ) . Full-batch gradient descent training scheme was commonly used in the earlier GNN research . While this is suitable for relatively small graphs , it requires storing all intermediate embeddings , which is not scalable for large graphs . The convergence can be slow since the parameters are updated only once per epoch . Hamilton et al . ( 2017 ) and Ying et al . ( 2018 ) proposed the training of GNNs with mini-batch stochastic gradient descent ( SGD ) methods . Mini-batch SGD training suffers from the neighborhood expansion and leads to time-complexity that grows exponential with respect to the GNN depth . To reduce the exponential complexity of receptive nodes , Chen et al . ( 2018 ) , Huang et al . ( 2018 ) , and Zou et al . ( 2019 ) proposed layer-wise sampling , where a fixed number of nodes are sampled in each layer . Importance sampling techniques were incorporated to reduce variance . Unfortunately the overhead of the iterative neighborhood sampling strategy is still significant and becomes worse as GNNs become progressively deeper . Chiang et al . ( 2019 ) and Zeng et al . ( 2020 ) proposed graph-wise sampling to further improve the sampling efficiency . This can be viewed as a special case of layer-wise sampling where the same set of nodes is sampled across all layers . Chen et al . ( 2017 ) and Cong et al . ( 2020 ) proposed variance reduction stochastic training frameworks that maintain a cache for the intermediate embeddings of all nodes . This can improve convergence but results in large memory requirements , stretching the capabilities of GPUs when training over large graphs . Due to this drawback , we do not consider such approaches in this paper , but it is an intrguing direction for future work . Most existing graph neural network papers do not clearly address how they set the batch size . Experimentally , we observe that batch size is a critical hyper-parameter and can significantly influence training time and test accuracy . The importance of the batch size has been recognized for non-graph deep learning models . Keskar et al . ( 2016 ) , He et al . ( 2019 ) and Masters & Luschi ( 2018 ) have shown that smaller batch sizes , in the range { 4 , 16 , . . . , 512 } , can achieve better generalization performance . The randomness of small batches proves beneficial . McCandlish et al . ( 2018 ) suggested that the batch size should be selected so that a balance is achieved between the “ noise ” and “ signal ” of the gradient . Radiuk ( 2017 ) showed that using larger batch sizes , of the order of 1024 , can be beneficial when training convolutional neural network models . Gower et al . ( 2019 ) , Alfarra et al . ( 2020 ) , and Smith ( 2018 ) introduced adaptive batch size approaches to further improve the convergence rate and generalization performance . To the best of our knowledge , no existing work has directly addressed the selection of the batch size for stochastic training for graph neural networks , and the objective of this paper is to fill that gap and provide guidelines for the GNN setting . 3 PRELIMINARIES . We represent a graph G = ( V , E ) with a set of nodes V = { v1 , . . . , vn } and set of edges E = { e1 , . . . , eM } by an adjacency matrix A ∈ Rn×n . For node v ∈ V , we let N ( v ) be the set of neighbors of v. In addition , we associate each node v to a feature vector xv ∈ R1×F , and let X ∈ Rn×F be the corresponding feature matrix . Let D be the degree matrix of the graph G , where Di , i = ∑ j Ai , j and Di , j = 0 if i 6= j . To ease the presentation , we use symbols such as R . T to denote that there exists an absolute constant c such that R ≤ c · T . 3.1 GRAPH NEURAL NETWORK MODELS . Graph neural networks ( GNNs ) can be applied to node prediction , link prediction and graph prediction tasks . In this work , we focus on the node prediction task . We are given labels of nodes from a training set and we need to predict the labels for nodes in a testing set . One paradigm for solving this problem is to learn representations for all the nodes and then map the representations to labels . Graph neural network aggregate the representations of neighbors into each node in order to integrate the graph structure into each node ’ s representation . Specifically , let H l ∈ RN×F l be the representation for layer l , where the ith row hli is the representation for node i at layer l and F l is the dimension of the representation ( H0 is set as the original node features X ) . The forward propagation for hidden states are defined as : H l+1 = σ ( H̃ l+1 ) and H̃ l+1 = ÃH lW l , ( 1 ) where W l ∈ RF l×F l+1 are trainable parameters , σ ( · ) is an activation function , and Ã is a normalization of adjacency matrix A , e.g. , the random walk normalization Ã = D−1A , or the symmetric normalization Ã = D−1/2AD−1/2 . The equations in ( 1 ) can be expressed for each node i as : hl+1i = σ ( h̃ l+1 i ) and h̃ l+1 i = ∑ j∈N ( i ) Ãijh l jW l. ( 2 ) 3.2 STOCHASTIC TRAINING FOR GNN MODELS . Sampling in training GNNs : Label sampling and neighbor sampling . In conventional deep learning models , every sample in a mini-batch contributes independently to the approximated gradient . Including more samples thus reduces the variance of the gradient estimate by statistical power . However , in GNN models , samples in a mini-batch are no longer independent . In fact , we have two different concepts of sampling . First , we sample a mini-batch of nodes in the training set and we call this label sampling . Since the representation of nodes also depends on neighbor nodes , the receptive field for each selected node grows exponentially as the number of layers increases . Neighbor sampling is adopted to constrain the number of receptive neighbor nodes . Existing frameworks use three main approaches to handle sampling for GNNs : node-wise , layer-wise and graph-wise sampling . Node-wise sampling . Hamilton et al . ( 2017 ) and Ying et al . ( 2018 ) adopt a uniformly random sampling of the labels . For neighbor sampling , they recursively sample a certain number of neighbors for each layer . Specifically , to evaluate the aggregation for layer l+1 , for each node i , a set of nodes Sli is sampled from the neighbors of node i , and equation ( 2 ) is evaluated as h̃l+1i = |N ( i ) | |Sli| ∑ j∈Sli Aijh l jW l. ( 3 ) |Sli| is predefined to limit the sampled nodes , but the size of the receptive field for each included mini-batch node still grows exponentially as the number of layers increases . Layer-wise sampling . Chen et al . ( 2018 ) , Huang et al . ( 2018 ) , and Zou et al . ( 2019 ) propose a different neighbor sampling method to reduce the number of receptive nodes . Nodes are sampled at the layer level instead of the node level . Importance sampling is adopted to reduce the variance of sampling and further improve convergence . Specifically , given the set of sampled nodes Sl+1 in layer l+ 1 , nodes in layer l are sampled with some probability distribution ql ( i|Sl+1 ) that is derived from minimizing variance of gradients Chen et al . ( 2018 ) , where i is the index of the node . For brevity , we denote the sampling distribution as ql ( i ) . From Chen et al . ( 2018 ) and Zou et al . ( 2019 ) , the forward propagation in equation ( 2 ) is defined as : h̃l+1i = 1 |Sl| ∑ j∈Sl Aij ql ( i ) hljW l. ( 4 ) Note that by controlling |Sl| , the number of receptive nodes only grow linearly with respect to layer size . In this framework , uniformly random sampling is used for labels . Graph-wise sampling . Zeng et al . ( 2020 ) and Chiang et al . ( 2019 ) propose graph-wise sampling . This can be regarded as a special case of layer-wise sampling , which uses the same set of nodes as both the sampled labels and sampled neighbors across all layers . Importance sampling is adopted in Zeng et al . ( 2020 ) and normalization is applied on both loss and neighbor aggregations to obtain an unbiased estimator for the gradients . Since in practice , layer-wise sampling and graph-wise sampling are much more efficient in practice than node-wise sampling , we focus on the influence of the batch size for layer-wise sampling and graph-wise sampling . For layer-wise sampling , the number of included label samples ( the batch size ) can differ from the number of neighbours sampled at each layer . In our analysis , we focus on the case where these values are equal ; see the supplementary material for further discussion .

Summary: In this paper, the authors studied the problem of batch size selection in graph neural networks. Since scaling up the batch size is the most efficient way to increase parallelism, this topic seems to be very important for making full use of modern computer architectures like GPUs or TPUs. The authors conduct a detailed analysis on the impact of batch size in graph neural networks. After a series of discussions, the authors conclude that an ideal batch size should be n/d where n is the total number of nodes and d is the average node degree. The authors suggested that this batch size should be able to give the best convergence/generalization performance for graph neural networks. Overall, the paper is properly written. In terms of presentation, the structure is clear and the idea is easy to understand.