Title: Tailoring Self-Attention for Graph via Rooted Subtrees
Venue: NeurIPS 2023
Accepted: yes
Keywords: Graph Based Learning
Abstract: Attention mechanisms have made significant strides in graph learning, yet they still exhibit notable limitations: local attention faces challenges in capturing long-range information due to the inherent problems of the message-passing scheme, while global attention cannot reflect the hierarchical neighborhood structure and fails to capture fine-grained local information. In this paper, we propose a novel multi-hop graph attention mechanism, named Subtree Attention (STA), to address the aforementioned issues. STA seamlessly bridges the fully-attentional structure and the rooted subtree, with theoretical proof that STA approximates the global attention under extreme settings. By allowing direct computation of attention weights among multi-hop neighbors, STA mitigates the inherent problems in existing graph attention mechanisms. Further we devise an efficient form for STA by employing kernelized softmax, which yields a linear time complexity. Our resulting GNN architecture, the STAGNN, presents a simple yet performant STA-based graph neural network leveraging a hop-aware attention strategy. Comprehensive evaluations on ten node classification datasets demonstrate that STA-based models outperform existing graph transformers and mainstream GNNs. The code
is available at https://github.com/LUMIA-Group/SubTree-Attention.

Tailoring Self-Attention for Graph via Rooted Subtrees

Abstract

Attention mechanisms have made significant strides in graph learning, yet they still exhibit notable limitations: local attention faces challenges in capturing long-range information due to the inherent problems of the message-passing scheme, while global attention cannot reflect the hierarchical neighborhood structure and fails to capture fine-grained local information. In this paper, we propose a novel multihop graph attention mechanism, named Subtree Attention (STA), to address the aforementioned issues. STA seamlessly bridges the fully-attentional structure and the rooted subtree, with theoretical proof that STA approximates the global attention under extreme settings. By allowing direct computation of attention weights among multi-hop neighbors, STA mitigates the inherent problems in existing graph attention mechanisms. Further we devise an efficient form for STA by employing kernelized softmax, which yields a linear time complexity. Our resulting GNN architecture, the STAGNN, presents a simple yet performant STA-based graph neural network leveraging a hop-aware attention strategy. Comprehensive evaluations on ten node classification datasets demonstrate that STA-based models outperform existing graph transformers and mainstream GNNs. The code is available at https://github.com/LUMIA-Group/SubTree-Attention .

1 Introduction

Graph Neural Networks (GNNs) have achieved remarkable performance in various tasks, such as drug discovery [35, 15], social networks [29, 31], traffic flow [11], and recommendation systems [44]. Most GNNs are based on the message-passing scheme [16], hierarchically aggregating information from multi-hop neighbors by stacking multiple layers. During this procedure, a rooted subtree can be generated for each node, representing the node's neighborhood structure. Nodes with similar neighborhood structures possess similar subtrees, which leads to similar node representations [13].

Local attention in graph learning can be seen as a natural combination of the message-passing scheme and the self-attention mechanism. By adaptively assigning weights among one-hop neighbors in a single layer, the local attention mechanism allows each node to focus on the most task-relevant neighbors [39, 36, 40, 3]. However, local attention limits the receptive field to one-hop neighbors. While stacking multiple local-attention layers to build a deep model can increase the receptive field, such message-passing-based deep architectures face challenges in capturing long-range dependencies [10] due to issues such as over-smoothing [5] and over-squashing [2].

On the other hand, the global attention mechanism originated from vanilla Transformer [38] has been widely adopted in graph learning domain [23, 12, 48, 28, 30], leveraging the fully-attentional architecture to address the aforementioned issues. However, the global attention employed by graph transformers cannot reflect the hierarchical neighborhood structure and fails to capture fine-grained local information, which is crucial in many real-world scenarios [25, 17, 46, 49]. To mitigate this deficiency, recent studies try to directly assemble global attention and message-passing-based

models by combining GNNs and Transformers, including simultaneously applying GNNs and Transformers [50] or building Transformers on top of GNNs [33, 27, 45].

Considering the limitations of both local and global attention, we propose a multi-hop graph attention mechanism, termed SubTree Attention (STA). It allows the root node to directly attend to further neighbors in the subtree, enabling the root node to gather information from the entire rooted subtree within one layer. It provides two main advantages: (i) Compared to deep architectures with multiple local attention layers, STA avoids issues associated with the message-passing scheme such as oversmoothing and over-squashing. (ii) Compared to global attention, STA can hierarchically capture the neighborhood structure by enabling each node to focus on its own rooted subtree.

Due to the exponential growth in neighborhood size with increased hops, directly calculating attention among multi-hop neighbors becomes impractical. Meanwhile, powers of the adjacency matrix have to be stored for the calculation among multi-hop neighbors. To address these issues, we employ kernelized softmax [19] to develop an algorithm that reduces the quadratic time complexity to linear while avoiding the need to store the powers of the adjacency matrix. This efficient algorithm can be viewed as keys and values performing a random walk on the graph and eventually landing on queries.

Furthermore, we provide a theoretical analysis of STA, demonstrating that under extreme settings, STA converges to the global self-attention. As a result, STA can be regarded as a bridge between local and global attention, effectively combining the benefits of both approaches. In addition, we present the STA module with multiple attention heads. We introduce a hop-wise gating mechanism, enabling attention heads to be specialized in capturing information from specific hops.

We then propose a simple yet performant multi-hop graph attention network, named STAGNN, which can leverage multi-hop information and in the meantime acts as a fully-attentional model. As for the evaluation, we test the performance of STAGNN on ten common node classification datasets. Despite its relatively simple architecture, STAGNN consistently outperforms existing GNNs and graph transformers. Furthermore, we demonstrate that STAGNN maintains a competitive performance with an extremely deep architecture. Additional ablation studies are conducted to show the effectiveness of subtree attention even in the presence of global attention.

2 Background

Let G = ( N , E ) be an undirected graph, with the associated nodes set N and edges set E . We use N = |N| to represent the number of nodes. X ∈ R N × f denotes the node feature, where f denotes the number of features. Let A be the adjacency matrix of G and let D be the diagonal degree matrix of A . A sym = D -1 / 2 AD -1 / 2 denotes the symmetric normalized adjacency matrix, while A rw = AD -1 denotes the random walk matrix. Let ˆ A be an arbitrary transition matrix, including A rw, A sym or other matrices representing message propagation. We use M i : and M : j to indicate the i th row and the j th column of the matrix M , respectively. And let [ [0 , K ] ] denote the set { 0 , 1 , . . . , K } .

2.1 Multi-Hop Representations

Many existing GNNs and diffusion-based models perform multi-hop message-passing in a single layer, taking advantage of multi-hop representations [8, 1, 47, 22, 51]. Among them, decoupled GCN [42] is a typical representative. One important reason for the over-smoothing problem in GCN is that neighborhood aggregation and feature transformation are coupled [4]. To address this issue, decoupled GCNs perform feature transformation and neighborhood aggregation, respectively. A general form of decoupled GCN can be described as follows [8]:

where { α k } k ∈ [ [0 ,K ] ] are the aggregation weights for different hops. K represents the number of propagation steps, which also corresponds to the height of the resulting rooted subtree. Although there exist various strategies to assign weights to different hops, these methods all apply a K -step propagation with transition matrix ˆ A , which inevitably results in over-smoothing when K is large.

There are relatively few attention-based methods that leverage multi-hop information. As a representative, Wang et al. (2021) [41] generalized GAT by Personalized PageRank [22]. Yet, this strategy

only employs attention scores from one-hop neighbors and necessarily acts as a low-frequency filter. In a recent work, Chen et al. (2022) [6] began with aggregating K -hop representations and subsequently processed them as a sequence of length K using a Transformer model. Despite the decent performance, this approach still adopts a K -step propagation to capture multi-hop information before taking the attention mechanism into account, which leads to the aforementioned issues. This observation inspires us to incorporate the attention mechanism into the propagation phase, rather than adhering to the K -step propagation with transition matrix ˆ A .

2.2 Global Self-Attention and Kernelized Softmax

In graph learning domain, the Global Self-Attention function SA ( · , · , · ) computes a weighted sum of all positions for each node. It first projects the node feature matrix into three subspaces:

