Source-linked AI summary
Graph Neural Tangent Kernel: Fusing Graph Neural Networks with Graph Kernels
Simon S. Du, Kangcheng Hou, Barnabás Póczos, Ruslan Salakhutdinov, Ruosong Wang, Keyulu Xu
TL;DR
Graph kernels are easy to train but can be limited by hand-crafted features, while GNNs extract stronger features yet are harder to train and analyze. The paper proposes GNTKs, kernels corresponding to infinitely wide gradient-trained GNNs, and reports theoretical learnability and strong benchmark performance.
Problem
Graph kernels have limited expressive power from hand-crafted features, whereas GNNs are harder to train and their theoretical guarantees are less understood.
Method
GNTKs translate infinitely wide GNNs trained by gradient descent into analytically computable pairwise graph kernels.
Results
GNTKs achieve state-of-the-art performance on four of seven graph-classification benchmarks, outperforming all baselines on those datasets.
Takeaways & Limitations
GNTKs combine the expressive power of GNNs with graph-kernel benefits and provide theoretical learning guarantees for broad classes of smooth graph functions.
Takeaways & Limitations
The theoretical analysis considers a simple GNN with a single BLOCK and assumes bounded graph sizes and an invertible kernel matrix.
Abstract
from arXiv · showhide
While graph kernels (GKs) are easy to train and enjoy provable theoretical guarantees, their practical performances are limited by their expressive power, as the kernel function often depends on hand-crafted combinatorial features of graphs. Compared to graph kernels, graph neural networks (GNNs) usually achieve better practical performance, as GNNs use multi-layer architectures and non-linear activation functions to extract high-order information of graphs as features. However, due to the large number of hyper-parameters and the non-convex nature of the training procedure, GNNs are harder to train. Theoretical guarantees of GNNs are also not well-understood. Furthermore, the expressive power of GNNs scales with the number of parameters, and thus it is hard to exploit the full power of GNNs when computing resources are limited. The current paper presents a new class of graph kernels, Graph Neural Tangent Kernels (GNTKs), which correspond to infinitely wide multi-layer GNNs trained by gradient descent. GNTKs enjoy the full expressive power of GNNs and inherit advantages of GKs. Theoretically, we show GNTKs provably learn a class of smooth functions on graphs. Empirically, we test GNTKs on graph classification datasets and show they achieve strong performance.
1 Introduction
The paper contrasts graph kernels’ trainability and guarantees with GNNs’ stronger feature extraction, then proposes GNTKs to combine these advantages. GNTKs translate infinitely wide GNNs into analytically computable graph kernels with theoretical and empirical support.
- Graph kernels use combinatorial graph properties, whereas GNNs aggregate through multiple layers and nonlinearities to extract higher-order features.
- GNNs can outperform graph kernels in practice but are harder to train because their objectives are highly non-convex.
- GNTKs correspond to infinitely wide GNNs trained by gradient descent and depend only on pairwise graph-kernel values computed analytically.
- GNTKs are designed to retain GNN expressive power while inheriting graph-kernel benefits.
- A general recipe covers GIN, GCN, and jumping-knowledge GNNs, while theory establishes polynomial-sample learning for broad smooth graph functions.
- Across seven graph-classification benchmarks, GNTK outperformed all baselines on four datasets and achieved state-of-the-art performance there.
2 Preliminaries
The preliminaries describe graph classification and the neighborhood-aggregation framework used to build GNNs. They formalize BLOCK transformations, READOUT pooling, and architecture-dependent aggregation choices.
- The paper studies graph classification: predicting labels for unseen graphs from labeled training graphs.
- GNNs recursively aggregate and transform neighboring node representations, then pool node features into a graph representation.
- The framework names neighborhood aggregation BLOCK and graph-level pooling READOUT operations.
- A BLOCK aggregates features from N(u) ∪ {u} and applies a nonlinear transformation such as an MLP or ReLU layer.
- The number R denotes fully connected layers within each BLOCK, while L denotes the number of BLOCK operations.
- Aggregation scaling distinguishes architectures: GCN averages neighbors with cu = 1/(|N(u)|+1), whereas GIN sums them with cu = 1.
- With R = 2, the BLOCK uses two hidden MLP layers after neighborhood aggregation.
- The standard READOUT sums node features, while alternatives include cross-layer jumping-knowledge aggregation.
3 GNTK Formulas
This section derives GNTKs by taking infinite-width limits of GNNs and recursively computing covariance and derivative quantities across graph-aggregation layers. The resulting pairwise graph kernel is assembled through READOUT formulas, with optional jumping knowledge.
- 3.1 Intuition of the Formulas: The GNTK construction translates a specified GNN architecture into an explicit kernel formula for pairs of input graphs.
- 3.1 Intuition of the Formulas: In sufficiently over-parameterized networks, nearly constant gradient-descent kernel dynamics become equivalent to kernel regression, motivating the NTK limit.
- 3.1 Intuition of the Formulas: As width m approaches infinity with Gaussian initialization, layerwise output covariances and derivative covariances can be computed by dynamic programming.
- 3.2 Formulas for Calculating GNTKs: For graphs G and G′, the formula defines a pairwise kernel value Θ(G, G′) for a GNN with L BLOCK operations and R fully connected layers per BLOCK.
- 3.2 Formulas for Calculating GNTKs: The construction initializes Σ(0)(G, G′) as the covariance matrix between input node features of the two graphs.
- 3.2 Formulas for Calculating GNTKs: Each BLOCK first performs neighborhood aggregation, then recursively applies R ReLU transformations to produce intermediate covariance and kernel quantities.
- 3.2 Formulas for Calculating GNTKs: The final GNTK output is computed from intermediate quantities through a READOUT formula.
- 3.2 Formulas for Calculating GNTKs: The worked example translates a GNN with L = 2, R = 1, and jumping knowledge into its corresponding GNTK.
4 Theoretical Analysis of GNTK
The theoretical analysis derives generalization bounds for GNTK predictors and proves polynomial-sample learnability for a class of graph labeling functions under boundedness conditions.
- Generalization framework: The analysis studies GNTK generalization in a supervised setting with i.i.d. graph-labeled training data and an invertible kernel matrix.The kernel matrix is defined by pairwise GNTK values between training graphs.
- Generalization framework: A standard Rademacher-complexity result gives a data-dependent population-loss bound for the GNTK predictor with probability at least 1 − δ.The bound depends on the kernel matrix and training labels.
- Sample complexity: Bounding y^TΘ^-1y and tr(Θ) yields concrete sample-complexity guarantees for GNTK learning.The analysis applies this strategy to graph labeling functions represented by the paper’s target function class.
- Sample complexity: Under the stated boundedness conditions, the GNTK corresponding to the simple GNN learns the specified labeling functions with polynomially many samples.The paper presents this as a sample-complexity analysis for a class of graph functions.
- Sample complexity: If graph sizes in the training set are bounded, the trace term satisfies tr(Θ) ≤ O(nV^2), where n is the number of training samples.The theorem assumes every training graph has at most V nodes.
5 Experiments
Experiments evaluate GNTKs on graph classification and study how corresponding GNN architecture choices affect performance. GNTKs achieve strong results while reflecting several architectural trends observed in GNNs.
- 5 Experiments: The experiments use 10-fold cross-validation and report mean and standard deviation of validation accuracies across bioinformatics and social-network benchmarks.
- 5.1 Results: GNTKs achieve state-of-the-art classification accuracy on most datasets and outperform all baselines on four datasets.They achieve the best performance on all social-network datasets.
- 5.1 Results: 83.6% accuracy on COLLAB and 67.9% on PTC exceed the best baseline results of 81.0% and 64.6%, respectively.
- 5.1 Results: GNTK is more computationally efficient than its GNN counterpart with the same architecture on the reported IMDB-B experiment.GIN takes 19 minutes on a TITAN X GPU, whereas GNTK takes 2 minutes.
- 5.2 Relation between GNTK Performance and the Corresponding GNN: On NCI, GNTKs with more BLOCK operations perform better, consistent with the value of additional global structural information for bioinformatics graphs.The authors associate this pattern with the usefulness of broader molecular and biological structure.
- 5.2 Relation between GNTK Performance and the Corresponding GNN: On social networks, sum aggregation with cu = 1 outperforms average aggregation with cu = 1/|N(u)|+1, matching the corresponding GNN pattern.The paper states that sum aggregation learns more graph structure information than average aggregation in this setting.
- 5.2 Relation between GNTK Performance and the Corresponding GNN: Applying jumping knowledge improves GNTK performance on NCI and IMDB, while increasing MLP layers improves performance by approximately 0.8%.
A.1 Proof of Theorem 4.2
This proof section develops the GNTK expression for a simple GNN and uses ReLU-related quantities and polynomial-kernel feature maps in its analysis.
- A.1 Proof of Theorem 4.2: For graphs G and G′, Section 3 provides the GNTK kernel function corresponding to the simple GNN.
- A.1 Proof of Theorem 4.2: The proof uses the ReLU activation σ(z) = max{0, z} and its derivative ˙σ(z) = 1[z ≥0].
- A.1 Proof of Theorem 4.2: The analysis introduces Φ^(2l)(·) as the feature map of a degree-2l polynomial kernel.
- A.1 Proof of Theorem 4.2: Writing Θ = Θ1 + Θ2, with Θ2 positive semidefinite, supports the proof's quadratic-form argument for any y ∈ R^n.
B Experiment Setup
The experiment setup computes GNTKs using the paper's formulas, tunes architecture-related hyperparameters, and evaluates graph classification with kernel regression or C-SVM.
- B Experiment Setup: GNTKs are calculated using the formulas from Section 3.2, with post-activation expectations computed using the approach of Arora et al.
- B Experiment Setup: The number of BLOCK operations is searched over candidate values from 1 through 14.
- B Experiment Setup: The number of fully connected layers per BLOCK is searched over 1, 2, and 3, alongside the parameter cu.
- B Experiment Setup: Graph classification uses kernel regression and C-SVM, with C selected by grid search over 120 logarithmically spaced values in [10^-2, 10^4].
- B Experiment Setup: GNTK uses fewer hyperparameters than GNN because it does not tune learning rate, momentum, weight decay, batch size, or MLP width.