Sketch-GNN: Scalable Graph Neural Networks with Sublinear Training Complexity
Abstract
Graph Neural Networks (GNNs) are widely applied to graph learning problemssuch as node classiﬁcation.
When scaling up the underlying graphs of GNNs to alarger size, we are forced to either train on the complete graph and keep the fullgraph adjacency and node embeddings in memory (which is often infeasible) ormini-batch sample the graph (which results in exponentially growing computationalcomplexities with respect to the number of GNN layers).
Various sampling-basedand historical-embedding-based methods are proposed to avoid this exponentialgrowth of complexities.
However, none of these solutions eliminates the lineardependence on graph size.
This paper proposes a sketch-based algorithm whosetraining time and memory grow sublinearly with respect to graph size by training
GNNs atop a few compact sketches of graph adjacency and node embeddings.
Based on polynomial tensor-sketch (PTS) theory, our framework provides a novelprotocol for sketching non-linear activations and graph convolution matrices in
GNNs, as opposed to existing methods that sketch linear weights or gradientsin neural networks.
In addition, we develop a locality sensitive hashing (LSH)technique that can be trained to improve the quality of sketches.
Experiments onlarge-graph benchmarks demonstrate the scalability and competitive performanceof our Sketch-GNNs versus their full-size GNN counterparts.
Introduction
Graph Neural Networks (GNNs) have achieved the state-of-the-art graph learning in numerousapplications, including classiﬁcation [26], clustering [3], recommendation systems [42], socialnetworks [16] and more, through representation learning of target nodes using information aggregatedfrom neighborhoods in the graph.
The manner in which GNNs utilize graph topology, however,makes it challenging to scale learning to larger graphs or deeper models with desirable computationaland memory efﬁciency.
Full-batch training that stores the Laplacian of the complete graph suffersfrom a memory complexity of O ( m + ndL + d 2 L ) on an n -node, m -edge graph with node featuresof dimension d when employing an L -layer graph convolutional network (GCN).
This linear memorycomplexity dependence on n and the limited memory capacity of GPUs make it difﬁcult to train onlarge graphs with millions of nodes or more.
As an example, the MAG240M-LSC dataset [21] is anode classiﬁcation benchmark with over 240 million nodes that takes over 202 GB of GPU memorywhen fully loaded.
To address the memory constraints, two major lines of research are proposed: (1) Sampling-basedapproaches [18, 11, 12, 14, 44] based on the idea of implementing message passing only betweenthe neighbors within a sampled mini-batch; (2) Historical-embedding based techniques, such as
GNNAutoScale [17] and VQ-GNN [15]), which maintain the expressive power of GNNs on sampledsubgraphs using historical embeddings.
However, all of these methods require the number of minibatches to be proportional to the size of the graph for ﬁxed memory consumption.
In other words, they signiﬁcantly increase computational time complexity in exchange for memory efﬁciency whenscaling up to large graphs.
For example, training a 4-layer GCN with just 333K parameters (1.3 MB)for 500 epochs on ogbn-papers100M can take more than 2 days on a powerful AWS p4d.24x largeinstance [21].
We seek to achieve efﬁcient training of GNNs with time and memory complexities sublinear ingraph size without signiﬁcant accuracy degradation.
Despite the difﬁculty of this goal, it should beachievable given that (1) the number of learnable parameters in GNNs is independent of the graphsize, and (2) training may not require a traversal of all local neighborhoods on a graph, but ratheronly the most representative ones (thus sublinear in graph size) as some neighborhoods may bevery similar.
In addition, commonly-used GNNs are typically small and shallow with limited modelcapacity and expressive power, indicating that a modest proportion of data may sufﬁce.
This paper presents Sketch-GNN , a framework for training GNNs with sublinear time and memorycomplexity with respect to graph size.
Using the idea of sketching, which maps high-dimensionaldata structures to a lower dimension through entry hashing, we sketch the n × n adjacency matrixand the n × d node feature matrix to a few c × c and c × d sketches respectively before training,where c is the sketch dimension.
While most existing literature focuses on sketching linear weightsor gradients, we introduce a method for sketching non-linear activation units using polynomial tensorsketch theory [19].
This preserves prediction accuracy while avoiding the need to “unsketch” back tothe original high dimensional graph-node space n , thereby eliminating the dependence of trainingcomplexity on the underlying graph size n .
Moreover, we propose to learn and update the sketches inan online manner using learnable locality sensitive hashing (LSH) [9].
This reduces the performanceloss by adaptively enhancing the sketch quality while incurring minor overhead sublinear in graph size.
In practice, we ﬁnd that the sketch-ratio c/n required to maintain “full-graph” model performancedrops as n increases; as a result, our Sketch-GNN enjoys sublinear training scalability.
Sketch-GNN applies sketching techniques to GNNs to achieve training complexity sublinear to thedata size.
This is fundamentally different from the few existing works which sketch the weights orgradients [29, 13, 25, 28, 36] to reduce the memory footprint of the model and speed up optimization.
Our approach is ﬂexible to architecture and has the potential to be generalized to other neural networksand data types, e.g., CNNs on gigapixel images.
To the best of our knowledge, Sketch-GNN is theﬁrst sub-linear complexity training algorithm for GNNs based on LSH and tensor sketching.
Thesublinear efﬁciency obtained applies not only to GNNs with a ﬁxed convolution matrix, such as GCN [26] and GraphSAGE [18], but also to GNNs with learnable convolution matrices, such as GAT [38].
Experiments on several large graph datasets, such as ogbn-product [21] with 2.45M nodes, demonstrate that Sketch-GNNs can match the performance of the standard model trained on the completegraph, while requiring signiﬁcantly reduced computations and memory for both ﬁxed (GCN, GraphSAGE) and learnable convolution (GAT) models.
For instance, SketchGCN on ogbn-arxiv [21]is 72% and 55% faster than the corresponding full-graph and sampling-based (GraphSAINT [44])baselines, while the pre-processing time is just 14% of overall reduction (when running 500 epochs).
Preliminaries
Basic Notations.
Consider a graph with n nodes and m edges.
Connectivity is given by the adjacencymatrix A ∈ { 0 , 1 } n × n and features on nodes are represented by the matrix X ∈ R n × d , where d is thenumber of features.
Given a matrix C , let C i,j , C i, : , and C : ,j denote its ( i, j ) -th entry, i th row, andj th , respectively.
denotes the element-wise (Hadamard) product, whereas C (cid:12) k representsthe k -th order element- power.
(cid:107) · (cid:107) F is the symbol for the Frobenius norm.
I n ∈ R n × n denotesthe identity matrix, whereas 1 n ∈ R n is the vector whose elements are all ones.
Med {·} representsthe element-wise median over a set of matrices.
Superscripts are used to indicate multiple instancesof the same kind of variable; for instance, X ( l ) ∈ R n × d l are the node representations on layer l .
Uniﬁed Framework of GNNs.
A Graph Neural Network (GNN) layer receives the node representation of the preceding layer X ( l ) ∈ R n × d as input and outputs a new representation X ( l +1) ∈ R n × d ,where X = X {+(
 ) ∈+} R n × d are the input features.