where W Q ∈ R d × d K , W K ∈ R d × d K and W V ∈ R d × d V are learnable projection matrices. Then the new representation of the i th node is computed as follows:

where sim ( · , · ) : R d × R d → R is a function used to evaluate the similarity between queries and keys. A common form of self-attention is called softmax attention, which applies the exponential of the dot

In fact, we can use an arbitrary positive-definite kernel κ to serve as sim ( · , · ) . Given a selected kernel κ and its corresponding feature map ϕ , we can rewrite sim ( · , · ) as: sim ( Q i : , K j : ) = ϕ ( Q i : ) ϕ ( K j : ) T . Thus, Equation 3 becomes:

There are many potential choices for the feature map ϕ . e.g., Tsai et al. (2019) [37] verified that RBF kernels perform on par with exponential kernels on neural machine translation and sequence prediction, and Choromanski et al. (2021) [9] opted for Positive Random Features (PRF).

The key advantage of Equation 4 is that all nodes share two identical summations ∑ N j =1 ϕ ( K j : ) T V j : and ∑ N j =1 ϕ ( K j : ) T , which only need to be computed once. By doing so, we can avoid computing the full attention matrix { sim ( Q i : , K j : ) } i ∈ N,j ∈ N and reduce the complexity from O ( N 2 ) to O ( N ) .

3 The Proposed Attention Mechanism: SubTree Attention

In this section, we present a detailed introduction to an efficient multi-hop attention mechanism called S ub T ree A ttention (STA). First, we give the definition of STA, followed by an efficient algorithm for computing STA based on kernelized softmax and the message-passing scheme. We then explain how multi-head STA makes attention heads hop-aware by incorporating a gate into each hop. Finally, we prove that STA approximates the global self-attention when the height of the subtree is O (log N ) .

3.1 SubTree Attention

In this subsection, we give the definition of our proposed multi-hop graph attention mechanism named subtree attention. Similar to the global self-attention function SA ( · , · , · ) , subtree attention takes queries, keys, and values as inputs and outputs new values.

We first define the method for computing attention weights among the k th hop neighbors, which we refer to as STA k . For the i th node, this process can be described as follows:

(a) SubTree Attention

STA k ( Q , K , V ) i : represents the i th node attending to its k th hop neighbors, which are also the k th level of its rooted subtree. SubTree Attention STA : ( R n × d k , R n × d k , R n × d v ) → R n × d o can then be calculated by aggregating the results from all levels of the rooted subtree:

AGGR can be any aggregation function, such as sum, concat [18], attention-based readout [6], or the GPR-like aggregation [8] which we present in section 4.

In simple terms, STA k allows each node to attend to the k th level of its own subtree. Meanwhile, STA aggregates the results from all levels to gather information from the entire subtree. Figure 1 shows the process of STA.

In contrast to global attention, subtree attention enables each node to compute attention weights on its own subtree, thus taking advantage of the key insight of Message-Passing GNNs (MP-GNNs). e.g., even if the majority of nodes within the receptive fields are the same for two different nodes, they will still gather different information through subtree attention if they have different subtree structures.

3.2 An Efficient Algorithm for SubTree Attention

Equation 5 gives a straightforward method to calculate STA k : we first compute the complete similarity matrix { sim ( Q i : , K j : ) } i ∈ N,j ∈ N and then calculate its Hadamard product with ˆ A k to obtain { ˆ A k ij sim ( Q i : , K j : ) } i ∈ N,j ∈ N . In short, we treat ˆ A k as a mask for the similarity matrix. This algorithm exhibits two primary disadvantages: (i) The computational cost associated with calculating the entire similarity matrix is O ( N 2 ) . (ii) ˆ A k quickly converges to a dense matrix. Storing ˆ A k in the GPU memory for computing the Hadamard product leads to considerable space complexity. Considering these two limitations, utilizing ˆ A k as a mask for the similarity matrix is suboptimal.

We now present an efficient algorithm for subtree attention. Considering the close relationship between rooted subtrees and MP-GNNs, we leverage the message-passing scheme to implement the computation of subtree attention. By permitting keys and values to propagate along the edges, we can achieve an algorithm that has linear time complexity and avoids the need to store ˆ A k .

Learning from Equation 4, we use a feature map ϕ to replace sim ( · , · ) . The choice of feature map is not the main focus of our work. Our model adopts a simple yet effective approach proposed by Katharopoulos et al. (2020) [19] that chooses ϕ ( x ) = elu ( x ) +1 as the feature map and demonstrates empirically that it performs on par with softmax attention. We can rewrite Equation 5 as follows:

Note that there are two summations ∑ N j =1 ˆ A k ij ϕ ( K j : ) T V j : and ∑ N j =1 ˆ A k ij ϕ ( K j : ) T in Equation 7. We can think of these two summations as a kind of message propagation . That is to say, we first compute ϕ ( K i : ) and ϕ ( K i : ) T V i : for each node. Then we let ϕ ( K i : ) and ϕ ( K i : ) T V i : undergo k steps message passing. Finally, we use the aggregated keys and values ∑ N j =1 ˆ A k ij ϕ ( K j : ) T V j : and ∑ N j =1 ˆ A k ij ϕ ( K j : ) T in conjunction with the node's own query ϕ ( Q i : ) to complete the computation of subtree attention. Figure 1 illustrates the whole process of this efficient algorithm. When we choose A rw as the transition matrix, this process can be regarded as keys and values performing a random walk on the graph, eventually landing on different queries. Note that message passing occurs on each edge, thus reducing the computational cost from O ( N 2 ) to O ( |E| ) . Additionally, message passing only requires the sparse adjacency matrix A , thereby circumventing the need to store ˆ A k . Furthermore, { STA i } i ∈ [ [1 ,K ] ] can be viewed as a nested process, calculated one after another.

Based on this algorithm, we can say that STA is an attempt to incorporate the message-passing scheme into the fully-attentional architecture. In fact, STA serves as a message-passing module for keys and values. In section 4, we design a novel multi-hop graph attention network employing STA for message propagation. We also provide a detailed complexity analysis in Appendix A.

3.3 SubTree Attention with Multiple Heads

In this subsection, we present STA with multiple attention heads. Kim et al. (2022) [20] discovered empirically that different attention heads tend to concentrate on neighbors at different hops. Certain attention heads can attend to remote nodes, while others consistently focus on nearby nodes, suggesting that attention heads can be specialized in capturing information from specific hops. To make better use of multiple attention heads in this context, we propose a hop-aware method of mixing them.

Suppose there are a total of H attention heads. STA with multiple attention heads, noted as MSTA, can be described as follows:

where [ ] denotes row-wise concatenation. Q h , K h , and V h represent the query, key, and value matrices for the h th head, respectively. W O denotes a linear projection matrix. g k ∈ R H is an H -dimensional vector and g h k is its h th element, representing the weight of the h th attention head at the k th hop. Compared to STA with a single attention head, we introduce in total H × K additional learnable parameters: { g i } i ∈ [ [1 ,K ] ] . We can regard g k as a hop-wise gate that determines the weight of each attention head at the k th hop.

In other words, we can reconsider the multi-hop attention mechanism in terms of multi-task learning. Different attention heads are seen as different experts, while aggregating information from different hops is seen as different task. g k signifies the process of selecting appropriate experts for each task.

3.4 Theoretical Analysis of SubTree Attention

MP-GNNs suffer from issues like over-smoothing or over-squashing when the height of the subtree increases. In this subsection, we theoretically demonstrate that STA avoids the issues associated with the message-passing scheme despite employing the same rooted subtree as MP-GNNs. Notice that we employ the random walk matrix A rw as the transition matrix in the STA module, i.e. , we have ˆ A = A rw in this subsection.

We first employ a slightly modified approach to rewrite the global self-attention module SA, which can be described as:

where π i = d ( i ) ∑ N j =1 d ( j ) and d ( j ) denotes the degree of the j th node. Note that Equation 9 is consistent with Equation 3.

Comparing Equation 5 and Equation 9, we find that the only difference between SA and STA k lies in the mask for the similarity matrix. STA k employs ˆ A k ij as the mask, whereas SA employs π i . We now demonstrate that when the height K of the rooted subtree takes the same order of magnitude as log( N ) , STA k can be approximately considered as SA.

