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: The goal of the paper is to propose a principled strategy to select batch size for training graph neural networks with SGD. Training (using GNNs) real world graphs with a large number of nodes/ edges may not always fit in CPU /GPU memory, hence constructing mini-batches is important. Specifically, the authors propose a strategy for the task of node classification - where they aim to select batch size based on number of nodes and average degree in a graph and show that their proposed guidelines have benefits in terms of training time as well as accuracy. The authors propose a metric - pseudo precision rate which is dependent on the computation cost and the variance of the gradients and derive a lower bound for this metric which factors into account the batch size.