Generalization in Graph Neural Networks: Improved PAC-Bayesian Bounds on
  Graph Diffusion
Graph neural networks are widely used tools for graph prediction tasks.
Motivated by their empirical performance, prior works have developed
generalization bounds for graph neural networks, which scale with graph
structures in terms of the maximum degree. In this paper, we present
generalization bounds that instead scale with the largest singular value of the
graph neural network's feature diffusion matrix. These bounds are numerically
much smaller than prior bounds for real-world graphs. We also construct a lower
bound of the generalization gap that matches our upper bound asymptotically. To
achieve these results, we analyze a unified model that includes prior works'
settings (i.e., convolutional and message-passing networks) and new settings
(i.e., graph isomorphism networks). Our key idea is to measure the stability of
graph neural networks against noise perturbations using Hessians. Empirically,
we find that Hessian-based measurements correlate with the observed
generalization gaps of graph neural networks accurately. Optimizing noise
stability properties for fine-tuning pretrained graph neural networks also
improves test performance on several graph-level classification tasks.
Haotian Ju, Dongyue Li, Aneesh Sharma, Hongyang R. Zhang
# Generalization in Graph Neural Networks: Improved

###### Abstract

Graph neural networks are widely used tools for graph prediction tasks. Motivated by their empirical performance, prior works have developed generalization bounds for graph neural networks, which scale with graph structures in terms of the maximum degree. In this paper, we present generalization bounds that instead scale with the largest singular value of the graph neural network's feature diffusion matrix. These bounds are numerically much smaller than prior bounds for real-world graphs. We also construct a lower bound of the generalization gap that matches our upper bound asymptotically. To achieve these results, we analyze a unified model that includes prior works' settings (i.e., convolutional and message-passing networks) and new settings (i.e., graph isomorphism networks). Our key idea is to measure the stability of graph neural networks against noise perturbations using Hessians. Empirically, we find that Hessian-based measurements correlate with the observed generalization gaps of graph neural networks accurately; Optimizing noise stability properties for fine-tuning pretrained graph neural networks also improves test performance on several graph-level classification tasks.

## 1 Introduction

A central measure of success for a machine learning model is the ability to generalize well from training data to test data. For linear and shallow models, the generalization gap between their training performance and test performance can be quantified via complexity notions such as the Vapnik-Chervonenkis dimension and Rademacher complexity. However, formally explaining the empirical generalization performance of deep models remains a challenging problem and an active research area (see, e.g., Hardt and Recht [25]). There are by now many studies for fully-connected and convolutional neural networks that provide an explanation for their superior empirical performance [4, 43]. Our work seeks to formally understand generalization in graph neural networks (GNN) [45], which are commonly used for learning on graphs [23].

As a concrete example for motivating the study of generalization performance, we consider the fine-tuning of pretrained graph neural networks [27]. Given a pretrained GNN learned on a diverse range of graphs, fine-tuning the pretrained GNN on a specific prediction task is a common approach for transfer learning on graphs. An empirical problem with fine-tuning is that, on one hand, pretrained GNNs use lots of parameters to ensure representational power. On the other hand, fine-tuning a large GNN would overfit the training data and suffer poor test performance without proper algorithmic intervention. Thus, a better understanding of generalization in graph neural networks can help us identify the cause of overfitting and, consequently, inspire designing robust fine-tuning methods for graph neural networks.
A naive application of the generalization bounds from fully-connected feedforward networks [4; 42] to GNNs would imply an extra term in the generalization bound that scales with \(n^{l-1}\), where \(n\) is the number of nodes in the graph, hence rendering the error bounds vacuous. Besides, Scarselli et al. [46] shows that the VC dimension of GNN scales with \(n\). Thus, although the VC dimension is a classical notion for deriving learning bounds, it is oblivious to the graph structure. Recent works have taken a step towards addressing this issue with better error analysis. Verma and Zhang (2019) find that one-layer graph neural networks satisfy uniform stability properties [49], following the work of Hardt et al. [26]. The generalization bound of Verma and Zhang [49] scales with the largest singular value of the graph diffusion matrix of the model. However, their analysis only applies to a single layer and node prediction. Garg et al. (2020) analyze an \(l\) layer message-passing neural network \(-\) with \(l-1\) graph diffusion layers and \(1\) pooling layer \(-\) for graph prediction tasks [16]. Their result scales with \(d^{l-1}\), where \(d\) is the maximum degree of the graph. Subsequently, Liao et al. (2021) develop a tighter bound but still scales with \(d^{l-1}\)[37]. For both results, the graph's maximum degree is used to quantify the complexity of node aggregation in each diffusion step.

In this paper, our major contribution is to show generalization bounds for graph neural networks by reducing the max degree to the spectral norm of the graph diffusion matrix. We analyze the stability of a graph neural network against noise injections. Denote an \(l\)-layer GNN by \(f\). Based on PAC-Bayesian analysis [39], the generalization gap of \(f\) will be small if \(f\) remains stable against noise injections; otherwise, the generalization gap of \(f\) will be large. By quantifying the noise stability of \(f\) via Lipschitz-continuity properties of its activation functions, one can get PAC-Bayes bounds for feedforward networks that correlate with their observed generalization gaps [12; 2; 31; 34]. We use a refined stability analysis of graph neural networks through Hessians and show tight generalization bounds on the graph diffusion matrix.

**Our Contributions.** The goal of this work is to improve the theoretical understanding of generalization in graph neural networks, and in that vein, we highlight two results below:

* First, we prove sharp generalization bounds for message-passing neural networks [9; 17; 33], graph convolutional networks [35], and graph isomorphism networks [54]. Our bounds scale with the spectral norm of \(P_{{}_{G}}^{l-1}\) for an \(l\)-layer network, where \(P_{{}_{G}}\) denotes a diffusion matrix on a graph \(G\) and varies between different models (see Theorem 3.1 for the full statement). We then show a matching lower bound instance where the generalization gap scales with the spectral norm of \(P_{{}_{G}}^{l-1}\) (see Theorem 3.2).