Theorem 1 Let ˆ A ∈ R N × N denote the random walk matrix of a connected and non-bipartite graph, i.e. , ˆ A = AD -1 . Let 1 = λ 1 ≥ · · · ≥ λ N be the eigenvalues of ˆ A . Then we have the following results:

And for a given ϵ , the smallest K 0 that satisfies the condition shown in Equation 10 is at most O ( log N ϵ 1 -max { λ 2 , | λ n |} ) . If we had more information about V , we could specify the convergence rate of STA k . e.g., if V is computed by V = σ ( XW V ) where σ is a non-negative activation function, then:

holds true when none of the denominators is equal to zero. And for a given η , the smallest K 1 that satisfies the condition shown in Equation 11 is at most O ( log N η 1 -max { λ 2 , | λ n |} ) .

Equation 10 demonstrates that ˆ A k ij converges to π i with logarithmic complexity, indicating that under general conditions, STA k quickly tends to SA. Notice that although STA k is a multi-hop graph attention module implemented using the message-passing mechanism, it displays the characteristics of global attention when the height of the subtree is O (log N ) . This property prevents STA from the issues associated with the message-passing scheme such as over-smoothing and over-squashing. From this perspective, subtree attention serves as a bridge connecting local and global attention. STA 1 plays the role of local attention, while STA O (log N ) acts as the global self-attention. Subtree attention achieves a hierarchical attention computation by a hop-wise aggregation of { STA k } k ∈ [ [1 , O (log N )] ] . A detailed proof of Theorem 1 is provided in Appendix B.

4 The Proposed Multi-Hop Graph Attention Network: STAGNN

In this section, we present a simple yet effective multi-hop graph attention network, named STAGNN. This model is built upon decoupled GCN, but employs STA as the message-passing module instead. STAGNN can be divided into two steps: first, we apply MLP to compute queries, keys, and values for each node; then, we use STA to propagate information. Formally, it can be described as:

86

87

88

89

90

91

70

72

74

76

78

CiteSeer

60

65

70

75

ROC-AUC

Deezer

25

30

35

40

Actor

GCN

GAT

MixHop

JKNet

LDS

IDGL

NodeFormer

STAGNN

where X is the input node feature and O is the learned representation for each node. We adopt the GPR-like aggregation [8] for HopAggregation in STA. To be precise, we assign learnable parameters { α k } k ∈ [ [0 ,K ] ] to each hop (initialized simply to 1), and then the nodes aggregate information from each hop based on these learned weights. Figure 2 shows the overall architecture of STAGNN.

Comparing Equation 1 and Equation 12, the only difference between decoupled GCN and STAGNN lies in the propagation method. Decoupled GCN relies on high powers of the normalized adjacency matrix to capture long-range information, which inevitably results in over-smoothing. In contrast, STAGNN utilizes subtree attention for message propagation, effectively learning more informative representations from multi-hop neighbors without suffering from the inherent problems associated with the message-passing scheme.

Wu et al. (2022) [43] have drawn attention to an issue named over-normalization : In the context of graphs with a large volume of nodes, the use of the global attention module may lead to a situation where the attention scores for the majority of nodes are nearly zero and thus resulting in gradient vanishing. Subtree attention can alleviate this problem by providing a hierarchical calculation focusing on each level of the rooted subtree instead of the whole graph.

5 Evaluation

We evaluate the performance of STAGNN on ten common node classification datasets, with detailed dataset information provided in Appendix C. We then verify the performance of STAGNN under extreme settings, empirically showing its capacity to tackle over-smoothing. Furthermore, We conduct an experiment that confirms the necessity of subtree attention even in the presence of global attention. Additional ablation studies are conducted for further discussions. For implementation, we fix the number of hidden channels at 64. More implementation details are presented in Appendix D. All experiments are conducted on an NVIDIA RTX4090 with 24 GB memory.

3

4

5

6

7

8

9

10

20 30 50 75100

Propagation steps

K

87.5

88.0

88.5

89.0

89.5

90.0

3

4

5

6

7

8

9

10

20 30 50 75100

Propagation steps

K

36.5

37.0

37.5

38.0

38.5

39.0

Actor

0

20

40

60

80

100

Hop

0

1

2

3

GPR weights when

K

=100

Actor

.

5.1 Experiments on Node Classification

To compare STAGNN with a wide variety of baselines, we select two recent works [43, 6] and follow their experiment settings and their choice of baselines. The metrics of the baselines are adopted from these works [43, 6]. The code for reproduction can be found in supplementary materials.

Comparison with Multi-Hop GNNs and Structure Learning Methods For a fair comparison, we strictly follow the experiment settings from Wu et al. (2022) [43]. We test the performance of STAGNNonfour common datasets: Cora, Citeseer, Deezer, and Actor. The first two are homogeneous graphs, while the latter two are heterophilic [24, 52, 34]. We compare STAGNN with two mainstream GNNs: GCN [21], GAT [39], two multi-hop GNNs: JKNet [47], MixHop [1], and three graph structure learning methods: LDS [14], IDGL [7], NodeFormer [43]. We apply the same random splits with train/valid/test ratios of 50%/25%/25% as [43]. Further details can be found in Appendix D.1. Figure 3 displays the experimental results, showing that STAGNN, which features a relatively simple model architecture, outperforms all the baselines on all four datasets. This result highlights the effectiveness of STAGNN in managing both homogeneous and heterophilic graphs.

Comparison with Decoupled GCNs and Graph Transformers For a fair comparison, we strictly follow the experiment settings from Chen et al. (2022) [6]. We test the performance of STAGNN on six common datasets: Pubmed, Corafull, Computer, Photo, CS and Physics. More information about these datasets can be found in Appendix C. We compare STAGNN with two mainstream GNNs: GCN[21], GAT [39], two decoupled GCNs: APPNP [22], GPRGNN [8], and five graph transformers: GT [12], Graphormer [48], SAN [23], GraphGPS [32], NAGphormer [6]. We apply the same random splits with train/valid/test ratios of 60%/20%/20% as [6]. The experimental results are shown in Table 1. STAGNN shows comparable or superior performance compared to all the baselines, which highlights the competitiveness of STAGNN when compared to existing GNNs and graph transformers.

5.2 Experiments on Deep STAGNN

We evaluate the performance of deep STAGNN on Cora and Actor, with the height of the subtree ranging from 3 to 100. The experimental results are presented in Figure 4. In contrast to MP-GNNs, STAGNN maintains robust performance even when the height of the subtree reaches 100. For Cora, the accuracy of STAGNN peaks at K = 10 , demonstrating its ability to effectively collect information from a large receptive field. We further visualize the GPR weights of STAGNN when the height of the subtree is set to K = 100 , and we observe distinct characteristics for Cora and Actor. In the case of Cora, the GPR weights exhibit a monotonic decrease, which aligns with the witnessed performance drop on Cora as the height of the subtree increases. Therefore, for Cora, we may consider keeping the height of the subtree within a reasonable range. In the case of Actor, the GPR weights eventually stabilize at a value close to 1. This finding suggests that the limiting state of STA k , i.e., SA π , is suitable for Actor, which is confirmed by the robust performance of deep STAGNN on Actor. More visualizations of GPR weights can be found in Appendix E. In summary, STAGNN can achieve impressive results even with an extremely deep architecture.

86

87

88

89

90

91

70

72

74

76

78

CiteSeer

60

65

70

75

ROC-AUC

Deezer

25

30

35

40

Actor

GAT

STAGNN-SUM

STAGNN-CONCAT

STAGNN-ATTN

STAGNN-GPR

5.3 Ablation Study

Study on the Necessity of SubTree Attention in the Presence of Global Attention We demonstrate that STA k converges to the global self-attention, which naturally leads to a critical question:

Is it still necessary to employ subtree attention when global attention is already present?

In this experiment, we seek to answer the question above by exploring two scenarios: (i) applying global attention independently and (ii) combining global attention with subtree attention. Formally, we extend STAGNN by replacing the STA module with global attention enhanced by 0, 1, 2, or 3 hop/hops of subtree attention. More details on the experiment settings can be found in Appendix D.2.

