SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Siyuan Huang Yunchong Song Jiayue Zhou Zhouhan Lin*
Shanghai Jiaotong University
siyuan_huang_sjtu@outlook.com ycsong@sjtu.edu.cn lin.zhouhan@gmail.com
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.
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
*Zhouhan Lin is the corresponding author
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.
Let  $\mathcal{G} = (\mathcal{N},\mathcal{E})$  be an undirected graph, with the associated nodes set  $\mathcal{N}$  and edges set  $\mathcal{E}$ . We use  $N = |\mathcal{N}|$  to represent the number of nodes.  $\mathbf{X}\in \mathbb{R}^{N\times f}$  denotes the node feature, where  $f$  denotes the number of features. Let  $\mathbf{A}$  be the adjacency matrix of  $\mathcal{G}$  and let  $\mathbf{D}$  be the diagonal degree matrix of  $\mathbf{A}$ .  $\mathbf{A}_{\mathrm{sym}} = \mathbf{D}^{-1 / 2}\mathbf{A}\mathbf{D}^{-1 / 2}$  denotes the symmetric normalized adjacency matrix, while  $\mathbf{A}_{\mathrm{rw}} = \mathbf{AD}^{-1}$  denotes the random walk matrix. Let  $\hat{\mathbf{A}}$  be an arbitrary transition matrix, including  $\mathbf{A}_{\mathrm{rw}}$ ,  $\mathbf{A}_{\mathrm{sym}}$  or other matrices representing message propagation. We use  $\mathbf{M}_i$ : and  $\mathbf{M}_{:j}$  to indicate the  $i^{\mathrm{th}}$  row and the  $j^{\mathrm{th}}$  column of the matrix  $\mathbf{M}$ , respectively. And let  $[0,K]$  denote the set  $\{0,1,\dots ,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]:
$$
\mathbf {O} = \sum_ {k = 0} ^ {K} \alpha_ {k} \operatorname {P R O P A G} _ {k} (\mathbf {H}), \operatorname {P R O P A G} _ {k} (\mathbf {H}) = \hat {\mathbf {A}} ^ {k} \mathbf {H}, \mathbf {H} = \operatorname {M L P} (\mathbf {X}) \tag {1}
$$
where  $\{\alpha_{k}\}_{k\in [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  $\hat{\mathbf{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  $\hat{\mathbf{A}}$ .
2.2 Global Self-Attention and Kernelized Softmax
In graph learning domain, the Global Self-Attention function  $\mathrm{SA}(\cdot ,\cdot ,\cdot)$  computes a weighted sum of all positions for each node. It first projects the node feature matrix into three subspaces:
$$
\mathbf {Q} = \mathbf {X} \mathbf {W} _ {Q}, \mathbf {K} = \mathbf {X} \mathbf {W} _ {K}, \mathbf {V} = \mathbf {X} \mathbf {W} _ {V} \tag {2}
$$
where  $\mathbf{W}_Q\in \mathbb{R}^{d\times d_K}$ ,  $\mathbf{W}_K\in \mathbb{R}^{d\times d_K}$  and  $\mathbf{W}_V\in \mathbb{R}^{d\times d_V}$  are learnable projection matrices. Then the new representation of the  $i^{\mathrm{th}}$  node is computed as follows:
$$
\operatorname {S A} (\mathbf {Q}, \mathbf {K}, \mathbf {V}) _ {i:} = \frac {\sum_ {j = 1} ^ {N} \operatorname {s i m} \left(\mathbf {Q} _ {i :} , \mathbf {K} _ {j :}\right) \mathbf {V} _ {j :}}{\sum_ {j = 1} ^ {N} \operatorname {s i m} \left(\mathbf {Q} _ {i :} , \mathbf {K} _ {j :}\right)} \tag {3}
$$
where  $\mathrm{sim}(\cdot ,\cdot):\mathbb{R}^d\times \mathbb{R}^d\to \mathbb{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 product to compute similarity:  $\mathrm{sim}(\mathbf{Q}_{i:},\mathbf{K}_{j:}) = \exp \left(\frac{\mathbf{Q}_{i:}\mathbf{K}_{j:}^{T}}{\sqrt{d_{K}}}\right)$ .
In fact, we can use an arbitrary positive-definite kernel  $\kappa$  to serve as  $\mathrm{sim}(\cdot ,\cdot)$ . Given a selected kernel  $\kappa$  and its corresponding feature map  $\phi$ , we can rewrite  $\mathrm{sim}(\cdot ,\cdot)$  as:  $\mathrm{sim}(\mathbf{Q}_{i:},\mathbf{K}_{j:}) = \phi (\mathbf{Q}_{i:})\phi (\mathbf{K}_{j:})^T$ . Thus, Equation 3 becomes:
$$
\operatorname {S A} (\mathbf {Q}, \mathbf {K}, \mathbf {V}) _ {i:} = \frac {\sum_ {j = 1} ^ {N} \phi \left(\mathbf {Q} _ {i :}\right) \phi \left(\mathbf {K} _ {j :}\right) ^ {T} \mathbf {V} _ {j :}}{\sum_ {j = 1} ^ {N} \phi \left(\mathbf {Q} _ {i :}\right) \phi \left(\mathbf {K} _ {j :}\right) ^ {T}} = \frac {\phi \left(\mathbf {Q} _ {i :}\right) \sum_ {j = 1} ^ {N} \phi \left(\mathbf {K} _ {j :}\right) ^ {T} \mathbf {V} _ {j :}}{\phi \left(\mathbf {Q} _ {i :}\right) \sum_ {j = 1} ^ {N} \phi \left(\mathbf {K} _ {j :}\right) ^ {T}} \tag {4}
$$
There are many potential choices for the feature map  $\phi$ . 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  $\sum_{j=1}^{N} \phi(\mathbf{K}_{j:})^T \mathbf{V}_{j}$  and  $\sum_{j=1}^{N} \phi(\mathbf{K}_{j:})^T$ , which only need to be computed once. By doing so, we can avoid computing the full attention matrix  $\{\mathrm{sim}(\mathbf{Q}_{i:}, \mathbf{K}_{j:})\}_{i \in N, j \in N}$  and reduce the complexity from  $\mathcal{O}(N^2)$  to  $\mathcal{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 SubTree Attention (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  $\mathcal{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  $\mathrm{SA}(\cdot ,\cdot ,\cdot)$ , 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^{\mathrm{th}}$  hop neighbors, which we refer to as  $\mathrm{STA}_k$ . For the  $i^{\mathrm{th}}$  node, this process can be described as follows:
$$
\begin{array}{l} \mathrm {S T A} _ {0} (\mathbf {Q}, \mathbf {K}, \mathbf {V}) _ {i}: = \mathbf {V} _ {i}; \\ \mathrm {S T A} _ {k} (\mathbf {Q}, \mathbf {K}, \mathbf {V}) _ {i:} = \frac {\sum_ {j = 1} ^ {N} \hat {\mathbf {A}} _ {i j} ^ {k} \operatorname {s i m} \left(\mathbf {Q} _ {i :} , \mathbf {K} _ {j :}\right) \mathbf {V} _ {j :}}{\sum_ {j = 1} ^ {N} \hat {\mathbf {A}} _ {i j} ^ {k} \operatorname {s i m} \left(\mathbf {Q} _ {i :} , \mathbf {K} _ {j :}\right)} \quad \forall k \in [ 1, K ] \tag {5} \\ \end{array}
$$
(a) SubTree Attention
(b) Efficient algorithm of SubTree Attention
Figure 1: (Left): Definition of SubTree Attention (STA).  $\mathrm{STA}_k$  represents that each node attends to the  $k^{\text{th}}$  level of the rooted subtree, and then STA aggregates information from the entire subtree. (Right): An efficient algorithm for computing subtree attention. We first compute  $\phi(\mathbf{K}_{i:})$  and  $\phi(\mathbf{K}_{i:})^T\mathbf{V}_{i:}$  for each node, and then let  $\phi(\mathbf{K}_{i:})$  and  $\phi(\mathbf{K}_{i:})^T\mathbf{V}_{i:}$  perform  $K$ -step random walk, respectively. After each step of random walk, we compute attention weights using  $\phi(\mathbf{Q}_{i:})$  and the aggregated keys and values. The computation of  $\{\mathrm{STA}_k\}_{k \in [1,K]}$  can be seen as a nested process.
$\mathrm{STA}_k(\mathbf{Q}, \mathbf{K}, \mathbf{V})_i$ : represents the  $i^{\text{th}}$  node attending to its  $k^{\text{th}}$  hop neighbors, which are also the  $k^{\text{th}}$  level of its rooted subtree. SubTree Attention STA:  $(\mathbb{R}^{n \times d_k}, \mathbb{R}^{n \times d_k}, \mathbb{R}^{n \times d_v}) \to \mathbb{R}^{n \times d_o}$  can then be calculated by aggregating the results from all levels of the rooted subtree:
$$
\operatorname {S T A} (\mathbf {Q}, \mathbf {K}, \mathbf {V}) _ {i:} = \operatorname {A G G R} \left(\left\{\operatorname {S T A} _ {k} (\mathbf {Q}, \mathbf {K}, \mathbf {V}) _ {i:} \mid k \in [ 0, K ] \right\}\right) \tag {6}
$$
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,  $\mathrm{STA}_k$  allows each node to attend to the  $k^{\mathrm{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  $\mathrm{STA}_k$ : we first compute the complete similarity matrix  $\{\mathrm{sim}(\mathbf{Q}_{i:},\mathbf{K}_{j:})\}_{i\in N,j\in N}$  and then calculate its Hadamard product with  $\hat{\mathbf{A}}^k$  to obtain  $\{\hat{\mathbf{A}}_{ij}^k\mathrm{sim}(\mathbf{Q}_{i:},\mathbf{K}_{j:})\}_{i\in N,j\in N}$ . In short, we treat  $\hat{\mathbf{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  $\mathcal{O}(N^2)$ . (ii)  $\hat{\mathbf{A}}^k$  quickly converges to a dense matrix. Storing  $\hat{\mathbf{A}}^k$  in the GPU memory for computing the Hadamard product leads to considerable space complexity. Considering these two limitations, utilizing  $\hat{\mathbf{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  $\hat{\mathbf{A}}^k$ .
Learning from Equation 4, we use a feature map  $\phi$  to replace  $\mathrm{sim}(\cdot ,\cdot)$ . 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  $\phi (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:
$$
\mathrm {S T A} _ {k} (\mathbf {Q}, \mathbf {K}, \mathbf {V}) _ {i:} = \frac {\sum_ {j = 1} ^ {N} \hat {\mathbf {A}} _ {i j} ^ {k} \phi (\mathbf {Q} _ {i :}) \phi (\mathbf {K} _ {j :}) ^ {T} \mathbf {V} _ {j :}}{\sum_ {j = 1} ^ {N} \hat {\mathbf {A}} _ {i j} ^ {k} \phi (\mathbf {Q} _ {i :}) \phi (\mathbf {K} _ {j :}) ^ {T}} = \frac {\phi (\mathbf {Q} _ {i :}) \sum_ {j = 1} ^ {N} \hat {\mathbf {A}} _ {i j} ^ {k} \phi (\mathbf {K} _ {j :}) ^ {T} \mathbf {V} _ {j :}}{\phi (\mathbf {Q} _ {i :}) \sum_ {j = 1} ^ {N} \hat {\mathbf {A}} _ {i j} ^ {k} \phi (\mathbf {K} _ {j :}) ^ {T}} \tag {7}
$$
Note that there are two summations  $\sum_{j=1}^{N} \hat{\mathbf{A}}_{ij}^{k} \phi(\mathbf{K}_{j:})^{T} \mathbf{V}_{j}$ : and  $\sum_{j=1}^{N} \hat{\mathbf{A}}_{ij}^{k} \phi(\mathbf{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  $\phi(\mathbf{K}_{i:})$  and  $\phi(\mathbf{K}_{i:})^{T} \mathbf{V}_{i}$ : for each node. Then we let  $\phi(\mathbf{K}_{i:})$  and  $\phi(\mathbf{K}_{i:})^{T} \mathbf{V}_{i}$ : undergo  $k$  steps message passing. Finally, we use the aggregated keys and values  $\sum_{j=1}^{N} \hat{\mathbf{A}}_{ij}^{k} \phi(\mathbf{K}_{j:})^{T} \mathbf{V}_{j}$ : and  $\sum_{j=1}^{N} \hat{\mathbf{A}}_{ij}^{k} \phi(\mathbf{K}_{j:})^{T}$  in conjunction with the node's own query  $\phi(\mathbf{Q}_{i:})$  to complete the computation of subtree attention. Figure 1 illustrates the whole process of this efficient algorithm. When we choose  $\mathbf{A}_{\mathrm{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  $\mathcal{O}(N^2)$  to  $\mathcal{O}(|\mathcal{E}|)$ . Additionally, message passing only requires the sparse adjacency matrix  $\mathbf{A}$ , thereby circumventing the need to store  $\hat{\mathbf{A}}^{k}$ . Furthermore,  $\{\mathrm{STA}_i\}_{i \in [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:
$$
\operatorname {M S T A} (\mathbf {Q}, \mathbf {K}, \mathbf {V}) = \operatorname {A G G R} \left(\left\{\operatorname {M S T A} _ {k} (\mathbf {Q}, \mathbf {K}, \mathbf {V}) \mid k \in [ 0, K ] \right\}\right)
$$
$$
\operatorname {M S T A} _ {k} (\mathbf {Q}, \mathbf {K}, \mathbf {V}) = \left[ \operatorname {h e a d} _ {k} ^ {1}, \dots , \operatorname {h e a d} _ {k} ^ {H} \right] \mathbf {W} _ {O} \quad \forall k \in [ 1, K ], \quad \operatorname {M S T A} _ {0} (\mathbf {Q}, \mathbf {K}, \mathbf {V}) = \mathbf {V} \tag {8}
$$
$$
\operatorname {h e a d} _ {k} ^ {h} = \hat {g} _ {k} ^ {h} \operatorname {S T A} _ {k} \left(\mathbf {Q} ^ {\mathbf {h}}, \mathbf {K} ^ {\mathbf {h}}, \mathbf {V} ^ {\mathbf {h}}\right) \quad \forall h \in [ 1, H ], \quad \hat {\boldsymbol {g}} _ {k} = \operatorname {s o f t m a x} \left(\boldsymbol {g} _ {k}\right)
$$
where  $\left[\right]$  denotes row-wise concatenation.  $\mathbf{Q}^{\mathrm{h}},\mathbf{K}^{\mathrm{h}}$  , and  $\mathbf{V}^{\mathrm{h}}$  represent the query, key, and value matrices for the  $h^{\mathrm{th}}$  head, respectively.  $\mathbf{W}_O$  denotes a linear projection matrix.  $\pmb {g}_k\in \mathbb{R}^H$  is an  $H$  -dimensional vector and  $g_{k}^{h}$  is its  $h^\mathrm{th}$  element, representing the weight of the  $h^\mathrm{th}$  attention head at the  $k^{\mathrm{th}}$  hop. Compared to STA with a single attention head, we introduce in total  $H\times K$  additional learnable parameters:  $\{g_i\}_{i\in [1,K]}$  . We can regard  $\pmb {g}_k$  as a hop-wise gate that determines the weight of each attention head at the  $k^{\mathrm{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.  $\pmb{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  $\mathbf{A}_{\mathrm{rw}}$  as the transition matrix in the STA module,i.e., we have  $\hat{\mathbf{A}} = \mathbf{A}_{\mathrm{rw}}$  in this subsection.
We first employ a slightly modified approach to rewrite the global self-attention module SA, which can be described as:
$$
\operatorname {S A} (\mathbf {Q}, \mathbf {K}, \mathbf {V}) _ {i:} = \frac {\sum_ {j = 1} ^ {N} \pi_ {i} \sin \left(\mathbf {Q} _ {i :} , \mathbf {K} _ {j :}\right) \mathbf {V} _ {j :}}{\sum_ {j = 1} ^ {N} \pi_ {i} \sin \left(\mathbf {Q} _ {i :} , \mathbf {K} _ {j :}\right)} \tag {9}
$$
where  $\pi_{i} = \frac{d(i)}{\sum_{j=1}^{N} d(j)}$  and  $d(j)$  denotes the degree of the  $j^{\text{th}}$  node. Note that Equation 9 is consistent with Equation 3.