Graph Neural Network Acceleration via Matrix Dimension Reduction
Authors: Shunhua Jiang
Yunze Man
Zhao Song
Danyang Zhuo
Venue:  (2021)
Keywords: Graph Neural Networks
Deep learning
Optimization
Kernel Method

Abstract: Graph Neural Networks (GNNs) have become the de facto method for machine learning on graph data (e.g., social networks, protein structures, code ASTs), but they require significant time and resource to train. One alternative method is Graph Neural Tangent Kernel (GNTK), a kernel method that corresponds to infinitely wide multi-layer GNNs. GNTK's parameters can be solved directly in a single step, avoiding time-consuming gradient descent. Today, GNTK is the state-of-the-art method to achieve high training speed without compromising accuracy. Unfortunately, solving for the kernel and searching for parameters can still take hours to days on real-world graphs. The current computation of GNTK has running time $O(N^4)$, where $N$ is the number of nodes in the graph. This prevents GNTK from scaling to datasets that contain large graphs. Theoretically, we present two techniques to speed up GNTK training while preserving the generalization error: (1) We use a novel matrix decoupling method to reduce matrix dimensions during the kernel solving. This allows us to reduce the dominated computation bottleneck term from $O(N^4)$ to $O(N^3)$. (2)  We apply sketching to further reduce the bottleneck term to $o(N^{\omega})$, where $\omega \approx 2.373$ is the exponent of current matrix multiplication. Experimentally, we demonstrate that our approaches speed up kernel learning by up to $19\times$ on real-world benchmark datasets.

Under review as a conference paper at ICLR 2021
GRAPH NEURAL NETWORK ACCELERATION VIA MATRIX DIMENSION REDUCTION
Anonymous authors Paper under double-blind review
A ABSTRACT
Graph Neural Networks (GNNs) have become the de facto method for ma- chinc learning on graph data (c.g., social networks, protein structures, code ASTs), but they require significant time and resource train. One alter- nalive method is Graph Neural Tangent Kernel (GNTK), a kernel method that corresponds to infinitely wide multi-layer GNNs. GNTK's parameters can be solved directly in a single step, avoiding time-consuming gradient de- scent. Today, GNTK is the state-of-the- speed without compromising accuracy. Unfortunately, solving for the ker- nel and scarching for parameters can still take hours to days on real-world graphs. The current computation of GNTK has rumning time O(N*), where the number of nodes in the graph. This prevents GNTK from scal- ing to datasets that contain large graphs. Theoretically, we present two techniques to speed up GNTK training while preserving the generalization error: (1) We a novel matrix decoupling method to reduce matrix di- mensions during the kernel solving. This allows us to reduce the dominated computation bottleneck term from O(N1) O(N3). (2) We apply sketch- ing to further reduce the bottleneck term to o(N"), where z 2.373 is the exponent of current matrix multiplication. Experimentally, we demonstrate that our approaches speed up kernel learning by up to 19x on real-world benchmark datasets.
method achieve high training
INTRODUCTION
Graph Neural Networks (GNNs) have quickly become the de facto method for machine learn- ing graph data. GNNs have delivered ground-breaking results in many important areas of AI, including social networking Yang et al. (2020a), bio-informatics Zitnik & Leskovec (2017); Yue al. (2020), recommendation systems Ying et al. (2018), and autonomous driv- ing Weng et al. (2020); Yang et. al. (2020b). However, efficient GNNs training has become a major challenge with the relentless increase in the complexity of GNN models and dataset sizes, both in terms of the number of graphs in dataset and the sizes of the graphs.
Recently, new direction for fast GNN training (GNTK). Solving for the kernel and searching for the parameters in GNTK is equivalent to using gradient descent to train an infinitely wide multi-layer GNN. GNTK is significantly faster than iterative gradient descent optimization because solving the parameters in GNTK just a single-step kernel learning process. addition, GNTK allows GNN training to scale with GNN model sizes because the training time grows only linearly with the complexity of GNN models. However, GNTK training can still take hours to days on typical GNN datasets today.
use Graph Neural Tangent Kernel
Our key observation is that, during the process of solving parameters in GNTK, most of the training time and resource is spent multiplications of large matrices. Let N be the maximum number of nodes in the graphs, these matrices can have sizes as large as N x 2 This means single matrix multiplication takes at least 2 time, and it prevents GNTK from scaling to larger graphs. Thus, in order to speed up GNTK training, we need to reduce matrix dimensions.

---

