Graph Neural Tangent Kernel: Fusing Graph Neural Networks with Graph Kernels

Introduction

Learning on graph-structured data such as social networks and biological networks requires one to design methods that effectively exploit the structure of graphs. Graph Kernels (GKs) and Graph Neural Networks (GNNs) are two major classes of methods for learning on graph-structured data. GKs, explicitly or implicitly, build feature vectors based on combinatorial properties of input graphs. Popular choices of GKs include Weisfeiler-Lehman subtree kernel (Shervashidze et al., 2011), graphlet kernel (Shervashidze et al., 2009) and random walk kernel (Vishwanathan et al., 2010; Gärtner et al., 2003). GKs inherit all benefits of kernel methods. GKs are easy to train, since the corresponding optimization problem is convex. Moreover, the kernel function often has explicit expressions, and thus we can analyze their theoretical guarantees using tools in learning theory. The downside of GKs, however, is that hand-crafted features may not be powerful enough to capture high-order information that involves complex interaction between nodes, which could lead to worse practical performance than GNNs.

GNNs, on the other hand, do not require explicitly hand-crafted feature maps. Similar to convolutional neural networks (CNNs) which are widely applied in computer vision, GNNs use multi-layer structures and convolutional operations to aggregate local information of nodes, together with non-linear activation functions to extract features from graphs. Various architectures have been proposed (Xu et al., 2019a, 2018). GNNs extract higher-order information of graphs, which lead to more powerful features compared to hand-crafted combinatorial features used by GKs. As a result, GNNs have achieved state-of-the-art performance on a large number of tasks on graph-structured data. Nevertheless, there are also disadvantages of using GNNs. The objective function of GNNs is highly non-convex, and thus it requires careful hyper-parameter tuning to stabilize the training procedure. Meanwhile, due to the non-convex nature of the training procedure, it is also hard to analyze the learned GNNs directly. For example, one may ask whether GNNs can provably learn certain class of functions. This question seems hard to answer given our limited theoretical understanding of GNNs. Another disadvantage of GNNs is that the expressive power of GNNs scales with the number of parameters. Thus, it is hard to learn a powerful GNN when computing resources are limited. Can we build a model that enjoys the best of both worlds, i.e., a model that extracts powerful features as GNNs and is easy to train and analyze like GKs?

In this paper, we give an affirmative answer to this question. Inspired by recent connections between kernel methods and over-parameterized neural networks (Arora et al., 2019b, a; Du et al., 2019, 2018; Jacot et al., 2018; Yang, 2019), we propose a class of new graph kernels, Graph Neural Tangent Kernels (GNTKs). GNTKs are equivalent to infinitely wide GNNs trained by gradient descent, where the word “tangent” corresponds to the training algorithm — gradient descent. While GNTKs are induced by infinitely wide GNNs, the prediction of GNTKs depends only on pairwise kernel values between graphs, for which we give an analytic formula to calculate efficiently. Therefore, GNTKs enjoy the full expressive power of GNNs, while inheriting benefits of GKs.

First, inspired by recent connections between over-parameterized neural networks and kernel methods Jacot et al. (2018); Arora et al. (2019a); Yang (2019), we present a general recipe which translates a GNN architecture to its corresponding GNTK. This recipe works for a wide range of GNNs, including graph isomorphism network (GIN) (Xu et al., 2019a), graph convolutional network (GCN) (Kipf and Welling, 2016), and GNN with jumping knowledge (Xu et al., 2018). Second, we conduct a theoretical analysis of GNTKs. Using the technique developed in Arora et al. (2019b), we show for a broad range of smooth functions over graphs, a certain GNTK can learn them with polynomial number of samples. To our knowledge, this is the first sample complexity analysis in the GK and GNN literature. Finally, we validate the performance of GNTKs on 7 standard benchmark graph classification datasets. On four of them, we find GNTK outperforms all baseline methods and achieves state-of-the-art performance. In particular, GNKs achieve 83.6% accuracy on COLLAB dataset and 67.9% accuracy on PTC dataset, compared to the best of baselines, 81.0% and 64.6% respectively. Moreover, in our experiments, we also observe that GNTK is more computationally efficient than its GNN counterpart.

