VQ-GNN: A Universal Framework to Scale up Graph Neural Networks using Vector Quantization
Authors: Mucong Ding
Kezhi Kong
Jingling Li
Chen Zhu
John P Dickerson
Furong Huang
Tom Goldstein
Venue: NeurIPS 2021 Poster (2021)
Keywords: Graph Neural Networks
Scalability
Vector Quantization

Abstract: Most state-of-the-art Graph Neural Networks (GNNs) can be defined as a form of graph convolution which can be realized by message passing between direct neighbors or beyond. To scale such GNNs to large graphs, various neighbor-, layer-, or subgraph-sampling techniques are proposed to alleviate the "neighbor explosion" problem by considering only a small subset of messages passed to the nodes in a mini-batch. However, sampling-based methods are difficult to apply to GNNs that utilize many-hops-away or global context each layer, show unstable performance for different tasks and datasets, and do not speed up model inference. We propose a principled and fundamentally different approach, VQ-GNN, a universal framework to scale up any convolution-based GNNs using Vector Quantization (VQ) without compromising the performance. In contrast to sampling-based techniques, our approach can effectively preserve all the messages passed to a mini-batch of nodes by learning and updating a small number of quantized reference vectors of global node representations, using VQ within each GNN layer. Our framework avoids the "neighbor explosion" problem of GNNs using quantized representations combined with a low-rank version of the graph convolution matrix. We show that such a compact low-rank version of the gigantic convolution matrix is sufficient both theoretically and experimentally. In company with VQ, we design a novel approximated message passing algorithm and a nontrivial back-propagation rule for our framework. Experiments on various types of GNN backbones demonstrate the scalability and competitive performance of our framework on large-graph node classification and link prediction benchmarks.

