Kernel Readout for Graph Neural Networks

1 Introduction

In recent years, graph neural networks (GNNs) have flourished in many fields such as machine learning [Yu and Jia, 2023; Wu et al., 2023b; Zhang et al., 2024], computer vision [Han et al., 2022a; Chen et al., 2023b; Cai et al., 2024b], and bioinformatics [Gilmer et al., 2017; Gasteiger et al., 2021]. Originating from the theoretical investigation of both spatial and spectral domains, GNN has become a general paradigm powerful for modeling non-Euclidean graph data. Typically, GNNs generate representations that preserve structural information by recursive neighborhood aggregation. Numerous studies [Welling and Kipf, 2017; Hamilton et al., 2017; Zheng et al., 2023; Jin et al., 2021] have been devoted to designing various GNN architectures for node-level tasks like node classification and clustering, and the most prevalent GNNs have been validated to generate expressive node representations.

∗Corresponding authors.

Beyond that, some critical real-world scenarios, e.g. protein interface prediction [Fout et al., 2017; Liu et al., 2020] and compounds functionality analysis [Kojima et al., 2020], require predicting for entire graphs, leading to graph-level tasks, where each graph, rather than a node, is a sample. A key challenge is the graph isomorphism problem, that is, how to recognize whether two graphs are isomorphic or not. Xu et al. [2018] first explored the connection between GNNs and the graph isomorphic problem, proving that GNNs are at least as powerful as the Weisfeiler-Lehman graph isomorphism test. Motivated by these, many researchers have proposed GNNs focusing on graph-level tasks [Gilmer et al., 2017]. Generally, these methods apply readout functions, also known as graph pooling functions, on node representations, and then obtain graph-level representations, i.e., a vector representing the entire graph.

Readout is a crucial technique bridging powerful GNNs and graph-level learning. The most widely adopted graphlevel GNN backbones [Xu et al., 2018] only used some simple readout functions like Sum, Max, and Mean. These functions typically collect first-order statistics only, which does not accurately characterize the distribution of graphs. Meanwhile, they are fixed and non-data-driven operations, so they may not fit different graphs well. Some recent work has put forward learnable readout functions to solve these issues. They are mainly categorized as attention-based [Li et al., 2016; Itoh et al., 2022; Fan et al., 2020] and sequencebased [Vinyals et al., 2016; Buterez et al., 2022] approaches. These methods have made some progress, but the gains are not steady, and most critically, no unified framework to guide the design of readout. [Sun and Fan, 2024] proposed a deep MMD graph kernel (MMD-GK) that integrates graph kernel learning with graph neural networks and has promising performance in graph classification and clustering. Although Deep MMD-GK does not require any readout operation, it has quadratic time and space complexities.

In this article, we propose a novel kernel-based readout function named KerRead, which utilizes the kernel technology in readout. The detailed process of KerRead is depicted in Figure 1. To be specific, we treat each dimension as a point with coordinates determined by the feature values of each

node in this dimension, then we define several adaptive centers and calculate kernels between these centers and dimension points, leading to several advantages over other methods: 1) KerRead implicitly augments the number of nodes in a graph to be large or even infinite through the Gaussian kernel; 2) Graph representations are generated by computing the kernel matrices so they are naturally in the same scale; 3) The centers are trainable and the number of centers can be flexibly defined. 4) KerRead adheres to the crucial property of permutation invariance, a vital characteristic for robust graph learning. In general, KerRead employs kernel technology to better reflect the latent distributions of diverse node representations, which is a more powerful and scalable method.

Our contributions are as follows.

• We propose a new paradigm to analyze and design readout functions from the perspective of kernel functions.
• We put forward KerRead which solves several drawbacks of existing methods.
• We conduct comprehensive experiments to evaluate the proposed KerRead, including graph classification and clustering with various GNN backbones.

2 RelatedWork

2.1 Graph Neural Networks