This paper is organized as follow. In Section 2, we provide necessary background and review operations in GNNs that we will use to derive GNTKs. In Section 3, we present our general recipe that translates a GNN to its corresponding GNTK. In Section 4, we give our theoretical analysis of GNTKs. In Section 5, we compare GNTK with state-of-the-art methods on graph classification datasets. We defer technical proofs to the supplementary material.

Preliminaries

Next, we formalize the GNN framework. We refer to the neighbor aggregation process as a BLOCK operation, and to graph-level pooling to as a READOUT operation.

A BLOCK operation aggregates features over a neighborhood $\mathcal{N}(u)\cup\{u\}$ via, e.g., summation, and transforms the aggregated features with non-linearity, e.g. multi-layer perceptron (MLP) or a fully-connected layer followed by ReLU. We denote the number of fully-connected layers in each BLOCK operation, i.e., the number of hidden layers of an MLP, by $R$.

When $R=1$, the BLOCK operation can be formulated as

When the number of fully-connected layers $R=2$, the BLOCK operation can be written as

READOUT Operation.

To get the representation of an entire graph $\bm{h}_{G}$ after $L$ steps of aggregation, we take the summation over all node features, i.e.,

There are more sophisticated READOUT operations than a simple summation Xu et al. (2018); Zhang et al. (2018a); Ying et al. (2018). Jumping Knowledge Network (JK-Net) Xu et al. (2018) considers graph structures of different granularity, and aggregates graph features across all layers as

Building GNNs using BLOCK and READOUT.

GNTK Formulas

In this section we present our general recipe which translates a GNN architecture to its corresponding GNTK. We first provide some intuitions on neural tangent kernels (NTKs). We refer readers to Jacot et al. (2018); Arora et al. (2019a) for more comprehensive descriptions.

Recent advances in optimization of neural networks have shown, for sufficiently over-parameterized neural networks, the matrix $\bm{H}(t)$ keeps almost unchanged during the training process Arora et al. (2019b, a); Du et al. (2019, 2018); Jacot et al. (2018), in which case the training dynamics is identical to that of kernel regression. Moreover, under a random initialization of parameters, the random matrix $\bm{H}(0)$ converges in probability to a certain deterministic kernel matrix, which is called Neural Tangent Kernel (NTK) Jacot et al. (2018) and corresponds to infinitely wide neural networks. See Figure 4 in the supplementary material for an illustration.

in the limit that $m\to\infty$ and $\theta$ are all Gaussian random variables, which can be viewed as a Gaussian process. For each layer in the GNN, we use $\bm{\Sigma}$ to denote the covariance matrix of outputs of that layer, and $\bm{\dot{\Sigma}}$ to denote the covariance matrix corresponds to the derivative of that layer. Due to the multi-layer structure of GNNs, these covariance matrices can be naturally calculated via dynamic programming.

2 Formulas for Calculating GNTKs

More specifically, we first perform a neighborhood aggregation operation

READOUT Operation.

Given these intermediate outputs, we can now calculate the final output of GNTK using the following formula.

To better illustrate our general recipe, in Figure 1 we give a concrete example in which we translate a GNN with $L=2$ BLOCK operations, $R=1$ fully-connection layer in each BLOCK operation, and jumping knowledge, to its corresponding GNTK.

Theoretical Analysis of GNTK

For a testing point $G_{te}$, the prediction of kernel regression using GNTK on this testing point is

The following result is a standard result for kernel regression proved using Rademacher complexity. For a proof, see Bartlett and Mendelson (2002).

The following two theorems guarantee that if labels are generated as described in (3), then the GNTK that corresponds to the simple GNN described above can learn this function with polynomial number of samples. We first give an upper bound on $\boldsymbol{y}^{\top}\boldsymbol{\Theta}^{-1}\boldsymbol{y}$.