Figure 1: **Left: Our spectral norm bounds for graph diffusion matrices are orders of magnitude smaller than maximum degree bounds on real-world graphs; In Figure 0(a), we measure the spectral norm and the max degree for five graph datasets. Right: In Figure 0(b), our Hessian-based generalization measure (plotted in green, scaled according to the _left axis_) matches the empirically observed generalization gaps of graph neural networks (plotted in yellow, scaled according to the _left axis_). The blue line shows a uniform-convergence bound (scaled according to the _right axis_) that is orders of magnitude larger than the observed gaps.**
* Second, our stability analysis of graph neural networks provides a practical tool for measuring generalization. Namely, we show that the noise stability of GNN can be measured by the trace of the loss Hessian matrix. The formal statement is given in Lemma 4.3, and our techniques, which include a uniform convergence of the Hessian matrix, may be of independent interest. We note that the proof applies to twice-differentiable and Lipschitz-continuous activation functions (e.g., tanh and sigmoid).

Taken together, these two results provide a sharp understanding of generalization in terms of the graph diffusion matrix for graph neural networks. We note that the numerical value of our bounds in their dependence on the graph is much smaller than prior results [16, 37], as is clear from Figure 0(a). Moreover, the same trend holds even after taking weight norms into account (see Figure 2, Section 3.3). Further, the Hessian-based bounds (see Lemma 4.3, Section 4) are non-vacuous, matching the scale of empirically observed generalization gaps in Figure 0(b).

Finally, motivated by the above analysis, we also present an algorithm that performs gradient updates on perturbed weight matrices of a graph neural network. The key insight is that minimizing the average loss of multiple perturbed models with independent noise injections is equivalent to regularizing \(f\)'s Hessian in expectation. We conduct experiments on several graph classification tasks with Molecular graphs that show the benefit of this algorithm in the fine-tuning setting.

## 2 Related Work

**Generalization Bounds:** An article by Zhang et al. (2017) finds that deep nets have enough parameters to memorize real images with random labels, yet they still generalize well if trained with true labels. This article highlights the overparametrized nature of modern deep nets (see also a recent article by Arora [1]), motivating the need for complexity measures beyond classical notions. In the case of two-layer ReLU networks, Neyshabur et al. (2019) show that (path) norm bounds better capture the "effective number of parameters" than VC dimension\(-\)which is the number of parameters for piecewise linear activations [5].

For multilayer networks, subsequent works have developed norm, and margin bounds, either via Rademacher complexities [4, 20, 38], or PAC-Bayesian bounds [2, 42, 36, 34]. All of these bounds apply to the fine-tuning setting following the distance from the initialization perspective. Our analysis approach builds on the work of Arora et al. (2018) and Ju et al. (2022). The latter work connects perturbed losses and Hessians for feedforward neural networks, with one limitation Hessians do not show any explicit dependence on the data. This is a critical issue for GNN as we need to incorporate the graph structure in the generalization bound. Our result instead shows an explicit dependence on the graph and applies to message-passing layers that involve additional nonlinear mappings. We will compare our analysis approach and prior analysis in more detail when we present the proofs in Section 4 (see Remark 4.4).

**Graph Representation Learning:** Most contemporary studies of learning on graphs consider either node-level or graph-level prediction tasks. Our result applies to graph prediction while permitting an extension to node prediction: see Remark 4.2 in Section 4. Most graph neural networks follow an information diffusion mechanism on graphs [45]. Early work takes inspiration from ConvNets and designs local convolution on graphs, e.g., spectral networks [6], GCN [35], and GraphSAGE [24] (among others). Subsequent works have designed new architectures with graphs attention [48] and isomorphism testing [54]. Gilmer et al. (2017) synthesize several models into a framework called message-passing neural networks. Besides, one could also leverage graph structure in the pooling layer (e.g., differentiable pooling and hierarchical pooling [59, 62]). It is conceivable that one can incorporate the model complexity of these approaches into our analysis. Recent work applies pretraining to large-scale graph datasets for learning graph representations [27]. Despite being an effective transfer learning approach, few works have examined the generalization of graph neural nets in the fine-tuning step.
Besides learning on graphs, GNNs are also used for combinatorial optimization [47] and causal reasoning [55]. There is another approach for graph prediction using kernels [50, 11]. There are also alternative graph diffusion processes besides GNN [18, 19, 61]. For references, see review articles [23, 7, 13, 53].

**Generalization in GNN:** Recent work explores generalization by formalizing the role of the algorithm, and the alignment between networks and tasks [56]. Esser et al. [14] finds that transductive Rademacher complexity-based bound provides insights into the behavior of GNNs (e.g., under stochastic block models). Besides, there are works about size generalization, which refer to performance degradation when models extrapolate to graphs of different sizes from the input [47, 58]. It is conceivable that the new tools we have developed may be useful for studying extrapolation.

**Expressivity of GNN:** The expressivity of GNN for graph classification can be related to graph isomorphism tests and has connections to one-dimensional Weisfeiler-Lehman testing of graph isomorphism [41, 54]. This implies limitations of GNN for expressing tasks such as counting cycles [44, 8, 3]. The expressiveness view seems orthogonal to generalization, which instead concerns the sample efficiency of learning. For further discussions and references, see a recent survey by Jegelka [30].

## 3 Sharp Generalization Bounds for Graph Neural Networks

We first introduce the problem setup for analyzing graph neural networks. Then, we state our generalization bounds for graph neural networks and compare them with the prior art. Lastly, we construct an example to argue that our bounds are tight.

### Problem setup

Consider a graph-level prediction task. Suppose we have \(N\) examples in the training set; each example is an independent sample from a distribution denoted as \(\mathcal{D}\), which is jointly supported on the feature space \(\mathcal{X}\) times the label space \(\mathcal{Y}\). In each example, we have an undirected graph denoted as \(G=(V,E)\), which describes the connection between \(n\) entities, represented by nodes in \(V\). For example, a node could represent a molecule, and an edge between two nodes is a bond between two molecules. Each node also has a list of \(d\) features. Denote all node features as an \(n\) by \(d\) matrix \(X\). For graph-level prediction tasks, the goal is to predict a graph label \(y\) for every example. We will describe a few examples of such tasks later in Section 5.2.