For two integers a S b, Under review as a conference paper at ICLR 2021
Our Contributions. We present two lechniques to speed up GNTK: (1) We use novel matrix decoupling method to reduce matrix dimensions during the training without harming the calculation results. This reduces the dominated computation bottleneck term from O(N*) to ((N²). (2) We propose a sketching method to further reduce the bottleneck term to oNN", where 2.373 is the exponent of current matrix multiplication.
We provide thcoretical results that the resulting randomized GNTK still has a good gener- alization bound. In experiments, we evaluate our method on standard graph classification benchmarks. Our method improves GNTK training time by up to 19x while maintaining the same level of accuracy.
BACKGROUND
Notations. For a positive integer n, we define [n] := ... n2.
we define [a, b] == {a, a + 1, 1,, ,6]. and (a. b) == {a + 1. b -1 Similarly we define la, b) and (a, b). For a full rank square matrix A, we use A-1 to denole its true inverse. We define the big O notation such that [(m) = ((g(n) means there exists E N4, and M E R denote its spectral norm. We use ||A]| to denote its Frobenius norm. We use A to denote the transpose of A. For a matrix and vector z, we define || Vx Ax. We use to denote the ReLU activation function, i.e. 6(z) max 0}. For function f : R - R, we use f' denote the derivative of
such that [(n) S M g(n) for all n 2 no- For a matrix A, we use ||||| or [All2
Graph neural network (GNN). A GNN has L levels of AGGREGATE operations, cach followed by a COMBINE operation. A COMBINE operation has R fully-connected layers with output dimension m, and uses ReLU non-linearity. In the end, the GNN has a READOUT operation that corresponds to the pooling operation of normal neural networks.
Consider graph G EV with |V| N. Each node V has feature vector h. Rd.
In GNN we will use vectors (1rr) such that 1 denotes the number of levels, and denotes the number of hidden layers. The size is h4.00)
(1.0) RRd and (1.0) E Rm for all / [2 L}. We
define the initial vector la (U,R)) hy E IR,,, Yu U.
For any I E [Z], the AGGREGATE operation aggregates the information from last level:
(1,0) == Cu
Σ
h R(-1.R),
(E)(())(w))
where cu E R is a scaling parameter, and N(u) is the set of neighbor nodes of u. The COMBINE operation then uses R mully-connected layers with ReLU activation: Vr e [R].
((.r) == (co/m)//2 ((W(.))) h(rr-1) R
where Co E R is scaling parameter, ((.1) E Rmdd, , and W(.r)) E mmxm for all (l,r) e [L] x [R] 1)). We let W == (W(()) he(4), Finally, the output of the GNN on graph G = (V, E) is computed by a READOUT operation:
Sgnn(W.G. = Σ (L.R)) E m.
uEV
For more details see Appendix Section B.1.
Neural tangent kernel. We briefly review the neural tangent kernel definition. Let Φ denote the non-linear activation function, e.g. ((2) max(2,0) the ReLU activation function.
Definition 2.1 (Neural tangent kernel Jacot et al. (2018)). For any input two data points 4, RRd, we define the kernel mapping Kitk Rd x Rd + R
Kntke z) = =>)==>/ssoooa (w 22)x zdw
2

---