Graph neural networks (GNNs) aim to learn low-dimensional and dense node representations from high-dimensional sparse graphs while preserving the original structural information. Initially, GNNs were operated in the spectral domain [Bruna et al., 2014; Defferrard et al., 2016; Welling and Kipf, 2017; Wang and Zhang, 2022; Wu et al., 2024] and transformed the feature signals of the graph into the frequency domain through Fourier transformation and subsequently designing various filters in the frequency domain, where the goal was to make the signal of target node more similar with its neighbors. Meanwhile, other graph neural network methods operated in the spatial domain [Hamilton et al., 2017;

Veliˇckovi´c et al., 2018; Xu et al., 2018; Cai et al., 2024a; Yu and Jia, 2024], which can acquire the node representation by message passing mechanism. For instance, GraphSAGE [Hamilton et al., 2017] aggregated the sampled neighborhoods of the target nodes to generate the target node representations. Graph Attention Network (GAT) [Veliˇckovi´c et al., 2018] introduced the multi-attentional module mechanism into GraphSAGE. Meanwhile, many widely developed domains, such as machine learning [Chen et al., 2023a; Fan et al., 2022; Chen et al., 2022; Fan, 2022] and data mining [Cai et al., 2022; Zheng et al., 2022; Fan, 2021; Cai and Fan, 2022], have extensively employed numerous GNN variants [Chen et al., 2023c; Chen et al., 2023d; Wu et al., 2023a; Zheng et al., 2021].

2.2 Graph-level Prediction

Graph-level tasks treat each entire graph as a sample, posing many new challenges for GNNs. For this purpose, many GNN methods were focusing on the graph-level tasks [Xu et al., 2018; Gilmer et al., 2017; Gao and Ji, 2019; Buterez et al., 2022; Ju et al., 2022; Sun and Fan, 2024; Han et al., 2022b]. For instance, [Gilmer et al., 2017] proposed a general framework for quantum chemistry representation called Message Passing Neural Networks. GUNet [Gao and Ji, 2019] added the graph pooling and unpooling layers during graph convolution. [Luo et al., 2023] explored the semi-supervised classification problem from a subgraph perspective. [Cai et al., 2024a] advanced graph clustering with a novel pseudolabel mechanism. [Liu et al., 2022] considered the long-tailed problem in the graph classification. [Sun et al., 2023] utilized the Lov´asz principle to build an unsupervised graph representation learning framework.

2.3 Readout Function

The readout function, a vital component in GNNs, pools the learned graph embedding matrix into a graph embedding vector for the subsequent graph classification task. The aforementioned graph-level methods use simple and naive readout functions, such as sum, mean, or max, to generate the final graph representation vector. [Li et al., 2016; Itoh et al., 2022; Fan et al., 2020] proposed attention-based readout functions to capture the importance of each node. [Lee et al., 2021] proposed learnable modules for readout functions to further improve their expressive power. To further enhance the expressiveness of readout functions, certain methods incorporate techniques designed for sequence data types into the readout process [Navarin et al., 2019; Xu et al., 2018; Buterez et al., 2022]. There are also some hierarchical readout methods, that is, reading out smaller and smaller graphs step by step as the graph is embedded until it is aggregated into a node [Zhang et al., 2018; Ying et al., 2018; Du et al., 2021], however, these hierarchical methods also need a readout function to gradually reduce the size of the graph, so in this article, we only focus the readout function which map the node embedding into graph vector directly.

3 ProposedMethod

3.1 Notation and ProblemFormulation

A graph is presented as G = (A, X), where A ∈ Rn×n is the adjacency matrix, where Aij = 1 denotes that vi is connected to vj and Aij = 0 otherwise. Let G = {G1, ..., GN} be the graph set. We aim to learn the graph representation vector zGi for each graph Gi to facilitate downstream tasks.

Encoding Graph GNNs are usually used to encode each graph and consist of some graph convolutional layers. Each layer is based on the message-passing mechanism, i.e., aggregating the neighborhood information for each target node, which is formed by two steps

where we have H(0) Gi = XGi . For convenience in expression, we will use HGi ∈ Rni×dL to denote the node embedding matrix of graph Gi at final layers in the following part.

Readout Function After convolution by GNNs and getting node emebdding matrix HGi , a readout function is utilized to map HGi ∈ Rni×dL into a graph vector zGi ∈ RdL for the following graph classification task:

where HGi [j, :] refers to the j-th row of Gi node embedding matrix, i.e., node vj in graph Gi.

3.2 DetailedMethod

In this part, we will introduce our proposed KerRead, which incorporates the principles of kernels into the realm of readout mechanisms. KerRead is a learnable readout function that matches GNNs' expressiveness well by solving the following challenges of existing methods:

• Challenge 1: How to 'readout' distinguishable graph representations in a data-driven manner?
• Challenge 2: How to fairly learn graph representations for all graphs in a unified size sample space?
• Challenge 3: How to capture latent distributions from graphs with scarce and varying numbers of nodes?

The following involves devising an adaptive center, and we leverage expressive kernel functions to calculate graph representations, e.g., Gaussian kernel. Moreover, we put forth the concept of multi-head centers to further enhance the model's capabilities. The overall framework is shown in Figure 1.

Adaptive Center In graph-level tasks, although graphs vary, nodes of different graphs are described by features with the same meaning and amount. Therefore, it makes more sense to think of readout in a dimension-wise perspective than in the sample-wise one, i.e., we describe graphs by considering how to express distributions in each dimension instead of directly considering fusing representations of all nodes.

More specifically, HGi [:, j] in the node representations can be viewed as a description or distribution for the distribution of graph Gi in the corresponding dimension, and these distributions of different graphs are naturally comparable. From this view, mapping the node representation matrix into a graph representation vector can be regarded as representing the distributions of each dimension by one value. Assume there is a central dimension called center in this paper, we can represent the distribution by calculating certain distances between each dimension and the center, and our key idea is employing the effective kernel method to achieve this. It is obvious that the design of the center is critical in our method. On the one hand, one single and fixed center can not adapt to different graphs; On the other hand, centers of different graphs can not be totally individual otherwise the graph representations may not be compared. For this purpose, we design an adaptive center which can adapt different graphs and share information among graphs. First we leverage a shared model fΘ : RdL → R to generate the adaptive center as

where s ∈ R is a learnable scale factor, and fΘ can be a certain neural network like Multi-Layer Perceptrons (MLP) performed on each row of HGi and Θ is the set of learnable parameters. Note that fΘ can also be a simple pooling function, in which Θ is the empty set. For example, the adaptive center cGi ∈ Rni can be

In this way, the learning of centers is data-driven and optimized in conjunction with GNNs. More importantly, the designed model can address graphs with different sizes as it adaptively generates a center for each graph. Meanwhile, all graphs share fΘ and the scale factor s, so that the training process incorporates information from various graphs. These provide bases to solve Challenge 1 and 2. After obtaining cGi with flexible size ni, i.e., the number of nodes in Gi, we can calculate the kernel value between each dimension of node embedding and the adaptive center to obtain the final graph vector in the following section.

Kernel Calculation Compared to node-level tasks, graphs commonly contain very scarce nodes in graph-level tasks (e.g., molecules), and the number of nodes in different graphs may be imbalanced. Thus, it is difficult to capture latent distributions and derive discriminative graph representations (Challenges 3). To tackle this, we further employ a transformation function ϕ to map each column of node embedding and the adaptive center into augmented sample space with more virtual nodes:

where HGi [:, j] denotes the j-th column of HGi , and ϕ : Rni → Rn′i maps a vector refers to ni nodes to a vector refers to n′i nodes, where n′i ≫ ni. In general, explicitly designing a ϕ is cost-prohibitive, particularly when n′i is significantly large or even infinite. Instead, we leverage kernel functions k : X × X → R capable of implicitly inducing ϕ:

Thus the final graph representation vector is:

There are numerous available kernel functions, each inducing a unique feature map ϕ. Perhaps the best-known kernel function is the Gaussian kernel, i.e.,

where cGi is the proposed learnable adaptive center corresponding to Gi, and the parameter σ ∈ R is also learnable. The mapping ϕ induced by a Gaussian kernel is infinitedimension. In our experiment, we evaluated various kernel functions, including the Laplace kernel, Sigmoid kernel, Polynomial kernel, and others. The performance results for these kernels can be found in Section 4.3.