**Message-passing neural networks (MPNN).** We study a model based on several prior works for graph-level prediction tasks [9, 17, 16, 37]. Let \(l\) be the number of layers: the first \(l-1\) layers are diffusion layers, and the last layer is a pooling layer. Let \(d_{t}\) denote the width of each layer for \(t\) from \(1\) up to \(l\). There are several nonlinear mappings in layer \(t\), denoted as \(\phi_{t},\rho_{t}\), and \(\psi_{t}\); further, they are all centered at zero. There is a weight matrix \(W^{(t)}\) of dimension \(d_{t-1}\) by \(d_{t}\) for transforming neighboring features, and another weight matrix \(U^{(t)}\) of dimension \(d\) by \(d_{t}\) for transforming the anchor node feature.

For the first \(l-1\) layers, we recursively compute the node embedding from the input features \(H^{(0)}=X\):

\[H^{(t)}=\phi_{t}\Big{(}XU^{(t)}+\rho_{t}\big{(}P_{{}_{G}}\psi_{t}(H^{(t-1)}) \big{)}W^{(t)}\Big{)}. \tag{1}\]

For the last layer \(l\), we aggregate the embedding of all nodes: let \(\mathbf{1}_{n}\) be a vector with \(n\) values of one:

\[H^{(l)}=\frac{1}{n}\mathbf{1}_{n}^{\top}H^{(l-1)}W^{(l)}. \tag{2}\]

Note that this setting subsumes many existing GNNs. Several common designs for the graph diffusion matrix \(P_{{}_{G}}\) would be the adjacency matrix of the graph (denoted as \(A\)). \(P_{{}_{G}}\) can also be the normalized adjacency matrix, \(D^{-1}A\), with \(D\) being the degree-diagonal matrix. Adding an identity matrix in \(A\) is equivalent to adding self-loops in \(G\). In GCN, we can set \(U^{(t)}\) as zero, \(\rho_{t}\) and \(\psi_{t}\) as identity mappings.
**Notations.** For any matrix \(X\), let \(\left\|X\right\|\) denote the largest singular value (or spectral norm) of \(X\). Let \(\left\|X\right\|_{F}\) denote the Frobenius norm of \(X\). We use the notation \(f(N)\lesssim g(N)\) to indicate that there exists a fixed constant \(c\) that does not grow with \(N\) such that \(f(N)\leq c\cdot g(N)\) for large enough values of \(N\). Let \(W\) and \(U\) denote the union of the \(W\) and \(U\) matrices in a model \(f\), respectively.

### Main results

Given a message-passing neural network denoted as \(f\), what can we say about its generalization performance? Let \(f(X,G)\) denote the output of \(f\), given input with graph \(G\), node feature matrix \(X\), and label \(y\). The loss of \(f\) for this input example is denoted as \(\ell(f(X,G),y)\). Let \(\hat{L}(f)\) denote the empirical loss of \(f\) over the training set. Let \(L(f)\) denote the expected loss of \(f\) over a random example of distribution \(\mathcal{D}\). We are interested in the generalization gap of \(f\), i.e., \(L(f)-\hat{L}(f)\). How would the graph diffusion matrix \(P_{{}_{G}}\) affect the generalization gap of graph neural networks?

To motivate our result, we examine the effect of incorporating graph diffusion in a one-layer linear neural network. That is, we consider \(f(X,G)\) to be \(\frac{1}{n}\mathbf{1}_{n}^{\top}P_{{}_{G}}XW^{(1)}\), which does not involve any nonlinear mapping for simplicity of our discussion. In this case, by standard spectral norm inequalities for matrices, the Euclidean norm of \(f\) (which is a vector) satisfies:

\[\left\|f(X,G)\right\| =\left\|\frac{1}{n}\mathbf{1}_{n}^{\top}P_{{}_{G}}XW^{(1)}\right\|\] \[\leq\left\|\frac{1}{n}\mathbf{1}_{n}^{\top}\right\|\cdot\left\|P_ {{}_{G}}\right\|\cdot\left\|X\right\|\cdot\left\|W^{(1)}\right\| \tag{3}\]

Thus, provided that the loss function \(\ell(\cdot,y)\) is Lipschitz-continuous, standard arguments imply that the generalization gap of \(f\) scales with the spectral norm of \(P_{{}_{G}}\) (divided by \(\sqrt{N}\)) [40]. Let us compare this statement with a fully-connected neural net that averages the node features, i.e., the graph diffusion matrix \(P_{{}_{G}}\) is the identity matrix. The spectral norm of \(P_{{}_{G}}\) becomes one. Together, we conclude that the graph structure affects the generalization bound of a single layer GNN by adding the spectral norm of \(P_{{}_{G}}\).

Our main result is that incorporating the spectral norm of the _graph diffusion matrix_\(P_{{}_{G}}^{l-1}\) is sufficient for any \(l\) layer MPNN. We note that the dependence is a power of \(l-1\) because there are \(l-1\) graph diffusion layers: see equation (1). Let \(f\) be an \(l\)-layer network whose weights \(\mathbf{W},\mathbf{U}\) are defined within a hypothesis set \(\mathcal{H}\): For every layer \(i\) from \(1\) up to \(l\), we have that

\[\left\|W^{(i)}\right\|\leq s_{i}, \left\|W^{(i)}\right\|_{F}\leq s_{i}r_{i},\] \[\left\|U^{(i)}\right\|\leq s_{i}, \left\|U^{(i)}\right\|_{F}\leq s_{i}r_{i}, \tag{4}\]

where \(s_{1},s_{2},\ldots,s_{l}\) and \(r_{1},r_{2},\ldots,r_{l}\) are bounds on the spectral norm and stable rank and are all greater than or equal to one, without loss of generality. We now present the full statement.