The experimental results are presented in Table 2. We find that incorporating subtree attention as an auxiliary to global attention significantly improves overall performance, implying that it is still necessary to employ subtree attention to capture the neighborhood structure even in the presence of global attention. This observation also inspires us that leveraging subtree attention to augment global attention could be an interesting approach for enhancing Graph Transformers.

Study on HopAggregation methods In this experiment, we investigate the influence of different HopAggregation functions within STAGNN. We compare GPR-like aggregation with sum, concat [18], and attention-based readout [6]. Details of the experiment settings can be found in Appendix D.3. The experimental results are shown in Figure 5. GPR-like aggregation outperforms the alternatives on all four datasets. This observation highlights the importance of enabling nodes to adaptively learn the weight of each hop, which is the main advantage of GPR-like aggregation.

6 Conclusion

We propose a novel multi-hop graph attention mechanism called Subtree Attention (STA), which effectively addresses the limitations of local and global attention in graph learning. The proposed STA mechanism bridges the gap between local and global attention, hierarchically capturing neighborhood structures while addressing issues associated with the message-passing scheme. The key algorithm for computing subtree attention, utilizing kernelized softmax and the message-passing scheme, reduces the computational cost from O ( N 2 ) to O ( |E| ) while avoiding the need to store high powers of the adjacency matrix. This process can be approximately viewed as keys and values performing a random walk on the graph. We further prove theoretically that under extreme settings, STA approximates the global self-attention. Based on the STA module, we propose a simple yet performant multi-hop graph attention network, coined STAGNN. Comprehensive evaluations on various node classification datasets demonstrate that STAGNN outperforms mainstream GNNs and graph transformers. Ablation studies further verify the effectiveness of subtree attention, even in the presence of global attention.

Current Limitations, Potential Impacts and Further Discussions In the present work, we mainly evaluate STAGNN, a novel multi-hop graph attention network incorporating STA. However, there are many other potential applications of STA, including combining STA with other graph learning methods that utilize self-attention mechanisms or supplementing global attention with subtree attention to enhance graph transformers. Furthermore, evaluating the robustness of STA and digging into its interpretability can also be part of future works. We provide a detailed discussion of potential impacts in Appendix G. And further analysis of the gate mechanism within the mixture of attention heads can be found in Appendix F.

References

S. Abu-El-Haija, B. Perozzi, A. Kapoor, N. Alipourfard, K. Lerman, H. Harutyunyan, G. V. Steeg, and A. Galstyan. Mixhop: Higher-order graph convolutional architectures via sparsified neighborhood mixing. In K. Chaudhuri and R. Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA , volume 97 of Proceedings of Machine Learning Research , pages 21-29. PMLR, 2019.

U. Alon and E. Yahav. On the bottleneck of graph neural networks and its practical implications. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021 . OpenReview.net, 2021.

S. Brody, U. Alon, and E. Yahav. How attentive are graph attention networks? In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022 . OpenReview.net, 2022.

D. Chen, Y. Lin, W. Li, P. Li, J. Zhou, and X. Sun. Measuring and relieving the over-smoothing problem for graph neural networks from the topological view. CoRR , abs/1909.03211, 2019.

D. Chen, Y. Lin, W. Li, P. Li, J. Zhou, and X. Sun. Measuring and relieving the over-smoothing problem for graph neural networks from the topological view. In The Thirty-Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, The Thirty-Second Innovative Applications of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2020, New York, NY, USA, February 7-12, 2020 , pages 3438-3445. AAAI Press, 2020.

J. Chen, K. Gao, G. Li, and K. He. Nagphormer: Neighborhood aggregation graph transformer for node classification in large graphs. CoRR , abs/2206.04910, 2022.

Y. Chen, L. Wu, and M. J. Zaki. Iterative deep graph learning for graph neural networks: Better and robust node embeddings. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual , 2020.

E. Chien, J. Peng, P. Li, and O. Milenkovic. Adaptive universal generalized pagerank graph neural network. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021 . OpenReview.net, 2021.

K. M. Choromanski, V. Likhosherstov, D. Dohan, X. Song, A. Gane, T. Sarlós, P. Hawkins, J. Q. Davis, A. Mohiuddin, L. Kaiser, D. B. Belanger, L. J. Colwell, and A. Weller. Rethinking attention with performers. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021 . OpenReview.net, 2021.

H. Dai, Z. Kozareva, B. Dai, A. J. Smola, and L. Song. Learning steady-states of iterative algorithms over graphs. In J. G. Dy and A. Krause, editors, Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässan, Stockholm, Sweden, July 10-15, 2018 , volume 80 of Proceedings of Machine Learning Research , pages 1114-1122. PMLR, 2018.

A. Derrow-Pinion, J. She, D. Wong, O. Lange, T. Hester, L. Perez, M. Nunkesser, S. Lee, X. Guo, P. W. Battaglia, V. Gupta, A. Li, Z. Xu, A. Sanchez-Gonzalez, Y. Li, and P. Veliˇ ckovi´ c. Traffic Prediction with Graph Neural Networks in Google Maps. 2021.

V. P. Dwivedi and X. Bresson. A generalization of transformer networks to graphs. CoRR , abs/2012.09699, 2020.

J. Feng, Y. Chen, F. Li, A. Sarkar, and M. Zhang. How powerful are k-hop message passing graph neural networks. In NeurIPS , 2022.

L. Franceschi, M. Niepert, M. Pontil, and X. He. Learning discrete structures for graph neural networks. In K. Chaudhuri and R. Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA , volume 97 of Proceedings of Machine Learning Research , pages 1972-1982. PMLR, 2019.

T. Gaudelet, B. Day, A. R. Jamasb, J. Soman, C. Regep, G. Liu, J. B. Hayter, R. Vickers, C. Roberts, J. Tang, et al. Utilising graph machine learning within drug discovery and development. arXiv preprint arXiv:2012.05716 , 2020.

J. Gilmer, S. S. Schoenholz, P. F. Riley, O. Vinyals, and G. E. Dahl. Neural message passing for quantum chemistry. In D. Precup and Y. W. Teh, editors, Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017 , volume 70 of Proceedings of Machine Learning Research , pages 1263-1272. PMLR, 2017.

A. Gulati, J. Qin, C. Chiu, N. Parmar, Y. Zhang, J. Yu, W. Han, S. Wang, Z. Zhang, Y. Wu, and R. Pang. Conformer: Convolution-augmented transformer for speech recognition. In H. Meng, B. Xu, and T. F. Zheng, editors, Interspeech 2020, 21st Annual Conference of the International Speech Communication Association, Virtual Event, Shanghai, China, 25-29 October 2020 , pages 5036-5040. ISCA, 2020.

W. L. Hamilton, Z. Ying, and J. Leskovec. Inductive representation learning on large graphs. In I. Guyon, U. von Luxburg, S. Bengio, H. M. Wallach, R. Fergus, S. V. N. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA , pages 1024-1034, 2017.

A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event , volume 119 of Proceedings of Machine Learning Research , pages 5156-5165. PMLR, 2020.

J. Kim, D. Nguyen, S. Min, S. Cho, M. Lee, H. Lee, and S. Hong. Pure transformers are powerful graph learners. In NeurIPS , 2022.

T. N. Kipf and M. Welling. Semi-supervised classification with graph convolutional networks. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings . OpenReview.net, 2017.

J. Klicpera, A. Bojchevski, and S. Günnemann. Predict then propagate: Graph neural networks meet personalized pagerank. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019 . OpenReview.net, 2019.

D. Kreuzer, D. Beaini, W. L. Hamilton, V. Létourneau, and P. Tossou. Rethinking graph transformers with spectral attention. In M. Ranzato, A. Beygelzimer, Y. N. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual , pages 21618-21629, 2021.

D. Lim, X. Li, F. Hohne, and S. Lim. New benchmarks for learning on non-homophilous graphs. CoRR , abs/2104.01404, 2021.

K. Lin, L. Wang, and Z. Liu. Mesh graphormer. In 2021 IEEE/CVF International Conference on Computer Vision, ICCV 2021, Montreal, QC, Canada, October 10-17, 2021 , pages 12919-12928. IEEE, 2021.

