Paper: How Graph Neural Networks Learn: Lessons from Training Dynamics in Function Space (ICLR-2024)

Reviewer comment (Theory):
The assumption of infinitely wide network is not realistic. It is better to analyze the evolution of the kernel.

Author response:
Though idealized, the infinite width limit or NTK regime is widely used for shedding insight into problems that are otherwise highly prohibitive to analyze. Previous studies for the NTK regime have led to a wealth of advances in numerous problems such as dataset distillation, matrix completion, generalization prediction, architectural search (see [1] and references therein). In our particular scenario, the insights drawn from the infinite width regime make predictions regarding behavior in more practical finite-width models. Importantly, these predictions can be largely verified based on extensive empirical verification spanning 17 datasets.  So indeed, although as the reviewer suggests infinitely wide models may be idealized, they still have practical value as an analysis tool, and our theoretical results still asymptotically hold for larger models in the finite width case.

**(Evolution of NTK)** As for evolution of the kernel in the finite width case, a general theoretical characterization is still largely infeasible [2]. Nevertheless, we have provided empirical analysis of the evolution of the NTK during training. As elaborated in Section 5.2, our observations on the kernel evolution include: 1) the NTK of GNNs tend to align with the graph; 2) the NTK also tends to align with the optimal kernel as it evolves during the training process, and heterophily hinders this process since alignment to both graph and the optimal target become contradictory. 

[1] Neural Tangent Kernel: A Survey, 2022

[2] Spectral Evolution and Invariance in Linear-width Neural Networks (ICML 2023 Workshop)

Reviewer comment (Theory):
The restriction to two-layer GNN or last-layer feature propagation is not realistic either.

Author response:
In Section 4.1, we have derived the explicit formula for the node-level GNTK, which is compatible with any GNN defined in Section 2. Such a formula demonstrates how the graph adjacency is integrated into the kernel similarity measure. The purpose of Section 4.2 is then to provide concrete examples for strengthening and illustrating these results, and thus architectural restrictions are necessary. 

Since there is typically no explicit, analytically-tractable expression of the NTK for very deep NNs (only a formula for recurrently computing it exists), we consider two-layer GNNs in many previous works on MLPs [3]. In addition to this however, we introduce a notable exception where there indeed exists an explicit NTK expression for deep GNNs with learnable embedding inputs and feature propagation at the last layer, which is why we consider this type of architecture. Moreover, feature propagation at the last layer is not uncommon; many popular GNNs follow this architecture such as APPNP [4] and GPR-GNN [5]. We hope this helps to clarify, but we are happy to provide further details if needed.

[3] Fine-Grained Analysis of Optimization and Generalization for Overparameterized Two-Layer Neural Networks (ICML 2019)

[4] Predict then propagate: Graph neural networks meet personalized pagerank (ICLR 2019)

[5] Adaptive Universal Generalized Pagerank Graph Neural Network (ICLR 2021)

Reviewer comment (Theory):
Is it possible to go beyond neural tangent kernel and two-layer GNN? What theoretical tools are needed? I assume there are methods developed for MLP.

Author response:
**(Beyond NTK Regime)** Going beyond the NTK regime may be possible in principle; however, the analysis either still requires certain degrees of overparameterization, leading to imprecise results based on bounds that are heavily dependent on the width (e.g. [5]), or uses other approximations such as second-order Taylor expansion of neural networks (e.g. [6]). Unfortunately, exact mathematical characterization of the optimization for finite width MLPs, even for those shallow ones, is still an open question in general neural network theory. In fact, there are quite a few recent papers just on analyzing single-neuron networks (e.g. [7,8]), indicating the difficulty of analyzing not-so-wide networks. 

**(Beyond Two-Layer)** Extending to deeper layers is somewhat more feasible. We have provided a formula for computing the NTK of arbitrarily deep GNNs in Section 4.1. However, the expression of the NTK for deeper neural networks is quite complex and not conducive to further analysis at this point. Still this is an interesting direction we hope to explore in future work if possible.

[6] Learning Over-Parameterized Two-Layer ReLU Neural Networks beyond NTK (COLT 2020)

[7] Beyond Linearization: On Quadratic and Higher-Order Approximation of Wide Neural Networks (ICLR 2020) 

[8] Agnostic learning of a single neuron with gradient descent (NeurIPS 2021)

[9] Learning a Neuron by a Shallow ReLU Network: Dynamics and Implicit Bias for Correlated Inputs (2023)

Reviewer comment (Theory):
In Section 4.2, the authors only consider two very special $\bar{\mathcal{X}}$, which makes the theory somehow limited.