Multi-head Centers In order to enhance the richness of information captured from node representation HGi , relying solely on a single center may be insufficient. To address this limitation, we further introduce the concept of multi-head center, which initializes m learnable scale factors denoted as s = [s1, ..., sm]. The adaptive center matrix CGi ∈ Rni×m is then defined as follows:

where Ψ = {s,Θ} is the parameter set containing scale factors and parameters of fΘ. Subsequently, the multi-head kernel matrix ZGi ∈ Rm×dL is expressed as:

where ZGi is only related to the number of heads m and hidden dimension dL. Finally, the fusion is performed on all graph representation vectors:

where hΥ : Rm → R is a fusion model and Υ is the parameter set. As a simple example, hΥ can be a weighted summation of {ZGi [j, :]}mj=1 , allowing for adaptive learning of contributions from each center to the final graph vector.

Overall Framework The entire pipeline of graph representation learning on graph Gi with KerRead is formulated as

where f GNN W is a certain GNN with parameter set W and we denote MΥ,Ψ,W the over-all graph representation learning model with KerRead for convenience. Recall that ϕ is the implicit mapping induced by the kernel function k. When k is a Gaussian kernel, then n′i = ∞.

Theorem 1. Suppose the kernel k is feature permutation invariant, then MΥ,Ψ,W is node permutation invariant-that is, for any graph G and its any permuted counterpart G ′ , MΥ,Ψ,W(G ′ ) = MΥ,Ψ,W(G) always holds.

The proof is deferred to Appendix A 1 . Note that the assumption in the theorem is very general. Many popular kernels such as Gaussian kernel, Laplacian kernel, and polynomial kernels are feature permutation invariant. Theorem 1 indicates the permutation invariance of our proposed framework, which has been proven to be very critical in graph-level learning [Buterez et al., 2022].

The training strategy of KerRead is based on the selected GNN backbone. For the graph classification task, we employ a supervised GNN algorithm, using cross-entropy as the loss function for supervised learning. Conversely, in handling graph clustering tasks, we opt for an unsupervised GNN backbone and utilize such as InfoNCE loss [Oord et al., 2018] for representation learning, and then input learned graph embeddings into clustering algorithms, such as K-means [Lloyd, 1982] or spectral clustering [Ng et al., 2001] to obtain the clustering results. This ensures that our model can be flexibly and effectively applied to different tasks.

For the complexity analysis of KerRead and MΥ,Ψ,W, we denote the number of graphs as N and the initial node feature dimension as d0, n = max({ni}Ni=1 ) is the max number of nodes and |E| = max({|Ei|}Ni=1 ) is the max number of edges among graphs, where Ei is the edge set of Gi. Considering an L-layer GNN, we define d = max({di}Li=1 ) is the max hidden dimension. For the forward pass of MΥ,Ψ,W in per iteration, the time complexity of KerRead is O(Nnmd), including adaptive centers generation and kernel calculation. The overall time complexity is O(N(nd+|E|)(Ld+d0)) with m ≪ d0. The space complexity of KerRead is O(m(n + d)), and that of the overall framework is O(N((n+d)(Ld+d0)+ L|E|)). Generally, the additional costs our approach imposes on the framework are negligible.

3.3 Connection to ExistingWork

For brief formulation, we adopt one center vector cGi for graph Gi, meaning a feature dimension in graph Gi, and compute the j-th element of graph representation zGi as

1 The appendix can be found at https://github.com/jiajunCAU/ KerRead/blob/master/Appendix.pdf

Based on this formula, the process of readout can be described as performing a kernel on the dimensions of node representations and a center vector:

For example, the widely adopted Sum can be formulated as

where 1Gi is an all-one vector and a linear kernel is employed, thus ϕ is an identity mapping. Similarly, we have

model to generate the attention score as the node weight. In this way, we innovatively explain several common readout methods from the perspective of the kernel, and find these traditional readout functions both are simplified versions and special instances of KerRead. Furthermore, this perspective also explains the causes of the above challenges: Challenge 1-fixed center; Challenge 2-non-adaptive center for all graphs; Challenge 3-linear kernel.

4 Experiments