L. László, L. Lov, and O. Erdos. Random walks on graphs: A survey. pages 1-46, 01 1996.

G. Mialon, D. Chen, M. Selosse, and J. Mairal. Graphit: Encoding graph structure in transformers. CoRR , abs/2106.05667, 2021.

E. Min, R. Chen, Y. Bian, T. Xu, K. Zhao, W. Huang, P. Zhao, J. Huang, S. Ananiadou, and Y. Rong. Transformer for graphs: An overview from architecture perspective. CoRR , abs/2202.08455, 2022.

F. Monti, F. Frasca, D. Eynard, D. Mannion, and M. M. Bronstein. Fake news detection on social media using geometric deep learning. arXiv preprint arXiv:1902.06673 , 2019.

L. Müller, M. Galkin, C. Morris, and L. Rampásek. Attending to graph transformers. CoRR , abs/2302.04181, 2023.

A. Pal, C. Eksombatchai, Y. Zhou, B. Zhao, C. Rosenberg, and J. Leskovec. Pinnersage: Multi-modal user embedding framework for recommendations at pinterest. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining , pages 2311-2320, 2020.

L. Rampásek, M. Galkin, V. P. Dwivedi, A. T. Luu, G. Wolf, and D. Beaini. Recipe for a general, powerful, scalable graph transformer. In NeurIPS , 2022.

Y. Rong, Y. Bian, T. Xu, W. Xie, Y. Wei, W. Huang, and J. Huang. Self-supervised graph transformer on large-scale molecular data. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual , 2020.

P. Sen, G. Namata, M. Bilgic, L. Getoor, B. Gallagher, and T. Eliassi-Rad. Collective classification in network data. AI Mag. , 29(3):93-106, 2008.

J. M. Stokes, K. Yang, K. Swanson, W. Jin, A. Cubillos-Ruiz, N. M. Donghia, C. R. MacNair, S. French, L. A. Carfrae, Z. Bloom-Ackermann, et al. A deep learning approach to antibiotic discovery. Cell , 180(4):688-702, 2020.

K. K. Thekumparampil, C. Wang, S. Oh, and L. Li. Attention-based graph neural network for semisupervised learning. CoRR , abs/1803.03735, 2018.

Y. H. Tsai, S. Bai, M. Yamada, L. Morency, and R. Salakhutdinov. Transformer dissection: An unified understanding for transformer's attention via the lens of kernel. In K. Inui, J. Jiang, V . Ng, and X. Wan, editors, Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing, EMNLP-IJCNLP 2019, Hong Kong, China, November 3-7, 2019 , pages 4343-4352. Association for Computational Linguistics, 2019.

A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. In I. Guyon, U. von Luxburg, S. Bengio, H. M. Wallach, R. Fergus, S. V. N. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA , pages 5998-6008, 2017.

P. Velickovic, G. Cucurull, A. Casanova, A. Romero, P. Liò, and Y. Bengio. Graph attention networks. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Conference Track Proceedings . OpenReview.net, 2018.

G. Wang, R. Ying, J. Huang, and J. Leskovec. Improving graph attention networks with large margin-based constraints. CoRR , abs/1910.11945, 2019.

G. Wang, R. Ying, J. Huang, and J. Leskovec. Multi-hop attention graph neural networks. In Z. Zhou, editor, Proceedings of the Thirtieth International Joint Conference on Artificial Intelligence, IJCAI 2021, Virtual Event / Montreal, Canada, 19-27 August 2021 , pages 3089-3096. ijcai.org, 2021.

F. Wu, A. H. S. Jr., T. Zhang, C. Fifty, T. Yu, and K. Q. Weinberger. Simplifying graph convolutional networks. In K. Chaudhuri and R. Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA , volume 97 of Proceedings of Machine Learning Research , pages 6861-6871. PMLR, 2019.

Q. Wu, W. Zhao, Z. Li, D. P. Wipf, and J. Yan. Nodeformer: A scalable graph structure learning transformer for node classification. In NeurIPS , 2022.

S. Wu, F. Sun, W. Zhang, X. Xie, and B. Cui. Graph neural networks in recommender systems: A survey. ACM Comput. Surv. , 55(5):97:1-97:37, 2023.

Z. Wu, P. Jain, M. A. Wright, A. Mirhoseini, J. E. Gonzalez, and I. Stoica. Representing long-range context for graph neural networks with global attention. In M. Ranzato, A. Beygelzimer, Y. N. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual , pages 13266-13279, 2021.

Z. Wu, Z. Liu, J. Lin, Y. Lin, and S. Han. Lite transformer with long-short range attention. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020 . OpenReview.net, 2020.

K. Xu, C. Li, Y. Tian, T. Sonobe, K. Kawarabayashi, and S. Jegelka. Representation learning on graphs with jumping knowledge networks. In J. G. Dy and A. Krause, editors, Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässan, Stockholm, Sweden, July 10-15, 2018 , volume 80 of Proceedings of Machine Learning Research , pages 5449-5458. PMLR, 2018.

C. Ying, T. Cai, S. Luo, S. Zheng, G. Ke, D. He, Y. Shen, and T. Liu. Do transformers really perform badly for graph representation? In M. Ranzato, A. Beygelzimer, Y. N. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual , pages 28877-28888, 2021.

A. W. Yu, D. Dohan, M. Luong, R. Zhao, K. Chen, M. Norouzi, and Q. V. Le. Qanet: Combining local convolution with global self-attention for reading comprehension. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 - May 3, 2018, Conference Track Proceedings . OpenReview.net, 2018.

J. Zhang, H. Zhang, C. Xia, and L. Sun. Graph-bert: Only attention is needed for learning graph representations. CoRR , abs/2001.05140, 2020.

J. Zhao, Y. Dong, M. Ding, E. Kharlamov, and J. Tang. Adaptive diffusion in graph neural networks. In M. Ranzato, A. Beygelzimer, Y. N. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual , pages 23321-23333, 2021.

J. Zhu, Y. Yan, L. Zhao, M. Heimann, L. Akoglu, and D. Koutra. Beyond homophily in graph neural networks: Current limitations and effective designs. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual , 2020.

A Complexity Analysis of STA

In this subsection, we analyze the time complexity of SubTree Attention (STA). STA has two key components: the feature map and the HopAggregation function. Both of these components offer a wide range of potential options for consideration. Different options will affect the time complexity of STA. In the following analysis, we will adopt the configuration used by STAGNN, i.e., we choose ϕ ( x ) = elu ( x ) + 1 as the feature map and use GPR-like aggregation as the HopAggregation function.

The computation of STA can be seen as an aggregation of { STA i } i ∈ [ [1 ,K ] ] , which refer to the attentionbased aggregation of each level of the rooted subtree. Therefore, we can start by analyzing the time complexity of STA k .

The calculation of STA k can be divided into three steps. In the first step, we compute ϕ ( K i : ) and ϕ ( K i : ) T V i : for each node. The time complexity of this step depends on the feature map. In our model, we chose ϕ ( x ) = elu ( x ) + 1 as the feature map. Thus, the time complexity of computing ϕ ( K i : ) is O ( Nd k ) . We also need to compute ϕ ( K i : ) T V i : for each node, the time complexity of this part is O ( Nd k d v ) . Therefore, the overall time complexity of the first step is O ( Nd k + Nd k d v ) .

In the second step, we let ϕ ( K i : ) and ϕ ( K i : ) T V i : propagate on the graph. For STA k , we need to propagate k times. The time complexity of propagating ϕ ( K i : ) once is O ( d k ) , and the time complexity of ϕ ( K i : ) T V i : propagating once is O ( d k d v ) . The message propagation occurs on each edge. Considering that there are in total |E| edges and k times propagation, the overall time complexity of this step is O ( k |E| d k + k |E| d k d v ) .