Author response:
In Section 4.1, we have given the general form of GNN training dynamics where the NTK can be derived based on our explicit formula. This part of the analysis is compatible with an arbitrary $\bar{\mathcal{X}}$ and demonstrates how the graph adjacency is naturally integrated into the kernel similarity measure. Section 4.2 complements 4.1 by providing concrete examples, for which purpose constraints on inputs are necessary. In comparison, previous work on GNN expressiveness often assume the graph is endowed with node-coloring [1], while other work assumes node features and graph structure are generated by a probabilistic model such as CSBM [2]. 

Moreover, we should also mention that the choices for $\bar{\mathcal{X}}$ are not arbitrary or contrived, but closely related to practice: $\bar{\mathcal{X}} = \mathbf I$ represents learnable node embeddings that are ubiquitous in GNN applications such as recommender systems and knowledge graphs, while $\bar{\mathcal{X}}$ from a graph spectral decomposition unifies many network embedding methods such as DeepWalk and Node2Vec, which themselves are being actively researched.

[1] Theory of Graph Neural Networks: Representation and Learning (2022)

[2] Understanding Non-linearity in Graph Neural Networks from the Perspective of Bayesian Inference (NeurIPS 2022)

Reviewer comment (Theory):
This question is related to the second point in the "Weakness". According to your derivations in Section B.4, I think you remove the nonlinearity or the dependence of the kernel on the parameters $W$ by taking the expectation for $W$. Please correct me if I misunderstood something. I am confused as to why you take the expectation -- in my opinion, training dynamics is the evaluation of residuals or parameters for any given/fixed initialization. If you want to take expectation over $W$, I think the equation (14) and (15) should be stated as something like expected residual. Please explain what is happening and why it is reasonable.

Author response:
The expectation stems from computing the NTK at its initialization where there are infinitely many parameters sampled from a Gaussian, the weighted sum of which induces expectation. To be more specific, as mentioned above, NTK $\mathbf \Theta_t$ at arbitrary time $t$ converges (w.r.t. model width) to a constant kernel function at its initialization $\mathbf \Theta_0$, and thus we need only to compute the $\mathbf \Theta_0$ where $\mathbf W^{(1)}$ is Gaussian distributed (as stipulated by the standard NTK parameterization stated in theorems). In the computation, expectation emerges from the fact that $\mathbf W^{(1)} = [w_1, w_2, \cdots, w_{\infty}] \in \mathbb R^{d\times \infty}$ is a concatenation of infinitely many vectors from $\mathcal N(0,\mathbf I_d)$, with a scaling factor $1/m$ (and $m\rightarrow \infty$) in the NTK for averaging. Similar results can be found for two-layer MLPs [7].

[7] Fine-Grained Analysis of Optimization and Generalization for Overparameterized Two-Layer Neural Networks (ICML 2019)

Reviewer comment (Experiments):
The empirical validation is limited in scope, focusing on smaller datasets and simpler models. Extensive testing involving state-of-the-art GNNs and more diverse benchmarks would be instrumental in corroborating the theoretical claims.

Author response:
**(Scale and Diversity of Datasets)** We remark that our empirical verification spans 17 datasets (Arxiv, Proteins, Products, Cora, Citeseer, Pubmed, Computers, Photo, CS, Physics, Roman-empire, Amazon-ratings, Minesweeper, Tolokers, Questions, Texas, Synthetic), including homophilic, heterophilic, and up-to-date (proposed in 2023) examples. Even so, we concede that some of these were deferred to the supplementary to save space and could easily be missed.  We could reorganize to include more in the main paper if the reviewer suggests and space allows. 

Additionally, with respect to size, the OGBN datasets in Table 1 are recognized to be relatively large-scale benchmarks, with up to 2,449,029 nodes and 61,859,140 edges.  Although there do exist a few larger benchmarks, dealing with these requires considerable engineering effort and compute, and scalability is not our primary focus. Overall then, we would argue that for an analysis-oriented paper in particular, our experimental coverage is considerable, especially so if we account for the supplementary.

**(SOTA Models)** As for baseline models for comparison, our experiments are specifically designed to complement and corroborate analytical results, not for challenging an ever-changing landscape of SOTA models.  That being said, in Appendix E.3, we have evaluated RP against an array of recent GNNs explicitly designed to handle heterophily, and it still performs competitively despite its far greater simplicity.

Reviewer comment (Novelty):
The theoretical contributions, while valuable, seem to offer only a modest advancement beyond existing studies on kernel alignment. Clarifying the distinctions from previous work would help to highlight the unique contributions of this study.