**Theorem 3.1**.: _Suppose all of the nonlinear activations in \(\left\{\phi_{t},\rho_{t},\psi_{t}:\forall\,t\right\}\) and the loss function \(\ell(\cdot,y)\) (for any fixed label \(y\in\mathcal{Y}\)) are twice-differentiable, Lipschitz-continuous, and their first-order and second-order derivatives are both Lipschitz-continuous._

_With probability at least \(1-\delta\) over the randomness of \(N\) independent samples from \(\mathcal{D}\), for any \(\delta>0\), and any \(\epsilon>0\) close to zero, any model \(f\) with weight matrices in the set \(\mathcal{H}\) satisfies:_

\[L(f)\leq(1+\epsilon)\hat{L}(f)+\mathrm{O}\left(\frac{\log(\delta^{-1})}{N^{3/4 }}\right)+\sum_{i=1}^{l}\sqrt{\frac{CBM_{i}\left(\max_{(X,G,y)\sim\mathcal{D}} \left\|X\right\|^{2}\left\|P_{G}\right\|^{2(l-1)}\right)\left(r_{i}^{2}\prod \limits_{j=1}^{l}s_{j}^{2}\right)}{N}}, \tag{5}\]

_where \(B\) is an upper bound on the value of the loss function \(\ell(x,y)\) for any \((x,y)\sim\mathcal{D}\), and \(C\) is a fixed Lipchitz constant depending on the activation's and the loss function's Lipschitz-continuity (see equation (43), Appendix A.2.4)._
As a remark, prior works by Garg et al. [16] and Liao et al. [37] consider an MPNN with \(W^{(t)}\) and \(U^{(t)}\) being the same for \(t\) from \(1\) up to \(l\), motivated by practical designs [17, 33]. Thus, their analysis needs to be conducted separately for GCN and MPNN with weight tying. By contrast, our result allows \(W^{(t)}\) and \(U^{(t)}\) to be arbitrarily different across different layers. This unifies GCN and MPNN without weight tying in the same framework so that we can unify their analysis. We defer the proof sketch of our result along with a discussion to Section 4.

### Comparison with prior art

In Table 1, we compare our result with prior results. We first illustrate the effects of graph properties on the generalization bounds. Then we will also show a numerical comparison to incorporate the other components of the bounds.

* Suppose \(P_{{}_{G}}\) is the adjacency matrix of \(G\). Then, one can show that for any undirected graph \(G\), the spectral norm of \(P_{{}_{G}}\) is less than the maximum degree \(d\) (cf. Fact A.1, Appendix A for a proof). This explains why our result is strictly less than prior results for MPNN in Table 1.
* Suppose \(P_{{}_{G}}\) is the normalized and symmetric adjacency matrix of \(G\): \(P_{{}_{G}}=\tilde{D}^{-1/2}\tilde{A}\tilde{D}^{-1/2}\), where \(\tilde{A}\) is \(A+\text{Id}\) and \(\tilde{D}\) is the degree-diagonal matrix of \(\tilde{A}\). Then, the spectral norm of \(P_{{}_{G}}\) is at most one (cf. Fact A.1, Appendix A for a proof). This fact explains why the graph dependence of our result for GCN is \(1\) in Table 1. Thus, we can see that this provides an exponential improvement compared to the prior results.

Thus, for the above diffusion matrices, we conclude that the spectral norm of \(P_{{}_{G}}\) is strictly smaller than the maximum degree of graph \(G\) (across all graphs in the distribution \(\mathcal{D}\)).

Numerical Comparison.Next, we conduct an empirical analysis to compare our results and prior results numerically. Following the setting of prior works, we use two types of models that share their weight matrices across different layers, including GCN [35] and the MPNN specified in Liao et al. [37]. For both models, we evaluate the generalization bounds by varying the network depth \(l\) between \(2,4,\) and \(6\).

We consider graph prediction tasks on three collaboration networks, including IMDB-B, IMDB-M, and COLLAB [57]. IMDB-B includes a collection of movie collaboration graphs. In each graph, a node represents an actor or an actress, and an edge denotes a collaboration in the same movie. The task is to classify each graph into the movie genre as Action or Romance. The IMDB-M is a multi-class extension with the movie graph label Comedy, Romance, or Sci-Fi. COLLAB includes a list of ego-networks of scientific researchers. Each graph includes a researcher and her collaborators as nodes. An edge in the graph indicates

\begin{table}
\begin{tabular}{|c|c|c|c|c|} \hline \hline Graph Dependence & **GCN** & **MPNN** & **GIN** & **GraphSAGE-Mean** \\ \hline Garg et al. (2020) & \(d^{l-1}\) & \(d^{l-1}\) & - & - \\ Liao et al. (2021) & \(d^{\frac{l-1}{2}}\) & \(d^{l-1}\) & - & - \\
**Ours (Theorems \(3.1\) and \(4.5\))** & \(1\) & \(\left\|A\right\|^{l-1}\) & \(\sum_{i=1}^{l-1}\frac{\left\|A\right\|^{l}}{l-1}\) & \(\left\|D^{-1}A\right\|^{l-1}\) \\ \hline \hline \end{tabular}
\end{table}
Table 1: How does the generalization gap of graph neural networks scale with graph properties? In this work, we show spectrally-normalized bounds on \(P_{{}_{G}}\) and compare our results with prior results in the following table. We let \(A\) denote the adjacency matrix, \(D\) be the degree-diagonal matrix of \(A\), and \(l\) be the depth of the GNN. Previous generalization bounds scale with the graph’s maximum degree denoted as \(d\). Our result instead scales with the spectral norm of \(P_{{}_{G}}\) and applies to new settings, including graph isomorphism networks (GIN) [54] and GraphSAGE with mean aggregation [24].
a collaboration between two researchers. The task is to classify each ego-network into the field of the researcher, including High Energy, Condensed Matter, and Astro Physics.

We report the numerical comparison in Figure 2. We report the averaged result over three random seeds. Our results are consistently smaller than previous results. As explained in Table 1, the improvement comes from the spectral norm bounds on graphs compared with the max degree bounds.

### A matching lower bound

Next, we show an instance with the same dependence on the graph diffusion matrix as our upper bound. In our example:

* The graph \(G\) is the complete graph with self-loops inserted in each node. Thus, the adjacency matrix of \(G\) is precisely a square matrix with all ones. We will set \(P_{{}_{G}}\) as the adjacency matrix of \(G\).
* In the first \(l-1\) graph diffusion layers, the activation functions \(\phi,\rho,\psi\) are all linear functions. Further, we fix all the parameters of \(\mathbf{U}\) as zero.
* The loss function \(\ell\) is the logistic loss.

Then, we demonstrate a data distribution such that there always exists some weight matrices within \(\mathcal{H}\) whose generalization gap must increase in proportion to the spectral norm of \(P_{{}_{G}}^{l-1}\) and the product of the spectral norm of every layer \(s_{1},s_{2},\ldots,s_{l}\).

**Theorem 3.2**.: _Let \(N_{0}\) be a sufficiently large value. For any norms \(s_{1},s_{2},\ldots,s_{n}\), there exists a data distribution \(\mathcal{D}\) on which with probability at least \(0.1\) over the randomness of \(N\) independent samples from \(\mathcal{D}\), for any \(N\geq N_{0}\), the generalization gap of \(f\) is greater than the following:_

\[\left|L(f)-\hat{L}(f)\right|\gtrsim\sqrt{\frac{\left(\max_{(\chi,G,\psi)\sim \mathcal{D}}\left\|P_{{}_{G}}\right\|^{2(l-1)}\right)\left(\prod\limits_{i=1}^{ l}s_{i}^{2}\right)}{N}}. \tag{6}\]

Figure 2: Comparing our result and prior results [16, 37] on three graph classification tasks. **Upper**: The experiments are conducted on GCNs. **Lower**: The experiments are conducted on MPNNs following the setup of Liao et al. [37].
Notice that the lower bound in (6) exhibits the same scaling in terms of \(G-\left\lVert P_{{}_{G}}\right\rVert^{l-1}\)--as our upper bound from equation (5). Therefore, we conclude that our spectral norm bound is tight for multilayer MPNN. The proof of the lower bound can be found in Appendix A.3.

**Remark 3.3**.: Our results from Theorem 3.1 and 3.2 together suggest the generalization error bound scales linearly in \(l\). To verify whether this is the case, we conducted an empirical study on three architectures (GCN, GIN-Mean, and GIN-Sum) that measured the growth of generalization errors as the network depth \(l\) varies. We find that the generalization error grows sublinearly with \(l\) to \(\left\lVert P_{{}_{G}}\right\rVert\). We also note that this sublinear growth trend has been captured by our Hessian-based generalization bound (cf. Figure 0(a)). It would be interesting to understand better why the sublinear trend happens and further provide insight into the behavior of GNN.

**Remark 3.4**.: Theorem 3.2 suggests that in the worst case, the generalization bound would have to scale with the spectral norms of the graph and the weight matrices. Although this is vacuous for large \(l\), later in Lemma 4.1, we show a data-dependent generalization bound using the trace of the Hessians, which is non-vacuous. As shown in Figure 0(a), Hessian-based measurements match the scale of actual generalization errors: the green line, calculated based on the trace of the loss Hessian matrix (cf. equation (7)), matches the scale of the actual generalization error plotted in the yellow line.

## 4 Proof Techniques and Extensions

Our analysis for dealing with the graph structure seems fundamentally different from the existing analysis. In the margin analysis of Liao et al. [37], the authors also incorporate the graph structure in the perturbation error. For bounding the perturbation error, the authors use a triangle inequality that results in a \((1,\infty)\) norm of the matrix \(P_{{}_{G}}\) (see Lemma 3.1 of Liao et al. [37] for GCN). We note that this norm can be larger than the spectral norm by a factor of \(\sqrt{n}\), where \(n\) is the number of nodes in \(G\): in the case of a star graph, this norm for the graph diffusion matrix of GCN is \(\sqrt{n}\). By comparison, the spectral norm of the same matrix is less than one (see Fact A.1, Appendix A).

How can we tighten the perturbation error analysis and the dependence on \(P_{{}_{G}}\) in the generalization bounds, then? Our proof involves two parts:

* **Part I:** By expanding the perturbed loss of a GNN, we prove a bound on the generalization gap using the trace of the Hessian matrix associated with the loss.
* **Part II:** Then, we explicitly bound the trace of the Hessian matrix with the spectral norm of the graph using the Lipschitzness of the activation functions.

**Part I:**_Measuring noise stability using the Hessian._ We first state an implicit generalization bound that measures the trace of the Hessian matrix. Let \(\mathbf{H}^{(i)}\) denote the Hessian matrix of the loss \(\ell(f(X,G),y)\) with respect to layer \(i\)'s parameters, for each \(i\) from \(1\) up to \(l\). Particularly, \(\mathbf{H}^{(i)}\) is a square matrix whose dimension depends on the number of variables within layer \(i\). Let \(\mathbf{H}\) denote the Hessian matrix of the loss \(\ell(f(X,G),y)\) over all parameters of \(f\).

**Lemma 4.1**.: _In the setting of Theorem 3.1, with probability at least \(1-\delta\) over the randomness of the \(N\) training examples, for any \(\delta>0\) and \(\epsilon\) close to \(0\), we get:_

\[L(f)\leq(1+\epsilon)\hat{L}(f)+(1+\epsilon)\sum_{i=1}^{l}\sqrt{\frac{B\cdot \left(\max_{(X,G,y)\sim\mathcal{D}}\operatorname{Tr}\left[\mathbf{H}^{(i)} \left[\ell(f(X,G),y)\right]\right]\right)s_{i}^{2}r_{i}^{2}}{N}}+\mathrm{O} \left(\frac{\log(\delta^{-1})}{N^{3/4}}\right). \tag{7}\]
Proof Sketch.At a high level, the above result follows from Taylor's expansion of the perturbed loss. Suppose each parameter of \(f\) is perturbed by an independent noise drawn from a Gaussian distribution with mean zero and variance \(\sigma^{2}\). Let \(\tilde{\ell}(f(X,G),y)\) be the perturbed loss value of an input example \(X,G\) with label \(y\). Let \(\mathcal{E}\) denote the noise injections organized in a vector. Using Taylor's expansion of the perturbed loss \(\tilde{\ell}\), we get:

\[\tilde{\ell}(f(X,G),y)-\ell(f(X,G),y)=\mathcal{E}^{\top}\nabla\ell(f(X,G),y)+ \frac{1}{2}\mathcal{E}^{\top}\mathbf{H}\big{[}\ell(f(X,G),y)\big{]}\mathcal{E} +\mathrm{O}(\sigma^{3}). \tag{8}\]

Notice that the expectation of the first-order expansion term above is equal to zero. The expectation of the second-order expansion term becomes \(\sigma^{2}\) times the trace of the loss Hessian. To derive equation (7), we use a PAC-Bayes bound of McAllester (39, Theorem 2). There are two parts to this PAC-Bayes bound:

* The expectation of the noise perturbation in equation (8), taken over the injected noise \(\mathcal{E}\);
* The KL divergence between the prior and the posterior, which is at most \(s_{i}^{2}r_{i}^{2}\) for layer \(i\), for \(i\) from \(1\) up to \(l\), within the hypothesis set \(\mathcal{H}\).

Thus, one can balance the two parts by adjusting the noise variance at each layer--this leads to the layerwise Hessian decomposition in equation (7).

A critical step is showing the uniform convergence of the Hessian matrix. We achieve this based on the Lipschitz-continuity of the first and twice derivatives of the nonlinear activation mappings. With these conditions, we prove the uniform convergence with a standard \(\epsilon\)-cover argument. The complete proof can be found in Appendix A.2.1.

**Remark 4.2**.: Our argument in Lemma 4.1 applies to graph-level prediction tasks, which assume an unknown distribution of graphs. A natural question is whether the analysis applies to node-level prediction tasks, which are often treated as semi-supervised learning problems. The issue with directly applying our analysis to semi-supervised learning is that the size of a graph is only finite. Instead, a natural extension would be to think about our graph as a random sample from some population and then argue about generalization in expectation of the random sample. It is conceivable that one can prove a similar spectral norm bound for node prediction in this extension. This would be an interesting question for future work.

**Part II:**_Spectral norm bounds of the trace of the Hessian._ Next, we explicitly analyze the trace of the Hessian at each layer. We bound the trace of the Hessian using the spectral norm of the weight matrices and the graph based on the Lipschitz-continuity conditions from Theorem 3.1. Notice that the last layer is a linear pooling layer, which can be deduced from layer \(l-1\). Hence, we consider the first \(l-1\) layers below.

**Lemma 4.3**.: _In the setting of Theorem 3.1, the trace of the loss Hessian matrix \(\mathbf{H}^{(i)}\) taken over \(W^{(i)}\) and \(U^{(i)}\) satisfies the following, for any \(i=1,2,\cdots,l-1\),_

\[\left|\mathrm{Tr}\left[\mathbf{H}^{(i)}\left[\ell(f(X,G),y)\right] \right]\right| \lesssim\ s_{l}^{2}\left(\sum_{p=1}^{d_{i-1}}\frac{d_{i}}{q=1} \left\|\frac{\partial^{2}H^{(l-1)}}{\partial\big{(}W^{(i)}_{p,q}\big{)}^{2}} \right\|_{F}+\sum_{p=1}^{d_{i}}\sum_{q=l}^{d_{i}}\left\|\frac{\partial^{2}H^{( l-1)}}{\partial\big{(}U^{(i)}_{p,q}\big{)}^{2}}\right\|_{F}+\left\|\frac{ \partial H^{(l-1)}}{\partial W^{(i)}}\right\|_{F}^{2}+\left\|\frac{\partial H ^{(l-1)}}{\partial U^{(i)}}\right\|_{F}^{2}\right) \tag{9}\] \[\lesssim\ \left\|X\right\|^{2}\left\|P_{c}\right\|^{2(l-1)}\prod_{j =1:\ j\neq i}^{l}s_{j}^{2}. \tag{10}\]

Proof Sketch.Equation (9) uses the chain rule to expand out the trace of the Hessian and then applies the Lipschitzness of the loss function. Based on this result, equation (10) then bounds the first and second derivatives of \(H^{(l-1)}\). This step is achieved via an induction of \(\partial H^{(j)}\) and \(\partial^{2}H^{(j)}\) over \(W^{(i)}\) and \(U^{(i)}\), for \(j=1,\ldots,l-1\) and \(i=1,\ldots,j\). The induction relies on the feedforward architecture and the Lipschitzness of the first and second derivatives. We leave out a few details, such as the constants in equations (9) and (10) that can be found in Appendix A.2.2 and A.2.3. Combining both parts together, we get equation (1).
**Remark 4.4**.: We compare our analysis with the approach of Liao et al. [37]. Both our analysis and Liao et al. [37] follow the PAC-Bayesian framework. But additionally, we explore Lipschitz-continuity properties of the first and second derivatives of the activation functions (e.g., examples of such activations include tanh and sigmoid). This allows us to measure the perturbation loss with Hessians, which captures data-dependent properties much more accurately than the margin analysis of Liao et al. [37]. It would be interesting to understand if one could still achieve spectral norm bounds on graphs under weaker smoothness conditions (see, e.g., Wei and Ma [51]). This is left for future work.

### Extensions

**Graph isomorphism networks.** This architecture concatenates every layer's embedding together for more expressiveness [54]. A classification layer is used after the layers. Let \(V^{(t)}\) denote a \(d_{i}\) by \(k\) matrix (recall \(k\) is the output dimension). Denote the set of these matrices by \(\mathcal{V}\). We average the loss of all of the classification layers. Let \(\hat{L}_{\text{\tiny{GIN}}}(f)\) denote the average loss of \(f\) over \(N\) independent samples of \(\mathcal{D}\). Let \(L_{\text{\tiny{GIN}}}(f)\) denote the expected loss of \(f\) over a random sample of \(\mathcal{D}\). See also equation (44) in Appendix A.4 for their precise definitions.