Under review as a conference paper at ICLR 2021
where Wh Ia) is d-dimensional mullivariate Gaussian distribulion, and Ø' is the derivalive of activation function Φ. If o is ReLU, then o (t) = 1 ift 20 0 and o' ('()) = if < 0. Given 21, 22. 2nn Rd, define kernel matriz K R as follows: Kijj K(z,,
The lower bound on smallest eigenvalue of neural tangent kernel matrix K (say À = Amin ( see Du et al. (2019b); Arora et al. (2019a;b); Song & Yang (2019); Brand et al. (2020) and separability of input data points {21, .. , } (6, see Li & Liang (2018); Allen-Zhu et al. (2019a;b)) play crucial role in deep learning theory. Due to Oymak & Soltanolkotabi (2020), À is at least 92(/n²) which unifies the two lines of research. The above work shows that as long as the neural network is sufficiently wide m Z poly(n, d, 1/6, 1/A) (where n is the number of input data points, d is the dimension of cach data). running (S)GD type algorithm is able minimize the training loss to zero.
Kernel regression and equivalence. Kernel method or Kernel regression is a fundamen- tal tool in machine learning Avron et al. (2017; 2019); Scholkopf Smola (2018). Recently, it has been shown that training an infinite-width neural network is equivalent to kernel regression Arora et al. (2019a). Further, the equivalence even holds for regularized neural network and kernel ridge regression Lee et al. (2020).
Let's consider the following neural nelwork, miuw 2 - Jun Jnn X)I2. Training this neu- ral network is equivalent to solving the following neural langent. kernel ridge regression problem: ming 3 (3, Note that fnu(3,, x) = (xx) 3 E R and (X = [mn 1)), fua (3, xn)]" E R" are the test data predictors. Here, & is the feature map corresponding the neural langent. kernel (NTK):
K
(w,
SSn (w,(W, (W,2)
OW |
(1 2)
T.. w
OW
where H, 2 Rd are any input data, and w,
i.dd
N I), r 1, m.
OUR GNTK FORMULATION
We show our GNTK formulation in this section. Our formulation builds upon the GNTK formulas of Du et al. (2019a). The descriptions this section presented in a simplified way. See Section B.2 and B.3 for more details.
3.1 EXACT GNTK FORMULAS
We consider GNN with L. AGGREGATE operations and COMBINE operations, and each COMBINE operation has R fully-connected layers. Let G = (U, F) and = (V, F) be two graphs with |U| = N and |V] N'. We use Ag and Aи to denote the adjacency matrix of G and H. We give the recursive formula to compute the kernel value Kgntk (G, H) e R induced by this GNN, which is defined as
K (oron (oonn (W, G) Df ")}
K Kgntk (G, = E
(=N((.1)
OW
Ow
where (0, I) is multivariate Gaussian distribution.
Rocall that the GNN uses scaling factors Cu for each node u E G. We define Cg to be the diagonal matrix such that (Ca)u C. for any E U. Similarly we define Cy. We will use intermediate matrices (((+)) e NxN' for each & e [0: L) and r re [0 R), where 1 denotes the level of AGGREGATE and COMBINE operations, and r denotes the level of fully-connected layers inside COMBINE operation.
Initially we define ( (G, E NNN' follows: Vu U,VEV,
[K(0. H)),,, ho))
where ha, hn E Rdd are the input features of u and v.
Next we recursively define (Gt,) H) for E [T] and r E [R].
3

---

Under review as a conference paper at ICLR 2021
Reference
Time
Du et al. (2019a)
O (Tmat (N, N, + NO LR- N One)) (Tmat (N, N, + + Tmatt (N, + LR N²)
Thm. 4.1 and 5.1
Table 1: When 1. O(1) and R 0(1); the dominate term in previous work is O(N*). We improve it to Tmaa (N, N,
Exact AGGREGATE. [K HG,,),,H,r H)u,, = Cu (E) bEA(())(())
The AGGREGATE operation gives the following formula:
Σ
Σ
[K H)
In the experiments the above equation is computed using Kronecker product:
veo(K(e.) (G, H)) ((CaAa) & (CnAn)) veo(K H)).
(1)
The dominating term of the final running time does not come from COMBINE and READOUT operations, thus we defer their details into Section B.2 in Appendix.
We briefly review the running time in previous work.
Theorem 3.1 (Running time of Du et al. (2019a), simplified version of Theorem D.2). Consider a GNN with L AGGREGATE operations and L COMBINE operations, and each COMBINE operation has R fully- connected layers. We compute the kernel matriz using graphs {G, = (V,, E with Vil S. N. Let E N, be the dimension of the feature vectors. The total runming time is O(n²) (Tmat N. L. N LR N
When using GNN, we usually use constant number of operations and fully-connedd connecccdeeeedd layers, ie., L (1), R O(1), and we have d o(N), while the size of the graphs can grow arbitrarily large. Thus it is easy to see that the dominated term in the above running time NA, the major contribution of this work is to reduce to o(Nw), where w Rs 2.373 is the exponent of current matrix multiplication.
3.2 APPROXIMATE GNTK FORMULAS
We follow the notations of previous section. Now the goal to compute an approximate version of the kernel value K(G, H) E R such that Kgatk( (G, H) x Kgntk(G H). We will use intermediate matrices K (G, H) NxN' for each & E [0 L. and r E [0 R). In the approximate version we add two random Gaussian matrices So R and Sy R//N in the AGGREGATE operation, where b < N and D' < N'.
A Approximate AGGREGATE operation. In the approximate version, we add two sketch- ing matrices SG bx and SH E R
k-(.0) H) CGAc (S) K (-(0-1.8) (G, H) (SH AnCH-
(2)
Not that for the special case ST.G SH I, the Eq. (2) degenerates to the the following case:
k(t H) CcAe k(e-1.m) (G, H) AnCu-
This equation is exactly the same as the equation Eq. (1) of the exact case. Soe Fact 4.2 for why they equivalent.
4 OUR TECHNIQUES: RUNNING TIME
The main contribution of our paper is to show that GNTK defined in Du et al. (2019a), while maintaining a similar generalization bound.
accelerate the computation of
In this section we present the techniques that to achieve faster running time. We will provide the generalization bound the next section. We first. present. our main running time theorem.
4

---

Under review as a conference paper at ICLR 2021
Theorem 4.1 (Main theorem, rumning time part, Theorem D.1). Consider GNN with AGGREGATE operations and COMBINE operations, and each COMBINE operation has R fully-connected layers. We compute the kernel matriz using graphs graphs {G, (V, E.)) with [Va S N. Let = o(N) be the sketch size. Let d E N+ be the dimension of the feature vectors. The total running time is
O(n²) (Tm N,, L.T.maa (N, LR. N2).
Note that we improve the dominating lerm from N to Tmat N, We achieve this improvement using two techniques:
1. We accelerate the multiplication of a Kronecker product with a vector by decoupling it into two matrix multiplications of smaller dimensions. In this way we improve the running time from N down to Tmat NN,N, N). We present more details in Section 4.2.
2. We further accelerate the two matrix multiplications by using two sketching matri- ces. In this way, we improve the running time from Tmai (N, N) to maa N, We present more details in Section 4.3.
4.1 NOTATIONS AND KNOWN FACTS
Fast matrix multiplication. We use the notation Time (n, d, m) to denote the time of multiplying an x d matrix with another d x m matrix. Let denote the exponent of matrix multiplication, i.e., Taae (n, n) = n² The first result shows w < 3 is Strassen 2.373, due to Williams (2012); Le Gall (2014).
(1969). The current best exponent is The common belief is w x 2 in the computational complexity community Cohn et al. (2005); Williams (2012); Jiang al. (2020). The following fact is well-known in the fast matrix multiplication literature Coppersmith (1982); Strassen (1991); Bürgisser et al. (1997) Taaa (a, b, c) = O(Tmat (a, c, b)) O(Tmar(T a, b)) for any positive integers a,b,c.
Kronecker product and vectorization. Given two matrices A E Baxxd and B E Re2xd2, We use @ to denote the Kronecker product, iee, for C = A x Reimuxdids, the (6 + (i2 1)) R1.J1 + (j2 1) di)-th entry of C' is Jn [di), E [d2]. For any give matrix H E Rddixda, use h vec(H) Rdida to denote the vector such that h = Har Vji E [di], ja E [d2].
Vi, e [n1], is [ng]
4.2 SPEEDUP VIA KRONECKER PRODUCT EQUIVALENCE
We make the following observation about kronecker product and vectorization. Proof is delayed to Section E.
(Equivalence between two matrix products and Kronecker product then ma- trix-vector multiplication). Given matrices E R xds BE Rnaxds and H E gddxdd we have vec(AHB = (A & B) vec(H).
Fact 4.2
In GNTK, when computing the 1-th AGGREGATE operation for I E [L], we need to compute product (A & B) veo( H)) with sizes A, B, H NNNN Note that A = CgAc, B AHAR, ((--1.)) (G, H) (see Eq. (1) in Section 3.1). Computing this product naively takes O(N4) time, since the Kronecker product A & B already has size N2 x N22 This is the O(N4) term of Du al. (2019a).
In our experiments, we instead compute AHBT, which takes O (N, N,N,N) time. And this is how we get our first improvement in running time.
4.3 SPEEDUP VIA SKETCHING MATRICES
The following lemma shows that the sketched version of matrix multiplication approximates the exact matrix multiplication. This justifies why we can use sketching matrices to speed up calculation.
5

---

Under review as a conference paper at ICLR 2021
Lemma 4.3 (Informal version of Lemiaa 5.4). Given n² malrices Hi,1.-... Hn,, E R and n matrices A1.... An E RNXN. Tet S E T denote a random malria where each Then with high probability, we have the
entry + Vo or - Vb' Vb each with probability 2. sollowing guaranlee: for all i,j E [n],
AT S SyA, STA A Hi
Note that the sizes of the matrices are A,, Ai, His RNNN and S,, Sg E RbxN, They corre- spond to A, = Ca, AG,, A3 = CG2 Aa,, Has = K (( Gj) (see Eq. (2) in Section 3.2).
Directly computing A, H, takes Tmaa (N,
time.
After adding two sketching matrices, using a certain ordering of computation, we can avoid the sine-consuming step of multiplying two N x N matrices. More specifically, we compute A, S S,S. A, in the following order:
Computing AT ST and SyA, both takes Tmai (N,, time.
Computing S: H.g takes Tmar(b, N) time.
Computing (5, ST takes Tina (b, N, time.
Computing (A S (S; takes Tiaa (N, b) time.
Computing (A) ST S,) (S,A,) takes Tmat (N, b, N) time.
Thus, we improve the rumning time from TN,
(N,
OUR TECHNIQUES! ERROR ANALYSIS
In this section, we prove that even though adding the sketching matrices in the GNTK formula will introduce some error, this error can be bounded, and we can still prove similar generalization bound as that of Du et al. (2019a).
Theorem 5.1 (Informal version of Theorem C.4). For each i E [ni), if the labels {y 1= satisfy y/4 = a EmEr (Teus 31) + Σt= 02 Enev(Tun (Tuu, (2))2, where Tu = cu EneN(Nu) hos I, BI, 12, Bt..... .2TT E Rd, and under the assumptions of Assump- tion C.6, and if we further have the conditions that
a1 02
4 (aml1B1112 + Σ 4VF a B2 o(n), = olvn),
then the generalization error of the approsimate GNTK be upper bounded by (G.v)D [[ffocc (G), O(1/n"), where constant E 10,,1).
LD(femek) = E
We use a standard generalization bound of kernel methods Bartlett & Mendelson (2002) (Theorem C.1) which shows that in order to prove generalization bound, it suffices to upper bound || " and tr[K]. We present our bound on ||y|| l The bound on tr[K] is simpler. For the full version of the prools, please see Section C.
Lemma 5.2 (Informal version of bound on |l R .). --1. Under Assumption C.6, we have
T
||vll=- 4/al- 1|3 + 4 1)/o2l 1|
E=
We provide a high-level proof sketch here. We first compute all the variables in the approxi- mate GNTK formula to get a close-form formula of K. Then combining with the assumption on the labels y, we show that Il R-- upper bounded by
| R 5 2 3 H H H 3)

---

Under review as a conference paper at ICLR 2021
where H, H E Rdxx are two matrices such that Vi,j E [n],
H Has - [HT H). = 1 (c, Ha, (S, Sa, Ca,
Ho. Ha, Co,
Next we show that H is PSD approximation of H H. Claim 5.3 (PSD approximation). We have (1 - H H H = H TH H 3 (1+ 1
Note that using this claim, we have
llull,, S (4a2. T H (HTH) H) H B) S
H
(H H) H B) s 4. al13112,
which finishes the proof. Now it remains to prove the claim. We prove it by using the following lemma which upper bounds the error of adding two sketching matrices. The proof of this lemma is deferred to Section E.
Lemma 5.4 (Frror bound of adding two sketching matrices). Let R E Ruixn, S E R be two independent AMS matrices Alon et al. (1999). Let B = Olog" 1.5 n). For any matrix A E IR and any vectors 9,h E IR", the following holds with probability
1/ 1/poly(n)
lg R RAS Sh - 9 Ah S
B Voi
11112/11111 + 3 Vo2
1 Alaal//l2
Alallllll +
vb, |
Using this lemma and the assumption of sketching sizes in the lemma statement, we can prove the following coordinate- wise upper bound:
- (H H)s-
|
[H Meal S
10"
Then we can upper bound ||A H H Hlla S I Hll,, which proves the claim. Thus we finish the proof.
6 EXPERIMENTS
In this section, we evaluate our proposed GNTK acceleration algorithm on various graph classification tasks. More details about the experiment setup can be found in Section of the supplementary material.
Datasets. We test our method on 7 benchmark graph classification datasets, including 3 social networking dataset (COLLAB, IMDBBINARY, LMDBMULTI) and bioinformatics datasets (PTC, NCL1, MUTAG and PROTEINS) Yanardag & Vishwanathan (2015). For bioinformatics datasct, each node has its categorical fcatures as input feature h to the algorithm. For social network dataset where nodes have input feature, each node as its feature better represent its structural information. The dataset statistics are shown in Table 2.
degree of
Baselines. We compare our proposed results with a number of state-of-the-art baselines for graph classification: (1) State-of-the-are deep graph neural networks architectures, in- cluding Graph Convolution Network (GCN) Kipf & Welling (2017), GraphSAGE Hamilton et al. (2017), PATCHY-SAN Nieport et al. (2016). Deep Graph CNN (DGCNN) Zhang et al. (2018a) and Graph Isomorphism Network (GIN) Xu et al. (2018a). (2) Kernel based methods, including the WL subtree kernel Shervashidze et al. (2011), Auonymous Walk Embeddings (AWL) Ivanov & Burnaev (2018), and RelGK Zhang et. al. (2018b). (3) Graph neural tangent kernel (GNTK) Du et al. (2019a). For deep learning methods, GNTK, RetGK and AWL, we report accuracy reported in the original papers. For WL subtree, report the accuracy of the implementation used in Xu et al. (2018a).
Results. We perform 10-fold cross validation and report the mean and standard deviation of accuracy. We show our performance by comparing with state-offthe-art Graph learning methods, including the original GNTK method. The accuracy is shown in Table 2 and
7

---

Under review as a conference paper at ICLR 2021
Datasets
COLLAB
IMDB-B
IMDB-M
PTC
NCN
MUTAG
PROTEINS
# of graphs
5000
1000
1500
344
4110
188
1113
# of classes
3
2
2
2
2
2
Avg # of nodes 74.5
19.8
13.0
25.5
29.8
17.9
39.1
GCN
79.0 + 1.8
74.0 = 3.4
51.9 + 3.8
64.2 + 4.3
80.2 + 2.0
85.6 + 5.8
76.0 + 3.2
GraphSAGE
72.3 = 5.3
50.9 2.2
63.9 = 7.7
77.7 + 1.5
85.1 7.6
75.9 3.2
PATCHY-SAN
72.6 2.2
2.2
45.2
60.0 + 4.8
78.6 1.9
d 4.2
+
DGONN
73.7
70.0
47.8
58.6
74.4
85.8
75.5
GIN
1.9
75.1 1.1+5.111
52.3 12.3+22.8 2.8
64.6 + 7.0
82.7 + 1.7
89.4 9..
76.2 2.8
WL Subtroe
78.9 1 1.9
73.8 3.9
50.9
59.9 1: 4.3
86.0
5.74
75.0
AWL
73.9
74.5
3.6
9.8
RetGK
0.3
1.0
47.7
62.5 1.6
84.3
90.3
75.8 4: 0.6
GNTK
83.6 + 1.0
3.6
62.8 4.6
67.9 4 6.9
84.2
9.5
75.6 15.6::4 4.2
Ours
83.6 = 1.0 76.9 10.9000
3.6
62.8
4.6
67.9 + 6.9
84.2
1.5
9.5
75.6 4.2
Table 2: Classification accuracy (%) for graph classification datasets with matrix de- coupling. We report the result of proposed method, oplimizing on original GNTK model.
Datasets
COLLAB
IMDEB
IMDR-M
PTC
NCI1
MUTAG
PROTEINS
GNTK
24 hrs
546.4
686.0
46.5
10,084.7
8.0
1,392.0
4,523.0 19x)
(6x) 112.5 (6.1x) 13.5 (3.4x)
(1.4x)
3.0 (2.7x)
782.7 (188x)
Running time analysis for our matrix decoupling method (in seconds). We report the kernel calculation time between the original GNTK method and our accelerated model.
Table 3:
running time is shown in Table 3. Our matrix decoupling method (MD) doesn't harm the result of GNTK while significantly accelerates the learning time of neural tangent kernel. Our proposed method achieves multiple times of improvements for all the datasets. In par- ticular, on COLLAB, our method achieves more than 19 times of learning time accelcration. We observe that the improvement of our method depends the sizes of the graphs. For large-scale dalaset like COLLAB, we achieve highest. acceleration because matrix multiplica- lion dominates the overall calculation time. And for bioinformatics datasets where number of nodes is relatively small, the improvement is not as prominent. Note that we only show the running time comparison between our method and the original GNTK method, because other state-of-the-al deep GNN methods takes significantly longer to learn via gradient descent Du al. (2019a). Analysis of our sketching method can be found in Section of the supplementary material.
7 CONCLUSION
Graph Neural Networks (GNNs) have become the most important method for machine learn- ing on graph data (e.g., social networks, protein structures), but training GNNs efficiently is major challenge. One alternative method is Graph Neural Tangent Kernel (GNTK), kernel method that is equivalent scent. GNTK's parameters can be solved directly in single step, avoiding time-consuming gradient descent. Because of this, GNTK has become the state-of-the-art method to achieve high training speed without compromising accuracy. Unfortunately, GNTK still takes hours to days to train on real-world graphs because it has a computation bottleneck of O(N*), where N denotes the number of nodes in the graph. We present two techniques to mitigate this bottleneck: (1) We use a novel matrix decoupling mothod to reduce matrix dimonsions during the kernel solving. This allows us to reduce this dominated computation bottleneck from ( to O(N)). (2) We apply sketching to further reduce the bottleneck to of ((N"), where Rs 2.373 is the exponent of current. matrix multiplication. We demonstrate that our approaches speed up kernel learning by up to 19x on real-world benchmark datasets.
train infinitely wide multi-layer GNNs using gradient de-
8

---

Under review as a conference paper at ICLR 2021
REFERENCES
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. On the convergence rate of training recurrent neural networks. In NewrIPS, 2019a.
Zeyuan Alle Yuanzhi Li, and Zhao Song. convergence theory for deep learning via wwee-parameterization. In ICML, 2019b.
Noga Alon, Yossi Matias, and Mario Szegedy. The space complexity of approximating the frequency moments. Journal of Computer and system sciences, 58(1):137-147, 1999.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of oplimization and generalization for overparamelerized two-layer neural net norksr.. In International Conference on Machine Learning (ICML), pp. 322 332. https: //atxi:/arxiv.. htggabb//990..0655 08584, 2019a.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, Ruslan Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. In NeurIPS. arXiv preprint arXiv:1904. 2019b.
Haim Avron, Michael Kapralov, Cameron Musco, Christopher Musco, Ameya Velingker, and Amir Zandieh. Random fourier features for kernel ridge regression: Approximation bounds and statistical guarantees. In ICML. https://arxiv. org/ /ptp/df///// /1804 09893,998..... pdf, 2017.
Haim Avron, Michael Kapralov, Cameron Musco, Christopher Musco, Ameya Velingker, and Amir Zandieh. A universal sampling method for reconstructing signals with simple fourier transforms. In STOC. pttpp://arxxie wwwp:/gggg/p///eeeeeeee 08723. pdf, 2019.
Peter Bartlett and Shahar Mendelson. Rademacher and gaussian complexities: Risk bounds and structural results. Journal of Machine Learning Research, 3(Nov):463 482-222, 482, 2002.
Jan den Brand, Binghui Peng, Zhao Song, and Omri Weinstein. Training (over- parametrized) neural networks in near-linear time. arXiv preprint 2rx::22006.110 2020.
Peter Bürgisser, Michael Clausen, and Mohammad A Shokrollahi. Algebraic complexity theory, volume 315. Springer Science & Business Media, 1997.
Henry Cohn, Robert Kleinberg, Balazs Szegedy, and Christopher Umans. Croup-theorettc algorithms for matrix multiplication. In 46th Annual IEEE Symposium on Foundations of Computer Science (FOCS'05) 05), pp. 379 388. IEFF, 2005.
Don Coppersmith. Rapid multiplication of rectangular matrices. SIAM Journal on Com- puling, 11 (33):467-477) 1982.
Simon Du, Kangcheng Hou, Russ Salakhutdinov, Barnabas Poczos, Ruosong Wang, and Keyulu Xu. Graph neural tangent kernel: Fusing graph neural networks with graph kernels. In Advances in Neural Information Processing Systems (NeurIPS), pp. 5723- 5733. https://arx:: wwwg/bs//g///eeeeeeeeee 13192, 2019a.
Simon Du, Xiyu Zhai, Baruabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. In ICLR. arXiv preprint arXiv:1810.02054, 2019b.
Will Hamilton, Zhitao Ying, and Jure Leskovec. Inductive representation learning on large graphs. In Advances neurol information processing systems (NeurIPS), pp. 1024-1034. https: pettps//aaze //arxi//.iiieee org/pdf/ 1100/ 02216, off,.. 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedinggs of the IEFF international conference computer vision, pp. 1026-1034 2015.

---

Under review as a conference paper at ICLR 2021
Sergey Ivanov and Evgeny Burnaev. Anonymous walk embeddings. aiv prepriat arXic:1805.11922 2018.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in ncural networks. In Advances in newral information processing systems (NewrIPS), pp. 8571 8580, 2018.
Shunhua Jiang, Zhao Song, Omri Weinstein, and Hengjie Zhang. Faster dynamic matrix inverse for faster lps. Xi preprint Xiv:2004.074 2020.
Thomas N. Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. In International Conference on Learning Representations (TCLR). https:// arxiv. org/ org/pdf/ 1609. www.11.099...... pdf, 2017.
François Le Gall. Powers of tensors and fast matrix multiplication. In Proceedings of the 39th
international symposium on symbolic and algebraic computation (ISSAC), pp. 296-303. ACM, 2014.
Jason D Lee, Ruoqi Shen, Zhao Song, Mengdi Wang, and Zheng Yu. Generalized leverage score sampling for neural networks. In NeurIPS, 2020.
Yuanzhi Li and Yingyu Liang. Learning overparameterized neural networks via stochastic gradient descent on structured data. In NewrIPS, 2018.
Mathias Niepert, Mohamed Ahmed, and Konstantin Kutzkov. Learning convolutional neural networks for graphs. In International conference on machine learning, pp. 2014-2023, 2016.
Samet Oymak and Mahdi Soltanolkotabi. Towards moderate overparametterizatio::: global convergence guarantees for training shallow neural networks. In IEEE Journal on Selected Areas in Information Theory. https: /taps://aree petps://arxivv wttp//f////////////////////e//// 190.. 04674. pdf, 2020.
Bernhard Scholkopf and Alexander Smola. Learning with kernels: support vector machines, regularization, opfimization, and beyond. Adaptive Computation and Machine Learning series, 2018
Nino Shervashidze, Pascal Schweitzer, Erik Jan Van Leeuwen, Kurt Mehlhorn, and Karsten M Borgwardt. Weisfeiler-lehmann graph kernels. Journal of Machine Learning Research, 12(9), 2011.
Zhao Song and Xin Yang. Quadratic suffices for over-paramettiiattion via matrix chernoff bound. In arxiv preprint. https://arxiv. hwtpg//dff1906666 1906. 03593, odf, 2019.
Volker Strassen. Gaussian climination is not optimal. 1969.
Volker Strassen. Degeneration and complexity of bilinear maps: some asymptotic spectra.
J. reine angew. Math, 413:127 180, 1991.
Xinshuo Weng, Yongxin Wang, Yunze Man, and Kris M Kitani. Gnn3dmot: Graph neural network for 3d multi-object tracking with 2d-3d multi-fcature learning. In Proceedings of the IEFE/ /CVF Conference Computer Vision and Pattern Recognition, pp. 6499-6500 2020.
Virginia Vassilevska Williams. Multiplying matrices faster than sospersmith-winngradd In Proceedings of the forty-fourth annual A ACM symposium on Theory of computing (STOC), pp. 887-898. ACM, 2012.
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? In International Conference on Learning Representations (ICLR). org/ abbs/ 00826, 2018a.
Keyulu Xu, Chengtao Li, Yonglong Tian, Tomohiro Sonobe, Ken-ichi Kawarabayashi, and Stefanie Jegelka. Representation learning on graphs with jumping knowledge works..s. In International Conference on Machine Learning, pp. 5453 5462, 2018b.
10

---

Under review as a conference paper at ICLR 2021
Pinar Yanardag and SVN Vishwanathan. Deep graph kernels. In Proceedings of the 21th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pp. 1365-13744 2015.
Carl Yang, Jieyu Zhang, Haonan Wang, Sha Li, Myungwan Kim, Matt Walker, Yiou Xiao, and Jiawei Han. Relation learning on social networks with multi-modal graph edge varia- tional antoencoders. In Proceedings of the 13th International Conference on Web Search and Data Mining, pp. 699-707, 2020a.
Zetong Yang, Yanan Sun, Shu Liu, and Jiaya Jia. 3dssd: Point-based 3d single stage object delector. In Proceedings of the IEEE/OK /CVF Conference on Computer Vision and Pollern Recognition, pp. 1.000-1108,, 2020b.
Rex Ying, Ruining Hc, Kaifeng Chen, Pong Eksombatchai, William L Hamilton, and Jure Leskovec. Graph convolutional neural networks for web-scale recommender systems. In Proceedings of the 24th ACM SIGKDD International Conference Knowledge Discovery & Data Mining, pp. 974-983, 2018.
Xiang Yuc, Zhen Wang, Jingong Huang, Srinivasan Parthasarathy, Soheil Moosavinasab, Yungui Huang, Simon M Lin, Wen Zhang, Ping Zhang, and Huan Sun. Graph embedding biomedical networks: methods, applications and evaluations. Bivinformalics, 36(4): 1241-1251, 2020.
Muhan Zhang, Zhicheng Cui, Marion Neumann, and Yixin Chen. An end-to-end deep learn- ing architecture for graph classification. In Pecroy-Second AAAI Conference on Artificial Inlelligence, 2018a.
Zhen Zhang, Mianzii Wang, Yijian Xiang, Yan Huang, and Arye Nehorai. Retgk: Graph kernels basod on return probabilitics of random walks. In Advances in Neural Information Processing Systems, pp. 3964 3974, 2018b.
Marinka Zitnik and Jure Leskovec. Predicting multicellular function through multi-layer tissuc networks. Bioinformatics, 13114):i190- 2017.
11


Reviews:
[{"review_id": "stUQ4WUMBPD", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 1, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "In this paper, the authors propose two techniques to speed up the training of the graph neural tangent kernel, by matrix decoupling and sketching. Experiments are convincing.\n\nWe find the title of the paper inappropriate, because the paper only consider the  graph neural tangent kernel.\n\nIn general, the paper is difficult to read because many important parts are only available in the appendices.\n\nThere are some spelling and grammatical errors that can be easily identified and corrected, such as \"The descriptions in this section is\".\n\n----------------\nFollowing the authors' response,  we have updated our rating accordingly considering the following facts: the authors didn't make any change to the paper within the rebuttal, while they had the possibility in response to several questions. They didn't address even our simplest concerns, about the title being inappropriate. They still want to keep the title too general, while the paper considers only graph neural tangent kernel. Even if the latter is equivalent to infinitely wide multi-layer GNN, it is still a very special case. Moreover, many of the major issues raised by the other reviews were not addressed.", "extra_scores": {}, "extra_text": {}}, {"review_id": "UC0VmMn2ARp", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes to apply Neural Tangent Kernel to GNN to improve training. This is a good initiative to apply state of the art methods to GNN training which would interest researchers and practitioners in the community. However, the novelty introduced in the paper low since the authors directly apply NTK to GNN.\n\nPros:\na.) A good point of the paper is the compressive experiments and theoretical guarantees which make the paper complete and useful to the community. The paper is also well written.\n\nb.) Good performances have been achieved from experiments with many benchmark datasets\n\n\nCons:\na.) My main issue with the paper is its lack of novelty. It is a good idea but the authors do not provide any new ideas for either NTK and GNN. The paper shows significant performance improvement with the proposed method, however, it is somewhat understandable. \n\nb.) More details about the experiments would be helpful. Can the authors give details about varying dimensions of the sketch matrices and its effect on speed and accuracy?\n\nc.) This paper focuses on using sketch matrices based on Alon et al.. It would be more informative if the authors compare using other sketching methods such as Gaussian matrices and sampling methods.", "extra_scores": {}, "extra_text": {}}, {"review_id": "GYhwJl3tXno", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 2, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\n#### **Post Rebuttal**\nI'm afraid the authors failed to answer my main question regarding the results and the applicability of their proposed approximation.\n\nTherefore, I decide to keep my score.\n\n---\n#### **Contribution Claims** \nThe paper vouches for the use of GNTK as a strong method for graph learning combining good properties of kernels and GNNs. However, it raises the issue of time complexity for calculating the kernel scaling with $O(N^4)$, and addresses it by suggesting two techniques to gain speed up. The proposed techniques are based on matrix decoupling, which simply improves complexity, and matrix sketching which approximates the correct kernel (in probability).\nThe authors further present a theoretical result that similar generalization guarantees, as attributed to the exact GNTK, are achieved with their approximation.\n\n**Strengths** \n- The paper gains a major speed up in GNTK calculation with no harm to performance (*up to a clarification regarding the reported results. See first item in weaknesses section) \n- The method generalization properties are analyzed and shown to be preserved.\n- Although the paper is packed with lots of details, the flow of the paper allows the reader to pick upon the general ideas presented in the paper.\n\n**Weaknesses** \n- *Results* - I have concerns regarding the reported performance in Table 2. I find it quite uncanny that the numbers are *exactly* the same as achieved by exact GNTK, although the proposed method is an approximation which also utilizes a stochastic component. In case I misunderstood, I would kindly ask the authors to clarify this point. \n- *Discussion compared to GNTK* - A discussion on how the exact GNTK generalization bound compares to the approximate one is missing.\n- *Related work* - The paper does not explain and mention previous and relevant uses of matrix sketching (either in deep learning or kernel methods in general) which makes it hard for someone who is not familiar with the term to understand solely from the paper what it means. It is also unclear how innovative the use of it is. \n- *Experimental setting* - Regarding the evaluation on social datasets, it is mentioned that a degree feature was added as an input. Is it true for the other baselines as well? \n- *Missing convincing use-cases* - As the paper claims to gain speed up making GNTK scalable to larger datasets, I would have liked to see the performance both in time and accuracy of the proposed method. \n\n#### **Decision Recommendation**\nFor now, I suggest to reject the paper.  \nI think that it is interesting and valuable but feels that it needs to be solidified. \nUpon clarification of the questionable results I would consider changing my decision. \n\n#### **Other Comments**\n- Typos - the paper contains many typos. e.g, Theorem 4.1 \"...n graphs n graphs...\" repeating twice.  \n\n\n\n\n\n\n\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: While the author response clarified some concerns, it could not convince the reviewers that the current version of the paper should be accepted for publication at ICLR. 

Decision: Reject