Paper: What Improves the Generalization of Graph Transformer? A Theoretical Dive into Self-attention and Positional Encoding (ICLR-2024)

Reviewer comment (Theory):
It is unclear whether the theoretical claims can be generalized to real-world scenarios and datasets given the restrictive assumption made under the presented results. In particular, Assumption 1 states that node labels follow a majority vote of discriminative patterns in core neighborhood, which is somewhat similar to network homophily. To verify this assumption on real-world datasets, Appendix E.1 provides additional analysis on the Cora network, which is known to be homophilic [A]. Thus, the question remains: "does the assumption hold for heterophilic networks as well?" Performing similar analyses on additional networks such as Actor and PascalVOC-SP-1G used in Section 5.2 would provide further support towards verifying this assumption.

Author response:
**Q1 ([W1] part 1 by the reviewer)**: "Does the assumption hold for heterophilic networks as well? Performing similar analyses on additional networks such as Actor and PascalVOC-SP-1G used in Section 5.2 would provide further support towards verifying this assumption.”

**A1**: Thank you for this question and the suggestions. **We would like to clarify that our Assumption 1 can be met by some heterophilous and long-range graph datasets. We also want to illustrate that our formulation of the core neighborhood is different from homophily, heterophily, and long-range dependency.** Homophily and heterophily are defined by the fraction of edges that connect nodes from the same class. If such fraction is close to $1$, then it is a homophilous graph. Otherwise, it is a heterophilous graph. Long range graphs are graphs where long-range information is needed for the learning tasks. Our definition of the core neighborhood is the neighborhood where the average winning margin between class-relevant and confusion nodes is the largest.

**We performed similar experiments on PubMed (homophilous graph), Actor (heterophilous graph), and PascalVOC-SP-1G (long range graph) in Section A.1 (It is Section E.1 in the first version).** The experiments include (1) the existence of discriminative nodes for each dataset (2) probing the core distance and justifying that the winning margin on the core distance $\Delta_n(z_m)>0$ (Assumption 1). 

