VQ-GNN: A Universal Framework to Scale up Graph Neural Networks using Vector Quantization

1 Introduction . The rise of Graph Neural Networks ( GNNs ) has brought the modeling of complex graph data into a new era . Using message-passing , GNNs iteratively share information between neighbors in a graph to make predictions of node labels , edge labels , or graph-level properties . A number of powerful GNN architectures [ 1–4 ] have been widely applied to solve down-stream tasks such as recommendation , social analysis , visual recognition , etc . With the soaring size of realistic graph datasets and the industrial need to model them efficiently , GNNs are hindered by a scalability problem . An L-layer GNN aggregates information from all L-hop neighbors , and standard training routines require these neighbors to all lie on the GPU at once . This prohibits full-batch training when facing a graph with millions of nodes [ 5 ] . ∗Equal contribution . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . A number of sampling-based methods have been proposed to accommodate large graphs with limited GPU resources . These techniques can be broadly classified into three categories : ( 1 ) Neighborsampling methods [ 2 , 6 ] sample a fixed-number of neighbors for each node ; ( 2 ) Layer-sampling methods [ 7 , 8 ] sample nodes in each layer independently with a constant sample size ; ( 3 ) Subgraphsampling methods [ 9 , 10 ] sample a subgraph for each mini-batch and perform forward and backpropagation on the same subgraph across all layers . Although these sampling-based methods may significantly speed up the training time of GNNs , they suffer from the following three major drawbacks : ( 1 ) At inference phase , sampling methods require all the neighbors to draw non-stochastic predictions , resulting in expensive predictions if the full graph can not be fit on the inference device ; ( 2 ) As reported in [ 5 ] and in Section 6 , state-of-the-art sampling-baselines fail to achieve satisfactory results consistently across various tasks and datasets ; ( 3 ) Sampling-based methods can not be universally applied to GNNs that utilize many-hop or global context in each layer , which hinders the application of more powerful GNNs to large graphs . This paper presents VQ-GNN , a GNN framework using vector quantization to scale most state-ofthe-art GNNs to large graphs through a principled and fundamentally different approach compared with the sampling-based methods . We explore the idea of using vector quantization ( VQ ) as a means of dimensionality reduction to learn and update a small number of quantized reference vectors ( codewords ) of global node representations . In VQ-GNN , mini-batch message passing in each GNN layer is approximated by a VQ codebook update and an approximated form of message passing between the mini-batch of nodes and codewords ; see Fig . 1 . Our approach avoids the “ neighbor explosion ” problem and enables mini-batch training and inference of GNNs . In contrast to samplingbased techniques , VQ-GNN can effectively preserve all the messages passed to a mini-batch of nodes . We theoretically and experimentally show that our approach is efficient in terms of memory usage , training/inference time , and convergence speed . Experiments on various GNN backbones demonstrate the competitive performance of our framework compared with the full-graph training baseline and sampling-based scalable algorithms . Paper organization . The remainder of this paper is organized as follows . Section 2 summarizes GNNs that can be re-formulated into a common framework of graph convolution . Section 3 defines the scalability challenge of GNNs and shows that dimensionality reduction is a potential solution . In Section 4 , we describe our approach , VQ-GNN , from theoretical framework to algorithm design and explain why it solves the scalability issue of most GNNs . Section 5 compares our approach to the sampling-based methods . Section 6 presents a series of experiments that validate the efficiency , robustness , and universality of VQ-GNN . Finally , Section 7 concludes this paper with a summary of limitations and broader impacts . 2 Preliminaries : GNNs defined as Graph Convolution . Notations . Consider a graph with n nodes and m edges ( average degree d = m/n ) . Connectivity is given by the adjacency matrix A ∈ { 0 , 1 } n×n and features are defined on nodes by X ∈ Rn×f0 with f0 the length of feature vectors . Given a matrix C , let Ci , j , Ci , : , and C : ,j denote its ( i , j ) -th entry , i-th row , j-th column , respectively . For a finite sequence 〈ib〉 : i1 , . . . , ib , we use C〈ib〉 , : to denote the matrix whose rows are the ib-th rows of matrix C. We use to denote the element-wise ( Hadamard ) product . ‖ · ‖p denotes the entry-wise ` p norm of a vector and ‖ · ‖F denotes the Frobenius norm . We use In ∈ Rn×n to denote the identity matrix , 1n ∈ Rn to denote the vector whose entries are all ones , and ein to denote the unit vector in Rn whose i-th entry is 1 . The 0-1 indicator function is 1 { · } . We use diag ( c ) to denote a diagonal matrix whose diagonal entries are from vector c. Andf represents concatenation along the last axis . We use superscripts to refer to different copies of same kind of variable . For example , X ( l ) ∈ Rn×fl denotes node representations on layer l. A Graph Neural Network ( GNN ) layer takes the node representation of a previous layer X ( l ) as input and produces a new representation X ( l+1 ) , where X = X ( 0 ) is the input features . A common framework for generalized graph convolution . Although many GNNs are designed following different guiding principles including neighborhood aggregation ( GraphSAGE [ 2 ] , PNA [ 11 ] ) , spatial convolution ( GCN [ 1 ] ) , spectral filtering ( ChebNet [ 12 ] , CayleyNet [ 13 ] , ARMA [ 14 ] ) , self-attention ( GAT [ 3 ] , Graph Transformers [ 15–17 ] ) , diffusion ( GDC [ 18 ] , DCNN [ 19 ] ) , WeisfeilerLehman ( WL ) alignment ( GIN [ 4 ] , 3WL-GNNs [ 20 , 21 ] ) , or other graph algorithms ( [ 22 , 23 ] ) . Despite these differences , nearly all GNNs can be interpreted as performing message passing on node features , followed by feature transformation and an activation function . As pointed out by Balcilar et al . [ 24 ] , GNNs can typically be written in the form X ( l+1 ) = σ ( ∑ s C ( s ) X ( l ) W ( l , s ) ) , ( 1 ) where C ( s ) ∈ Rn×n denotes the s-th convolution matrix that defines the message passing operator , s ∈ Z+ denotes index of convolution , and σ ( · ) denotes the non-linearity . W ( l , s ) ∈ Rfl×fl+1 is the learnable linear weight matrix for the l-th layer and s-th filter . Within this common framework , GNNs differ from each other by choice of convolution matrices C ( s ) , which can be either fixed or learnable . A learnable convolution matrix relies on the inputs and learnable parameters and can be different in each layer ( thus denoted as C ( l , s ) ) : C ( l , s ) i , j = C ( s ) i , j︸︷︷︸ fixed ·h ( s ) θ ( l , s ) ( X ( l ) i , : , X ( l ) j , : ) ︸ ︷︷ ︸ learnable ( 2 ) where C ( s ) denotes the fixed mask of the s-th learnable convolution , which may depend on the adjacency matrix A and input edge features Ei , j . While h ( s ) ( · , · ) : Rfl × Rfl → R can be any learnable model parametrized by θ ( l , s ) . Sometimes a learnable convolution matrix may be further row-wise normalized as C ( l , s ) i , j ← C ( l , s ) i , j / ∑ j C ( l , s ) i , j , for example in GAT [ 3 ] . We stick to Eq . ( 2 ) in the main paper and discuss row-wise normalization in Appendices A and E. The receptive field of a layer of graph convolution ( Eq . ( 1 ) ) is defined as a set of nodesR1i whose features { X ( l ) j , : | j ∈ Ri } determines X ( l+1 ) i , : . We re-formulate some popular GNNs into this generalized graph convolution framework ; see Table 1 and Appendix A for more . The back-propagation rule of GNNs defined by Eq . ( 1 ) is as follows , ∇X ( l ) ` = ∑ s ( C ( l , s ) ) T ( ∇X ( l+1 ) ` σ′ ( σ−1 ( X ( l+1 ) ) ) ) ( W ( l , s ) ) T , ( 3 ) which can also be understood as a form of message passing . σ′ and σ−1 are the derivative and inverse of σ respectively and ∇X ( l+1 ) ` σ′ ( σ−1 ( X ( l+1 ) ) ) is the gradients back-propagated through the non-linearity . 3 Scalability Problem and Theoretical Framework . When a graph is large , we are forced to mini-batch the graph by sampling a subset of b n nodes in each iteration . Say the node indices are i1 , . . . , ib and a mini-batch of node features is denoted by XB = X〈ib〉 , : . To mini-batch efficiently for any model , we hope to fetch Θ ( b ) information to the training device , spend Θ ( Lb ) training time per iteration while taking ( n/b ) iterations to traverse through the entire dataset . However , it is intrinsically difficult for most of the GNNs to meet these three scalability requirements at the same time . The receptive field of L layers of graph convolution ( Eq . ( 1 ) ) is recursively given by RLi = ⋃ j∈R1i RL−1j ( starting with R1i ⊇ { i } ∪ Ni ) , and its size grows exponentially with L. Thus , to optimize on a mini-batch of b nodes , we require Ω ( bdL ) inputs and training time per iteration . Sampling a subset of neighbors [ 2 , 6 ] for each node in each layer does not change the exponential dependence on L. Although layer- [ 7 , 25 ] and subgraph-sampling [ 9 , 10 ] may require only Ω ( b ) inputs and Ω ( Lb ) training time per iteration , they are only able to consider an exponentially small proportion of messages compared with full-graph training . Most importantly , all existing sampling methods do not support dense convolution matrices with O ( n2 ) non-zero terms . Please see Section 5 for a detailed comparison with sampling-based scalable methods after we introduce our framework . Idea of dimensionality reduction . We aim to develop a scalable algorithm for any GNN models that can be re-formulated as Eq . ( 1 ) , where the convolution matrix can be either fixed or learnable , and either sparse or dense . The major obstacle to scalability is that , for each layer of graph convolution , to compute a mini-batch of forward-passed features X ( l+1 ) B = X ( l+1 ) 〈ib〉 , : , we need O ( n ) entries of C ( l , s ) B = C ( l , s ) 〈ib〉 , : and X ( l ) , which will not fit in device memory . Our goal is to apply a dimensionality reduction to both convolution and node feature matrices , and then apply convolution using compressed “ sketches ” of C ( l , s ) B and X ( l ) . More specifically , we look for a projection matrix R ∈ Rn×k with k n , such that the product of low-dimensional sketches C̃ ( l , s ) B = C ( l , s ) B R ∈ Rb×k and X̃ ( l ) = RTX ( l ) ∈ Rk×fl is approximately the same as C ( l , s ) B X ( l ) . The approximated product ( of all nodes ) C̃ ( l , s ) X̃ ( l ) = C ( l , s ) RRTX ( l ) can also be regarded as the result of using a low-rank approximation C ( l , s ) RRT ∈ Rn×n of the convolution matrix such that rank ( C ( l , s ) RRT ) ≤ k. The distributional Johnson–Lindenstrauss lemma [ 26 ] ( JL for short ) shows the existence of such projectionR withm = Θ ( log ( n ) ) , and the following result by Kane and Nelson [ 27 ] shows that R can be chosen to quite sparse : Theorem 1 . For any convolution matrix C ∈ Rn×n , any column vector X : ,a ∈ Rn of the node feature matrix X ∈ Rn×f ( where a = 1 , . . . , f ) and ε > 0 , there exists a projection matrix R ∈ Rn×k ( drawn from a distribution ) with only an O ( ε ) -fraction of entries non-zero , such that Pr ( ‖CRRTX : ,a − CX : ,a‖2 < ε‖CX : ,a‖2 ) > 1− δ , ( 4 ) with k = Θ ( log ( n ) /ε2 ) and δ = O ( 1/n ) . Now , the sketches C̃ ( l , s ) B and X̃ ( l ) take up O ( b log ( n ) ) and Θ ( fl log ( n ) ) memory respectively and can fit into the training and inference device . The sparsity of projection matrix R is favorable because : ( 1 ) if the convolution matrix C ( l , s ) is sparse ( e.g. , direct-neighbor message passing where only O ( d/n ) -fraction of entries are non-zero ) , only an O ( εd ) -fraction of entries are non-zero in the sketch C̃ ( l , s ) ; ( 2 ) During training , X̃ ( l ) is updated in a “ streaming ” fashion using each minibatch ’ s inputs X ( l ) B , and a sparse R reduces the computation time by a factor of O ( ε ) . However , the projection R produced following the sparse JL-lemma [ 27 ] is randomized and requires O ( log2 ( n ) ) uniform random bits to sample . It is difficult to combine this with the deterministic feed-forward and back-propagation rules of neural networks , and there is no clue when and how we should update the projection matrix . Moreover , randomized projections destroy the “ identity ” of each node , and for learnable convolution matrices ( Eq . ( 2 ) ) , it is impossible to compute the convolution matrix only using the sketch of features X̃ ( l ) . For this idea to be useful , we need a deterministic and identity-preserving construction of the projection matrix R ∈ Rn×k to avoid these added complexities .

Summary: This paper is about using vector quantization to scale up GNN. The main idea is to use quantized representations combined with a low-rank projection of graph convolution matrix to avoid the "neighbor explosion" problem of GNNs. The proposed VQ-GNN is applied for node classification and link prediction, and shows comparable results.