In this section, we first introduce the benchmark datasets, the backbone, the baseline readout function, and parameter settings, and then we present the comparison results of our methods with other readout functions in different GNN backbones on graph classification tasks and graph clustering tasks.

4.1 Experimental Setting

Datasets This study employs a total of 8 graph-level datasets, comprising 6 chemical molecule datasets (MUTAG, DD, PROTEINS, NCI1, Mutagenicity and OGBG-Molhiv) and 2 social network graph datasets (IMDB-B, IMDB-M), these datasets are collected in TU datasets 2 and open graph benchmark 3 . The description and statistics of the dataset are in Appendix C.

Backbone and Baseline We use 6 widely adopted supervised GNN methods, including three baseline GNNs (GCN [Welling and Kipf, 2017], GraphSAGE [Hamilton et al., 2017], SGC [Wu et al., 2019]), three more expressive GNNs (GAT [Veliˇckovi´c et al., 2018], GIN [Xu et al., 2018] and GUNet [Gao and Ji, 2019]), and two unsupervised GNNs (InfoGraph [Sun et al., 2019] and GraphCL [You et al., 2020]) are also used in this article. To ensure a fair comparison, we select 10 readout functions as our baseline for comparison. These encompass simple readout functions (Sum, Mean, Max), along with Attention [Li et al., 2016], Set2set [Vinyals et al., 2016], Deep Sets [Zaheer et al., 2017], SRead [Lee et al., 2021], Set Transformer, Janossy MLP, and Janossy GRU [Buterez et al., 2022].

2 https://chrsmrrs.github.io/datasets/

3 https://ogb.stanford.edu/

Parameter Settings We employ the officially released source code and adopt the parameters suggested in the respective papers for all baseline readout functions and GNN backbones. Additional information on experimental settings can be found in Appendix C.

4.2 Performance Analysis

In this section, we analyze the performance of graph classification and graph clustering.

Graph Classifications We conducted a comprehensive analysis comparing the graph classification performance of KerRead with ten other readout functions across six GNN backbone models on eight datasets, as detailed in Table 1, Table 2, and Table 3. We obtain the following observations from this comparison:

• (i) KerRead consistently outperforms other readout methods across various GNN backbones. KerRead surpasses the

most powerful baseline Set2set readout by 2.3% and 2.4% on GCN and GAT backbones on average.

• (ii) KerRead has achieved the best performance on datasets with different sizes. Compared to baseline readout functions, we observe an AUC increase ranging from 0.8% to 19.2% on large-scale dataset OGBG-Molhiv on average.
• (iii) The average performance of different readouts in SGC was significantly lower than GUNet and GraphSAGE, which indicates that SGC does not possess a substantial advantage in graph-level tasks similar to its superiority in node-level tasks.
• (iv) Sequence-based readout methods, such as Janossy GRU, Set2Set, and Set Transformer, demonstrated better average performance on SGC. We speculate that incorporating complex sequential information can enhance the expressive power of final graph embedding when the backbone cannot capture enough information.

In particular, KerRead outperforms Sum, Mean, Max, and Attention in all comparisons, which aligns with earlier analysis since these readout functions are special cases of KerRead, and their expressivity is weaker than KerRead. Furthermore, KerRead's performance is more stable than some sequencebased methods, we consider KerRead satisfies permutation invariance, so it is not influenced by the order of input nodes, thus avoiding any impact on performance.

Graph Representation Learning & Clustering To provide a more comprehensive evaluation of KerRead, we conducted graph clustering experiments on seven benchmark datasets using two classic unsupervised graph backbone networks, as shown in Table 4. In general, KerRead continues to demonstrate highly competitive performance in graph clustering tasks. Compared to other readout functions, we observed an increase in ACC from 2.6% to 20.7% and an improvement in NMI from 2.9% to 24.9% on MUTAG when using InfoGraph. On GraphCL, ACC increased from 2.1% to 12.8%, and NMI increased from 0.3% to 24.3%. However, KerRead did not achieve optimal performance on the IMDBB and IMDB-M datasets. We speculate that this may be attributed to insufficient initial features in these two datasets, leading to the inability to capture enough information when computing the kernel values between each dimension of node embedding and the adaptive center.