To show the existence of discriminative nodes, we first generate Figures [10](https://imgur.com/z2VKqpA), [11](https://imgur.com/3qbJ3zE), [12](https://imgur.com/ySSCagL), and [13](https://imgur.com/W8pagMX) on Cora, PubMed, Actor, and PascalVOC-SP-1G to illustrate that node features of each class can be represented in a low dimension. Then, we use the averaged low-dimensional representation of nodes in each class to identify one discriminative pattern for each class. **Tables [2](https://imgur.com/IddhDqN), [3](https://imgur.com/BpC9B98), [4](https://imgur.com/htl1yWQ), and [5](https://imgur.com/AvM7Bay) indicate that high fractions of nodes are close to each discriminative pattern, which verifies the assumption on the discriminative/non-discriminative nodes.** 

To investigate the core distance of each dataset, we compute the fraction of nodes of which the label is aligned with the majority vote of the discriminative nodes in the distance-$z$ neighborhood. To extend the definition from binary classification in our formulation to multi-classification tasks, we use the average number of confusion nodes per class in the distance-$z$ neighborhood as $|\mathcal{D}\_{hash}^n\cap\mathcal{N}_z^n|$ ("hash" stands for #), the number of confusion nodes in the distance-$z$ neighborhood of node $n$. Figure [14](https://imgur.com/A4uhaZ0) shows the value of normalized $\bar{\Delta}(z)$ for $z=1,2,\cdots,12$, where $\bar{\Delta}(z)$ is divided by $|\mathcal{N}_z^n|$ to control the gap of different numbers of nodes in different neighborhoods. **The empirical result indicates that the core distances of the four datasets are all 1 given the largest $\bar{\Delta}(z)$ at $z=1$.**

 
**We also find the difference of $\bar{\Delta}(z)$ value in distance-$z$ neighborhoods with different graph properties.** (1) For the heterophilous graph Actor, the difference between different $\bar{\Delta}(z)$ is very small. (2) For the long-range graph PascalVOC-SP-1G, the value of $\bar{\Delta}(z)$ when $z=12$ is also remarkable. (3) for the homophilous graph PubMed and Cora, the value of $\bar{\Delta}(z)$ when $z=1$ is much larger than the value of $\bar{\Delta}(z)$ when $z$ is large. These show the difference between the four datasets. 

We then compute the fraction of nodes that satisfy the winning majority vote, i.e., $\Delta_n(z)>0$. **Table [6](https://imgur.com/c7giuaj) shows that at least 85% nodes of the whole graph satisfy that the majority vote wins, which verifies Assumption 1.**

Reviewer comment (Theory):
The data model assumes that the data is balanced where discriminative and non-discriminative nodes are distributed uniformly across the graph. It is unclear if this is also a reasonable assumption to make in various real-world networks.

Author response:
**Q2 ([W1] part 2 by the reviewer)**: “Similarly, the data model assumes that the data is balanced where discriminative and non-discriminative nodes are distributed uniformly across the graph. Is this also a reasonable assumption to make in various real-world networks?”

**A2**: We thank the reviewer for asking this question. After checking the paper, we noticed that the explanation that “discriminative and non-discriminative nodes are distributed uniformly across the graph” is confusing and does not accurately describe our model. We apologize for this error. **What we want to define for the “balanced dataset” is that the numbers of positive and negative LABELS in the graph are close to each other.** This serves as a  commonly used assumption in the theoretical analysis of neural networks [Shi et al., 2022, Brutzkus et al., 2022, Karp et al., 2021, Li et al., 2023] using the feature learning framework. This condition is also assumed in a recent work [Zhang et al., 2023]  in ICLR 2023 on GNN training using feature learning analysis. In our revision, we rewrite this sentence into “​​**We assume the dataset is balanced, i.e., the gap between the numbers of positive and negative labels is at most $O(\sqrt{N})$.**” Here, $N$ is the total number of nodes in the graph. We are sorry that the previous footnote here may be confusing, and we have moved it to another place.

Brutzkus et al., “An Optimization and Generalization Analysis for Max-Pooling Networks
”, UAI 2020.

Karp et al., “Local Signal Adaptivity: Provable Feature Learning in Neural Networks Beyond Kernels”, Neurips 2022

Shi et al., “A theoretical analysis on feature learning in neural networks: Emergence from inputs and advantage over fixed features”, ICLR 2022.

Li et al., “A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity”, ICLR 2023.

Zhang et al., “Joint Edge-Model Sparse Learning is Provably Efficient for Graph Neural Networks”, ICLR 2023

Reviewer comment (Theory):
The proof sketch in Subsection 4.6 is not so intuitive and hard to follow. Considering that the technical contribution of the paper is mostly theoretical, it would be great if the proof sketch stated what each lemma is arguing for intuitively, a rough picture of how each is proven, and how the Lemmas altogether culminate towards Theorem 4.1. The current proof sketch seems to be missing a significant portion that connects the lemmas to the main theorem, which makes the claims presented above somewhat unconvincing on their own.

Author response:
**Q3 ([W2] part 1 by the reviewer)**: “it would be great if the proof sketch stated what each lemma is arguing for intuitively, a rough picture of how each is proven, and how the Lemmas altogether culminate towards Theorem 4.1.”

**A3**: We apologize for the inconvenience of understanding the proof sketch and thank the reviewer for the suggestion. We have made revisions to Section 4.6 in our submission. 

**For what each lemma conveys**, we state that “By supportive Lemmas 8 and 9, we first characterize two groups of neurons that respectively activate the self-attention layer output of $\mu_1$ and $\mu_2$ nodes from initialization. Then, Lemma 4 shows that the neurons of $W_O$ in these two groups grow along the two directions of the discriminative pattern embeddings. Lemma 7 indicates that the updates of $W_V$ consist of neuron weights from these two groups. Meanwhile, Lemma 5 states that $W_Q$ and $W_K$ evolve to promote the magnitude of query and key embeddings of discriminative nodes. Lemma 6 depicts the training trajectory of the learning parameter of PE that emphasizes the core neighborhood.”

**For how each lemma is proven**, we briefly state that “we make the proof of each lemma tractable by studying gradient growth per distance-$z$ neighborhood for each $z$ rather than directly characterizing the gradient growth over the whole graph. Such a technique enables a dynamic tracking of per-parameter gradient updates. As a novel aspect, we prove Lemma 6 by showing that its most significant gradient component is proportional to the average winning margin in the core neighborhood.” in Section 4.6 due to limited space. We add a proof sketch of each lemma at the beginning of Appendix D.

**For how these lemmas culminate towards Theorem 4.1**, we state that “First, Lemma 5 and 6 collaborate to illustrate that attention weights correlated with class-relevant nodes become dominant when $\eta t=\Theta(1)$. Second, we compute the network output by Lemmas 4 and 7. By enforcing the output to be either $>1$ or $<-1$ for zero Hinge loss, we derive the sample complexity bound and the required number of iterations by concentration bounds.”

Reviewer comment (Theory):
Several questions in the proof sketch could be clarified: "Why do we use a two-phase-mini-batch SGD where we first solely train \( W_O \)?" and "How do the gradient updates of the weights \( W_O \), \( W_Q \), \( W_K \), \( W_V \), and \( b \) roughly lead to the generalization guarantees in Theorems 4.1–4.3?"

Author response:
**Q4 ([W2] part 2 by the reviewer)**: “"Why do we use a two-phase-mini-batch SGD where we first solely train W_O
?" and "How to the gradient updates of the weights W_O, W_V, W_Q, W_V, and b roughly lead to the generalization guarantees in Theorems 4.1-4.3?"”

**A4**: **For the first question on two-phase-mini-batch SGD, we have deleted such a requirement in our revision**. Thank you for asking this inspiring question. Now, we find we can improve the proof, and this condition is no longer necessary. This is because we theoretically verify that the vanilla SGD on all parameters, instead of training $W_O$ solely first, can already direct $W_O$ lucky neurons to activate the MLP layer along the training due to a $\gamma_d=\Theta(1)$ fraction of discriminative nodes in the graph (This is also the goal of the extra first phase of the two-phase training in the previous version). We characterize the updates of the lucky neurons to the desired directions in all iterations in the revised Lemma 9.

**For the second question on how the updates of weights lead to Theorem 4.1-4.3**, we have the following answer. **First, in A3, we have discussed how the conclusion in Theorem 4.1 is developed**. The main idea is to use Lemmas 5 and 6 to infer dominant attention weights on class-relevant nodes and then use Lemmas 4 and 7 to characterize the required number of samples and iterations for the desired network output of zero Hinge loss. **Second, for the reasoning of Theorem 4.2 and 4.3, we added the discussion in Section 4.6 in the revision**. We state that “When the self-attention layer weights are fixed at $0$ in Theorem 4.2, since that a given core neighborhood and a $\gamma_d=\Theta(1)$ fraction of discriminative nodes still ensures non-trivial attention weights correlated with class-relevant nodes along the training, the updates of $W_O$ and $W_V$ are order-wise the same as Lemmas 4 and 7. Then, we can apply Lemmas 4 and 7 to derive the required number of samples and iterations for zero generalization. Likewise, given a known core neighborhood in Theorem 4.3, the remaining parameters follow the same order-wise update as Lemmas 4, 5, and 7. Hence, Theorems 4.2 and 4.3 can be proved.”

Reviewer comment (Experiments):
In Figure 8, it is unclear why the "Uniform sampling" method that samples nodes across the whole graph for feature aggregation was chosen as a baseline. The reviewer does not fully understand what this baseline is hoping to convey.

Author response:
**Q5 ([Q1] by the reviewer)**:  “In Figure 8, why was the "Uniform sampling" method that samples nodes across the whole graph for feature aggregation chosen as a baseline? I don't think I fully understand what this baseline is hoping to convey.”

**A5**: We initially hope to show that the performance of the best PE-based sampling strategy is better than using uniform sampling. To avoid any confusion, we deleted this curve in our revision.

Reviewer comment (Presentation):
Also in Figure 8, it is unclear where in the plot we find that SPD-based PE "correctly reflects the homophily, heterophily, and long-range dependency of these three datasets." Is it from where the performance peak of PE-based sampling is located along the x-axis?

Author response:
**Q6 ([Q2] by the reviewer)**: “Also in Figure 8, where in the plot do we find that SPD-based PE "correctly reflects the homophily, heterophily, and long-range dependency of these three datasets"? Is it from where the performance peak of PE-based sampling is located along the x-axis?”

**A6**: Thank you for the question. We realized that our statement that “SPD-based PE also correctly reflects the homophily, heterophily, and long-range dependency of these three datasets” is misleading . Therefore, we deleted this sentence in the revision. **The main message from the performance of PE-based sampling is that (1) the red curves (PE-based sampling) in Figure 8 show different preferences on neighborhoods of the three datasets, indicating the existence of the core neighborhood (2) as $z$ increase, the trend of the red curves matches the trend of entries of the trainable positional encoding parameter $b$, indicating the success of the position encoding.**

What we meant to have additionally commented is that these three graphs have different patterns of the distance-$z$ neighborhood, and this difference might result from the fact that these graphs are  homophily, heterophily, and long-range dependency, respectively. By different patterns, we mean that the core neighborhood is $z=1$ in Figure 8 left, the impact of all distance $z$ is similar in Figure 8 middle, and both distance 1 and distance 12 have signficant impact in Figure 8 right.

Reviewer comment (Writing):
There are a number of typos that could use some additional proofreading:  
- Page 4: "We assme \( \kappa\geq 4\tau \)" → "We assume \( \kappa\geq 4\tau \)"  
- Figure 9 (x-axis title): "Fraction of labeld nodes" → "labeled nodes"  
- The legend on the Left plot of Figure 9 is different from the other two.  
- Middle plot of Figure 9 has an additional y-axis.

Author response:
**Q7**: There are a number of typos that could use some additional proofreading:

**A7**: Thank you for mentioning this problem. We have proofread the paper and made revisions to the submission. 

We would like to clarify that in Figure 9 left, the legend should be that “Red: GT with PE; Green: GT without PE; Blue: GCN”, which is the same as the other two. This shows that the conclusion that GT with PE has a better performance than that without PE and is also better than GCN holds for PubMed dataset. 

The reason that the middle plot of Figure 9 has two y-axis is that GCN performance has a large gap to the Graph Transformer performance with/without PE. If we use one y-axis, the difference between the Graph Transformer performance with and without PE will look very small. To avoid this problem, we draw two y-axis.

Reviewer comment (Experiments):
Is the proposed analytical framework applicable to node classifications for heterophily graphs? How about other graph tasks such as graph classification?

Author response:
**A3**: **The proposed analytical framework can be applied to node classification for heterophilous graphs that satisfy our proposed data assumption in Section 4.2.** To be specific, as long as the heterophilous graph has a neighrborhood that a majority vote of nodes in this neighborhood can decide each label, even if the winning margin is very small, such a heterophilous graph satisfies our data assumptions. We then verify this holds in practice and use the dataset Actor in Appendix A1 as an example. We first use the averaged low-dimensional representation (by principal component analysis) of nodes in each class to identify one discriminative pattern for each class. Table 4 indicates that high fractions of nodes are close to each discriminative pattern, which verifies the assumption on the discriminative/non-discriminative nodes. Figure 14 shows the value of normalized $\bar{\Delta}(z)$ for $z=1,2,\cdots,12$, indicating that the core distance $z_m=1$ for Actor. We then compute the fraction of nodes that satisfying the winning majority vote, i.e., $\Delta_n(z)>0$. Table 6 shows that at least $85%$ nodes of the whole graph satisfy that the majority vote wins, which verifies Assumption 1. By the findings above, we can conclude that some common practical datasets satisfy our data assumption, and the proposed analytical framework is generic. 


For the extension to graph classification tasks, **we provide a discussion in Section E.4 in the Appendix.** Our analytical framework can be applied to study the generalization of Graph Transformers on graph classification tasks, where the graph labels are determined by a non-trivial mechanism. To be more specific, we assume the existence of the core neighborhood of some central nodes in the graph to decide the ground truth graph label. This is motivated by graph classification on social networks where the connections between the central person and other people in the graph decide the graph label. **Then, we can formulate the graph classification problem as the node classification based on subgraphs.** Our technique can be used to handle this problem.

Reviewer comment (Writing):
In the introduction, the article states that "where nodes with discriminative patterns in some so-called core neighborhoods determine the labels by a majority vote, while non-discriminative nodes do not affect the labels." This sentence is a bit confusing for me. What are nodes with discriminative patterns, and how are so-called core neighborhoods defined, since a larger fraction of discriminative nodes improves generalization performance is an important contribution of this paper, perhaps a more detailed explanation or a citation is needed here.

Author response:
**A1**: We first apologize for the confusion by this sentence. The definition of discriminative patterns is introduced at the beginning of Section 4.2. We assign each node of the graph a pattern $\mu_j, j=1,2,\cdots, M$ with a noise. The nodes of $\mu_1$ or $\mu_2$ are discriminative nodes, while nodes of $\mu_j, 3\leq j\leq M$ are non-discriminative nodes. We assume the existence of a distance-z neighborhood where the ground truth label of every node is decided by such neighborhood. The node is class 1 if there are more nodes with $\mu_1$ patterns than that with $\mu_2$ patterns in the distance-z neighborhood, and the node is class 2 if there are more nodes with $\mu_2$ patterns than that with $\mu_1$ patterns. Such a neighborhood is called the core neighborhood. 

For the referred sentence in the Introduction section, to make the presentation concise, we have revised it to “**We focus on a semi-supervised binary node classification problem on structured graph data, where each node feature corresponds to either a discriminative or a non-discriminative pattern, and each ground truth node label is determined by the dominate discriminative pattern in the core neighborhood. **”

Reviewer comment (Novelty):
It seems to be common knowledge that a smaller fraction of erroneous labels, and smaller errors in the initial model mentioned in the paper affect performance, and there doesn't seem to be a strong incentive to adopt a theory to analyze the mechanism in depth.

Author response:
**A2**: Thank you for raising this question.  **We realize this statement can be confusing. What we should have meant is that this paper provides a formal theoretical characterization of how these factors affect generalization quantitatively.** We show that our sample complexity bound in equation 9 is proportional to $1/\gamma_d^2$ and $(1-\epsilon_S)^2$ where $\gamma_d$ is the fraction of discriminative nodes, and $\epsilon_S$ is the confusion rate. This indicates that the graphs with a larger fraction of discriminative nodes or a clearer-cutting vote among discriminative nodes can have a better generalization performance. When $\gamma_d$ is close to 0 and $\epsilon_S$ is close to $1/2$, i.e., the winning margin of the majority vote in the core neighborhood is small, the sample complexity bound proportional to $1/(\gamma_d-(\sigma+\tau))^2$ is large, which means the generalization is sensitive to the error in the initial model. A good pre-trained model is desired in this case. Our result also characterizes that the fraction of erroneous labels $p$ makes the sample complexity and the required number of iterations increase rapidly in an order of $(1-2p)^{-1/2}$ when $p$ is very close to $1/2$.

To address such concern, we have revised the related sentences in the abstract as “**This paper provides the quantitative characterization of the sample complexity and number of iterations for convergence dependent on the fraction of discriminative nodes, the dominant patterns, the fraction of erroneous labels, and the initial model errors.** “

Reviewer comment (Experiments):
The article focuses on the factors affecting the performance of the shallow Graph Transformer on the semi-supervised node classification task, and it is unclear whether the conclusions can be generalized to more complex frameworks, or to other downstream tasks.

Author response:
**A3**: This is a great question. **The reason why we focus on the theoretical analysis of one-layer Graph Transformers is due to the theoretical challenges of studying Transformers.** We would like to clarify that, **first, our work is the first to theoretically analyze the generalization of Graph Transformers.** As far as we know, existing works on the generalization or training dynamics of GNNs [Maskey et al., 2022, Tang & Liu 2023, Zhang et al., 2023] do not involve architectures with the attention structure.  **Second, existing works on the generalization analysis of other Transformers cannot be easily extended to study Graph Transformers.** Compared to our work, the state-of-the-art generalization analysis on other Transformers [Li et al., 2023, Tarzanagh et al., 2023, Oymak et al., 2023, Tian et al., 2023] don’t consider any graph-based labelling function and trainable positional encoding, which are crucial and necessary for node classification tasks. **Third, despite one layer, the architecture we formulate has captured major components in GTs, including trainable self-attention, MLP, and positional encoding. The resulting learning problem is already highly nonconvex.** Since we would like to study the training dynamics of the neural network, the complexity of the optimization and generalization problem scales up with the number of trainable parameters in the network. Given multiple trainable parameters, theoretically studying shallow Graph Transformers is already non-trivial. 

**Although our theoretical analysis is built on one-layer Graph Transformers, empirical results on multi-layer Graph Transformers are aligned with the theory.** Note that in recent works [Zhang et al., 2022, Wu et al., 2022, Chen et al., 2023, Cai et al., 2023, Shirzad et al., 2023, Luo et al., 2023], the number of Graph Transformer layers to learn datasets PubMed, Actor and PascalVOC-SP ranges from 2 to 6. In our empirical experiments on real datasets, we use four-layer Graph Transformers, which are approximately as deep as Graph Transformers proposed recently for node classification. 

By the way, we provide an extension to other positional encodings and multi-classification tasks in Sections E.2 and E.5 in the Appendix to show the versatility of our analytical framework.

**For the extension to other downstream tasks, we provide a discussion on graph classification in Section E.4 in the Appendix.** Our analytical framework can be applied to study the generalization of Graph Transformers on graph classification tasks. To be more specific, we assume the existence of the core neighborhood of some central nodes in each graph to decide the ground truth graph label. This is motivated by graph classification on social networks where the connections between the central person and other people in the graph decide the graph label. Then, we can apply the generalization analysis of node classification based on the core neighborhood to the graph classification problem. Our technique can be used to handle this problem.

Maskey et al., “Generalization analysis of message passing neural networks on large random graphs”, Neurips 2022.

Tang & Liu, “Towards understanding the generalization of graph neural networks”, ICML 2023.

Zhang et al., “Joint edge-model sparse learning is provably efficient for graph neural networks”, ICLR 2023.

Li et al., “A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity”, ICLR 2023.

Oymak et al., “On the role of attention in prompt-tuning”, ICML 2023.

Tarzanagh et al., “Max-margin token selection in attention mechanism”, Neurips 20
23.
Tian et al., “Scan and snap: Understanding training dynamics and token composition in 1-layer transformer”, Neurips 2023.

Zhang et al., “Hierarchical Graph Transformer with Adaptive Node Sampling”, Neurips 2022.

Wu et al., “NodeFormer: A Scalable Graph Structure Learning Transformer for Node Classification”, Neurips 2022.

Chen et al., “NAGphormer: A Tokenized Graph Transformer for Node Classification in Large Graphs”, ICLR 2023.

Cai et al., “On the Connection Between MPNN and Graph Transformer”, ICML 2023.

Shirzad et al., “EXPHORMER: Sparse Transformers for Graphs”, ICML 2023.

Luo et al., “Transformers over Directed Acyclic Graphs”, Neurips 2023.