Next, we state a generalization bound for graph isomorphism networks. Let \(f\) be any \(l\)-layer MPNN with weights defined in a hypothesis space \(\mathcal{H}\): the parameters of \(f\) reside within the constraints from equation (4); further, for every \(i\) from \(1\) up to \(l\), the spectral norm of \(V^{(i)}\) is less than \(s_{l}\). Building on Lemma 4.3, we show a bound that scales with the spectral norm generalization of the averaged graph diffusion matrices. Let \(P_{\text{\tiny{GIN}}}\) denote the average of \(l-1\) matrices: \(P_{\text{\tiny{G}}},P_{\text{\tiny{G}}}^{2},\ldots,P_{\text{\tiny{G}}}^{l-1}\). We state the result below.

**Corollary 4.5**.: _Suppose the nonlinear activation mappings and the loss function satisfy the conditions stated in Theorem 3.1. With probability at least \(1-\delta\) for any \(\delta\geq 0\), and any \(\epsilon\) close to zero, any \(f\) in \(\mathcal{H}\) satisfies:_

\[L_{\text{\tiny{GIN}}}(f)\leq(1+\epsilon)\hat{L}_{\text{\tiny{GIN}}}(f)+\sum_{i =1}^{l}\sqrt{\frac{BCd_{i}\cdot\left(\max_{(X,G,y)\sim\mathcal{D}}\|X\|^{2} \left\|P_{\text{\tiny{GIN}}}\right\|^{2}\right)\left(r_{i}^{2}\prod_{j=1}^{l} s_{j}^{2}\right)}{N}}+\text{O}\left(\frac{\log(\delta^{-1})}{N^{3/4}}\right), \tag{11}\]

_where \(B\) is an upper bound on the value of the loss function \(\ell\) across the data distribution \(\mathcal{D}\), and \(C\) is a fixed constant that only depends on the Lipschitz-continuity of the activation mappings and the loss._

The proof can be found in Appendix A.4. In particular, we apply the trace norm bound over the model output of every layer. The classification layer, which only uses a linear transformation, can also be incorporated.

Fine-tuned Graph Neural Networks.We note that all of our bounds can be applied to the fine-tuning setting, where a graph neural network is initialized with pretrained weights and then fine-tuned on the target task. The results can be extended to this setting by setting the norm bounds within equation (4) as the distance between the pretrained and fine-tuned model.

## 5 Noise Stability Optimization for Fine-tuning GNNs

A common practice to apply graph and deep learning is to adopt a pretrained GNN and fine-tune it for a target problem. Typically, only a small amount of data is available for fine-tuning. Thus, the fine-tuned model may overfit the training data, incurring a large generalization gap. A central insight from our analysis is that maintaining a small perturbed loss ensures lower generalization gaps. Motivated by this observation, we present an optimization algorithm to minimize the perturbed loss of a model.
Let \(f\) denote a model and \(\tilde{\ell}(f)\) be the perturbed loss of \(f\), with noise injected inside \(f\)'s weight matrices. Recall from step (8) that \(\tilde{\ell}(f)\) is equal to \(\ell(f)\) plus several expansion terms. In particular, minimizing the expectation of \(\tilde{\ell}(f)\) is equivalent to minimizing \(\hat{L}(f)\) plus the trace of the Hessian matrix. To estimate this expectation, we sample several noise perturbations independently. Because Taylor's expansion of \(\tilde{\ell}(f)\) also involves the gradient, we cancel this out by computing the perturbed loss with the negated perturbation. Algorithm 1 describes the complete procedure.

```
0: A training dataset \(\{(X_{i},G_{i},y_{i})\}_{i=1}^{N}\) with node feature \(X_{i}\), graph \(G_{i}\), and graph-level label \(y_{i}\), for \(i=1,\dots,N\).
0: Number of perturbations \(m\), noise variance \(\sigma^{2}\), learning rate \(\eta\), and number of epochs \(T\).
0: A trained model \(f^{(T)}\).
1: At \(t=0\), initialize the parameters of \(f^{(0)}\) with pretrained GNN weight matrices.
2:for\(1\leq t\leq T\)do
3:for\(1\leq i\leq m\)do
4: Add perturbation \(\mathcal{E}_{i}\) drawn from a normal distribution with mean zero and variance \(\sigma^{2}\).
5: Let \(\tilde{L}_{i}(f^{(t-1)})\) be the training loss of the model \(f^{(t-1)}\) with weight matrix perturbed by \(\mathcal{E}_{i}\).
6: Let \(\tilde{L}_{i}^{{}^{\prime}}(f^{(t-1)})\) be the training loss of the model \(f^{(t-1)}\) with weight matrix perturbed by \(-\mathcal{E}_{i}\).
7:endfor
8: Use stochastic gradient descent to update \(f^{(t)}\) as \(f^{(t-1)}-\frac{\eta}{2m}\sum_{i=1}^{m}\big{(}\nabla\tilde{L}_{i}\big{(}f^{(t -1)}\big{)}+\nabla\tilde{L}_{i}^{{}^{\prime}}\big{(}f^{(t-1)}\big{)}\big{)}\).
9:endfor
```

**Algorithm 1** Noise stability optimization for fine-tuning graph neural networks

We evaluate the above algorithm for fine-tuning pretrained GNNs. Empirical results reveal that this algorithm achieves better test performance compared with existing regularization methods for five graph classification tasks.

### Experimental setup

We focus on graph classification tasks, including five datasets from the MoleculeNet benchmark [52]. In each dataset, the goal is to predict whether a molecule has a certain chemical property given its graph representation. We use pretrained GINs from Hu et al. [27] and fine-tune the model on each downstream task. Following their experimental setup, we use the scaffold split for the dataset, and the model architecture is fixed for all five datasets. Each model has 5 layers; each layer has 300 hidden units and uses average pooling in the readout layer. We set the parameters, such as the learning rate and the number of epochs following their setup.

We compare our algorithm with previous regularization methods that serve as benchmark approaches for improving generalization. This includes early stopping, weight decay, dropout, weight averaging [29], and distance-based regularization [21]. For implementing our algorithm, we set the number of perturbations as 10 and choose the noise standard deviation \(\sigma\) with a grid search in \(\{0.01,0.02,0.05,0.1,0.2,0.5\}\).

