Sketch-GNN: Scalable Graph Neural Networks with Sublinear Training Complexity
Abstract
Graph Neural Networks (GNNs) are widely applied to graph learning problems such as node classiﬁcation.
When scaling up the underlying graphs of GNNs to a larger size, we are forced to either train on the complete graph and keep the full graph adjacency and node embeddings in memory (which is often infeasible) or mini-batch sample the graph (which results in exponentially growing computational complexities with respect to the number of GNN layers).
Various sampling-based and historical-embedding-based methods are proposed to avoid this exponential growth of complexities.
However, none of these solutions eliminates the linear dependence on graph size.
This paper proposes a sketch-based algorithm whose training 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 novel protocol for sketching non-linear activations and graph convolution matrices in
GNNs, as opposed to existing methods that sketch linear weights or gradients in neural networks.
In addition, we develop a locality sensitive hashing (LSH) technique that can be trained to improve the quality of sketches.
Experiments on large-graph benchmarks demonstrate the scalability and competitive performance of our Sketch-GNNs versus their full-size GNN counterparts.
Introduction
Graph Neural Networks (GNNs) have achieved state-of-the-art graph learning in numerous applications, including classiﬁcation [27], clustering [3], recommendation systems [43], social networks [16] and more, through representation learning of target nodes using information aggregated from 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 computational and memory efﬁciency.
Full-batch training that stores the Laplacian of the complete graph suffers from a memory complexity of O ( m + ndL + d 2 L ) on an n -node, m -edge graph with node features of dimension d when employing an L -layer graph convolutional network (GCN).
This linear memory complexity dependence on n and the limited memory capacity of GPUs make it difﬁcult to train on large graphs with millions of nodes or more.
As an example, the MAG240M-LSC dataset [21] is a node classiﬁcation benchmark with over 240 million nodes that takes over 202 GB of GPU memory when fully loaded.
To address the memory constraints, two major lines of research are proposed: (1) Sampling-based approaches [18, 11, 12, 14, 45] based on the idea of implementing message passing only between the 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 sampled subgraphs using historical embeddings.
However, all of these methods require the number of mini- batches 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 when scaling 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 large instance [21].
We seek to achieve efﬁcient training of GNNs with time and memory complexities sublinear in graph size without signiﬁcant accuracy degradation.
Despite the difﬁculty of this goal, it should be achievable given that (1) the number of learnable parameters in GNNs is independent of the graph size, and (2) training may not require a traversal of all local neighborhoods on a graph but rather only the most representative ones (thus sublinear in graph size) as some neighborhoods may be very similar.
In addition, commonly-used GNNs are typically small and shallow with limited model capacity 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 memory complexity with respect to graph size.
Using the idea of sketching, which maps high-dimensional data structures to a lower dimension through entry hashing, we sketch the n × n adjacency matrix and 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 weights or gradients, we introduce a method for sketching non-linear activation units using polynomial tensor sketch theory [19].
This preserves prediction accuracy while avoiding the need to “unsketch” back to the original high dimensional graph-node space n , thereby eliminating the dependence of training complexity on the underlying graph size n .
Moreover, we propose to learn and update the sketches in an online manner using learnable locality-sensitive hashing (LSH) [9].
This reduces the performance loss by adaptively enhancing the sketch quality while incurring minor overhead sublinear in the graph size.
In practice, we ﬁnd that the sketch-ratio c/n required to maintain “full-graph” model performance drops 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 the data size.
This is fundamentally different from the few existing works which sketch the weights or gradients [30, 13, 26, 29, 37] to reduce the memory footprint of the model and speed up optimization.
To the best of our knowledge, Sketch-GNN is the ﬁrst sub-linear complexity training algorithm for GNNs, based on LSH and tensor sketching.
The sublinear efﬁciency obtained applies to various types of GNNs, including GCN [27] and GraphSAGE [18].
Compared to the data compression approach [22, 23], which compresses the input graph to a smaller one with fewer nodes and edges before training, our Sketch-GNN is advantageous since it does not suffer from an extremely long preprocessing time (which renders the training speedups meaningless) and performs much better across GNN types/architectures. The remainder of this paper is organized as follows.
Section 2 summarizes the notions and preliminaries of GNNs and sketching. Section 3 describes how to approximate the GNN operations on the full graph topology with sketches. Section 3.3 introduces potential drawbacks of using ﬁxed sketches and develops algorithms for updating sketches using learnable LSHs. In Section 4, we compare our approach to the graph compression approach and other GNN scalability methods.
In Section 5, we report the performance and efﬁciency of Sketch-GNNs as well as several proof-of-concept and ablation experiments. Finally, Section 6 concludes this paper with a summary of limitations, future directions, and broader impacts.  Preliminaries
Basic Notations.
Consider a graph with n nodes and m edges.
Connectivity is given by the adjacency matrix A ∈ { 0 , 1 } n × n and features on nodes are represented by the matrix X ∈ R n × d , where d is the number of features.
Given a matrix C , let C i,j , C i, : , and C : ,j denote its ( i, j ) -th entry, i -th row, and j -th column, respectively. (cid:12) denotes the element-wise (Hadamard) product, whereas C (cid:12) k represents the k -th order element-wise power.
(cid:107) · (cid:107) F is the symbol for the Frobenius norm.
I n ∈ R n × n denotes the identity matrix, whereas 1 n ∈ R n is the vector whose elements are all ones.
Med {·} represents the element-wise median over a set of matrices.
Superscripts are used to indicate multiple instances of 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 different guiding principles, such as neighborhood aggregation (GraphSAGE), spatial convolution (GCN), self-attention (GAT), and Weisfeiler-Lehman (WL) alignment (GIN [44]), 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 this paradigm 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 of well-known GNNs under this framework.
Unless otherwise speciﬁed, we assume q = 1 and d = d l for every layer l ∈
[ L ] for notational convenience.
Count Sketch and Tensor Sketch.
(1) Count sketch [7, 41] 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 a linear transformation of u , i.e., CS ( u ) =
R u .
Here, R ∈ R c × n denotes the so-called count sketch matrix , which has exactly one non-zero element per column.
(2) Tensor sketch [32] is proposed as a generalization of count sketch to the tensor product of vectors.
Given z ∈ R n and 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 vector z ∈ 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 , where h ( j 1 , · · · , j k ) =
( h (1) ( j 1 ) + · · · + h ( k ) ( j k ))
mod c .
By deﬁnition, a tensor sketch of order k =degenerates to count sketch; TS 1 ( · ) =
CS ( · ) .
(3) We deﬁne count sketch of a matrix U ∈ R d × n as the count sketch of each row vector
individually, i.e., CS ( U ) ∈
R
d × c where [ CS ( U )]
i, : =
CS ( U i, : ) .
The tensor sketch of a matrix is deﬁned in the same way.
Pham and Pagh [32] devise a fast computation of tensor sketch of U ∈ R d × n (sketch dimension c and order k ) using count sketches and the Fast
Fourier Transform (FFT):
[Equation]
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 hash table(s) that only requires a data independent uniformity, i.e., with high probability the hash-buckets are of similar size.
In contrast, locality sensitive hashing (LSH) is a hashing scheme that uses locality-sensitive hash function H : R d → [ c ] to ensure that nearby vectors are hashed into the same bucket (out of c buckets in total) with high probability while distant ones are not.
SimHash achieves 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
[Equation]
where [ · (cid:107) · ] denotes concatenation of two vectors and arg max returns the index of the largest element.
SimHash is efﬁcient for large batches of vectors [1].
In this paper, we apply a learnable version of SimHash that is proposed by Chen et al.
[9], in which the projection matrix P is updated using 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 training is based solely on (dimensionality-reduced) compact sketches of the convolution and node feature matrices, the sizes of which can be set independently of the graph size n .
In each layer, Sketch-GNN receives some sketches of the convolution matrix C and node representation matrix X ( l ) and outputs some sketches of the node representations X ( l +1) .
As a result, the memory and time complexities are inherently independent of n .
The bottleneck of this problem is estimating the nonlinear activated product σ ( 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-rank decompositions.
As a direct corollary of the (distributional) Johnson–Lindenstrauss (JL) lemma [25], 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 the estimators 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 discussed in [19].
In this paper, we apply PTS to sketch the message passing of GNNs, including the nonlinear activations.
3.1 Sketch-GNN: Approximated Update Rules
Polynomial Tensor Sketch.
Our goal is to approximate the update rule of GNNs (Eq.
(1)) in each layer.
We ﬁrst expand the element-wise non-linearity σ as a power series, and then approximate the powers 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 are introduced to enable learning or data-driven selection of the weights when combing the terms of different 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 the coefﬁcients { c k } rk =1 .
To identify the coefﬁcients, Han et al.
[19] design a coreset-based regression algorithm, which requires at least O ( n ) additional time and memory.
Since the coefﬁcients { c k } rk = that achieve the best performance for the classiﬁcation tasks do not necessarily approximate a known activation, we propose learning the coefﬁcients { c k } rk =1 to optimize the classiﬁcation loss directly using gradient descent with simple L 2 regularization.
Experiments indicate that the learned coefﬁcients can approximate the sigmoid activation with relative errors comparable to those of the coreset-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 r pairwise-independent count sketches { CS ( k ) ( · ) } rk =1 with sketch dimension c , associated with hash tables 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 construct tensor sketches { TS k ( · ) } rk =2 up to the maximum order r .
In Sketch-GNN , sketches of node representations (instead of the O ( n ) standard representation) are propagated between layers.
To get rid of the dependence on n , we count sketch both sides of Eq.
[Equation]
l
[Equation]
[Equation]
k
[Equation]
) X
[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 of row-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 expand the 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 the original GNN and learns the sketches of representations.
Looking at the both ends of Eq.
(5), we obtain a formula that approximates the sketches of X ( l +1) using the sketches of X ( l ) and C , with learnable 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 ) } rk =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 layer representation (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 the smallest 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 )) , which is 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 as n increases; see Section 5.
As for the number of independent sketches r , we know from Lemmathat the dependence of r on n is r = Ω(3 log c n ) which is negligible when n is not too small; thus, in practice 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.
We defer the challenging problem of bounding the error of sketches and coefﬁcients learned by gradients to future studies.
Although the error bound is in expectation, we do not train over different sketches per iteration due to the instability caused by randomness.
Instead, we introduce learnable locality sensitive 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) to approximate 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) The performance is limited by the quality of the initial sketches.
For example, if the randomly-generated hash tables { h ( k ) } rk =1 have unevenly distributed buckets, there will be more hash collisions and consequently worse sketch representations.
The performance will suffer because only sketches are used in training.
(2) More importantly, when multiple Sketch-GNN layers are stacked, the input representation X ( l ) changes during training (starting from the second layer).
Fixed hash tables are not tailored to the “changing” hidden representations.
We seek a method for efﬁciently constructing high-quality hash tables tailored for each hidden embedding.
Locality sensitive hashing (LSH) is a suitable tool since it is data-dependent and preserves data similarity by hashing similar vectors into the same bucket.
This can signiﬁcantly improve 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 ] is replaced 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 the embedding vector of node i .
As a result, we deﬁne a data-dependent hash table
[Equation]
that can be used for computing the sketches of S ( l,k ) X and S ( l,k,k (cid:48)
[Equation]
C .
This LSH-based sketching can be directly applied to sketch the ﬁxed convolution matrix and the input feature matrix.
If SimHash is 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 multiplications that 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 we explicitly unsketch in each layer, the estimated hidden representations (cid:101) X ( l ) ( l = 2 , . . .
, L − 1) cannot be accessed and used to compute the hash tables.
However, unsketching any hidden representation, i.e., (cid:101) X (
[Equation]
) =
[Equation]
{ 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 complete representation.
(2) It’s unclear how to change the underlying hash table of a sketch across layers without unsketching to the n -dimensional space, even if we know the most up-to-date hash tables suited to each layer.
The challenge (2) , i.e., changing the underlying hash table of across layers, can be solved by maintaining 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 without unsketching, 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.
This algorithm, which we term online learning of sketches , is made up of two key parts: (Part 1) select a subset 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 network training, 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ﬁed in Fig.
1b (left) in Section 5.
Based on this insight, we only need to update a small fraction of the hash 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 distant vectors 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 largest gradients, i.e., B = { i | h ( l,k ) ( i ) =
arg max j [ S ( l,k ) X ] j, : for some k } .
The memory and overhead required 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 a sampled 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 gradient descent, as described in [9], with a O ( c | B | d + | B | 2 ) overhead.
Experimental validation of this LSH update 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 r sketches { 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 all nodes for node classiﬁcation (or some node pairs for link prediction).
However, the complexity of loss 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, we employ the locality sensitive hashing (LSH) technique again for loss calculation so that only a subset of node losses are evaluated based on a set of hash tables.
Speciﬁcally, we construct an LSH hash table for each class in a node classiﬁcation problem, which indexes all of the labeled nodes of this class 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 we defer the descriptions to Appendix C.
One-time Preprocessing.
If the convolution matrix C is ﬁxed (GCN, GraphSAGE), the “two-sided sketch” 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 } rk (cid:48)
=1 can be computed during the preprocessing phase.
If the convolution matrix C is sparse (which is true for most GNNs following Eq.
(1) on a sparse graph), we can use the sparse matrix representations for the sketches {{ S ( k,k (cid:48) )
C
[Equation]
(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 count sketches 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 compact graph-size independent sketches (i.e., O ( c ) memory).
Although the preprocessing time required to compute these sketches is O ( n ) , it is a one-time cost prior to training, and it is widely known that sketching is practically very fast.
Complexities of Sketch-GCN.
The theoretical complexities of Sketch-GNN is summarized as follows, 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 the theoretical complexities of Sketch-GNN to Appendix F.
We generalize Sketch-GNN to more GNN models in Appendix D and the pseudo-code which outlines the complete workﬂow of Sketch-GNN can be ﬁnd in Appendix E.
Related Work
Towards sublinear GNNs. Nearly all existing scalable methods focus on mini-batching the large graph and resolving the memory bottleneck of GNNs, without reducing the epoch training time.
Few recent work focus on graph compression [22, 24] can also achieve sublinear training time by coarsening (e.g., using [31]) the graph during preprocessing or condensing the graph with dataset condensation techniques like gradient-matching [46], so that we can train GNNs on the coarsened/condensed graph with fewer nodes and edges.
Nevertheless, these strategies suffer from two issues : (1) Although graph coarsening/condensation 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 coarsening algorithm used by [22] takes more than 68 minutes to process the 233K-node Reddit graph [45].
The long preprocessing time renders any training speedups meaningless.
The test performance of a model trained on the coarsened graph highly depends on the GNN type.
For graph condensation, if we do not carefully choose the GNN architecture used during condensation, the test performance of downstream GNNs can suffer from a 9.5% accuracy drop on the Cora graph [23].
For graph coarsening, although the performance of [22] on GCN is good, signiﬁcant performance degradations are observed on GraphSAGE and GAT; see Section 5.
Other scalable methods for GNNs can be categorized into four classes, all of them still require linear training 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 from the complete graph following three schemes: (1) node-wisely sample a subset of neighbors in each layer to reduce the neighborhood size; (2) layer-wisely sample a set of nodes independently in 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 ignoring messages from other neighbors.
O ( br L ) nodes are sampled in a mini-batch (where b is the mini-batch size), and the epoch time is O ( ndr L ) ; therefore, GraphSAGE is impractical for deep GNNs on a large graph.
FastGCN [12] and LADIES [48] are layer-sampling methods that apply importance sampling to reduce variance.
(B.2) The subgraph-wise scheme has the best performance and is most prevalent.
Cluster-GCN [14] partitions the graph into many densely connected subgraphs and samples a subset of subgraphs (with edges between subgraphs added back) for training per iteration.
GraphSAINT [45] 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 from sampling strategies, historical-embedding-based methods propose mitigating sampling errors and improving performance using some stored embeddings.
GNNAutoScale [17] keeps a snapshot of all 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 independent of n .
(D)
Linearized GNNs [42, 4, 33] replace the message passing operation in each layer with a one-time message passing during preprocessing.
They are practically efﬁcient, but the theoretical complexities remain O ( n ) .
Linearized models usually over-simplify the corresponding GNN and limit its expressive power.
We defer discussion of more scalable GNN papers and the broad literature of sketching and LHS for neural networks to Appendix G
Experiments
[Figure]
l a ti v e E rr o r [Figure] [Figure]
[Table] [Table]
In this section
[Table]
we evaluate the proposed Sketch-GNN algorithm and compare it with the (oracle) “full-graph” training baseline, a graph-coarsening based approach ( Coarsening [22]) and a dataset condensation based approach ( GCond [23]) which enjoy sublinear training time
[Table]
and other scalable
[Table]
[Table]
[Figure]
methods including: a sampling-based method ( GraphSAINT [45]), a historical-embedding based method ( VQ-GNN [15]), and a linearized GNN ( SGC [42]).
We test on two small graph benchmarks including Cora , Citeseer and several large graph benchmarks including ogbn-arxiv (169K nodes, 1.2M edges), Reddit (233K nodes, 11.6M edges), and ogbn-products (2.4M nodes, 61.9M edges) from [20, 45].
See Appendix H for 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 of the 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 for each epoch, characterized by the Hamming distances between consecutive updates.
The changes drop rapidly as training progresses, indicating that apart from the beginning of training, the hash codes 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 as the learned hash table gradually approaches the “ground truth”, i.e., the hash table constructed from the unsketched hidden representation.
Performance of Sketch-GNNs.
We ﬁrst compare the performance of Sketch-GNN with the other sublinear training methods, i.e., graph coarsening [22] and graph condensation [23] under various sketch ratios to understand how their performance is affected by the memory bottleneck.
Since graph condensation (GCond) requires learning the condensed graph from scratch and cannot be scaled to large graphs with a large sketch ratio [23], we ﬁrst compare with GCond and Coarsening on the two small graphs using a 2-layer GCN in Table 1.
We see GCond and Sketch-GNN can outperform graph coarsening by a large margin and can roughly match the full-graph training performance.
However,
GCond suffers from a processing time that is longer than the training time (see below) and generalizes poorly across GNN architectures.
In Table 2, we compare the performance of Sketch-GNN and
GCond across two GNN architectures (GCN and GraphSAGE).
While graph condensation (GCond) relies on a “reference architecture” during condensation, Sketch-GNN does not require preprocessing, and the sublinear complexity is granted by sketching “on the ﬂy”.
In Table 2, we see the performance of GCond is signiﬁcantly degraded when generalized across architectures, while Sketch-GNNs’ performance is always close to that of full-graph training.
In Table 3, 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 are signiﬁcant performance degradations when applying Coarsening to GraphSAGE and GAT, even under sketch 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 small across all architectures, even when the sketch ratio is 0 .
Therefore, our approach generalizes to more 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 4, 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 SketchGNN with appropriate sketch ratios ( 0 . 4 on ogbn-arxiv , 0 . 3 on Reddit , and 0 .
2 on ogbn-product ).
From Table 4, we conﬁrm that, with an appropriate sketch ratio, the performance of Sketch-GNN is usually close to the “full-graph” oracle and competitive with the other scalable approaches.
The needed sketch ratio c/n for Sketch-GNN to achieve competitive performance reduces as graph size grows.
This further illustrates that, as previously indicated, the required training complexities (to get acceptable performance) are sublinear to the graph size.
Efﬁciency of Sketch-GNNs.
For efﬁciency measures, we are interested in the comparison to
Coarsening and GCond, since these two approaches achieve sublinear training time at the cost of some preprocessing overheads.
Firstly, we want to address that both Coarsening and GCond suffer from an extremely long preprocessing time.
On ogbn-arxiv , Coarsening and GCond require 358 and 494 seconds on average, respectively, to compress the original graph.
In contrast, our Sketch-GNN sketch the input graph “on the ﬂy” and does not suffer from a preprocessing overhead.
On ogbnarxiv with a learning rate of 0 .
001 , full-graph training of GCN for 300 epochs is more than enough for convergence, which only takes 96 seconds on average.
The preprocessing time of Coarsening and GCond is much longer than the convergence time of full-graph training, which renders their training speedups meaningless.
However, Sketch-GNN often requires more training memory than Coarsening and GCond to maintain the copies of sketches and additional data structures, although these memory overheads are small, e.g., only 16.6 MB more than Coarsening on ogbn-arxiv with sketch ratio c/n = 0 .
All three sublinear methods (Corasening, GCond, Sketch-GNN) lead to a denser adjacency/convolution matrix and thus increased memory per node.
However, this overhead is small for Sketch-GNN because although we sketched the adjacency, its sparsity is still preserved to some extent, as sketching is a linear/multi-linear operation.
Ablation Studies: (1) Dependence of sketch dimension c on graph size n .
Although the theoretical approximation error increases under smaller sketch ratio c/n , we observe competitive experimental results with smaller c/n , especially on large graphs.
In practice, the sketch ratio required to maintain “full-graph” model performance decreases with n .
(2) Learned Sketches versus Fixed Sketches.
We ﬁnd that learned sketches can improve the performance of all models and on all datasets.
Under sketch-ratio c/n = 0 .
2 , the Sketch-GCN with learned sketches achieves 0 .
7004 ± 0 .
0096 accuracy on ogbn-arxiv while ﬁxed randomized sketches degrade performance to 0 .
6649 ± 0 . 0106 .
Conclusion
We present Sketch-GNN , a sketch-based GNN training framework with sublinear training time and memory complexities.
Our main contributions are (1) approximating nonlinear operations in GNNs using polynomial tensor sketch (PTS) and (2) updating sketches using learnable locality-sensitive hashing (LSH).
Our novel framework has the potential to be applied to other architectures and applications where the amount of data makes training even simple models impractical.
The major limitation of Sketch-GNN is that the sketched nonlinear activations are less expressive than the original activation functions, and the accumulated error of sketching makes it challenging to sketch much deeper GNNs.
We expect future research to tackle the above-mentioned issues and apply the proposed neural network sketching techniques to other types of data and neural networks.
Considering broader impacts, we view our work mainly as a methodological and theoretical contribution, and there is no obviously foreseeable negative social impact.
Acknowledgments and Disclosure of Funding
This work is supported by National Science Foundation NSF-IIS-FAI program, DOD-ONR-Ofﬁce of Naval Research,
DOD-DARPA-Defense Advanced Research Projects Agency
Guaranteeing AI
Robustness against Deception (GARD) and Adobe, Capital One and JP Morgan faculty fellowships.
Checklist
For all authors..
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
[Yes] See Section 1.
(b) Did you describe the limitations of your work?
[Yes] Currently, our work has two major limitations: (1) our theoretical assumptions and results may not perfectly correspond to the reality; see the theoretical remarks in Section 3, and (2) our implementation is not fully-optimized with the more advanced libraries; see the efﬁciency discussions in Section 5.
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] We see our work as a theoretical and methodological contribution toward more resource-efﬁcient graph representation learning.
Our methodological advances may enable larger-scale network analysis for societal good.
However, progress in graph embedding learning may potentially inspire other hostile social network studies, such as monitoring ﬁnegrained user interactions.
(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 all theoretical results?
[Yes] See Lemmaand Theorem 1.
(b) Did you include complete proofs of all theoretical results?
[Yes] See Appendix B. 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] See Appendix H. (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] See Appendix H. (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] See Section 5.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] See Appendix H. 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] See Appendix H. (b) Did you mention the license of the assets?
[Yes] See Appendix H. (c) Did you include any new assets either in the supplemental material or as a URL?
[No] (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/curating?
[No] (e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?
[No] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]