How Graph Neural Networks Learn: Lessons from Training Dynamics in Function Space

Thanks for the constructive feedback. We address each point as follows.

Weakness.1: "The assumption of infinitely wide network is not realistic. It is better to analyze the evolution of the kernel."

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)

Weakness.2: "The restriction to two-layer GNN or last-layer feature propagation is not realistic either."

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)

Question: "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."

(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)

Thanks for the constructive feedback. We address each point as follows.

Weakness.1: "In Section 4.2, the authors only consider two very special $\bar{\mathcal{X}}$, which makes the theory somehow limited."

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)

Weakness.2: "The training dynamics of GNNs should be highly nonlinear. More explicitly, $\mathbf \Theta_t$ depends on $\mathbf W$. However, the derived dynamics in theorems are linear."

This is an important question. The answer is related to NTK theory, which we leverage to translate an intractable problem into a tractable one. While the training dynamics of NN (neural network) models are generally nonlinear, recent analysis has shown that they asymptotically behave as linear models as width increases. For example, Theorem 2 in the seminal NTK work [3] states that NTK $\mathbf \Theta_t$ at arbitrary time $t$ converges (w.r.t. width) to a constant kernel function at initialization $\mathbf \Theta_0$. This result holds universally for other architectures [4]. Proofs of such results are usually mathematically intense, and counter-intuitive as the reviewer pointed out. Informally, it can be understood as follows: if the width goes to infinity, the movement of parameters caused by optimization is such that the function $f(x)$ itself changes, but perhaps counter-intuitively, the Jacobian and NTK do not change appreciably. This explains the dissociation of $\mathbf \Theta_t$ and $\mathbf W$.

For the linearization, it essentially stems from the first order Taylor expansion of a model around the initialization [5], i.e. $f(x) \approx f_0(x) + \nabla_{\mathbf{W}} f_0(x)^\top (\mathbf W - \mathbf W_0)$, which is indeed a linear model w.r.t. $\mathbf W$ that approximates the original $f(x)$. Such an approximation becomes precise for infinite width NNs, since the Jacobian stays constant and the Hessian can be proven to approach zero [6]. And because both $\nabla_{\mathbf{W}} f_0(x)^\top$ and $\mathbf W$ are in $\mathbf R^{\infty}$, NN optimization becomes equivalent to kernel regression [3]. That being said, it is important to note that the nonlinearity in the model is preserved in the derivative, and thus the benefits that GNNs can obtain from nonlinear activations are also preserved in our analysis.

[3] Neural Tangent Kernel: Convergence and Generalization in Neural Networks (NeurIPS 2018)

[4] Tensor programs iib: Architectural universality of neural tangent kernel training dynamics (ICML 2021)

[5] Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent (NeurIPS 2019)

[6] On the linearity of large non-linear models: when and why the tangent kernel is constant (NeurIPS 2020)

Question.1: "I think you remove the nonlinearity or the dependence of the kernel on the parameters $\mathbf W$, by taking the expectation for $\mathbf W$"; "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."

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)

Question.2: "In Theorem 5, do you have the training dynamics for optimizing $\mathbf W^{(2)}$ ? In Theorem 6, is the training dynamics for optimizing all parameters in the GNN?"

Yes, Theorem 5 can be extended straight-forwardly to incorporate $\mathbf W^{(2)}$, i.e. $\nabla_{\mathbf{W}} \mathcal{F}^{\top} \nabla_{\mathbf{W}} \mathcal{F} = \nabla_{\mathbf{W^{(1)}}} \mathcal{F}^{\top} \nabla_{\mathbf{W^{(1)}}} \mathcal{F} + \nabla_{\mathbf{W^{(2)}}} \mathcal{F}^{\top} \nabla_{\mathbf{W^{(2)}}} \mathcal{F}$. The latter one is relatively easy to derive as $\mathbf{W^{(2)}}$ is applied to the last layer. Consequently, the propagation matrix in Theorem 5 should be revised to $\mathbf{A}\left(\mathbf{A}^2 \odot \mathbf{S} + \mathbf A^\prime \right) \mathbf{A}$, where $\mathbf A^\prime_{ij} = \mathbb{E}_{{w} \sim\mathcal{N}(0,\mathbf I)}[\sigma(w^\top \mathbf A_i) \sigma(w^\top \mathbf A_j)]$ measures similarity of neighborhood patterns for two nodes. We did not originally consider $\mathbf{W^{(2)}}$ because it is often omitted by convention. All other analyses (including Theorem 6) consider optimizing all parameters.

Thanks for the constructive feedback. As the reviewer concerns are spread across strengths (S), weaknesses (W), and questions (Q), we will organize and address them grouped by topic.

Q.1: "How does the theoretical analysis diverge from previous studies on kernel alignment?"; S.Significance: "The insights on generalization are incremental on existing theory on kernel alignment."; W.Significance in theory: "The theoretical contributions, while valuable, seem to offer only a modest advancement beyond existing studies on kernel alignment."

(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)

Q.2: "Are there any strategies or recommendations proposed by the authors to tackle this issue [heterophily] beyond the current analysis?"; S.Significance: "The limitations of GNNs on heterophily are already well-known."

(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.