IGLU: Efficient GCN Training via Lazy Updates

1 INTRODUCTION . The Graph Convolution Network ( GCN ) model has received much attention as an effective graph representation learning technique . It can exploit network topology while embedding data points enabling superior performance in several applications such as node classification on graphs ( Kipf & Welling , 2017 ) , recommendation systems ( Ying et al. , 2020 ) and program repair ( Yasunaga & Liang , 2020 ) . Their success notwithstanding , training GCNs at scale remains challenging especially on large and dense graphs and multiple convolution layers . This is mainly due to the aggregation operation that enables GCNs to adapt to graph topology – a node ’ s output layer embedding in a GCN is influenced by embeddings of its neighbors in the previous layer which recursively depend on embeddings of their own neighbors in the still previous layer , and so on . Even in GCNs with 2-3 layers , when processing loss terms corresponding to a small sample of nodes in a mini-batch , this causes back propagation to update a large multi-hop neighborhood , causing mini-batch SGD-based techniques to scale poorly especially in dense graphs . Efforts to overcome this problem try to limit the number of nodes that receive updates as a result of a back-propagation step Chiang et al . ( 2019 ) ; Hamilton et al . ( 2018 ) ; Zeng et al . ( 2020 ) . This is done either by sub-sampling the neighborhood or clustering ( note the distinction between nodes sampled to create a mini-batch and neighborhood sampling done to limit the neighborhood of the mini-batch that receives updates ) . Variance reduction techniques have also been studied Chen et al . ( 2018a ) to reduce additional variance introduced by neighborhood sampling . However , these techniques often require large graphs to be heavily subsampled resulting in poor accuracy due to insufficient aggregation . They also do not guarantee unbiased gradients or rigorous convergence guarantees . See Section 2 for a more detailed discussion about state-of-the-art methods for GCN training . Our Contributions : This paper presents IGLU , an efficient technique for training GCNs based on lazy updates . An analysis of the gradient structure in GCNs reveals the most expensive component of the back-propagation step initiated at a node to be ( re- ) computation of forward-pass embeddings for its vast multi-hop neighborhood . Based on this observation , IGLU performs back-propagation with significantly reduced complexity using intermediate computations that are cached at regular intervals . This completely avoids neighborhood sampling and is a stark departure from the state-of-the-art . IGLU is architecture-agnostic and can be readily implemented on a wide range of GCN architectures . Avoiding neighborhood sampling also allows IGLU to completely avoid variance artifacts and offer provable convergence to a first-order stationary point under standard assumptions . In experiments , IGLU offered superior accuracies and accelerated convergence on a range of benchmark datasets . 2 RELATED WORKS . ( Bruna et al. , 2013 ; Defferrard et al. , 2016 ; Kipf & Welling , 2017 ) introduced the GCN architecture for transductive learning on graphs . Later works extended to inductive settings and explored architectural variants such as the GIN ( Xu et al. , 2019 ) . Much effort has focused on speeding-up GCN training . Sampling Based Approaches : The simplest neighborhood sampling strategy adopted by GraphSAGE ( Hamilton et al. , 2018 ) iteratively sub-samples the multi-hop neighborhood of a node before proceeding to initiate back-propagation at that node . Only the sub-sampled neighbors participate in the back-propagation step limiting the amount of computation . Layer sampling strategies such as FastGCN ( Chen et al. , 2018b ) , LADIES ( Zou et al. , 2019 ) and ASGCN ( Huang et al. , 2018 ) decrease computation by sampling nodes at each GCN layer , using importance sampling to reduce variance and improve connectivity among sampled nodes . FastGCN used the same sampling distribution for all layers and struggled to maintain connectivity unless large batch-sizes were used . LADIES used a per-layer distribution conditioned on nodes sampled for the layer above allowing more efficient use of samples under a budget . ASGCN used a linear model to jointly infer node importance weights . Recent works such as Cluster-GCN ( Chiang et al. , 2019 ) and GraphSAINT ( Zeng et al. , 2020 ) propose subgraph sampling creating mini-batches out of subgraphs and restricting back-propagation to nodes within the subgraph . To avoid losing too many edges , large mini-batch sizes are used . Cluster-GCN performs graph clustering and chooses multiple clusters per mini-batch while reinserting cross-cluster edges whereas GraphSAINT samples large subgraphs directly using random walks . Bias and Variance : Sampling techniques face bias as non-linear activations in the GCN architecture make it difficult to offer unbiased estimates of the loss function e.g . Zeng et al . ( 2020 ) offer un-biased estimates only if non-linearities are discarded . Sampling techniques also face increased variance for which variance-reduction techniques have been proposed such as VR-GCN ( Chen et al. , 2018a ) , MVS-GNN ( Cong et al. , 2020 ) and AS-GCN ( Huang et al. , 2018 ) . VR-GCN samples nodes whose embeddings are to be updated at each layer and uses stale embeddings for the rest and offers variance elimination in the limit under suitable conditions . MVS-GNN handles variance due to mini-batch creation by performing importance weighted sampling to construct mini-batches . Bandit Sampler ( Liu et al. , 2020 ) formulates variance reduction as an adversarial bandit problem . Other Approaches : Recent approaches decouple propagation from prediction as a pre-processing step e.g . PPRGo ( Bojchevski et al. , 2020 ) , APPNP ( Klicpera et al. , 2018 ) and SIGN ( Frasca et al. , 2020 ) . APPNP makes use of the relationship between the GCNs and PageRank to construct improved propagation schemes via personalized PageRank . PPRGo extends APPNP by approximating the dense propagation matrix via the push-flow algorithm . SIGN proposes inception style pre-computation of graph convolutional filters to speed up training and inference . GNNAutoScale ( Fey et al. , 2021 ) builds on VR-GCN and makes use of historical embeddings for scaling GNN training to large graphs . IGLU in Context of Related Work : IGLU avoids neighborhood sampling entirely and instead speeds-up learning using stale computations . Intermediate computations are cached and lazily updated at regular intervals e.g . once per epoch . We note that IGLU ’ s caching is distinct and much more aggressive ( e.g . lasting an entire epoch ) than the internal caching performed by popular frameworks such as TensorFlow and PyTorch ( where caches last only a single iteration ) . Recomputing these values in bulk offers IGLU economies of scale . IGLU faces zero sampling variance but incurs bias due to the use of stale computations . Fortunately , this bias is provably bounded , and can be made arbitrarily small by adjusting the step length and refresh frequency of the stale computations . 3 IGLU : EFFICIENT GCN TRAINING VIA LAZY UPDATES . Problem Statement : Consider the problem of learning a GCN architecture on an undirected graph G ( V , E ) with each of the N nodes endowed with an initial feature vector x0i ∈ Rd0 , i ∈ V . X0 ∈ Rn×d0 denotes the matrix of these initial features stacked together . N ( i ) ⊂ V denotes the set of neighbors of node i . A denotes the ( normalized ) adjacency matrix of the graph . A multi-layer GCN architecture uses a parameterized function at each layer to construct a node ’ s embedding for the next layer using embeddings of that node as well as those of its neighbors . Specifically xki = f ( x k−1 j , j ∈ { i } ∪ N ( i ) ; E k ) , where Ek denotes the parameters of k-th layer . For example , a standard GCN layer is given by , xki = σ ∑ j∈V Aij ( W k ) > xk−1j , where Ek is simply the matrix W k ∈ Rdk−1×dk and dk is the embedding dimensionality after the kth layer . More involved architectures exist that incorporate operations such as layer normalization , batch normalization , etc . In this paper , Ek will always denote the collection of all parameters of the kth layer e.g . offset and scale parameters used in a layer norm operation . Xk ∈ Rn×dk will denote the matrix of kth layer embeddings stacked together , giving us the handy shorthand Xk = f ( Xk−1 ; Ek ) . Given aK-layer GCN and a multi-label/multi-class task with C labels/classes , a fully-connected layer WK+1 ∈ RdK×C and appropriate activation function such as softmax is applied to get predictions which are then fed into the task loss . We note that IGLU does not require the task loss to decompose over the classes . The convergence proofs only require a smooth training objective function . Neighborhood Explosion : To understand the reason behind neighborhood explosion and cost of mini-batch based SGD training , consider a univariate regression problem with a no-frills 2-layer GCN with unidimensional features and sigmoidal activation within the hidden layers i.e . K = 2 and C = 1 = d0 = . . . = dk . This GCN is parameterized by w1 , w2 , w3 ∈ R and offers the output ŷi = w3σ ( z2i ) where z2i = ∑ j∈V Aijw 2x1j ∈ R. In turn , we have x1j = σ ( z1i ) where z1i = ∑ j′∈V Ajj′w 1x0j′ ∈ R and x0j′ ∈ R are the initial features of the nodes . Given a task loss ` : RN × RN e.g . least squares , denoting ` ′i = ` ′ ( ŷi , yi ) gives us ∂ ` ( ŷi , yi ) ∂w1 = ` ′i · ∂ŷi ∂z2i · ∂z 2 i ∂w1 = ` ′i · w3σ′ ( z2i ) · ∑ j∈V Aijw 2 ∂x1j ∂w1 = ` ′i · w3σ′ ( z2i ) · ∑ j∈V Aijw 2σ′ ( z1j ) · ∑ j′∈V Ajj′x 0 j′ . The nesting of the summations is conspicuous and establishes the neighborhood explosion : when seeking gradients in a K-layer GCN on a graph with average degree m , up to an mK−1-sized neighborhood of a node may be involved in the back-propagation update initiated at that node . Note that the above expression involves terms such as σ′ ( z2i ) , σ ′ ( z1j ) and that the values of z 2 i , z 1 j etc change whenever the model receives updates . Consequently , for a fresh mini-batch of nodes chosen for a particular iteration , terms such as σ′ ( z2i ) , σ ′ ( z1j ) need to be computed afresh if the gradient is to be computed exactly . Performing these computations amounts to doing forward pass operations that frequently involve a large neighborhood of the nodes of the mini-batch . Sampling strategies try to limit this cost by directly restricting the neighborhood over which such forward passes are computed . However , this introduces both bias and variance into the gradient updates as discussed in Section 2 . IGLU instead lazily updates various incomplete gradient ( defined below ) and node embedding terms that participate in the above expression . This completely eliminates sampling variance but introduce a bias due to the use of stale terms . However , this bias is not just provably bounded , but can be made arbitrarily small by adjusting the step length and frequency of refreshing these terms . Lazy Updates for GCN Training : Consider an arbitrary GCN architecture with the following structure : for some parameterized layer functions we have Xk = f ( Xk−1 ; Ek ) where Ek denotes the collection of all parameters of the kth layer e.g . weight matrices , offset and scale parameters used in layer norm operations , etc . Xk ∈ RN×dk denotes the matrix of kth layer embeddings stacked together and X0 ∈ RN×d0 are the initial features . For a K-layer GCN on a multi-label/multi-class task with C labels/classes , a fully-connected layer WK+1 ∈ RdK×C is additionally included to offer ŷi = ( W K+1 ) > xKi ∈ RC . Let us use the shorthand Ŷ ∈ RN×C to denote the matrix where the predicted outputs ŷi over all the nodes are stacked . We assume a task loss function ` : RC×RC×R+ that need not decompose over the classes . We will use the abbreviation ` i : = ` ( ŷi , yi ) . For sake of simplicity , we assume that the loss function itself includes any activation such as softmax that needs to be applied over the predictions ŷi . Let L = ∑ i∈V ` i denote the total loss function . Motivation : We define the the loss derivative matrix G = [ gic ] = RN×C as gic : = ∂ ` i∂ŷic . As the proof of Lemma 1 ( See Appendix C ) shows , the loss derivative with respect to parameters Ek at any layer has the form ∂L ∂Ek = ∑N j=1 ∑dk p=1 ( ∑ i∈V ∑ c∈ [ C ] gic · ∂ŷic ∂Xkjp ) ∂Xkjp ∂Ek . Note that the loss derivative is expressed in terms of partial derivatives ∂Xkjp ∂Ek that can be computed for any node using only embeddings of its neighbors in the ( k − 1 ) th layer i.e . Xk−1 thus avoiding any neighborhood explosion . This means that neighborhood explosion must be encountered while computing the terms encapsulated in the round brackets . Let us first formally recognize these terms as incomplete gradients . The notation ∂P∂Q ∣∣∣ R denotes the partial derivative of P w.r.t Q while keeping R fixed i.e . treated as a constant . Definition 1 . For any layer k ≤ K , define its incomplete task gradient to be αk = [ αkjp ] ∈ RN×dk , αkjp : = ∂ ( G Ŷ ) ∂Xkjp ∣∣∣∣∣ G = ∑ i∈V ∑ c∈ [ C ] gic · ∂ŷic ∂Xkjp The following lemma completely characterizes the loss gradients of the GCN but also shows that the incomplete gradient terms αk , k ∈ [ K ] can be efficiently computed using a recursive formulation that also does not involve any neighborhood explosion . Lemma 1 . The following results hold whenever the task loss L is differentiable : 1 . For the final fully-connected layer we have ∂L ∂WK+1 = ( XK ) > G as well as for any k ∈ [ K ] and any parameter Ek in the kth layer , ∂L ∂Ek = ∂ ( α k Xk ) ∂Ek ∣∣∣ αk = ∑ i∈V ∑dk p=1 α k ip · ∂Xkip ∂Ek . 2 . For the final layer , we have αK = G ( WK+1 ) > as well as for any k < K , we have αk = ∂ ( α k+1 Xk+1 ) ∂Xk ∣∣∣ αk+1 i.e . αkjp = ∑ i∈V ∑dk+1 q=1 α k+1 iq · ∂Xk+1iq ∂Xkjp . Lemma 1 establishes a recursive definition of the incomplete gradients using terms such as ∂Xk+1iq ∂Xkjp that only concern a single layer . Thus , computing αk for any k ∈ [ K ] does not involve any neighborhood explosion since only the immediate neighbors of a node need be consulted . Lemma 1 also shows that if αk are computed and frozen , the loss derivatives ∂L ∂Ek only involve additional computation of terms such as ∂Xkip ∂Ek which yet again involve a single layer and do not cause neighborhood explosion . This motivates lazy updates to αk , Xk values in order to accelerate back-propagation . However , performing lazy updates to both αk , Xk offers suboptimal performance . Hence IGLU adopts two variants described in Algorithms 1 and 2 . The backprop variant * keeps embeddings Xk stale for an entire epoch but performs eager updates to αk . The inverted variant on the other hand keeps the incomplete gradients αk stale for an entire epoch but performs eager updates to Xk . Algorithm 1 IGLU : backprop order Input : GCN G , initial features X0 , task loss L 1 : Initialize model parameters Ek , k ∈ [ K ] , WK+1 2 : while not converged do 3 : Do a forward pass to compute Xk for all k ∈ [ K ] as well as Ŷ 4 : Compute G then ∂L ∂WK+1 using Lemma 1 ( 1 ) and update WK+1 ←WK+1 − η · ∂L ∂WK+1 5 : Compute αK using G , WK+1 , Lemma 1 ( 2 ) 6 : for k = K . . . 2 do 7 : Compute ∂L ∂Ek using αk , Xk , Lemma 1 ( 1 ) 8 : Update Ek ← Ek − η · ∂L ∂Ek 9 : Update αk using αk+1 using Lemma 1 ( 2 ) 10 : end for 11 : end while Algorithm 2 IGLU : inverted order Input : GCN G , initial features X0 , task loss L 1 : Initialize model parameters Ek , k ∈ [ K ] , WK+1 2 : Do an initial forward pass to computeXk , k ∈ [ K ] 3 : while not converged do 4 : Compute Ŷ , G and αk for all k ∈ [ K ] using Lemma 1 ( 2 ) 5 : for k = 1 . . .K do 6 : Compute ∂L ∂Ek using αk , Xk , Lemma 1 ( 1 ) 7 : Update Ek ← Ek − η · ∂L ∂Ek 8 : Update Xk ← f ( Xk−1 ; Ek ) 9 : end for 10 : Compute ∂L ∂WK+1 using Lemma 1 ( 1 ) and use it to update WK+1 ←WK+1 − η · ∂L ∂WK+1 11 : end while SGD Implementation : Update steps in the algorithms ( steps 4 , 8 in Algorithm 1 and steps 7 , 10 in Algorithm 2 ) are described as a single gradient step over the entire graph to simplify exposition – in practice these steps are implemented using mini-batch SGD . This is done as usual by sampling a mini-batch of nodes S and computing task gradients only w.r.t L̂S = ∑ i∈S ` i instead of L. Contribution : As noted in Section 2 , IGLU uses caching in a manner fundamentally different from popular frameworks such as PyTorch or TensorFlow . Their caches are short-lived and always seek to compute exact gradients unlike IGLU that computes gradients much faster but with bounded bias . IGLU also extends other caching-based techniques such as VR-GCN which choose to cache node embeddings . In contrast , IGLU offers two variants and the variant that uses inverted order of updates ( Algorithm 2 ) and caches incomplete gradients outperforms the backprop variant of IGLU ( Algorithm 1 ) that caches node embeddings instead . Theoretical Analysis : Note that , conditioned on the stale parameters ( either αk or Xk depending on which variant is being executed ) , the gradients used by IGLU to perform model updates ( steps 4 , 8 in Algorithm 1 and steps 7 , 10 in Algorithm 2 ) do not exhibit any sampling bias . However , the staleness itself introduces bias in the training process . By controlling the step length η and the frequency with which the stale parameters are updated , this bias can be provably controlled resulting in guaranteed convergence to a first-order stationary point . Due to lack of space , we postpone the detailed statement and proof of the convergence guarantee to Appendix C. Theorem 2 ( IGLU Convergence ( Informal ) ) . Suppose the task loss function L has O ( 1 ) -Lipschitz gradients and IGLU is executed with small enough step lengths η with model updates ( steps 4 , 8 in Algorithm 1 and steps 7 , 10 in Algorithm 2 ) being carried out either in a full-batch or mini-batch SGD manner , then within T iterations , IGLU converges to a model iterate satisfying : 1 . ‖∇L‖22 ≤ O ( 1/T 2 3 ) if update steps are carried out on the entire graph in a full-batch . 2 . ‖∇L‖22 ≤ O ( 1/ √ T ) if update steps are carried out using mini-batch SGD . This result holds under minimal assumptions of objective smoothness and boundedness that are standard Chen et al . ( 2018a ) ; Cong et al . ( 2020 ) , yet offers convergence rates comparable to those offered by standard mini-batch SGD . However , whereas works such as ( Chen et al. , 2018a ) assume bounds on the sup-norm i.e . L∞ norm of the gradients , Theorem 2 only requires an L2 norm bound . Note that objective smoothness requires the architecture to use smooth activation functions . However , IGLU offers similar performance whether using non-smooth activations e.g . ReLU or smooth ones e.g . GELU ( see Appendix B.7 ) as is also observed by other works ( Hendrycks & Gimpel , 2020 ) . * The backprop variant is named so since it updates model parameters in the order back-propagation would have updated them i.e . WK+1 followed by EK , EK−1 , . . . whereas the inverted variant performs updates in the reverse order i.e . starting from E1 , E2 all the way to WK+1 .

Summary: The work proposes IGLU, an algorithm to scale-up GNNs using stale computations instead of traditional neighborhood sampling. IGLU has bounded bias if the loss and activation functions are smooth. Results on large-scale benchmarks show IGLU achieves SOTA and scales better than previous methods.