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 proposes a technique to improve the scalability of GNNs for large graphs. In many existing works, neighbor nodes, GNN layers, or subgraphs are subsampled to reduce the burden on GPU memory. In this paper, instead of subsampling, the vector quantization (VQ) on the embedding vector of each node is introduced to reduce the memory and ensure scalability. Technically, this paper develops a new approximation for message passing on codebook-represented graphs. In theoretical evaluation, the paper shows that the prediction errors due to VQ can be upper-bounded by using Lipshitz regularization. In the experimental evaluation, the generalization performance of the proposed VW method is comparable to existing scalable GNNs.