4.3 Sensitivity Analysis

In this section, we will focus on conducting a sensitivity analysis of the parameters of our proposed KerRead. This includes examining the impact of varying the number of heads, different kernel functions, distinct update strategies on the final results, and the efficiency analysis.

Kernel Function Selection As depicted in Figure 2, we analyze two common kernel functions, namely the Laplace kernel function and the Gaussian kernel, on the DD and PROTEINS datasets. Overall, the performance of the Gaussian kernel surpasses that of the Laplace kernel. Besides, we have further compared the graph classification performance of five commonly used kernel functions on seven datasets. The experimental results are shown in Table 5, revealing the following observations: (i) Overall, KerRead outperforms most baselines regardless of the kernel function used. (ii) The use of non-linear kernel functions yields better performance compared to linear kernel functions. Specifically, the Gaussian kernel outperforms the Linear kernel by an average accuracy increase of 2.9%.

The Number of Heads To explore the impact of the number of heads on the results, we present the graph classification results on the DD and PROTEINS datasets under different numbers of heads as shown in Figure 2. Overall, optimal performance is observed when the number of heads ranges from 2 to 6. However, as the number of heads increases, there may be a negative impact on performance. Specifically, the Laplace kernel achieves optimal performance on the DD dataset when the number of heads is 5, surpassing the performance at 1 and 10 heads by 2.04% and 2.51%, respectively. Similarly, the Gaussian kernel attains its best performance on the DD dataset with 5 heads, exhibiting improvements of 3.46% and 3.88% compared to 1 and 10 heads, respectively. These findings suggest that an appropriate increase in the number of heads can enhance performance, but an excessive number may lead to detrimental effects on the model.

Training Time We conducted an efficiency analysis comparing our model with a baseline readout. As illustrated in Figure 3, we depict the training time (measured in milliseconds per epoch) variation of KerRead under 1, 2, and 5 heads, as well as other methods. Generally, KerRead requires more training time as the number of heads increases. However, its training time remains closely aligned with the linear readout function and consistently stays below the complex readout functions (Set Transformer, Janossy GRU), however, KerRead achieves significant performance improvements compared to these linear readout functions. Furthermore, KerRead demonstrates optimal performance on the DD dataset, where each epoch takes approximately 4.1% longer than the linear readout function, which indicates that the gains in accuracy achieved by KerRead outweigh the marginal pressure introduced by the training time.

Learning Strategies We mentioned that the parameters c and σ in the model are learnable as shown in Eq (9) when employing Gaussian kernel. However, what impact would ensue if these parameters were fixed? As illustrated in Figure 4, we present the performance comparison of these four update strategies across seven datasets. Our observations indicate that, on average, the optimal performance is attained when both parameters are learnable. Across the seven datasets, U1 outperforms U4 by 1.67%, 3.48%, 2.33%, 0.24%, 0.07%, 1.6%, and 1.47%, respectively. This suggests that concurrently updating both parameters indeed facilitates the learning of a superior model.

4.4 Supplementary Experiment

Due to space constraints, we have included the other curial experiments in Appendix C, which include the t-SNE visualizations.

5 Conclusions

This work proposed a novel graph pooling strategy called kernel readout. The main idea is to transfer the node embedding from the sample space with limited nodes to the augmented sample space with infinite virtual nodes, and then generate the graph vector by kernel methods. The strategy has been applied to six supervised GNN backbones and two unsupervised GNN backbones and outperformed previous methods of graph pooling, on eight graph datasets.

Acknowledgements

This work was partially supported by the National Natural Science Foundation of China under Grant No.62376236, the General Program JCYJ20210324130208022 of Shenzhen Fundamental Research, the Guangdong Key Lab of Mathematical Foundations for Artificial Intelligence (PPA00003), the funding UDF01001770 of The Chinese University of Hong Kong, Shenzhen, and the National Natural Science Foundation for Young Scholars of China under Grants No.61502500 and No.61602500.

Contribution Statement

Jiajun Yu and Zhihao Wu contributed equally as co-first authors.