Although GNNs are designed following differentguiding principles, such as neighborhood aggregation (GraphSAGE), spatial convolution (GCN),self-attention (GAT), and Weisfeiler-Lehman (WL) alignment (GIN [43]), the great majority of
GNNs can be interpreted as performing message passing on node features, followed by feature transformation and an activation function.
The update rule of these GNNs can be summarized as [15]
[Equation]
Where C ( q ) ∈ R n × n denotes the q -th convolution matrix that deﬁnes the message passing operator,q ∈ Z + is index of convolution, σ ( · ) is some choice of nonlinear activation function, and W ( l,q ) [-
-]∈
R d l × d l +1 denotes the learnable linear weight matrix for the l -th layer and q -th ﬁlter.
GNNs under thisparadigm differ from each other by their choice of convolution matrices C ( q ) , which can be eitherﬁxed (GCN and GraphSAGE) or learnable (GAT).
In Appendix A.1, we re-formulate a number ofwell-known GNNs under this framework.
Unless otherwise speciﬁed, we assume q = 1 and d = d lfor every layer l ∈ [ L ] for notational convenience.
Count Sketch and Tensor Sketch.
(1) Count sketch [7, 40] is an efﬁcient dimensionality reduction method that projects an n -dimensional vector u into a smaller c -dimensional space using a random hash table h : [ n ] → [ c ] and a binary Rademacher variable s : [ n ] → {± 1 } ,where [ n ] = { 1 , . . . , n } .
Count sketch is deﬁned as CS ( u ) i
= (cid:80) h ( j )=
i s ( j ) u j , which is alinear transformation of u , i.e., CS ( u ) =
R u .
Here, R ∈ R c × n denotes the so-called countsketch matrix , which has exactly one non-zero element per column.
(2) Tensor sketch [31] isproposed as a generalization of count sketch to the tensor product of vectors.
Given z ∈ R nand an order k , consider a k number of i.i.d.
hash tables h (1) , . . .
, h ( k ) :
[ n ] →
[ c ] and i.i.d.binary Rademacher variables s (1) , . . .
, s ( k ) :
[ n ] → {± 1 } .
Tensor sketch also projects vectorz ∈ R n into R c , and is deﬁned as TS k ( z )
i = (cid:80) h ( j 1 , ··· ,j k ) =
i s (1) ( j 1 ) · · · s ( k ) ( j k )
z j 1 · · · z j k , whereh ( j 1 , · · · , j k ) =
( h (1) ( j 1 ) + · · · + h ( k ) ( j k ))
mod c .
By deﬁnition, a tensor sketch of order k = 1degenerates to count sketch; TS 1 ( · ) =
CS ( · ) .
(3) We deﬁne count sketch of a matrix U ∈ R d × n as thecount sketch of each row vector individually
, i.e., CS ( U ) ∈ R d × c where [ CS ( U )] i, : =
CS ( U i, : ) .
Thetensor sketch of a matrix is deﬁned in the same way.
Pham and Pagh [31] devise a fast computationof tensor sketch of U ∈ R d × n (sketch dimension c and order k ) using count sketches and the Fast
Fourier Transform (FFT):
[Equation]
(2) where CS ( p ) ( · ) is the count sketch with hash function h ( p ) and Rademacher variable s ( p ) .
FFT ( · )and FFT − 1 ( · ) are the FFT and its inverse applied to each row of a matrix.
Locality Sensitive Hashing.
The deﬁnition of count sketch and tensor sketch is based on hashtable(s) that only requires a data independent uniformity, i.e., with high probability the hash-bucketsare of similar size.
In contrast, locality sensitive hashing (LSH) is a hashing scheme that useslocality-sensitive hash function H : R d → [ c ] to ensure that nearby vectors are hashed into thesame bucket (out of c buckets in total) with high probability while distant ones are not.
SimHashachieves the locality-sensitive property by employing random projections [8].
Given a random matrix
P ∈ R c/ {+
 ×+} d , SimHash deﬁnes a locality-sensitive hash function H (
[Equation]
= [Equation] (3) where [ · (cid:107) · ] denotes concatenation of two vectors and arg max returns the index of the largestelement.
SimHash is efﬁcient for large batches of vectors [1].
In this paper, we apply a learnableversion of SimHash that is proposed by Chen et al.
[9], in which the projection matrix P is updatedusing gradient descent; see Section 3.3 for details.
Sketch-GNN Framework via Polynomial Tensor Sketch
Problem and Insights.
We intend to develop a “sketched counterpart” of GNNs, where trainingis based solely on (dimensionality-reduced) compact sketches of the convolution and node featurematrices, the sizes of which can be set independently of the graph size n .
In each layer, Sketch-GNNreceives some sketches of the convolution matrix C and node representation matrix X ( l ) and outputssome sketches of the node representations X ( l +1) .
As a result, the memory and time complexitiesare inherently independent of n .
The bottleneck of this problem is estimating the nonlinear activatedproduct σ ( CX ( l ) W ( l ) ) , where W ( l ) is the learnable weight of the l -th layer.
Before considering the nonlinear activation, as a ﬁrst step, we approximate the linear product
CX ( l ) W ( l ) , using dimensionality reduction techniques such as random projections and low-rankdecompositions.
As a direct corollary of the (distributional) Johnson–Lindenstrauss (JL) lemma [24],there exists a projection matrix R ∈ R c × n such that CX ( l ) W ( l ) ≈ (cid:0) CR T (cid:1) (cid:0) RX ( l ) W ( l ) (cid:1) [15].
Tensor sketch is one of the techniques that can achieve the JL bound [2]; for an error bound, see
Lemma 1 in Appendix B.
Count sketch offers a good estimation of a matrix product, CX ( l ) W ( l ) ≈ CS ( C ) CS (( X ( l ) W ( l ) ) T ) T .
While tensor sketch can be used to approximate the power of matrix product, i.e., ( CX ( l ) W ( l ) ) (cid:12) k ≈
TS k ( C )
TS k
(( X ( l ) W ( l ) ) T ) T , where ( · ) (cid:12) k is the k -th order element-wise power.
If we combine theestimators of element-wise powers of CX ( l ) W ( l ) , we can approximate the (element-wise) activationσ ( · ) on CX ( l ) W ( l ) .
This technique is known as a polynomial tensor sketch (PTS) and is discussedin [19].
In this paper, we apply PTS to sketch the message passing of GNNs, including the nonlinearactivations.
3.1 Sketch-GNN: Approximated Update Rules
Polynomial Tensor Sketch.
Our goal is to approximate the update rule of GNNs (Eq.
(1)) in eachlayer.
We ﬁrst expand the element-wise non-linearity σ as a power series, and then approximate thepowers using count/tensor sketch, i.e.,
[Equation]
where the k = 0 term always evaluates to zero as σ (0) = 0 .
In Eq.
(4), coefﬁcients c k areintroduced to enable learning or data-driven selection of the weights when combing the terms ofdifferent order k .
This allows for the approximation of a variety of nonlinear activation functions,such as sigmoid and ReLU.
The error of this approximation relies on the precise estimation of thecoefﬁcients { c k } rk =1 .
To identify the coefﬁcients, Han et al.
[19] design a coreset-based regressionalgorithm, which requires at least O ( n ) additional time and memory.
Since the coefﬁcients { c k } rk =1that achieve the best performance for the classiﬁcation tasks do not necessarily approximate aknown activation, we propose learning the coefﬁcients { c k } rk =1 to optimize the classiﬁcation lossdirectly using gradient descent with simple L 2 regularization.
Experiments indicate that the learnedcoefﬁcients can approximate the sigmoid activation with relative errors comparable to those of thecoreset-based method; see Fig.
1a in Section 5.
Approximated Update Rules.
The remaining step is to approximate the operations of GNNs using
PTS (Eq.
(4)) on sketches of convolution matrix C and node representation matrix X ( l ) .
Consider rpairwise-independent count sketches { CS ( k ) ( · ) } rk =1 with sketch dimension c , associated with hashtables h (1) , . .
, h ( r ) and binary Rademacher variables s (1) , . . .
, s ( r ) , deﬁned prior to training an
L -layer Sketch-GNN .
Using these hash tables and Rademacher variables, we may also constructtensor sketches { TS k ( · ) } rk =
up to the maximum order r .
In Sketch-GNN , sketches of node representations (instead of the O ( n ) standard representation) arepropagated between layers.
To get rid of the dependence on n , we count sketch both sides of Eq.
[Equation]
where S ( l +1 ,k (cid:48) )
X
=
CS ( k (cid:48) )
(( X ( l +1) ) T ) ∈ R
d × c is the transpose of column-wise count sketch of
X ( l +1) , and the superscripts of S ( l +1 ,k (cid:48) )
X indicate that it is the k (cid:48) -th count sketch of X ( l +1) (i.e.,sketched by CS ( k ) ( · ) ).
In the second line of Eq.
(5), we can move the matrix, c ( l )
k TS k
(( X ( l ) W ( l ) ) T ) ,multiplied on the left to TS k ( C ) T out of the count sketch function CS ( k (cid:48) ) ( · ) , since the operation ofrow-wise count sketch CS ( k (cid:48) ) ( · ) is equivalent to multiplying the associated count sketch matrix R ( k (cid:48) )on the right, i.e., for any U ∈ R n × n , CS ( k (cid:48) ) ( U ) =
UR ( k (cid:48) ) .
In the third line of Eq.
(5), we denote the“two-sided sketch” of the convolution matrix as S ( l,k,k (cid:48) ) C :
=
CS ( k (cid:48) ) ( TS k ( C ) T ) ∈ R c × c and expandthe tensor sketch
TS k
(( X ( l ) W ( l ) ) T ) using the FFT-based formula (Eq.
Eq.
(5) is the (recursive) update rule of Sketch-GNN , which approximates the operation of theoriginal GNN and learns the sketches of representations.
Looking at the both ends of Eq.
(5), weobtain a formula that approximates the sketches of X ( l +1) using the sketches of X ( l ) and C , withlearnable weights W ( l ) ∈ R d × d and coefﬁcients { c ( l ) k ∈ R } rk =1 .
The forward-pass and backwardpropagation between the input sketches { S (0 ,k ) X } rk =1 and the sketches of the ﬁnal layer representations{ S ( L,k ) } r k =1 take O ( c ) time and memory (see Section 3.3 for complexity details).
3.2 Error Bound on Estimated Representation
Based on Lemma 1 and the results in [19], we establish an error bound on the estimated ﬁnal layerrepresentation (cid:101) X
( L ) for GCN; see Appendix B for the proof and discussions.
Theorem 1.
For a Sketch-GNN with L layers, the estimated ﬁnal layer representation is(cid:101) X ( L )
= Med { R ( k ) S
( L,k ) X | k = 1 , · · · , r } , where the sketches are recursively computed using Eq.
For Γ ( l ) = max { 5 (cid:107) X ( l ) W ( l ) (cid:107) 2 F , (2 + 3 r ) (cid:80) i ( (cid:80) j
[ X ( l ) W ( l ) ] i,j ) r } , it holds that E ( (cid:107) X
{+( L ) − (cid:101) X ( L ) (cid:107) 2 F ) / (cid:107) X ( L ) (cid:107) 2 F ≤ (cid:81) Ll =
+}
(1 + 2 / (1 + cλ ( l )2 /nr Γ ( l ) )) − 1 , where λ ( l ) ≥ 0 is thesmallest singular value of the matrix Z ∈ R nd × r and Z :
,k is the vectorization of ( CX ( l ) W ( l ) ) (cid:12) k .
Moreover, if ( c ( λ ( l ) ) 2 /nr Γ ( l ) ) (cid:29) 1 holds true for every layer, the relative error is O
( L ( n/c )) , whichis proportional to the depth of the model, and inversely proportional to the sketch ratio ( c/n ) .
Remarks.
Despite the fact that in Theorem 1 the error bound grows for smaller sketch ratios c/n ,we observe in experiments that the sketch-ratio required for competitive performance decreases asn increases; see Section 5.
As for the number of independent sketches r , we know from Lemma 1that
the dependence of r on n is r = Ω(3 log c n ) which is negligible when n is not too small; thus, inpractice r = 3 is used.
The theoretical framework may not completely correspond to reality.
Experimentally, the coefﬁcients{{ c ( l ) k } rk =1 } Ll =1 with the highest performance do not necessarily approximate a known activation.
Wedefer the challenging problem of bounding the error of sketches and coefﬁcients learned by gradientsto future studies.
Although the error bound is in expectation, we do not train over different sketchesper iteration due to the instability caused by randomness.
Instead, we introduce learnable localitysensitive hashing (LSH) in the next section to counteract the approximation limitations caused by theﬁxed number of sketches.
3.3 A Practical Implementation: Learning Sketches using LSH
Motivations of Learnable Sketches.
In Section 3, we apply polynomial tensor sketch (PTS) toapproximate the operations of GNNs on sketches of the convolution and feature matrices.
Nonetheless,the pre-computed sketches are ﬁxed during training, resulting in two major drawbacks : (1) Theperformance is limited by the quality of the initial sketches.
For example, if the randomly-generatedhash tables { h ( k ) } rk =1 have unevenly distributed buckets, there will be more hash collisions andconsequently worse sketch representations.
The performance will suffer because only sketches areused in training.
(2) More importantly, when multiple Sketch-GNN layers are stacked, the inputrepresentation X ( l ) changes during training (starting from the second layer).
Fixed hash tables arenot tailored to the “changing” hidden representations.
We seek a method for efﬁciently constructing high-quality hash tables tailored for each hiddenembedding.
Locality sensitive hashing (LSH) is a suitable tool since it is data-dependent andpreserves data similarity by hashing similar vectors into the same bucket.
This can signiﬁcantlyimprove the quality of sketches by reducing the errors due to hash collisions.
Combining LSH with Sketching.
At the time of sketching, the hash table h ( k ) : [ n ] → [ c ] isreplaced with an LSH function H ( k ) :
R d → [ c ] , for any k ∈ [ r ] .
Speciﬁcally, in the l -th layer of a
Sketch-GNN, we hash the i -th node to the H ( k ) ( X ( l ) i, : ) -th bucket for every i ∈ [ n ] , where X ( l ) i, : is theembedding vector of node i .
As a result, we deﬁne a data-dependent hash table
[Equation]
6) that can be used for computing the sketches of S ( l,k ) X and S ( l,k,k (cid:48) ) C .
This LSH-based sketching canbe directly applied to sketch the ﬁxed convolution matrix and the input feature matrix.
If SimHashis used, i.e., H ( k ) ( u ) = arg max (cid:0)(cid:2) P ( k ) u (cid:107) − P ( k ) u (cid:3)(cid:1) (Eq.
(3)), an additional O ( ncr (log c + d ))computational overhead is introduced to hash the n nodes for the r hash tables during preprocessing;see Appendix F more information.
SimHash(es) are implemented as simple matrix multiplicationsthat are practically very fast.
In order to employ LSH-based hash functions customized to each layer to sketch the hidden representations of a Sketch-GNN (i.e., l = 2 , . . . , L − 1 ), we face two major challenges : (1) Unless weexplicitly unsketch in each layer, the estimated hidden representations (cid:101) X ( l ) ( l = 2 , . . .
, L − 1) cannotbe accessed and used to compute the hash tables.
However, unsketching any hidden representation,i.e., (cid:101) X ( l ) =
Med { R ( k ) S ( l,k ) X
| k = 1 , · · · , r } , requires O ( n ) memory and time.
We need to come up with an efﬁcient algorithm that updates the hash tables without having to unsketch the completerepresentation.
(2) It’s unclear how to change the underlying hash table of a sketch across layerswithout unsketching to the n -dimensional space, even if we know the most up-to-date hash tablessuited to each layer.
The challenge (2) , i.e., changing the underlying hash table of across layers, can be solved bymaintaining a sparse c × c matrix T ( l,k ) :
= R ( l,k ) ( R ( l +1 ,k ) )
T for each k ∈ [ r ] , which only requires
O ( cr ) memory and time overhead; see Appendix C for more information and detailed discussions.
We focus on challenge (1) for the remainder of this section.
Online Learning of Sketches.
To learn a hash table tailored for a hidden layer using LSH withoutunsketching, we develop an efﬁcient algorithm to update the LSH function using only a size- | B |subset of the length- n unsketched representations, where B denotes a subset of nodes we select.
Thisalgorithm, which we term online learning of sketches , is made up of two key parts: (Part 1) select asubset of nodes B ⊆ [ n ] to effectively update the hash table, and (Part 2) update the LSH function
H ( · ) with a triplet loss computed using this subset.(1) Selection of subset B : Because model parameters are updated slowly during neural networktraining, the data-dependent LSH hash tables also changes slowly (this behavior was detailed in [9]).
The amount of updates to the hash table drops very fast along with training, empirically veriﬁedin Fig.
1b (left) in Section 5.
Based on this insight, we only need to update a small fraction of thehash table during training.
To identify this subset B ∈ [ n ] of nodes, gradient signals can be used.
Intuitively, a node representation vector hashed into the wrong bucket will be aggregated with distantvectors and lead to larger errors and subsequently larger gradient signals.
Speciﬁcally, we proposeﬁnding the candidate set B of nodes by taking the union of the several buckets with the largestgradients, i.e., B = { i | h ( l,k ) ( i ) =
arg max j [ S ( l,k ) X ] j, : for some k } . The memory and overheadrequired to update the entries in B in the hash table is O ( | B | ) .(2) Update of LSH function:
In order to update the projection matrix P that deﬁnes a SimHash
H ( k ) : R d → [ c ] (Eq.
(3)), instead of the O ( n ) full triplet loss introduced by [9], we consider asampled version of the triplet loss on the candidate set B with O ( | B | ) complexity, namely
[Equation]
where P + = { ( (cid:101) X i, : , (cid:101)
X j, : ) | i, j ∈ B, (cid:104) (cid:101)
X i, : , (cid:101)
X j, : (cid:105) > t + } and P − = { ( (cid:101) X : ,i , (cid:101) X : ,j ) | i, j [-
-]∈
B, (cid:104) (cid:101) X : ,i , (cid:101) X : ,j (cid:105) <
t − } are the similar and dissimilar node-pairs in the subset B ; t + > t − andα > 0 are hyper-parameters.
This triplet loss L ( H, P + , P − ) is used to update P using gradientdescent, as described in [9], with a O ( c | B | d + | B | 2 ) overhead.
Experimental validation of this LSHupdate mechanism can be found in Fig.
1b in Section 5.
Avoiding O ( n ) in Loss Evaluation.
We can estimate the ﬁnal layer representation using the rsketches { S ( L,k ) } rk
=1 , i.e., (cid:101)
X ( L ) =
Med { R ( k ) S ( L,k ) X | k = 1 , · · · , r } and compute the losses of allnodes for node classiﬁcation (or some node pairs for link prediction).
However, the complexity ofloss evaluation is O ( n ) , proportional to the number of ground-truth labels.
In order to avoid O ( n )complexity completely, rather than un-sketching the node representation for all labeled nodes, weemploy the locality sensitive hashing (LSH) technique again for loss calculation so that only a subsetof node losses are evaluated based on a set of hash tables.
Speciﬁcally, we construct an LSH hashtable for each class in a node classiﬁcation problem, which indexes all of the labeled nodes of thisclass and can be utilized to choose the nodes with poor predictions by leveraging the locality property.
This technique, introduced in [10], is known as sparse forward-pass and back-propagation, and wedefer the descriptions to Appendix C.
One-time Preprocessing.
If the convolution matrix C is ﬁxed (GCN, GraphSAGE), the “two-sidedsketch” S ( l,k,k (cid:48) )
C = CS ( k (cid:48) )
( TS k ( C ) T ) ∈ R c × c is the same in each layer and may be denoted as
S ( k,k (cid:48) ) C .
In addition, all of the r 2 sketches of C , i.e., {{ S ( k,k (cid:48) )
C ∈ R c × c } rk =1 } [-r k 
-]rk (cid:48)
=1 can be computedduring the preprocessing phase.
If the convolution matrix C is sparse (which is true for most GNNsfollowing Eq.
(1) on a sparse graph), we can use the sparse matrix representations for the sketches{{ S ( k,k (cid:48) )
C ∈ R c × c } rk =1 } rk [-
-](cid:48)
=1 , and the total memory taken by the r 2 sketches is O ( r 2 c ( m/n ))
where(2 m/n ) is the average node degree (see Appendix F for details).
We also need to compute the r countsketches of the input node feature matrix X
= X (0) , i.e., { S (0 ,k ) X }
rk =1 during preprocessing, which requires O ( rcd ) memory in total.
In this regard, we have substituted the input data with compactgraph-size independent sketches (i.e., O ( c ) memory).
Although the preprocessing time required tocompute these sketches is O ( n ) , it is a one-time cost prior to training, and it is widely known thatsketching is practically very fast.
Complexities of Sketch-GCN.
The theoretical complexities of Sketch-GNN is summarized asfollows, where for simplicity we assume bounded maximum node degree, i.e., m = O ( n ) .
Training Complexity : (1a) Forward and backward propagation :
O ( Lcrd (log( c ) + d + m/n )) =
O ( c ) time and O ( Lr ( cd + rm/n ))
=
O ( c ) memory.
(1b)
Hash and sketch update : O ( Lr ( c +| B | d ))
=
O ( c ) time and memory.
(2) Preprocessing : O ( r ( rm + n + c ))
=
O ( n ) time and
O ( rc ( d + rm/n ))
=
O ( c ) memory.
(3) Inference : O ( Ld ( m + nd ))
=
O ( n ) time and O ( m +
Ld ( n + d ))
=
O ( n ) memory (the same as a standard GCN).
We defer a detailed summary of thetheoretical complexities of Sketch-GNN to Appendix F.
We generalize Sketch-GNN to more GNN models in Appendix D and the pseudo-code which outlinesthe complete workﬂow of Sketch-GNN can be ﬁnd in Appendix E.
Related Work
Scalable methods for GNNs can be categorized into four classes, all of them still require lineartraining complexities.
(A)
On a large sparse graph with n nodes and m edges, the “full-graph”training of a L -layer GCN with d -dimensional (hidden) features per layer requires O ( m + ndL + d 2 L )memory and O ( mdL + nd 2 L ) epoch time.
(B) Sampling-based methods sample mini-batches fromthe complete graph following three schemes: (1) node-wisely sample a subset of neighbors ineach layer to reduce the neighborhood size; (2) layer-wisely sample a set of nodes independentlyin each layer; (3) subgraph-wisely sample a subgraph directly and simply forward-pass and backpropagate on that subgraph. (B.1) GraphSAGE [18] samples r neighbors for each node while ignoringmessages from other neighbors. O ( br L ) nodes are sampled in a mini-batch (where b is the mini-batchsize), and the epoch time is O ( ndr L ) ; therefore, GraphSAGE is impractical for deep GNNs on alarge graph. FastGCN [12] and LADIES [46] are layer-sampling methods that apply importancesampling to reduce variance.
(B.2) The subgraph-wise scheme has the best performance and ismost prevalent. Cluster-GCN [14] partitions the graph into many densely connected subgraphs andsamples a subset of subgraphs (with edges between subgraphs added back) for training per iteration.
GraphSAINT [44] samples a set of nodes and uses the induced subgraph for mini-batch training. Both Cluster-GCN and GraphSAINT require O ( mdL + nd 2 L ) epoch time, which is the same as“full-graph” training, although Cluster-GCN also needs O ( m ) pre-processing time. (C)
Apart fromsampling strategies, historical-embedding-based methods propose mitigating sampling errors andimproving performance using some stored embeddings. GNNAutoScale [17] keeps a snapshot ofall embeddings in CPU memory, leading to a large O ( ndL ) memory overhead. VQ-GNN [15]maintains a vector quantized data structure for the historical embeddings, whose size is independentof n . (D)
Linearized GNNs [41, 4, 32] replace the message passing operation in each layer with aone-time message passing during preprocessing.
They are practically efﬁcient, but the theoreticalcomplexities remain O ( n ) .
Linearized models usually over-simplify the corresponding GNN andlimit its expressive power.
Towards sublinear GNNs.
Nearly all existing scalable methods focus on mini-batching the largegraph and resolving the memory bottleneck of GNNs, without reducing the epoch training time.
Few recent work focus on graph compression [22, 23] can also achieve sublinear training time bycoarsening (e.g., using [30]) the graph during preprocessing and training GNNs on the coarsenedgraph with fewer nodes and edges.
Nevertheless, this strategy suffers from two issues : (1) Althoughgraph coarsening is a one-time cost, the memory and time overheads are often worse than O ( n )and can be prohibitively large on graphs with over 100K nodes.
Even the fastest graph coarseningalgorithm used by [22] takes more than 68 minutes to process the 233K-node Reddit graph [44];see Table 1.
The long preprocessing time renders any training speedups meaningless.
The testperformance of a model trained on the coarsened graph highly depends on the GNN type.
Althoughthe performance of [22] on GCN is good, signiﬁcant performance degradations are observed on
GraphSAGE and GAT; see Section 5.
We defer discussion of more scalable GNN papers and the broad literature of sketching and LHS forneural networks to Appendix G.
Experiments
[Figure]
[Figure]
[Figure]
SketchRatio
[Figure]
c
[Figure]
R e l a ti v e E rr o r
CoresetTaylorLearnedviaGD(Ours) (a) Error of PTS     0. Epoch ∆ H a mm i n g 5 10 15 20 250.
Epoch H a mm i n g ( h l e a r n e d , h g r o und − t r u t h )
[Figure]
[Figure]
[Table]
[Table] [Table]
In this section, we evaluate the proposed Sketch-GNN algorithm and compare it with the (oracle) “fullgraph” training baseline, a graph-coarsening based method ( Coarsening [22]) which has sublineartraining time, and other scalable methods including: a sampling-based method ( GraphSAINT [44]),a historical-embedding based method ( VQ-GNN [15]), and a linearized GNN ( SGC [41]).
We teston several large graph benchmarks including ogbn-arxiv (169K nodes, 1.2M edges), Reddit (233Knodes, 11.6M edges), and ogbn-products (2.4M nodes, 61.9M edges) from [20, 44].
See Appendix Hfor the implementation details.
Proof-of-Concept Experiments: (1) Errors of gradient-learned PTS coefﬁcients : In Fig.
1a,we train the PTS coefﬁcients to approximate the sigmoid activated σ ( CXW ) to evaluate its approximation power to the ground-truth activation. The relative errors are comparable to those ofthe coreset-based method.
(2) Slow-change phenomenon of LSH hash tables :
In Fig. 1b (left),we count the changes of the hash table constructed from an unsketched hidden representation foreach epoch, characterized by the Hamming distances between consecutive updates.
The changesdrop rapidly as training progresses, indicating that apart from the beginning of training, the hashcodes of most nodes do not change at each update.
(3) Sampled triplet loss for learnable LSH :
In Fig.
1b (right), we verify the effectiveness of our update mechanism for LSH hash functions, asthe learned hash table gradually approaches the “ground truth”, i.e., the hash table constructed fromthe unsketched hidden representation.
Efﬁciency of Sketch-GNNs.
For efﬁciency measures, we are interested in the comparison to
Coarsening, as both approaches achieve sublinear training time at the cost of some preprocessingoverheads.
We use a 3-layer GCN as the backbone and set the sketch ratios ( c/n , ratio of sketchdimension c to graph size n ) of both algorithms to c/n = 0 . 1 , meaning that the coarsened graphcontains n/ 10 nodes.
We measure their preprocessing time, average epoch training time, and peaktraining memory, as reported in Table 1.
Although not rigorously comparable, we also set the minibatch size of GraphSAINT and VQ-GNN to b = n/ 10 .
We report the average epoch training time andpeak training memory for each method and the "full-graph" training baseline.
In addition to Table 1,the following are also recorded: (1) Coarsening requires 980 MB to preprocess ogbn-arxiv , whereas
Sketch-GNN only requires 539 MB.
Our preprocessing on the largest dataset, ogbn-product (2.4Mnodes), takes only 414s.
The wallclock time for the validation accuracy to reach 99% of its best is88 ± 8 s for SketchGCN, which is shorter than VQ-GNN’s 103 ± 11 s and GraphSAINT’s 120 ± 4 s.
From Table 1 and the aforementioned results, we can draw four important conclusions : (1) SketchGNN achieves the fastest average epoch time.
The coarsened graph is typically much denser andincreases the time required for message passing.
(2) Sketch-GNN usually converges faster than
GraphSAINT and VQ-GNN.
(3) Our preprocessing time is signiﬁcantly less than that of Coarsening.
Coarsening suffers from an extremely long preprocessing time, rendering the training speed-upsmeaningless.
Moreover, our preprocessing time scales well with graph size and sparsity.
(4) We alsorequire less preprocessing memory as sketching is linear/multi-linear operation and usually preservessparsity.
(5) Sketch-GNN often requires more training memory than Coarsening in order to maintainthe copies of sketches and additional data structures, although these memory overheads are small.
Performance of Sketch-GNNs.
We ﬁrst compare the performance of Sketch-GNN with Coarseningunder various sketch ratios to understand how their performance is affected by the memory bottleneck.
In Table 2, we report the test accuracy of both approaches on ogbn-arxiv , with a 3-layer GCN,
GraphSAGE, or GAT as the backbone and a sketch ratio of 0 . 1 , 0 . 2 , or 0 . 4 .
We see there aresigniﬁcant performance degradations when applying Coarsening to GraphSAGE and GAT, even undersketch ratio 0 .
4 , indicating that Coarsening may be compatible only with speciﬁc GNNs (GCN and
APPNP as explained in [22]).
In contrast, the performance drops of Sketch-GNN are always smallacross all architectures, even when the sketch ratio is 0 .
Therefore, our approach generalizes tomore GNN architectures and consistently outperforms the Coarsening method.
We move on to compare Sketch-GNN with linearized GNNs (SGC), sampling-based (GraphSAINT),and historical-embedding-based (VQ-GNN) methods.
In Table 3, we report the performance of
SGC, the “full-graph” training (oracle), GraphSAINT and VQ-GNN with mini-batch size 50K(their performance is not affected by the choice of mini-batch size if it is not too small), and
Sketch-GNN with appropriate sketch ratios ( 0 . 4 on ogbn-arxiv , 0 . 3 on Reddit , and 0 .
2 on ogbnproduct ).
From Table 3, we conﬁrm that, with an appropriate sketch ratio, the performance of
Sketch-GNN is always close to the “full-graph” oracle and competitive with the other scalableapproaches.
Impressively, the needed sketch ratio c/n for Sketch-GNN to achieve competitiveperformance reduces as graph size grows.
This further illustrates that, as previously indicated, therequired training complexities (to get acceptable performance) are sublinear to the graph size.
Ablation Studies: (1) Dependence of sketch dimension c on graph size n .
Although the theoreticalapproximation error increases under smaller sketch ratio c/n , we observe competitive experimentalresults with smaller c/n especially on large graphs.
In practice, the sketch-ratio required to maintain“full-graph” model performance decreases with n , as veriﬁed in Table 3: c/n = 0 .
4 is needed on ogbnarxiv with 169K nodes but c/n = 0 .
2 is adequate on ogbn-product with 2.45M nodes.
(2) Learned
Sketches versus Fixed Sketches.
We ﬁnd that learned sketches can improve the performanceof all models and on all datasets.
Under sketch-ratio c/n = 0 .
2 , the Sketch-GCN with learnedsketches achieves 0 .
7004 ± 0 .
0096 accuracy on ogbn-arxiv while ﬁxed randomized sketches degradeperformance to 0 .
6649 ± 0 . 0106 .
Conclusion
We present Sketch-GNN , a sketch-based GNN training framework with sublinear training time andmemory complexities.
Our main contributions are (1) approximating nonlinear operations in GNNsusing polynomial tensor sketch (PTS) and (2) updating sketches using learnable locality sensitivehashing (LSH).
Our novel framework has the potential to be applied to other architectures andapplications where the amount of data makes training even simple models impractical.