In the third step, we use the information ∑ N j =1 ˆ A k ij ϕ ( K j : ) T V j : and ∑ N j =1 ˆ A k ij ϕ ( K j : ) T aggregated by each node, along with the node's own query ϕ ( Q i : ) , to complete the computation of STA. For each node, we need to calculate ϕ ( Q i : ) ∑ N j =1 ˆ A k ij ϕ ( K j : ) T V j : , the time complexity of this part is O ( Nd k d v ) . At the same time, for each node, we need to calculate ϕ ( Q i : ) ∑ N j =1 ˆ A k ij ϕ ( K j : ) T , the time complexity of this part is O ( Nd k ) . So the total time complexity of this step is O ( Nd k + Nd k d v ) .

In summary, the total time complexity of STA k is O (2 Nd k +2 Nd k d v + k |E| d k + k |E| d k d v ) .

Next, we analyze the time complexity of STA when the height of the rooted subtree is K . It should be noted that { STA i } i ∈ [ [1 ,K ] ] can be viewed as a nested process, calculated one after another. Therefore, the first two steps of the above-mentioned calculation of STA k do not need to be repeated. We only need to complete the full calculation of STA K and perform the third step mentioned above K times. Therefore, the time complexity of calculating STA is O (( K +1) Nd k +( K +1) Nd k d v + K |E| d k + K |E| d k d v ) . In general, we can think of the time complexity of STA as O ( K |E| d k d v ) .

B Proof for Theorem 1

B.1 Proof for Equation 10

Let ˆ A denote the random walk matrix of a connected and non-bipartite graph, and let A sym denote the symmetric normalized adjacency matrix. Let 1 = λ 1 ≥ · · · ≥ λ N be the eigenvalues of ˆ A , which are also the eigenvalues of A sym [26]. Let v 1 , v 2 , . . . , v n be the corresponding orthonormal eigenvectors ( v 1 , v 2 , . . . , v n here are column vectors ) . Let π j = d ( j ) ∑ N i =1 d ( i ) and d ( i ) denotes the degree of the i th node. ˆ λ = 1 -max { λ 2 , | λ n |} denotes the corresponding spectral gap, and let D be the diagonal degree matrix. - → 1 denotes an all-ones column vector.

In this subsection, we prove the following results:

And for a given ϵ , the smallest K 0 that satisfies the condition above is at most O ( log N ϵ 1 -max { λ 2 , | λ n |} ) . We begin by considering an arbitrary distribution p i ∈ R N , which is a column vector and ∥ p i ∥ 2 = 1 . Notice that v 1 , v 2 , . . . , v n form an orthonormal basis, we can rewrite D -1 2 p i as:

We next consider the new distribution obtained when p i undergoes k -step random walk. Notice that ˆ A = AD -1 = D 1 2 ( D -1 2 AD -1 2 ) D -1 2 = D 1 2 A sym D -1 2 . Thus we have:

We now consider c 1 D 1 2 A k sym v 1 . As we know that D 1 2 - → 1 is an eigenvector of A sym with eigenvalue 1. We then have:

Notice that ∥ D 1 2 - → 1 ∥ 2 = √ ∑ N i =1 d ( i ) . Using the fact that v 1 , v 2 , . . . , v n form an orthonormal basis and D -1 2 p i = ∑ N i =1 c i v i , we then have:

Notice that A k sym v 1 = λ k 1 v 1 and λ 1 = 1 . Thus we have:

Considering Equation 14 and Equation 17, we have:

where ||| D 1 2 ||| p = sup x ∈ R N ∥ D 1 2 x ∥ 2 ∥ x ∥ 2 = √ d max. Thus we have:

and immediately:

And using the fact that v 1 , v 2 , . . . , v n are orthonormal and 1 -ˆ λ = max { λ 2 , | λ N |} = max {| λ 2 | , | λ 3 | , . . . , | λ N |} , we then have:

Notice that ∥ D -1 2 p i ∥ 2 2 ≤ ||| D -1 2 ||| 2 p ∥ p i ∥ 2 2 = 1 d min . Therefore:

Using Cauchy-Schwarz, we then have:

In conclusion, given an arbitrarily small positive number ϵ , for all k 0 greater than or equal to 1 ˆ λ log N ϵ , the L1 norm of the difference between ˆ A k p i and the vector π is less than or equal to ϵ . This result establishes that 1 ˆ λ log N ϵ indeed serves as an upper bound.

Notice that the vector p i is an arbitrary distribution. Thus, we may consider p i to be one of the i th unit basis vector in the N -dimensional space: { p 1 , p 2 , . . . , p N } , where each vector has only one element equal to 1 (the i th element) and all other elements equal to 0. Thus we have ∥ ˆ A k p j -π ∥ 1 = ∑ N i =1 | ˆ A k ij -π i | . Then given ϵ > 0 , we have that:

which demonstrates immediately the first part of Theorem 1:

B.2 Proof for Equation 11

In this subsection, we prove the following results: if V is computed by V = σ ( XW V ) where σ is a non-negative activation function, then:

holds true when none of the denominators is equal to zero. And for a given η , the smallest K 1 that satisfies the condition shown in Equation 11 is at most O ( log N η 1 -max { λ 2 , | λ n |} ) .

This result can, indeed, be viewed as a straightforward corollary of Equation (10). The crucial prerequisite is that all elements of the vector V must be positive. Importantly, there are no specific and immediately:

requirements imposed on the non-negative activation function σ , meaning it can be any function that ensures non-negativity.

STA k and SA are defined as follows:

Equation 27 shows their form as row vectors. Their j th elements are:

Hence, we have:

For clarity, we proceed under the assumption that none of the denominators equal zero, which is reasonable considering the context. Let δ it represent the difference between ˆ A k it and π i : ˆ A k it = π i + δ it . Given η ∈ ]0 , 1[ , we aim to determine an upper bound of the convergence rate between STA k and SA.

Using Equation 10, we take ϵ = η N 2 and we have immediately:

We can rewrite Equation 29 as:

Assuming that k ≥ 2 log N η 1 -ˆ λ . Considering the fraction ∑ N t =1 δ it sim ( Q i : , K t : ) V tj ∑ N t =1 π i sim ( Q i : , K t : ) V tj in the first part of Equation 31. Using Equation 30 and the fact that sim ( Q i : , K t : ) and V tj are all positive, we have:

Notice that ∀ t ∈ [ [1 , N ] ] , π i ≥ 1 N 2 . Hence, we have:

Therefore:

Thus we have

and

Considering the second part ∑ N t =1 π i sim ( Q i : , K t : ) ∑ N t =1 ( π i + δ it ) sim ( Q i : , K t : ) of Equation 31. Utilizing the same line of reasoning, we can obtain:

Considering Equation 35, Equation 37 and Equation 31, we finally prove that:

which proves the second part of Theorem 1.

C Dataset Information

In this section, we present the datasets used in our experiments. These different types of data provide a robust platform to evaluate the performance of our methods.

The detailed information for each dataset is presented in Table 3. These datasets are drawn from the areas of citation networks, co-purchasing networks, co-authorship networks, and social networks: · Citation Networks : The citation networks datasets include Cora, Citeseer, Pubmed, and CoraFull. Nodes in these networks correspond to scientific publications, while the edges represent citations between these documents. In addition to the topological structure, each node carries a binary attribute vector, encoding the presence or absence of specific words from a pre-determined dictionary. The dimensionality of these attribute vectors varies from 1,433 in Cora to 8,710 in CoraFull. Moreover, each document node is associated with a unique class label, signifying the document's overarching scientific discipline. · Co-authorship Networks : We utilize the CoauthorCSDataset and CoauthorPhysicsDataset that capture co-authorship relationships in Computer Science and Physics domains, respectively. Nodes represent individual authors and edges encode co-authorship relations, thus creating an undirected graph. · Co-purchasing Networks : We utilize the AmazonCoBuyComputerDataset and AmazonCoBuyPhotoDataset, derived from Amazon's co-purchasing network. Nodes denote products and edges symbolize frequent co-purchase incidents. Moreover, the nodes can carry diverse product-specific information. · Social Networks : The Deezer-Europe dataset is a dataset representing a social network of Deezer users collected via the public API in March 2020. The nodes in this network symbolize Deezer users hailing from various European countries, while the edges represent reciprocal follower relationships between these users. The features of each node are derived from the preferences of the users, specifically, the artists they have expressed an interest in. The task associated with this graph involves binary node classification, wherein the objective is to predict the user's gender. · Co-occurrence Networks : We utilize the Actor dataset, a type of co-occurrence network based on the Microsoft Academic Graph. Nodes represent actors, and an edge signifies their co-appearance on the same Wikipedia page.

D Implementation Details

Positional Encoding We use Laplacian positional encoding to capture the structural information. As positional encoding is not the focus of our work, we use a simple approach to combine positional encoding with the original features of the nodes, which is also applied by [6]. Formally, we first calculate the eigenvectors corresponding to the smallest m eigenvalues of the Laplace matrix to construct the matrix P ∈ R n × m . Then we take X ′ = [ X , P ] as the new input, where [ ] denotes row-wise concatenation. For all the datasets, we set m = 3 .

D.1 Node Classification

Training Details We choose two recent studies [43, 6] and we adhere to their experimental configurations. The metrics for the baselines are also derived from these works [43, 6]. For Cora, Citeseer, Deezer and Actor, we apply the same random splits with train/valid/test ratios of 50%/25%/25% as [43]. We conduct 5 runs with different splits and take the mean accuracy and standard deviation for comparison. For Pubmed, Corafull, Computer, Photo, CS and Physics, we apply the same random splits with train/valid/test ratios of 60%/20%/20% as [6]. We conduct 10 runs with different splits and take the mean accuracy and standard deviation for comparison. Specifically, we utilize the ROC-AUC measure for binary classification on the Deezer dataset. For other datasets containing more than two classes, we opt for Accuracy as the metric. We employ the Adam optimizer for gradient-based optimization. The training procedure can at most repeat until a given budget of 3000 epochs and we set the patience of early stop to 200 epochs. We report the test accuracy of the epoch which has the highest accuracy on the validation set.

Hyperparameters For the model configuration of STAGNN, we fix the number of hidden channels at 64. We use grid search for hyper-parameter settings. The learning rate is searched within {0.001,0.01}, dropout probability searched within {0.0,0.2,0.4,0.6}, weight decay searched within {0.0001,0.0005,0.001,0.005}, height of the rooted subtree K searched within {3,5,10}, number of attention heads searched within {1,2,4,6,8}. The best hyper-parameters are provided in supplementary materials.

D.2 Study on the Necessity of SubTree Attention in the Presence of Global Attention

In this experiment, we extend STAGNN by replacing the STA module with global attention enhanced by 0, 1, 2, or 3 hop/hops subtree attention. We now present a detailed mathematical description of the experimental configurations. Formally, we compare the performance of the STAGNN-based model equipped with four different attention strategies: Global Attn Only, 1-hop STA + GA, 2-hops STA + GA and 3-hops STA + GA on six datasets: Pubmed, Corafull, Computer, Photo, CS and Physics, with the same experiment setting described in subsection D.1.

First, we calculate keys, queries and values.

Next, the output of the four different models (equipped with global attention enhanced by subtree attention of different heights) can be described as:

Global Attn Only :

1-hop STA + GA :

2-hops STA + GA :

3-hops STA + GA :

STAGNN-SUM :

STAGNN-CONCAT :

where W O is a linear projection matrix.

STAGNN-ATTN :

where W a is a linear projection matrix and [ ] denotes row-wise concatenation.

E More Visualizations of GPR Weights

We conduct more visualizations of the GPR weights on Cora and Actor, with heights K of the rooted subtrees ranging from 3 to 75. The results are shown in Figure 7.

In the case of Cora, we observe that as the depth K of the rooted subtree increases, STA keeps increasing the GPR weights of the local neighborhood in order to preserve the local information from being covered up by the global information.

α T here represents the coefficient of teleportation, because we can regard the global attention enhanced by subtree attention here as the random walk with teleportation. The only difference between these models is that they use subtree attention of different heights as an auxiliary to global attention. As shown in Table 2, we can observe that 2-hops STA + GA and 3-hops STA + GA outperform Global Attn Only by a large margin.

D.3 Study on HopAggregation Methods

In this experiment, we investigate different choices of the HopAggregation functions within the STA module. We compare GPR-like aggregation with sum, concat [18], and attention-based readout [6]. We now present a detailed mathematical description of the experimental configurations. Formally, we compare the performance of the following four models: STAGNN-GPR (origin STAGNN), STAGNNSUM, STAGNN-CONCAT and STAGNN-ATTN on four datasets: Cora, Citeseer, Deezer-Europe and Actor, with the same experiment setting described in subsection D.1.

First, we calculate keys, queries and values.

Next, the output of the four different models (STAGNN with different HopAggregation methods) can be described as:

STAGNN-GPR (origin STAGNN) :

0

1

2

3

1.0

1.5

2.0

2.5

3.0

GPR weights when

K

=3

0

1

2

3

4

1.0

1.5

2.0

2.5

GPR weights when

K

=4

0

1

2

3

4

5

1.0

1.5

2.0

2.5

GPR weights when

K

=5

0

2

4

6

1.0

1.5

2.0

2.5

GPR weights when

K

=6

cora

film

0

2

4

6

1.0

1.5

2.0

GPR weights when

K

=7

0

2

4

6

8

1.0

1.5

2.0

GPR weights when

K

=8

0

2

4

6

8

1.0

1.5

2.0

GPR weights when

K

=9

0

2

4

6

8

10

1.0

1.5

2.0

GPR weights when

K

=10

0

5

10

15

20

0.5

1.0

1.5

2.0

GPR weights when

K

=20

0

10

20

30

0.5

1.0

1.5

2.0

2.5

GPR weights when

K

=30

0

10

20

30

40

50

0

1

2

3

GPR weights when

K

=50

0

20

40

60

0

1

2

3

GPR weights when

K

=75

F Further discussion of the Gate Mechanism within the Mixture of Attention Heads

In this subsection, we conduct an ablation study of the gate mechanism within the mixture of attention heads. The sub-tree attention module with multiple attention heads is defined as follows:

MSTA w/ gate vector g , w/ softmax (origin STAGNN) :

The hop-wise gate vector here g k ∈ R H is an H -dimensional vector and g h k is its h th element. Compared to STA with a single attention head, we introduce in total H × K additional learnable parameters: { g i } i ∈ [ [1 ,K ] ] .

For comparison, we consider two variants.

MSTA w/ gate vector g k , w/o softmax : MSTA ( Q , K , V ) = AGGR ( { MSTA k ( Q , K , V ) | k ∈ [ [0 , K ] ] } ) MSTA k ( Q , K , V ) = [ head 1 k , . . . , head H k ] W O ∀ k ∈ [ [1 , K ] ] , MSTA 0 ( Q , K , V ) = V head h k = g h k STA k ( Q h , K h , V h ) ∀ h ∈ [ [1 , H ] ] (50)

MSTA w/o gate vector g k :

The experimental results are shown in Table 4. We find that the performance of MSTA w/ gate vector g k , w/o softmax and MSTA w/o gate vector g k are almost the same, which means that using the gate

vector without softmax is approximately equivalent to not using the gate vector. In fact, on closer examination, we find that without softmax, the learned gate vector would be a vector with all equal elements, which means that it is difficult for the model to learn different weights of attention heads at each hop without the help of softmax. Additionally, we observe that for most datasets, using the gating mechanism leads to improvement of the overall performance.

G Potential Impacts

Besides learning better node representations, our proposed Subtree Attention (STA) has potential impacts on various aspects of graph learning. Compared to global attention, STA can help the model to better learn the hierarchical structure of the graph. Therefore, STA can be utilized as a plug-in module for designing local-aware Transformers on graph, acting as a competitor of all the GNNassisted Transformers. STA opens new avenues for model design by combining the message-passing scheme with fully-attentional architectures, which can significantly enhance both the computational efficiency and expressive power of fully-attentional models on graph data. Furthermore, STA bridges the gap between local and global graph attention methods. This opens up possibilities for the design and application of hierarchical attention models that can leverage both local neighborhood and global structural information from graph data.

Peer reviews:
comments: SUMMARY:
The paper introduces Subtree Attention (STA), a new graph attention mechanism that overcomes limitations of existing mechanisms in graph learning. STA combines fully-attentional structure with rooted subtrees, approximating masked global attention under extreme settings. By computing attention weights among multi-hop neighbors directly, STA addresses previous mechanism problems. The authors present STAGNN, a graph neural network that utilizes STA and a hop-aware attention strategy. Extensive evaluations on ten node classification datasets show that STA-based models outperform existing graph transformers and mainstream GNNs.

STRENGTHS:
1, This paper proposes a novel multi-hop graph attention mechanism, Subtree Attention (STA), which bridges the fully-attentional structure and the rooted subtree. 
2, By incorporating kernelized softmax, they develop an optimized approach for STA, resulting in a linear time complexity.
3, Experiments have shown that STAGNN achieves good performance on Node classification tasks.

WEAKNESSES:
The performance improvement of NAGphormer is limited, as demonstrated in Table 1. For instance, when evaluating photos, NAGphormer achieved a score of 95.49±0.11, while STAGNN achieved 95.64±0.27. The difference between 95.49 and 95.64 is only 0.15, which is smaller than 0.27. A similar trend is observed in the case of CS. Furthermore, in Physics, STAGNN performs worse than NAGphormer. Additionally, it is worth noting that the variance of STAGNN appears to be much larger than that of NAGphormer.

QUESTIONS:
Why do you only evaluate STA on the node classification task? What about Graph Classification?

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-05

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a novel multi-hop graph attention mechanism called SubTree Attention (STA) to address the limitations of both local and global attention in Graph Neural Networks (GNNs). STA allows the root node to attend directly to further neighbors in the subtree, enabling it to gather information from the entire rooted subtree within one layer. This mechanism avoids issues associated with deep architectures using local attention layers and captures the hierarchical neighborhood structure more effectively than global attention. The authors propose an efficient algorithm based on kernelized softmax to calculate attention among multi-hop neighbors. They also provide a theoretical analysis of STA, demonstrating its convergence to a degree-masked global self-attention. The paper introduces STAGNN, a multi-hop graph attention network that integrates the STA module into decoupled GCNs, achieving superior performance compared to existing GNNs and graph transformers. Experimental evaluations on node classification datasets confirm the effectiveness of STA and the competitive performance of STAGNN.\n\nSTRENGTHS:\n* The authors propose a strategy to incorporate attention in multi-hop propagation GNNs and get good performances.\n* Good theoretical analysis is provided.\n\nWEAKNESSES:\n* Lack of baselines (BernNet,...). The authors ignore some important baselines of propagation-based GNNs, including BernNet and successive work.\n* The kernalized attention is a method from Transformer variants. Need citation.\n* Datasets are not large, should include OGB-datasets, at least ogbn-arxiv\n\nQUESTIONS:\n* How is Graph Transformer applied to node-classification? Is it the pipeline described? Because graph for node-classification is larger than graphs for graph-level tasks? Maybe subgraph sampling is needed?\n* Two questions about novelty. The STAGNN is an attention-based method to combine multi-hop information of graphs?  The difference of using vanilla transformer and STA is that STA has kernalized attention? If not, can you briefly describe the novelty here?\n* Is the multi-head attention applied in STAGNN?\n* If the attention will converge to degree-based global attention, why will STA work? Is the convergence achieved in training?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe present manuscript proposes a novel graph attention layer, that lies in the middle of the local aggregation scheme of message passing and the global (non-structured) nature of a full-attention. The newly proposed Subtree Attention constructs for each node the similarity pairs of key and query matrices from its k-hop neighborhood. \n\nGiven the transformed matrices, they perform an efficient aggregation throughout the neighborhood hops by considering the random walk adjacency matrix as the transition matrix that is used in the graph attention. They show that given the utilization of the random walk matrix combined with a kernelized softmax function, they reduce the computational cost to $\\mathcal{O}(\\mathcal{E})$, where $\\mathcal{E}$ the number of graph edges.\n\nThe authors prove theoretically that given a set of assumptions, the Subtree attention can approximate the global attention, showing that the proposed layer is able to tackle over-smoothing and over-squashing problems that come with MPNNs, while bridging the local with the global attention.\n\nTheir experimental study shows that STAGNN (their method) has a superior performance with respect to the state-of-the-art graph transformer models.\n\nSTRENGTHS:\n- The Subtree attention is a novel and interesting idea on bridging the notions of the local neighborhood attention with the global graph attention.\n- The theoretical study shows that subtree attention can approximate global attention, which can be very useful for treating over-smoothing in MPNNs.\n- Empirically, the model STAGNN seems to outperform its competitor baselines. \n- The paper is well-written and easy to follow.\n\nWEAKNESSES:\n- The experimental study is limited to six benchmark datasets, which it is unclear whether long-range interactions play any role, or instead shallower architectures are, also, effective. \n- It would be very useful testing STAGNN, also, in other datasets (e.g. from OGB), that reportedly depend on long-range interactions, or in synthetic datasets, so that the impact of the proposed model can be shown better.\n\nQUESTIONS:\n- Can the authors show any empirical evidence that the examined datasets can show a need for deeper models? In other words, do these datasets consist of any long-range interactions, or do they consist only of shallow interactions?\n\nLIMITATIONS:\nThe authors make a discussion on the supplementary material over possible extensions and ways to improve the STAGNN by incorporating it to more models.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors propose a novel Graph Transformer called STA-GNN (SubTree-attention-GNN) to address the over-smoothing and over-squashing in the message-passing scheme. Different from the previous Graph Transformer, STA-GNN's attention mechanism is called SubTree Attention (STA), which computes a node's attention according to different levels of its subtree. To overcome the exponentially growing computation cost STA brings, the authors leverage the linearized attention technique raised by [19]. Experiments show that STAGNN outperforms SOTA GNN on node classification tasks.\n\nSTRENGTHS:\n1. The authors provide a thorough demonstration of STA's workflow with a theoretical analysis of how STA addresses over-smoothing and over-squishing. Besides. The authors propose an effective computing algorithm, which decreases STA time complexity.\n\n2. The experiment result outperforms SOTA GCN and Graph Transformer.\n\n3. Ablation studies show that STA is beneficial to global attention in graph learning.\n\nWEAKNESSES:\n1. STA's timing performance is unclear. The authors leverage kernelized Softmax to design an efficient algorithm with time complexity O(epsilon) (the number of edges). However, it is hard to compare STA with Graph transformers' (O(N) where N is the number of nodes) from the time complexity aspect, since the number of nodes and edges are independent in an undirected graph. This paper lacks experiments on STA's and Graph Transformers' timing performance. Without evidence, people will worry about if STA can achieve superior performance within a reasonable time limit.\n\n2. Analysis of STA's space complexity is missing. STA efficient algorithm reduces the number of calculations by storing and re-using previous calculation results for future calculation. But there is no discussion on the potential problem of STA's memory cost. A memory cost comparison can better show the efficiency of STA.\n\n3. Hop2Token used in NAGphormer [1] also learns the graph at the level of hops like STA. The experiment results against with NAGphormer does not seem very superior, especially when the experiment is not comprehensive to include a comparison of time and memory cost.\n\n4. Equation (5) doesn't match the context in Figure 1 (a). N is defined as the number of nodes at the start of Section 2. However, figure 1 (a) says the STA_{k} should only attend the k-th hop neighbors, which is conflict with Equation (5).\n\t\n5. Lack of explanation when introducing the rewriting from equation (3) to equation (4) (in Sec. 2.2). It is not obvious how to rewrite a similarity function with a selected kernel and feature map. [2] is about how to rewrite the similarity function using a linear kernel, which should be cited here.\n\n[1] J. Chen, K. Gao, G. Li, and K. He. Nagphormer: Neighborhood aggregation graph transformer for node classification in large graphs. CoRR, abs/2206.04910, 2022.\n\n[2] A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret. Transformers are rnns: Fast autoregressive\ntransformers with linear attention. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 5156–5165. PMLR, 2020.\n\nQUESTIONS:\n1. In section 5.2, the authors claim \"In contrast to MP-GNNs, STAGNN maintains robust performance even when the height of the subtree reaches 100\". Where is the source of \"the performance of MP-GNNs doesn't maintain robust when the height of the subtree reaches 100\"?\n\nLIMITATIONS:\nNo limitations are mentioned in this paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}