For each $i\in[n]$, if the labels $\{y_{i}\}_{i=1}^{n}$ satisfy

Combining Theorem 4.2 and Theorem 4.3 with Theorem 4.1, we know if

is bounded, and $|V_{i}|$ is bounded for all graphs $G_{i}=(V_{i},E_{i})$ in the training set, then the GNTK that corresponds to the simple GNN described above can learn functions of forms in (3), with polynomial number of samples. To our knowledge, this is the first sample complexity analysis in the GK and GNN literature.

Experiments

In this section, we demonstrate the effectiveness of GNTKs using experiments on graph classification tasks. For ablation study, we investigate how the performance varies with the architecture of the corresponding GNN. Following common practices of evaluating performance of graph classification models Yanardag and Vishwanathan (2015), we perform 10-fold cross validation and report the mean and standard deviation of validation accuracies. More details about the experiment setup can be found in Section B of the supplementary material.

The benchmark datasets include four bioinformatics datasets MUTAG, PTC, NCI1, PROTEINS and three social network datasets COLLAB, IMDB-BINARY, IMDB-MULTI. For each graph, we transform the categorical input features to one-hot encoding representations. For datasets where the graphs have no node features, i.e. only graph structure matters, we use degrees as input node features.

1 Results

We compare GNTK with various state-of-the-art graph classification algorithms: (1) the WL subtree kernel Shervashidze et al. (2011); (2) state-of-the-art deep learning architectures, including Graph Convolutional Network (GCN) Kipf and Welling (2016), GraphSAGE Hamilton et al. (2017), Graph Isomorphism Network(GIN) Xu et al. (2019a), PATCHY-SANNiepert et al. (2016) and Deep Graph CNN (DGCNN) Zhang et al. (2018a); (3) Graph kernels based on random walks, i.e., Anonymous Walk Embeddings Ivanov and Burnaev (2018) and RetGK Zhang et al. (2018b). For deep learning methods and random walk graph kernels, we report the accuracies reported in the original papers. The experiment setup is deferred to Section B.

The graph classification results are shown in Table 1. The best results are highlighted as bold. Our proposed GNTKs are powerful and achieve state-of-the-art classification accuracy on most datasets. In four of them, we find GNTKs outperform all baseline methods. In particular, GNTKs achieve 83.6% accuracy on COLLAB dataset and 67.9% accuracy on PTC dataset, compared to the best of baselines, 81.0% and 64.6% respectively. Notably, GNTKs give the best performance on all social network datasets. Moreover, In our experiments, we also observe that with the same architecture, GNTK is more computational efficient that its GNN counterpart. On IMDB-B dataset, running GIN with the default setup (official implementation of Xu et al. (2019a)) takes 19 minutes on a TITAN X GPU and running GNTK only takes 2 minutes.

2 Relation between GNTK Performance and the Corresponding GNN

We conduct ablation study to investigate how the performance of GNTK varies as we change the architecture of the corresponding GNN. We select two representative datasets, one social network dataset IMDBBINARY, and another bioinformatics dataset NCI1. For IMDBBINARY, we vary the number of BLOCK operations in $\{2,3,4,5,6\}$. For NCI1, we vary the number of BLOCK operations in $\{8,10,12,14,16\}$. For both datasets, we vary the number of MLP layers in $\{1,2,3\}$.

We investigate how the performance of GNTKs is correlated with number of BLOCK operations and the scaling factor $c_{u}$. First, on the bioinformatics dataset (NCI), we observe that GNTKs with more layers perform better. This is perhaps because, for molecules and bio graphs, more global structural information is helpful, as they provide important information about the chemical/bio entity. On such graphs, GNTKs are particularly effective because GNTKs can easily scale to many layers, whereas the number of layers in GNNs may be restricted by computing resources.