### Experimental results

Table 2 reports the test ROC-AUC performance averaged over multiple binary prediction tasks in each dataset. Comparing the average ranks of methods across datasets, our algorithm outperforms baselines on all five molecular property prediction datasets. The results support our theoretical analysis that the noise stability property of GNN is a strong measure of empirical generalization performance. Next, we provide details insights from applying our algorithm.

First, we hypothesize that our algorithm is particularly effective when the empirical generalization gap is large. To test the hypothesis, we vary the size of the training set in the BACE dataset; we compare the performance of our algorithm with early stopping until epoch 100. We plot the generalization gap between
the training and test losses during training, shown in Figure 2(a)-2(b). As the trend shows, our algorithm consistently reduces the generalization gap, particularly when the training set size \(N\) is 600.

Second, we hypothesize that our algorithm helps reduce the trace of the Hessian matrix (associated with the loss). We validate this by plotting the trace of the Hessian as the number of epochs progresses during training, again using the BACE dataset as an example. Specifically, we average the trace over the training dataset. Figure 2(c) shows the averaged trace values during the fine-tuning process. The results confirm that noise stability optimization reduces the trace of the Hessian matrix (more significantly than early stopping). We note that noise stability optimization also reduces the largest eigenvalue of the Hessian matrix, along with reducing the trace. This can be seen in Figure 2(d).

Lastly, we study the number of perturbations used in our algorithm. While more perturbations would lead to a better estimation of the noisy stability, we observe that using 10 perturbations is sufficient for getting the most gain. We also validate that using negated perturbations consistently performs better than not using them across five datasets. This is because the negated perturbation cancels out the first-order term in Taylor's expansion. In our ablation study, we find that adding the negated perturbation performs better than not using it by 1% on average over the five datasets.

**Remark.** We note that noise stability optimization is closely related to sharpness-aware minimization (SAM) [15]. Noise stability optimization differs in two aspects compared with SAM. First, SAM requires solving constrained minimax optimization, which may not even be differentiable [10]. Our objective remains the same after perturbation. Second, SAM reduces the largest eigenvalue of the Hessian matrix, which can be seen from Taylor's expansion of \(\tilde{\ell}(f)\). We reduce the trace of the Hessian matrix, which includes reducing the largest eigenvalue as part of the trace. There is another related work that regularizes noise stability in NLP [28]. Their approach adds noise perturbation in the input and regularizes the loss change in the output. Our approach directly adds the perturbation in the weight matrices.

\begin{table}
\begin{tabular}{l c c c c c c} \hline \hline Dataset & SIDER & ClinTox & BACE & BBBP & Tox21 & \\ \# Molecule Graphs & 1,427 & 1,478 & 1,513 & 2,039 & 7,831 & Avg. Rank \\ \# Binary Prediction Tasks & 27 & 2 & 1 & 1 & 12 & \\ \hline Early Stopping & 61.06\(\pm\)1.48 & 68.25\(\pm\)2.63 & 82.86\(\pm\)0.95 & 67.80\(\pm\)1.05 & 77.52\(\pm\)0.23 & 5.8 \\ Weight Decay & 61.30\(\pm\)0.21 & 67.43\(\pm\)2.88 & 83.72\(\pm\)0.99 & 67.98\(\pm\)2.41 & 78.23\(\pm\)0.35 & 5.0 \\ Dropout & 63.90\(\pm\)0.90 & 73.70\(\pm\)2.80 & 84.50\(\pm\)0.70 & 68.07\(\pm\)1.30 & 78.30\(\pm\)0.30 & 3.6 \\ Weight Averaging & 63.67\(\pm\)0.34 & 78.78\(\pm\)1.49 & 83.93\(\pm\)0.36 & 70.26\(\pm\)0.24 & 77.59\(\pm\)0.11 & 3.4 \\ Distance-based Reg. & 64.36\(\pm\)0.48 & 76.68\(\pm\)1.19 & 84.65\(\pm\)0.48 & 70.37\(\pm\)0.44 & 78.62\(\pm\)0.24 & 2.2 \\ \hline
**Ours (Alg. 1)** & **65.13\(\pm\)0.18** & **80.18\(\pm\)0.82** & **85.07\(\pm\)0.43** & **71.22\(\pm\)0.36** & **79.31\(\pm\)0.24** & **1.0** \\ \hline \hline \end{tabular}
\end{table}
Table 2: Test ROC-AUC (%) score for five molecular property prediction datasets with different regularization methods. The reported results are averaged over five random seeds.

Figure 3: In Figures 2(a) and 2(b), we show that our algorithm is particularly effective at reducing the generalization gap for small training dataset sizes \(N\). In Figures 2(c) and 2(d), we find that both the trace and the largest eigenvalue of the loss Hessian matrix decreased during training.
Conclusion

This work develops generalization bounds for graph neural networks with a sharp dependence on the graph diffusion matrix. The results are achieved within a unified setting that significantly extends prior works. In particular, we answer an open question mentioned in Liao et al. [37]: a refined PAC-Bayesian analysis can improve the generalization bounds for message-passing neural networks. These bounds are obtained by analyzing the trace of the Hessian matrix with the Lipschitz-continuity of the activation functions. Empirical findings suggest that the Hessian-based bound matches observed gaps on real-world graphs. Thus, our work also develops a practical tool to measure the generalization performance of graph neural networks. The algorithmic results with noise stability optimization further demonstrate the practical implication of our findings.

Our work opens up many interesting questions for future work. Could the new tools we have developed be used to study generalization in graph attention networks [48]? Could Hessians be used for measuring out-of-distribution generalization gaps of graph neural networks?

## Acknowledgement

Thanks to Renjie Liao, Haoyu He, and the anonymous referees for providing constructive feedback on our work. Thanks to Yang Yuan for the helpful discussions. HJ and DL acknowledge financial support from the startup fund of Khoury College of Computer Sciences, Northeastern University.