Author response:
**(Pior Work on Kernel Alignment)** Kernel-target alignment [1] is a well-established concept, with application to kernel methods, e.g., as an objective for learning kernels [2]. However, the connection with generalization is mostly restricted to classic models such as Parzen window estimators [3]. Overall, this research field has been active well before the era of deep learning, and subsequently, established connections to deep learning models are loose, e.g., prior work applying kernel alignment to deep model generalization does not exist to the best of our knowledge. Importantly though, we do not aim to make general advancements in this classic field itself.

**(Ours)** Instead, we repurpose the generic notion of kernal alignment within a completely new use-case, namely, quantifying and demystifying the generalization performance of neural network architectures, and GNNs in particular. To this end, we introduce the novel construct of kernel-graph alignment, which is quite distinct from classical kernel alignment.  We then leverage this construct to generate a series of significant insights regarding the evolution of GNN functions during optimizaiton, informing when and why the resulting learned functions can generalize. 

Beyond this, we are unaware of existing work that elucidates neural network architectural designs (GNNs or otherwise) through the lens of kernel alignment as we have done.  That being said, we are happy to contextualize with respect to any additional references the reviewer might refer us to. 

[1] On kernel-target alignment (NeurIPS 2001)

[2] Learning the kernel matrix with semidefinite programming (JMLR 2004)

[3] Remarks on Some Nonparametric Estimates of a Density Function (1956)

Reviewer comment (Novelty):
How does the theoretical analysis diverge from previous studies on kernel alignment? Clarification of the novel insights would be appreciated.

Author response:
**(Pior Work on Kernel Alignment)** Kernel-target alignment [1] is a well-established concept, with application to kernel methods, e.g., as an objective for learning kernels [2]. However, the connection with generalization is mostly restricted to classic models such as Parzen window estimators [3]. Overall, this research field has been active well before the era of deep learning, and subsequently, established connections to deep learning models are loose, e.g., prior work applying kernel alignment to deep model generalization does not exist to the best of our knowledge. Importantly though, we do not aim to make general advancements in this classic field itself.

**(Ours)** Instead, we repurpose the generic notion of kernal alignment within a completely new use-case, namely, quantifying and demystifying the generalization performance of neural network architectures, and GNNs in particular. To this end, we introduce the novel construct of kernel-graph alignment, which is quite distinct from classical kernel alignment.  We then leverage this construct to generate a series of significant insights regarding the evolution of GNN functions during optimizaiton, informing when and why the resulting learned functions can generalize. 

Beyond this, we are unaware of existing work that elucidates neural network architectural designs (GNNs or otherwise) through the lens of kernel alignment as we have done.  That being said, we are happy to contextualize with respect to any additional references the reviewer might refer us to. 

[1] On kernel-target alignment (NeurIPS 2001)

[2] Learning the kernel matrix with semidefinite programming (JMLR 2004)

[3] Remarks on Some Nonparametric Estimates of a Density Function (1956)

Reviewer comment (Novelty):
Given the recognized challenges of GNNs in dealing with heterophily, are there any strategies or recommendations proposed by the authors to tackle this issue beyond the current analysis?

Author response:
**(Recommendation for Heterophily)** While recommending specific GNN architectures for tackling heterophily is an active research field mostly beyond our scope, we can nonetheless recommend more high-level design principles motivated by our analysis: 
1. Increasing $A(\mathbf A, \mathbf \Theta^*)$ by designing propagation matrices that deviate from the original adjacency matrix, such that the resulting GNN NTK could better align with the target, can lead to better generalization performance.  Moreover, this perspective helps to substantiate certain existing heuristics, such as assigning negative weights to edges. 
2. Decreasing $A(\mathbf \Theta_t, \mathbf \Theta^*)$ by designing non-standard GNN architectures, i.e., that deviate from our defintion in Section 2, can reduce the degradation of NTK-target alignment caused by heterophily; in doing so, the negative effects of heterophilic graphs on generalization can also be mitigated. Incidently, this also helps to quantify regimes where non-GNN architectures such as MLPs can sometimes perform better.
 

We should also mention that even the extremely simple RP algorithm that naturally emerges from our analysis works well in practice handling heterophily.  See Section E.3 of the supplementary for empirical examples.

**(Significance of Connection with Heterophily)** Although it is already well-known from empirical observations that GNN performance may degrade at times on heterophilic graphs, our analysis is still quite relevant in providing principled, complementary explanations for why and how this happens.