Moreover, the performance of GNTK is correlated with that of the corresponding GNN. For example, in social networks, GNTKs with sum aggregation $c_{u}=1$ work better than average aggregation $c_{u}=\frac{1}{\left|\mathcal{N}(u)\right|+1}$. The similar pattern holds in GNNs, because sum aggregation learns more graph structure information than average aggregation Xu et al. (2019a). This suggests GNTK can indeed inherit the properties and advantages of the corresponding GNN, while also gaining the benefits of graph kernels.

Effects of Jumping Knowledge and Number of MLP Layers

In the GNN literature, jumping knowledge network (JK) is expected to improve performance Xu et al. (2018); Fey (2019). In Figure 3, we observe that a similar trend holds for GNTK. The performance of GNTK is improved on both NCI and IMDB datasets when jumping knowledge is applied. Moreover, increasing the number of MLP layers can increase the performance by $\sim 0.8\%$. These empirical findings further confirm that GNTKs can inherit the benefits of GNNs, since improvements on GNN architectures are reflected in the improvements GNTKs.

We conclude that GNTKs are attractive for graph representation learning because they can combine the advantages of both GNNs and GKs.

Acknowledgments

S. S. Du and B. Póczos acknowledge support from AFRL grant FA8750-17-2-0212 and DARPA D17AP0000. R. Salakhutdinov and R. Wang are supported in part by NSF IIS-1763562, Office of Naval Research grant N000141812861, and Nvidia NVAIL award. K. Xu is supported by NSF CAREER award 1553284 and a Chevron-MIT Energy Fellowship. This work was performed while S. S. Du was a Ph.D. student at Carnegie Mellon University and K. Hou was visiting Carnegie Mellon University.

References

Appendix A Missing Proofs

By Section 3, for two graph $G$ and $G^{\prime}$, the GNTK kernel function that corresponds to the simple GNN can be described as

Let $\bm{\Phi}^{(2l)}(\cdot)$ be the feature map of the polynomial kernel of degree $2l$, i.e.,

A.2 Proof of Theorem 4.3

Since for each node $u$, $\bm{\overline{h}}_{u}=c_{u}\sum_{v\in\mathcal{N}(u)\cup\{u\}}\bm{h}_{v}$, and $c_{u}=\left(\left\|\sum_{v\in\mathcal{N}(u)\cup\{u\}}\bm{h}_{v}\right\|_{2}\right)^{-1}$, we have $\|\bm{\overline{h}}_{u}\|_{2}=1$. Moreover,

Appendix B Experiment Setup

To calculate GNTKs, we adopt the formulas provided in Section 3.2. To calculate the expectation of the post-activation output, i.e., (1) and (2), we use the same approach as in Arora et al. [2019a] (cf. Section 4.3 in Arora et al. [2019a]).

For GNTKs, we tune the following hyperparameters.

The number of BLOCK operations. We search from candidate values $\{1,2,\ldots,14\}$.

The number of fully-connected layers in each BLOCK operation. We search from candidate values $\{1,2,3\}$.

The parameter $c_{u}$. We search from candidate values $\left\{1,\frac{1}{\left|\mathcal{N}(u)\right|+1}\right\}$.

To utilize the GNTKs we compute to perform graph classification, we test with kernel regression and $C$-SVM as the final classifier. In our experiments, the regularization parameter $C$ in $C$-SVM is determined using grid search from 120 values evenly chosen from $[10^{-2},10^{4}]$, in log scale.

We would like to remark that GNTK has strictly smaller number of hyper-parameters than GNN since we do not need to tune the learning rate, momentum, weight decay, batch size and the width of the MLP layers for GNTK. Furthermore, we find on bioinformatics datasets, we get consistently good results by setting the number of BLOCK operations to be $10$, the number of MLP layers to be $1$ and $c_{u}$ to be $1/|\mathcal{N}(u)|$. We get 75.3% accuracy on PROTEINS, 67.9% on PTC, and 83.6% on NCI1. For social network datasets, by setting the number of BLOCK operations to be $2$, the number of MLP layers to be $2$ and $c_{u}$ to be $1$, we get 76.7% accuracy on IMDB-B, 52.8% on IMDB-M, and 83.3% on COLLAB.