VQ-GNN: A Universal Framework to Scale-up Graph Neural Networks using Vector Quantization
Mucong Ding", Kezhi Kong", , Jingling Li, Chen Zhu, John Dickerson, Furong Huang, Tom Goldstein
Department of Computer Science, University of Maryland (mcding, kong, jingling, chenzhu, john, furongh, tomg) @ cs. umd.edu
Abstract
Most state-of-the-art Graph Neural Networks (GNNs) can be defined as a form of graph convolution which can be realized by message passing between direct neighbors or beyond. To scale such GNNs to large graphs, various neighbor- layer, subgraph-sampling techniques problem by considering only a small subset of messages passed to the nodes in a mini-batch. However, sampling-base or global context each layer, show unstable performance
proposed to alleviate the "neighbor explosion"
methods are difficult to apply to GNNs that
utilize many-hops--ae for different tasks and datasets, and do not speed up model inference. We propose principled and fundamentally different approach, VQ-GNN, universal framework scale up any convolution-based GNNs using Vector Quantization (VQ) without compromising the performance. contrast to sampling- techniques, our approach effectively preserve all the messages passed mini-batch of nodes by learning and updating a small number of quantized reference vectors of global node representations, using VQ within each GNN layer. Our framework avoids the "neighbor explosion" problem of GNNs using quantized representations combined with low-rank version of the graph convolution matrix. We show that such compact low-rank version of the gigantic convolution matrix is sufficient both theoretically and experimentally. approximated message passing algorithm and a nontrivial back-propagation rule for our framework. Experiments on various types of GNN backbones demonstrate the scalability and competitive performance of our framework on large-graph node classification and link prediction benchmarks.
company with VQ, we design a novel
Introduction
The rise of Graph Neural Networks (GNNs) has brought the modeling of complex graph data into new era. Using cesssage-passing GNNs iteratively share information between neighbors in graph to make predictions of node labels, edge labels, or graph-level properties. A number of powerful GNN architectures [14 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 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 [
"Equal contribution.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Mini-Batch Message Passing
approx. by
VQ, Codebook Update
and
Approx. Message Passing
Nodes Mini Batch (before code re-assign)
Nodes Mini-Batch
Messages passed Mini-Batch
(after code re re-assign)
Other Graph Edges
Codeword Assignment
Other Graph Nodes
Codewords
Codeword Moving Average
Figure 1: In framework, VQ-GNN, each mini-batch message passing (left) is approximated by VQ codebook update (middle) and an approximated message passing (right). All the messages passed the nodes the current mini-batch are effectively preserved. Circles are nodes, and rectangles are VQ codewords. double circle indicates nodes in the current mini-bate batch. Color represents codeword assignment. During VQ codebook update, codeword assignment of nodes in the mini-batch is refreshed (node 1), and codewords are updated using the assigned nodes. During approximated message passing, messages from out-of-mini-bee nodes are approximated by messages from the corresponding codewords, messages from nodes assigned to same codeword are merged (a and b), and tatra-mini-bate messages are not changed (c and d).
A number of sampling-hase methods have been proposed to accommodate large graphs with limited GPU resources. These techniques can be broadly classified into three categories: (1) Neighbor-
sampling methods [2] 6 sample fixed-number of neighbors for each node; (2) Layer-sampling methods 77 ( sample nodes in each layer independently with constant sample size; (3) Subgraph- sampling methods 99, 10 sample subgraph for each mini-batch and perform forward and back- propagation on the same subgraph across all layers. Although these sampling-bas based methods may significantly speed up the training time of GNNs, they suffer from the following three major draw- backs: (1) At inference phase, sampling methods require all the neighbors to draw non nto-stocc predictions, resulting in expensive predictions if the full graph cannot be fit the inference device; (2) As reported in [5] and in Section 6 state-of-the-ar sampling-baselie baseliess fail to achieve satisfactory results consistently across various tasks and datasets; (3) Sampling-based methods cannot be uni- versally 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-of- the-art GNNs to large graphs through principled and fundamentally different approach compared with the sampling- campling-based based methods. We explore the idea of using vector quantization (VQ) as a means of dimensionality reduction to learn and update 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 VQ codebook update and approximated form of message passing between the mini-batch of nodes and codewords; see Fig. [ Our approach avoids the "neighbor explosion" problem and enables mini-hatch training and inference of GNNs. In contrast sampling- based techniques, VQ-GNN can effectively preserve all the messages passed to mini-batch of nodes. We theoretically and experimentally show that our approach 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 aall-graph training baseline and sampling based scalable algorithms.
Paper organization. The remainder of this paper is organized as follows. Section summarizes GNNs that can be re-formulated into common framework of graph convolution. Section defines the scalability challenge GNNs and shows that dimensionality reduction is potential solution. In Section 4 we describe our approach, VQ-GNN, from theoretical framework to algorithm design and explain why solves the scalability issue of most GNNs. Section Section5]comme commpree our approach to methods. 6ection
the sampling-base robustness, and universality of VQ-GNN, Finally, Section 7 concludes this paper with summary of limitations and broader impacts.
series of experiments that validate the efficiency,

---

Table 1: Summary of GNNs re-formulated as generalized graph convolution.
Model Name
Design Idea
Conv. Matrix Type
# of Conv.
Convolution Matrix
GCN
L
Spatial Conv.
Fixed
C= B-1/2AD-1/2 cu) = In
SAGE-Mean
Message Passing
Fixed
(Ce) - D-1A
(ts) = A+ In and
GAT
3
Self-Attention
Learnable
# of heads
(X f,2 x
exp (LeakyReLU (x.- || xx wihe)) ( =l!"ا)))
1 Where Á A I.. D D In. 'C represents mean aggregator. Weight matrix 2 W wh IWG.
Need row-wise normalization. C is non-zero if and only if A., 1. thus GAT follows diroct-neighbor aggregation.
2 Preliminaries: GNNs defined as Graph Convolution
Notations. Consider graph with n nodes and m edges (average degree d = m/n). Connectivity is given by the adjacency matrix A {0, 1 j and features defined on nodes by X RR* with So the length of feature vectors. Given matrix C, let Ct.j. Ci., and C., denote its (i,j)-th entry, C(as). denote the
i-th row, j-th column, respectively. For finite sequence (i6) :i,...,, 1b. matrix whose rows are the ig-th rows of matrix C. We use @ to denote the element-wise (Hadamard) product. || Ipp denotes the entry-wise ( norm of vector and || || denotes the Frobenius norm. We use In E RT*x to denote the identity matrix, 1, E IR" to denote the vector whose entries are all ones, and e,, to denote the unit vector in 1² 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. And || represents concatenation along the last axis. We use superscripts to refer to different copies of same kind of variable. For example, ( E R=* fi denotes node representations on layer Graph Neural Network (GNN) layer takes the node representation of previous layer X input and produces representation X(+11,, where X (0) is the input features.
common framework for generalized graph convolution. Although many GNNs are designed fol- lowing different guiding principles including neighborhood aggregation (GraphSAGE (2). PNA [I],, spatial convolution (GCN D), spectral filtering (ChebNet [12, CayleyNet [3, ARMA 14),
self-attention (GAT 3. Graph Transformers (15177). diffusion (GDC [18] DCNN (99). Weisfeiler- Lehman (WL) alignment (GIN 44, 3WL-GNNs (20. 21).. or other graph algorithms ([22) 231).
Despite these differences, nearly all GNNs can be interpreted 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(+1)=0 =a Σ cl x(x)))))
(1)
where (+)) E R denotes the s-th convolution matrix that defines the message passing operator, E Z_ denotes index of convolution. and (-)) denotes the non-linearit.. W((,)) E Rixfi the learnable linear weight matrix for the 1-th layer and s-th filter.
Within this common framework, GNNs differ from each other by choice of convolution matrices ((ss), which can be either fixed or learnable. A learnable convolution matrix relies on the inputs and learaalee parameters and can be different in each layer (thus denoted as ((l,))
(1) ge.. !X
fixed
leamable
where ((s) denotes the fixed mask of the s-th learnable convolution, which may depend on the adjacency matrix and input edge features Ex,, While ((-)) R x Rí + R can be any learnable model parametrized by g(l.e) Sometimes a learnable convolution matrix may be further row-wise normalized as C(S,)) + ///e/// C.l.a)) for example in GAT (3. We stick to Eq. (2) in the main paper and discuss row-wise normalization in Appendices A and The receptive field of
3

---

layer of graph convolution (Eq. (1) is defined as set of nodes R, whose features { | E R4}
determines X (1+1) We re-formulate some popular GNNs into this generalized graph convolution framework; see Table | and Appendix [A for more.
The baoo-proongatioon rule of GNNs defined by Eq. (I is as follows,
V Vx = Σ (CCW))) ( (xxanne @ o (x(+)))) (x (Wil.))"
(
(3)
which can also be understood form of message passing, o' and o are the derivative and inverse of a respectively and V Vxatnl @ o' d' (x(((t))) is the gradients pack-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 < n nodes in each iteration. Say the node indices are ......,99 and a mini-batch of node features is denoted by XB = (i)..- To mini-batch efficiently for any model, we hope to fetch (b)) information to the training device, spend ((bb) 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. 4 is recursively given by Ri USER; R2 R..-1 (starting with R2 2 {i} U.N,), and its size grows exponentially with J.. Thus, optimize on mini-hatch of nodes,
require ((bdd) inputs and training time per iteration. Sampling subset of neighbors [2, 6 for each node in each layer does change the exponential dependence on L. Although layer- 7 25 and Sabgraph-sampling [0] may require only ((6) inputs and (bbb) training time per iteration, they are only able to consider an exponentially small proportion of messages compared with full full-graph training. Most importantly, all existing sampling methods do not support dense convolution matrices with (n?²) 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. ( 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 mini-batch of nowward-passe features ((+1) ((+1) (16).
need ((n) entries of
cff.*) c(l.s) and ((),, which will not fit in device memory. (bs)::
Our goal is to apply a dimensionality reduction to both convolution and node feature matrices, and then apply convolution using compressed "sketches" of CHL and ((), More specifically, we look for a projection matrix R E R= with k < such that the product of fow-dimensional sketches R. E b and ()) R (!) E -R fi is approximately the same C(l.s)) (!) O(l.a) c((++)
The approximated product (of all nodes) (1l,,) (1) = (+L.s) RR (!) can also be regarded as the result of using a low-rank approximation (ttss) RR Ruxn of the convolution matrix such that rank (C(l.) RRR) < k. The distributional Aechnson-Lindensternn
lemma [26) (JL for short) shows the existence of such projection R with (llo(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 E R='x, any column vector X.,a E RR of the node feature matrix X E Raxf (where a = 1.....f) and E > O, there exists a projection matrix R E Raak (drawn from a distribution) with only ((E)-fraction of entries non-zero, such that
Pr (ICRR" X., - CX < xea( > 1
with & = @(log(n)/s²) and ó O(1/n).
Now, the sketches )(I,a) and (1) take up Olb (())))))) and O(i 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.. strect-nnighhon message passing where only ((dd/))-faacionn of entries are non-zero), only an (( (Eed)-fraction of entries are non-zero in the sketch Ca.s;; (2) During training, X updated in a "streaming" fashion using each mini- batch's inputs X and sparse R. reduces the computation time by factor of O(e). However, the projection Rt produced following the sparse JL-lemma 27 is randomized and requires Ologg (n)) uniform random bits to sample. It is difficult to combine this with the deterministic feed-forward and pack-propagation rules of neural networks, and there is 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), impossible to compute the convolution matrix only using the sketch of features ( For this idea to be useful, we need a deterministic and construction of the projection matrix R E Rn* to avoid these added complexities.
Proposed Method: Vector Quantized GNN
Dimensionality reduction using Vector Quantization (VQ). A natural and widely-used method to reduce the dimensionality of data in a deterministic and Shentity-preserving manner is Vector Quantization (28) (VQ), a classical data compression algorithm that can be formulated as the following optimization problem:
min
- RR s.t. Re: E (e ex)
(5)
(E XERVex XEREX
which is classically solved via k-means
28 Here the sketch of features X is called the feature "codewords." R is called the codeword assignment matrix, whose rows are unit vectors in R$ i.e., Rs.. if and only if the i-th node is assigned to the w-th cluster in k-means. The objective Cluster Sum Squares (WCSS), and we can define the relative error of VQ as E = || - RX || [||]|||]]]]] x| The rows X are the & codewords (i.e., centroids in k-means), and can be computed as = diag ( 1, R X, which is slightly different from the definition in Section as a row-wise normalization of RT is required. The sketch of the convolution matrix C can still be computed as C = CR. In general, VQ provides us principled framework to learn the low-dimensional sketches X and C, in a deterministic and mnde-idennityy-peeervngg manner. However, enable mini-batch training and inference of GNNs using VQ, three more questions need to be answered:
in Eq. (5) called
How to approximate the forward-passed mini-batch features of nodes using the learned codewords?
How back-propagate through VQ and estimate the mini-batch gradients of nodes?
How to update the codewords and assignment matrix along with the training of GNN?
In the following part of this section, we introduce the VQ-GNN algorithm by answering all the three questions and presenting a scalability analysis.
Approximated forward and backward message passing, To approximate the forward pass through GNN layer (Eq. ( with mini-batch of nodes (is). we can divide the messages into two categories: intra-mini-batch messages, and messages from out-of-mini-batch nodes; see the right figure of Fig. Intra-mini-batch messages C!(.,)) x can always be computed exactly, where (1.ss) = (.(()) Rgix,, because they only rely on the previous layer's node features of the current mini-batch. Equipped with the codewords (1) and the codeword assignment of all nodes R(D, we can approximate the messages from (i(l.e) (1,, where (1) = diag 1T) T (0) as defined above and Ci(l.s) = ((lle)) R. Here, out Gll.e) is the remaining part of the convolution matrix after removing the intra-mini-batch mes- sages, thus (C(l.e) ).j
nodes as
) E (i)) for any j E {l...., J.,, and Ö(.ss) is the sketch of all.s) In general, we can easily approximate the forward-paseer mini- batch features x-(+1) by X ((++)) o Σ, (Cl..e) x ((l..) x (l)....))).
(1)
5

---

I-th Layer
Min-fuech Feuturcs
Forward Message-Passingg
FormardSOweed Miri-forch Feeture
Codswords
Grid
Used only furwwrr
Min Basch aradlerts
Backward Massaggn-Pssiingg
Min-baich
tradient
only back propagation Both fonword and backward
Figure 2: Three types messages contribute to the mini-batch features and gradients. We only need "red" and "green" messages for the forward- pass. However, "blue" messages are required for prok-propagation. The "red", "blue", and "green" messages character- ized Cout. (CT)ous Jout. and Cin respectively (Eqs. (6) and
Figure 3: each layer. VQ-GNN estimates bowwardppassd mini-batch features using the previ- ous layer's mini-batch features and the feature code- words through approximated forward message- (Eq. (6). The hack-propagated mini-batch gradients are estimated a symmetric manner with the help of gradient codewords (Eq.
However, the above construction of (1+1) does not allow us to pack-prop through VQ straight- forwardly using chain rules. During back-propagation, we aim at approximating the previous layer's mini-batch gradients V x given the gradients of the (approximated) output V g (Eq. (3)
Firstly, we do know how compute the partial derivative of Č(L.A) and ( with respect to X because the learning and updating of VQ codewords and assignment are data dependent and are usually realized by an iterative optimization algorithm. Thus, we need to go through an iterative
computation graph to evaluate the partial derivative of R()) with respect to x-() access to many historical features and gradients, thus violating the scalability constraints. Secondly, even if we apply some very rough approximation during back proppappaaatioon as in [29], that is, assuming
B which requires
that the partial derivative of R(!) with respect X² can be ignored (i.e., the codeword assignment matrix is detached from the computation graph, known as "straight through" back- poock-ropaggtion) we not able to evaluate the derivatives of codewords X (1) because they rely some node features out of the current mini-batch and are not the training device. Generally speaking, designing back nack-propagation rule for VQ under the mini-batch training setup is challenging new problem.
is helpful to re-examine what is happening when we back-propagate on the full graph. In Section 2 we see that back-proe propaaatioannn of a layer of convolution boasdedto----edd GNN can also be realized by message passing (Eq. (3. In Fig. 2 we show the messages related into three types. The "green" and "red" messages are the intra- mntr-mini--btth messages and the messages from out-of-ming cif-miniibatch nodes, respectively. Apart from them, although the "blue" messages to out-of-n nut-of-mini-- batch nodes do not contribute to the forward- passed mini-bated batch features, they are used during back proppaaaaattioo and are an important part the mack-propagated mini-batch gradients. Since both forward- waassd-pass pass and back-pne pock-propagation can be realized by message passing, can we approximate the back- propppppp mini-batch gradients V x in a symmetric manner? We can introduce a set
mini-batch of nodes can be classified
of gradient codewords (((+1) d (R² (,, 1,, RT ((++1) using the same assignment matrix, where ((++1 V & O o' (x is the gradients back-propagated through non-linearity. Each gradient codeword corresponds one-to-one with a feature codeword since we want to use only assignment matrix R. Each pair of codewords concatenated together during VQ updates. Following this idea, we define the approximated forward and backward message passing as follows:
((+1))
((lis) (C(Ca)
(ll out O
Σ
B
wilhaa)
(6)
Jou
passing weight matrix (8.)
mini batch features and fest. codewards
E xin
Σ (Brl))))
(wrll))
mini butch gradients codewords

---

where ((1..)) E (bb+t))((++m) is the approximated message passing weight matrix and is shared during the forward-pass and back prop-ppopagttion process. The lower halves of the left-hand side vectors of Eqs. (6) and 7 are used in neither the forward nor the backward calculations and are never calculated during training or inference. The approximated forward and backward message passing enables the end-to-end mini- training and inference of GNNs and is the core of our VQ-GNN framework.
Error- on estimated features and gradients. We can effectively upper bound the estima- tion errors of mini-bater features and gradients using the relative error of VQ under some mild conditions. For ease of presentation, assume the GNN has only one convolution matrix in the following theorems.
Theorem 2. If the VQ relative error of l-th layer is e0) the convolution matrix ( is either fixed or learnable with the Lipschitz constant of hages (-) 182 + upper-bounded by Lip (hp() and the Lipschitz constant of the non-linearity is Lip(a), then the estimation error of forward-passed mini-batch features satisfies,
| -
(1
(8)
Corollary 3. the conditions in Theorem Si hold and the non-linearity satisfies [=)] any IR, then the estimation error of back- propaapopd mini-batch gradients satisfies,
S for
xxyelle xx X Vl 5 (1+
max 1/
(/ (9)
Note that the error bounds rely on the Lipschitz constant of h(-) when the convolution matrix is learnable. In practice, we can Lipshitz regularize GNNs like GAT 3 without affecting their performance; see Appendix E
the complete algorithm and analysis of scalability. The only remaining question is how to update the learned codewords and assignments during training? In this paper, we use the VQ update rule proposed in [29], which updates the codewords as exponential moving averages of the mSeini-batch inputs; see Appendix [E for the detailed algorithm. We find such an exponential moving average technique suits us well for the mini-batch training of GNNs and resembles the online k-means algorithm. See Fig. for the schematic diagram of VQ-GNN, and the complete pseudo-code is in Appendix
VQ-GNN:
With VQ-GNN, we can mini-batch train and perform inference on large graphs using GNNs, just like regular neural network (e.g., MLP). We have to maintain a small codebook of k codewords and update it for each iteration, which takes an extra (Tk/') memory and O [[mnk]] f') training time per epoch, where I. and f are the numbers of layers and (hidden) features of the GNN respectively. We can effectively preserve all messages related to a mini-batch while randomly sampling nodes from the graph. The number of intra-mini-bate Thus we only need to pass dbbdd/n + 6k) messages per iteration and O + nk) per epoch. In practice, when combined with techniques including product VQ and implicit whitening (see Ap- pendix we can further improve the stability and performance of VQ-GNN. These theoretical and experimental analyses justify the efficiency of the proposed VQ-GNN framework.
messages is Obbdd/n) when the nodes are sampled randomly.
5 Related Work
In this section, we review some of the recent scalable GNN methods and analyze their theoretical memory and time complexities, with a focus on scalable algorithms that can be universally applied to a variety of GNN models (like our VQ-GNN framework), including NS- SAGE [2] Cluster- [9] and GraphSAINT [IO]. We consider GCN here the simplest benchmark. For GCN with L layers and f-dimensional (hidden) features in each layer, when applied to a sparse graph with nodes and m edges (i.e., average degree d = m/n) for ffull-grapht training and inference: the memory usage O + 2) and the training/inference time is O [Lmm + Lnf22). We further assume the graph is large and consider the training and inference device memory is ((b) where is the mini-batch
We call the neighbor sampling method in 2 NS-SAGE and the GNN model in the same paper BAGE-Mean avoid ambiguity.
7

---

Table 2: Memory and time complexities of sampling-basee baaeddd methods and our approach; see Section for details.
Scalable Method
Memory Usage
Pre-computation Time
Training Time
Inference Time
NS-SAGE
O(br"++
O(1r - ---rr-1-p
Cluster-GCN
((LS++If)
O(m)
O(LmS + LAfa)
O(na" nd! fa)
GraphSAINT-RW
O +.f9)
O(T2nf+ IAnf2)
VQ-GNN (Ours)
OLLbf Lf Lkf)
O[Lbdf Lnf Lnkf)
O, Ln(g))
size (i.e., the memory bottleneck limits the mini-batch size), and generally d < b < < m holds. We divide sampling baselines into three categories, and the complexities of selected methods are in Table 2 See Appendix D for more related work discussions.
Neighbor-sampling. Neighbor sampling scheme chooses a subset of neighbors in each layer to reduce the amount of message passing required. NS-SAGE 2 samples neighbors for each node and only aggregates the messages from the sampled node. For a GNN with L layers, (brr) nodes sampled in mini-batch, which leads the complexities growing exponentially with the number of layers L; see Table Therefore, NS-SAGE is not scalable on large graphs for a model with an arbitrary number of layers. NS-SAGE requires all the neighbors to draw non-stochastic predictions in the inference phase, resulting in (dd) inference time since we cannot fit O(n) nodes all at once to the device. VR-GCN 66] proposes a variance reduction technique to further reduce the size r of sampled neighbors. However, VR-GCN requires a O ((Lnf) side memory of all the nodes' hidden features and suffers from this added memory complexity.
Layer-sampling. These methods perform node sampling independently in each layer, which results in a constant sample size across all layers and limits the exponential expansion of neighbor size. FastGCN (7] applies importance sampling to reduce variance. Adapt [25) improves FastGCN by an additional sampling network but also incurs the significant overhead of the sampling algorithm.
Subgraph-sampling, Some proposed schemes sample subgraph for each mini-batch and perform
forward and backward passes on the same subgraph across layers. Cluster-GCA 9 partitions large graph into several densely connected subgraphs and samples subset of subgraphs (with edges between clusters added back) to train in each mini-batch. Cluster-GCN requires ((m) pre- computation time and (bdd) time to recover the intra-cluster edges when loading each mini-batch. GraphSAINT
[0] samples a set of nodes and takes the induced subgraph for mini-batch training. We consider the best- pest-performing variant, GraphSAINT-RW, which uses L steps of random walk to induce subgraph from randomly sampled nodes. ((L)) nodes and edges are covered in each of the n/b mini-batches. Although ((L)) nodes are sampled with some repetition in epoch, the number of edges covered (i.e., messages considered in each layer of a GNN) is also ((L)) which randomly samples nodes for each
and is usually much smaller than m. GraphSAINT- mini-batch, does not suffer from this L factor in the complexities. However, its performance is worse than GraphSAINT-RW's Like NS-SAGE and some other sampling methods, Cluster-GCh and Grap cannot draw predictions on a randomly sampled subgraph in the inference phase. Thus they suffer from the same (dd4) inference time complexity NS-SAGE; see Table
Experiments
In this section, we verify the efficiency, robustness, and universality of VQ-GNN using a series of experiments. See Appendix(Fjicc implementation details and Appendix G for ablation studies and more experiments.
Scalability and efficiency: memory usage, convergence, training and inference time. We sum- marize the memory usage of scalable methods and our VQ-GNN framework in Table Based on the implementations of the PyG library (30), memory consumption of GNN models usually grows linearly with respect to both the number of nodes and the number of edges in mini-batch. On the ogbn-arxiv benehmark, we fix the number of wradient-descendes nodes and the number of messages passed in a mini-batch to be 85K and 1.5M respectively for fair comparisons among the sampling
8

---

Table 3: Peak memory usage, Evaluated when fixing the number of gradient-descended nodes the number of messages passed per mini-hatch be the same. Tested for GCN and SAGE-Mean the ogbn-arxiv henchmark.
Fixed
85K nodes per batch
1.5M messages pussed per butch
GNN Model
GCN
SAGE Mean
GCN
SAGE-Mean
NS-SAGE
1140.3
953.7
Cluster-GCN
501.5 MB
514.1 XB
757.4 NB
769.3 MB
GraphSAINT-RW
526.5 MB
519.2 MB
661.6 MB
650.4 MB
VQ-GNN (Ours)
758.0 MB
801.8
485.5
508.5
0.70
0.70
0.85
Valldation souuoec
0.50
0.35
0:55
0.50
0.50
0.45
0.45
VQ-GRE (Cuns)
0.40
VO-SNN (Durs) GraphssINT RW ChatenCio
GraphSMIKT-NG RW
Cluster GCN
NS-SAGE
20
00
80
100
120
140
20
40
80
1000
120
140
Training time (sec)
Training time (swc)
(a) GCN
(b) SAGE-Mean
Figure 4: Convergence curves (validation accuracy vs. training time). Mini-batch size and learning rate are kept the same. Tested for GCN and SAGE-Mcan on the ogbn-arxiv benchmark.
methods and our approach. VQ-GNN might require some small extra memory when provided with the same amount of nodes per batch, which is the cost to retain all the edges from the original graph. However, our VQ-GNN framework can effectively preserve all the edges connected to a mini-bati batch of nodes (i.e., never drop edges); see Fig. 1 batch, our method can show significant memory efficiency compared with the sampling baselines.
Thus when we fix the number of messages passed per
Fig. 4ssows the convergence comparison of various scalability methods, where we see VQ-GNN is superior in terms of the convergence speed with respect the training time. When training GCN and SAGE-Mean on the ogbn-arxiv benchmark for a specific amount of time (e.g.. 100 s), the validation performance VQ-GNN is always the highest. The training time Fig. eexcludes the time for data loading, pre- pre-processing and validation set evaluation.
Our VQ-GNN approach also leads to compelling inference speed-ups. Despite the training-efficieng issues of GNNs, conducting inference on large-scale graphs suffers from some unique challenges. According to our discussions in Section [5] and following the standard implementations provided by the Open Graph Benchmark (OGB) [5), the three sampling-bas inference procedure) require all of the L-hopp neighbors of the mini-batch nodes lie on the device once during the inference phase. The inference time of SAGE-Mean trained with sampling-methon on the ogbn-arxiv benchmark is 1.61 s, while our method can accelerate inference by order of magnitude and reduce the inference time to 0.40 s.
baselines (which share the same
Performance comparison across various datasets, settings, and tasks. We validate the efficacy of our method on various benchmarks in Table The four representative benchmarks are selected because they have very different types of datasets, settings, and tasks. The ogbn-arxiv benchmark is a common citation network of aiXiv papers, while Reddit is very dense social network of Reddit posts, which has much more features per node and larger average node degree; see Table Appendix F for detailed statistics of datasets. PPI is node classification benchmark under the inductive learning setting, i.c., neither attributes nor connections of test nodes are present during training, while the other benchmarks are all transductive. VQ-GNN can be applied under the inductive setting with only one extra step: during the inference stage, we now need to find the codeword assignments (i.e., the nearest codeword) of the test nodes before making predictions since have access to the test nodes during training. Neither the learned codewords nor the GNN parameters are updated during inference. ogbl-collab is link prediction benchmark where the labels and loss are intrinsically different.
is very challenging for a scalable method perform well on all benchmarks. Table we confirm that VQ-GNN is more robust than the three sampling- based methods. Across the four benchmarks,

---

Table 4: Performance comparison between sampling-based haselines and our approach, VQ-GNN.
Task
Node Classification (Transductive)
Node Clussification (Transductive)
Benchmark
ogbi-arxxiv (Acc.istd)
Reddit (Acc.=std.)
GNN Model
GCN
SAGE-Mean
GAT
GCN
SAGE-Mean
GAT
"Full-Graph"
.7029 1. .0036
.6982 .0038
.7097
.0055
OOMF
OOM
OOM?
NS-SAGE.
NA
.7094 + 00000
.7123
.0044
NA
.9615 + .0089
.9426
.0043
Cluster-GCN
6805 + .0774
.976 + 00049
.6660 + .0062
.9664 + .0034
.4566 + .0061
.980 + .0055
GraphSAINT-RW
.7079 + .0377
.6987 = 0039
.7117 = .0032
.9255 = .0057
.9581 = 0007
.9431
.0067
VQ-GNN (Ours)
.0555 + .0033
.7028 = 00047
.7033 + .0034
.3999 = 0001
.4499 = 0004 .9438
.0059
Task
Node Classification (Induetive)
Link Prediction (Trunsductive) ogbl-collab (Hits@502sstd)
Benchmark
PPI (FI tstre
GNN Model
GCN
SAGE-Mean
GAT
GCN
SAGE-Mean
GAT
"Full-Graph"
9173 1 .009
.3558
.0046
.9222 1: .0035
.475 1 .0107
.810 1 00081
.4048
.0125
NS-SAGE.
NA
.9121
.0033
.9407
.0055
NA
.4766 + .041
.3499
.0142
Cluster-GCN
.852 1. .066
.8810 1 .0001
.9051 1: .077
.4068 + .0066
3486 1 0216
.3905 1 .0152
GraphSAINT-RW
9110 # .0557
.382 = .0744
.9612 + .0442
.4368 = .0109
3359 + .0128
.389- = .0114
VQ-GNN (Ours)
.9499 1 .0588
.9788 1 .0019
.937 1 0033
.416 + 00134
4673 1 .0644
.4102 1 .0090
1 "NS-AGGE sampling method not compatible with the GCN backbone. "OOM"" refers "out memory". The
"full-graph" training on the Reddit benchmark requires more than 11 GB of memory. multiple labels per node, and the evaluation metric is Fl score instcad of accuracy.
"The PPI benchmark comes with
VQ-GNN can always achieve performance similar with or better than the oracle "full-graph" training performance, while the other scalable algorithms may suffer from performance drop in some cases. For example, NS-SAGE fails when training GAT on ogbl-collab, Cluster-GCN consistently falls behind on PPI, and GraphSAINT-ROE Mean and GAT backbones. We think
performance drops on the ogbl-collat when using SAGE- the robust performance of VQ-GNN is its unique value among the many other scalable solutions. VQ-GNN framework is robust because it provides bounded approximations of gralhh" training (Theorem 2 and Corollary
B while most the other scalable algorithms do not enjoy such theoretical guarantee. VQ-GNN is also universal to various backbone models, including but not limited to GCN, LAGE-Mean and GAT shown here; see Appendix G for more experiments on GNNs that utilize multi-hop neighborhoods and global context, e.g., graph transformers.
7 Conclusion
Summary of our framework: strengths, weaknesses, future directions, and broader impacts. This paper introduced the proposed VQ-GNN framework, which can scale most of GNNs to large graphs through a principled and fundamentally different approach compared with sampling- based methods. We have shown both theoretically and experimentally that our approach is efficient in memory usage, training and inference time, and convergence speed. VQ-GNN can be universally applied to most GNN models and different graph learning tasks and can equivalently scale-up GNNs utilizing many-hops-away or global context for each layer. However, the performance of VQ-GNN relies on the quality of approximation provided by VQ. In practice, for VQ to work adequately in GNN, a set of techniques are necessary. Because the limited time, we did not heuristically explore all possible techniques or optimize the VQ design. Given that our preliminary design of VQ in GNN already achieved competitive performance compared with the state-of-the hypothesize that further optimization of VQ design could improve performance. We hope our work opens up promising new avenues of research for scaling up GNNs, which also has the potential to be applied to other data domains wherever the size of a single sample is large, e.g., long time-series or videos. Considering broader impacts, we view our work mainly as a methodological and theoretical contribution, which paves the way for more resource-efficient graph representation learning. We envision our methodological innovations can enable more scalable ways to do large-network analysis for social good. However, progress graph embedding learning might also trigger other hostile social network analyses, e.g., extracting fine-grained user interactions for social tracking.
sampling baselines, we
10

---

Acknowledgments and Disclosure of Funding
Goldstein, Kong, and Chen were supported by the Office of Naval Research, AFOSR MURI program, the DARPA Young Faculty Award, and the National Science Foundation Division of Mathematical Sciences. Additional support was provided by Capital One Bank and JP Morgan Chase. Huang and Ding were supported by startup fund from the Department of Computer Science of the University of Maryland, National Science Foundation IIS-1850220 CRII Award 030742-00001, NOD-DARPA-Defente Advanced Research Projects Agency Guaranteeing AI Robustness against Deception (GARD), Air Force Material Command, and Adobe, Capital One and JP Morgan faculty fellowships. Li and Dickerson were supported in part by NSF CAREER Award IIS- 1846237, NSF D-ISN Award #2039862, NSF Award CCF- 1852352. NIH RO1 Award WLM-013039-M NIST MSE Award #20126334, DARPA GARD 12020001, DoD WHS Award #HQ003420F0035, ARPA-E Award #4334192 and a Google Faculty Research Award.
References
[1] Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. In ICLR, 2016.
[2] William L Hamilton, Zhitao Ying, and Jure Leskovec. Inductive representation learning on large graphs. In NeurIPS, 2017.
[3] Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. In ICLR, 2018
[4] Keyulu Xu. Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful graph neural networks? In ICLR,, 2019.
[5] Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. arXiv preprint (rXiv:2005.00683: 2020.
[6] Jianfei Chen, Jun Zhu, and Le Song. Stochastic training of graph convolutional networks with variance reduction. In JCML, 2018.
[7] Jie Chen, Tengfei Ma. and Cao Xiao. Fastgen: Fast learning with graph convolutional networks via importance sampling. In ICLR, 2018.
[8] Difan Zou, Ziniu Hu, Yewen Wang, Song Jiang, Yizhou Sun, and Quanquan Gu. Layer-dependent impor- tance sampling for training deep and large graph convolutional networks, arXiv preprint arXiv: 191.00333,, 2019.
[9] Wei-Lin Chiang, Xuanqing Liu, Si Si, Yang Li, Samy Bengio, and Cho-Jui Hsich. Cluster-gen: An efficient algorithm for training deep and large graph convolutional networks. In Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pages 257-266. 2019.
[10] Hanqing Zeng, Hongkuan Zhou, Ajitesh Srivastava, Rajgopal Kannan, and Viktor Prasanna. Graphsaint: Graph sampling based inductive learning method. In ICLR, 2019.
[11] Gabriele Corso, Luca Cavalleri, Dominique Beaini, Pietro Liò, and Petar Veličkovid. Principal neighbour- hood aggregation for graph nets. NeurlPS, 33, 2020.
[12] Michaël Defferard, Xavier Bresson, and Pierre Vandergheynst. Convolutional neural networks on graphs with fast localized spectral filtering. In NewrIPS, 2016.
[13] Ron Levic, Federico Monti, Xavier Bresson, and Michael M Bronstein. Cayleynets: Graph convolutional neural networks with complex rational spectral filters. JEEE Transactions on Signal Processing, 2018.
[14] Filippo Maria Bianchi, Daniele Grattarola, Lorenzo Livi, and Cesare Alippi. Graph neural networks with convolutional arma filters. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2021.
[15] OmriPuny Heli YaronLipman. Global attention improves graph networks generalization. arXiv preprint 2xXiv:2006607844, 2020.
[16] Yu Rong. Yatao Bian, Tingyang Xu, Weiyang Xie, Ying Wei, Wenbing Huang, and Junzhou Huang. Self Self-supervised graph transformer large-scale molecular data. NeurIPS, 33, 2020.
11

---

[17] Jiawei Zhang, Haopeng Zhang, Congying Xia, and Li Sun. Graph-bert: Only attention is needed for learning graph representations. arXiv preprint 4rX1v:2001.05440, 2020.
[18] Johannes Klicpera, Stefan Weißenberger, and Stephan Günnemann. Diffusion improves graph learning. In NeurIPS, 2019.
[19] James Atwood and Don Towsley, Diffusion-convolationa neural networks. In NewrlPS, pages 1993-2001, 2016.
[20] Christopher Morris, Martin Ritzert, Matthias Fey. William L Hamilton, Jan Eric Lenssen, Gaurav Rattan, and Martin Grohe. Weisfeiler and leman go neural: Higher-order graph neural networks. In AAAI, 2019.
[21] Haggai Maron, Heli Ben-Hamu, Hadar Serviansky, and Yaron Lipman. Provably powerful graph networks. In NeurIPS, 2019.
[22] Keyulu Xu, Jingling Li, Mozhi Zhang. Simon S Du, Ken-ichi Kawarabayashi, and Stefanie Jegelka. What can neural networks reason about? In ICLR, 2020.
[23] Andreas Loukas. What graph neural networks cannot learn: depth width. In ICLR, 2020.
[24] Muhammet Balcilar, Renton Guillaume, Pierre Héroux, Benoit Gaüzère, Sébastien Adam, and Paul Honeine. Analyzing the expressive power of graph neural networks in a spectral perspective. ICLR, 2021.
[25] Wenbing Huang, Tong Zhang. Yu Rong, and Junzhou Huang. Adaptive sampling towards fast graph representation learning. NeurIPS, 2018.
[26] William Johnson and Joram Lindenstrauss. Extensions of lipschitz mappings into a hilbert space. Contemporary mathematics, (66189-20066 1984.
[27] Daniel Kane and Jelani Nelson. Sparser johnson lindenstrauss transforms. Journal of the ACM (JACM). 61(1):1-23,200 2014.
[28] Robert M. Gray and David L. Neuhoff. Quantization, IEEE transactions on information theory, 44(6): 2325-2383, 1998.
[29] Aaron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learning. arXiv preprint 1:Xiv:/711.0090 2017.
[30] Matthias Fey and Jan Eric Lenssen. Fast graph representation learning with pytorch geometric. arXiv preprint arXiv: orXiv:/903.022 023002,,,, 2019.
[31] Yujia Li, Daniel Tarlow, Mare Brockschmidt, and Richard Zemel. Gated graph sequence neural networks. In ICLR, 2015.
[32] Chris Ding and Xiaofeng He. K-means clustering via principal component analysis. In Proceedings of the twenty-first international conference on Machine learning. page 29, 2004.
[33] Felix Wu, Amauri Souza, Tianyi Zhang, Christopher Fifty, Tao Yu, and Kilian Weinberger. Simplifying graph convolutional networks. In International conference machine learning, pages 6861-6871. PMLR, 2019.
[34] Aleksandar Bojehevski, Johannes Klicpera, Bryan Perozzi, Amol Kapoor, Martin Blais, Benedek Rózem- berezki, Michal Lukasik, and Stephan Günnemann. Scaling graph neural networks with approximate pagerank. In Proceedings of 26th ACM SIGKDD International Conference on Knowledge Discovery Data Mining, pages 2464-2473, 2020.
[35] Emanuele Rossi, Fabrizio Frasca, Ben Chamherlain, Davide Eynard, Michael Bronstein, and Federico Monti. Sign: Scalable inception graph neural networks. arXiv preprint arXiv:2004.11193 2020.
[36] Shyam A Tailor, Javier Fernandez- Fermandez-Marques, and Nicholas D Lane. Degree- Quantization-awwe aware training for graph neural networks. arXiv preprint arXiv:2008.05000 2020.
[37] Boyuan Feng, Yuke Wang, Xu Li, Shu Yang, Xueqiao Peng, and Yufei Ding. Sgquant: Squeezing the last bit on graph neural networks with specialized quantization. In 2020 JEEE 32nd International Conference Tools with Artificial Intelligence (ICTAI), pages 1044-1052. IEEE, 2020.
[38] Ali Razavi, Aaron van den Oord, and Oriol Vinyals. Generating diverse high-fidelit images with vq-vae-2. arXiv preprint 1906 1906.00446. 2019.
12

---

[39] Lars Maaloe, Marco Fraccaro, Valentin Liévin, and Ole Winther. Biva: A very deep hierarchy of latent variables for generative modeling. arXiv preprint 1rXiv:190..02122, 2019.
[40] Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsu- pervised learning of visual features by contrasting cluster assignments. arXiv preprint arXiv:2006.0988. 2020.
[41] Alexci Baevski, Steffen Schneider, and Michael Auli vq-wav2vec: Sclf-supervised learning discrete speech representations. arXiv preprint 1rXiv:1910.05553, 2019.
[42] Alexei Baevski, Henry Zhou, Abdelrahman Mohamed, and Michael Auli. wav2vec 2.0; A framework for self-supervised learning of speech representations. arXiv preprint arXiv:2006. 2020.
[43] Sinong Wang. Belinda Li, Madian Khabsa, Han Fang. and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04760 2020.
[44] Zhengyang Geng, Meng- Hao Guo, Hongxu Chen, Xia Li, Wei, and Zhouchen Lin. Is attention better than matrix decomposition. In International Conference on Learning Representations, 2021.
[45] Hanwei Wu and Markus Flierl. Learning product codebooks using vector-quantized autoencoders for image retrieval. In 2019 IEEE Global Conference on Signal and Information Processing (GlobalSIP), pages 1-5. IEEE, 2019.
[46] David Berthelot, Colin Raffel, Aurko Roy, and Ian Goodfellow, Understanding and improving interpolation in autoencoders via adversarial regularizer. arXiv preprint arXiv: 00543, 2018.
[47] George Dasoulas, Kevin Scaman, and Aladin Virmaux. Lipschitz normalization for self-attention layers with application graph neural networks. arXiv preprint arXiv:2 04886, 2021.
[48] Matthias Fey, Jan E Lenssen, Frank Weichert, and Jure Leskovec. Gnnautoscale: Scalable and expressive graph neural networks via historical embeddings. arXiv preprint 0rX6; 2021.
Checklist
The checklist follows the references. Please read the checklist guidelines carefully for information on how to answer these questions. For each question, change the default [TODO] [Yes], [No], or [N/A]. You are strongly encouraged to include justification to your answer, either by referencing the appropriate section of your paper or providing brief inline description. For example:
Did you include the license to the code and datasets? [Yes] See Section.
Did you include the license to the code and datasets? [No] The code and the data are proprietary.
Did you include the license to the code and datasets? [N/A]
Please do not modify the questions and only use the provided macros for your answers. Note that the Checklist section does not count towards the page limit. In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questions/answers below.
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See Section 7
(c) Did you discuss any potential negative societal impacts of your work? [Yes] See Sec- tion 7
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of theoretical results? [Yes] See Theorems]
and 2 and Corollary 3
(b) Did you include complete proofs of all theoretical results? [Yes] See Appendix B
13

---

3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material
URL)? [Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] See Appendix and the code.
(c) Did you report error bars (c.g., with respect the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, cloud provider)? [Yes] See Appendix FF
4. If you using existing assets (e.g., code, data, models) curating/releasing assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [Yes]
(c) Did you include any new assets either in the supplemental material URL? [No]
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [N/A]
(e) Did you discuss whether the data you using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
14


Reviews:
[{"review_id": "N8gp86i_DR9", "reviewer": "Reviewer_f3BS", "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.\nTechnically, this paper develops a new approximation for message passing on codebook-represented graphs.\nIn theoretical evaluation, the paper shows that the prediction errors due to VQ can be upper-bounded by using Lipshitz regularization.\nIn the experimental evaluation, the generalization performance of the proposed VW method is comparable to existing scalable GNNs.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Originality: \nVQ + DNNs has already been studied in the famous VQ-VAE [29] paper and its followers. However, the combination of GNNs and VQ is new to me. Reducing the dimensionality of the embedding vectors is certainly a natural approach to memory compression. If there has been no previous attempt pursuing the scalability of GNNs in this direction, then I think this paper proposes a novel and valuable idea to the GNN community.\n\nThe discussion in Sec. 3 shows that if a good dimensionality reduction matrix R can be found, we can approximate well the node feature vectors. In this paper, VQ is chosen as one of the dimensionality reduction methods because it is a ''natural and widely-used method'', but are there any other ''natural and widely-used'' dimensionality reduction method choices? For example, if we don't need to use a fixed codebook, aren't PCA and FLDA also candidates?　\nI expect more discussions on this issue to justify the choice of VQ. \n\nThe appendix acknowledges that the quantized representation of Degree-Quant [35] improves the memory efficiency of GNNs, but [35] does not consider scalability and therefore does not conflict with the contribution of this paper. \nI'm not fully sure regarding this point. \nThe main objective of both this paper and [35] is to reduce the memory burden of the computational device while avoiding subsampling. This paper and [35] try to achieve the same objective with different approaches: vector quantization in this paper and numerical quantization in [35]. Then these are not orthogonal, I think. \nIt is very possible that I misunderstand the essences, and I would appreciate clarification on this point in the rebuttal discussion.\n\nI think that the new algorithm for learning correctly on VQ codebook-represented graph mini-batches is a great job since it is technically non-trivial and the outcome seems reasonable. \n\n## Quality: \nExcept for the details of some of the formulas, I find no major problems or concerns with the technical discussion in the main manuscript.\n\nI did not follow the details of the proofs.\n\n## Clarity:\nIn the left-hand side of equations (6.7), the lower halves of the vectors are omitted. \nWhat variables are being addressed here? Is it a variable that is not used in either the forward or backward calculations?\n\nTable 3 is difficult to understand for me. \nIn particular, the comparison of ``memory usage when the number of message passing is fixed'' What is the main message here? \nI read the paper several times, but in the end, I cannot figure it out. It's a concept I haven't seen in the context of many GNN papers, so unless you can explain it more clearly, I can't read any advantage of the proposed method from this table.\n\nL302-303 ''We summarize the training and inference time........ in Table 3.\" should be Table 4. \n\n## Significance: \n\nThe main weakness of this paper is that the experimental results.\n\nThe main purpose of this paper is, in my understanding, to propose a GNN technique with scalability to large graphs. However, Table 2 says that the memory usage of the proposed method is comparable to that of existing subsampling methods. This implies that there is no large graph dataset such that the proposed method is the first and the only solution for scalable GNNs. \n\nAnother issue is the generalization performance. Unlike existing subsampling methods, the proposed method is unique in that it can retain all nodes and layers (L52-53). This paper examines the advantage of this characteristic in terms of numerical performance in Table 5. Among the four scalable GNNs, the proposed method obtains the highest performance in only one out of six cases. The best performer is NS-SAGE, which takes first place in three cases, and GraphSAINT is the runner-up, best in two cases.\n\nThis means that the proposed method does not show outstanding results or unique values among several solutions that provide scalability for GNNs.\n\nSo, why should we prepare a newly proposed method instead of using existing methods that have already been evaluated and implemented (e.g., GraphSAGE)? It is difficult to answer this question clearly, therefore it is difficult for me to recommend for acceptance in the current manuscript status.\n\nIn terms of computation time and time complexity, the proposed method is clearly superior to others (Table 2, 4). Therefore, I suggest that it would be better to focus on the point of fast computation rather than the general scalability.\n\n\n(+) VQ for dimensionality reduction for GNN is (probably) new\n(+) Approximate message passing for VQ codebook developed\n(+) Theoretical guarantee for bounded prediction error\n(-) More discussions for the choice of VQ dimensionality reduction are expected. \n(-) [35] is on the totally different lines of researchers, I'm not fully sure about that. \n(--) Experimental results for scalability and generalization are weak. \n(-) What is the main message of Table 3? \n\n## After feedbacks\nI found the author feedbacks effectively resolve many of my misunderstandings and concerns. \nI raise my evaluations to 6, weak accept. Good luck!", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Many VQ-related techniques are required to make the proposed method work satisfactorily. Discussed in Sec. 7. ", "review_text": "", "extra_scores": {"ethical_concerns": "No specific concerns", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "LnXvxwbpdz", "reviewer": "Reviewer_UTjD", "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.", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "1) Although this paper is about scaling up the GNN training, the experiment is only on small-scale datasets, only one dataset for node classification and one for link prediction. It would be interesting to see larger scale experiments with nodes size in millions.\n\n2) Seems this paper is about scaling up the training, it would be interesting to check the real training time vs accuracy. Why in table 4, per epoch time is faster than Cluster-GCN and GraphSAINT?   It seems to me that per epoch the proposed method will be slower as it needs additional time to update the codebook.\n\n3) How is the performance w.r.t. number of layers? how about inductive and transductive settings?\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "NA", "review_text": "", "extra_scores": {"ethical_concerns": "NA", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3 hours"}, "extra_text": {}}, {"review_id": "73tbtJA1hWa", "reviewer": "Reviewer_ai2z", "summary": "\nThe paper presents a vector quantization method to reduce the memory consumption of the convolutional operations for scalable GNNs. The method avoids subsampling nodes in the mini batches to prevent all the messages could be kept in the message passing operations. It quantizes the convolutional matrices leveraging the dimensionality reduction method. The approximated message passing with help of the quantization codewords is shown to have low relative error theoretically under certain conditions and have comparable accuracy with the original algorithm empirically.\n", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Positive points:\n1. The motivation and challenges are well-presented and related work seems to be adequate.\n2. I like the idea of applying dimensionality reduction technique to quantize the feature and convolutional matrices in the GNNs. The experimental results seem convincing, though limited experiments are provided.\n\nComments:\n1.\tThe notations and explanations are not very clear, especially in the ``approximated forward and backward message passing`` part of section 4.\n2.\tI would prefer to see more experiments and discussion on the hyperparameters selection and their affects to the algorithm such as the selection of the reduced dimension the mini-batch sizes, and the effect of different mini-batch sampling strategies.\n\nDuring the rebuttal phase, the authors add additional experiments following the suggestions and I'm satisfied with the feedback and additional clarification given by the author. Considering the theoretical/experimental impact of the paper, I'd keep my rating towards marginal acceptance.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3 hours"}, "extra_text": {}}, {"review_id": "6-M9bfXtxE", "reviewer": "Reviewer_UGM2", "summary": "This paper presents VQ-GNN: a scalable framework for training graph neural networks that approximates full-graph message passing with a memory of global notes that is updated during training. The authors derive approximate forward and backward passes and show that the network can preserve full-batch results while reducing computational and memory costs.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I think the general idea of approximating the general form of convolutional matrices with a low-rank decomposition is good. Is there a connection to the Reformer architecture (Kitaev, Kaiser et al 2020) and locality-sensity-hashing approximations to full attention matrices in Transformers?\n\nThe proposed method demonstrates substantial gains on link-prediction tasks relative to alternative scalable GNN training methods. The performance (accuracy / hits) is similar - and in some cases better - than the full-graph algorithm. Combined with the memory and training / inference speed gains, it makes for a compelling set of results.\n\nI like the analysis of the fundamental scalability limitations of other ‘scalable’ GNN training algorithms. The highlight is the futility of managing the exponential receptive field dependence on the number of layers by dropping nodes or layers. In general the paper contains extensive theoretical results, and analyses of the method’s complexity relative to alternatives.\n\nIt would be useful to see some experimental analysis on codebook / minibatch size? For instance a plot of model performance / computational costs as a function of the minibatch or codebook size. Otherwise it’s not clear how sensitive the framework is to these factors.\n\nIt seems like a potentially important bit of experimental detail is added in the appendix [L641-646]:\n\n“To mitigate the error induced by VQ in the high-dimensional space, we split feature vectors into small pieces. In practice, we find that when the dimension of each piece is 4, our algorithm generally works well. When the split dimension is 4 we have 32 separate branches each layer to do the VQ. These branches are independent and can be paralleled. At the end of each layer, separated feature vectors are concatenated together to restore the original hidden dimension, and the restored feature is input to the next layer\n\nI found it a bit hard to understand this. Are the 4-dimensional vector pieces independently assigned to codebooks? If this is an important component then it would be useful to report some analysis in the paper. And the same for the other VQ-details: whitening and Lipshitz regularization. If the model is highly sensitive to these elements then that is a limitation. \n \nOverall, the proposed VQ-GNN framework is well motivated, well-executed and has promising results. Some additional analysis would be welcome, and potentially a restructuring so that some of the important practical details (product VQ etc) are made more apparent in the main body.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors have addressed the limitations around the tricks used to make the VQ work well. I think it is reasonable to explore these design choices in greater detail in future work. The authors briefly discuss societal risks relating to the use of large-scale GNNs in exploitative analysis of social networks.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}]

Author rebuttal:

Meta-review: 3 of 4 ratings were "accept". The concerns from the reviewer who gave the lowest rating (a 5) were not that clear to me through the extended  conversation with the authors, and I'm not swayed by their concerns.

The authors did extensive further experiments in response to the feedback, which reinforced their results.

Overall the believe the paper is a borderline-to-clear accept, given the novelty and performance.
Decision: Accept (Poster)