SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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  $\pmb{g}_k$ , w/ softmax (origin STAGNN):
$$
\mathrm {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 \mathrm {M S T A} _ {0} (\mathbf {Q}, \mathbf {K}, \mathbf {V}) = \mathbf {V} \tag {49}
$$
$$
\mathrm {h e a d} _ {k} ^ {h} = \hat {g} _ {k} ^ {h} \operatorname {S T A} _ {k} (\mathbf {Q} ^ {\mathbf {h}}, \mathbf {K} ^ {\mathbf {h}}, \mathbf {V} ^ {\mathbf {h}}) \quad \forall h \in [   [ 1, H ]   ], \quad \hat {\boldsymbol {g}} _ {k} = \operatorname {s o f t m a x} (\boldsymbol {g} _ {k})
$$
The hop-wise gate vector here  $g_{k} \in \mathbb{R}^{H}$  is an  $H$ -dimensional vector and  $g_{k}^{h}$  is its  $h^{\mathrm{th}}$  element. Compared to STA with a single attention head, we introduce in total  $H \times K$  additional learnable parameters:  $\{g_{i}\}_{i \in [1,K]}$ .
For comparison, we consider two variants.
- MSTA  $w/$  gate vector  ${\mathbf{g}}_{k}$  ,  $w/o$  softmax:
$$
\mathrm {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 \mathrm {M S T A} _ {0} (\mathbf {Q}, \mathbf {K}, \mathbf {V}) = \mathbf {V} \tag {50}
$$
$$
\operatorname {h e a d} _ {k} ^ {h} = 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 ]
$$
- MSTA w/o gate vector  $\pmb{g}_k$ :
$$
\mathrm {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 \mathrm {M S T A} _ {0} (\mathbf {Q}, \mathbf {K}, \mathbf {V}) = \mathbf {V} \tag {51}
$$
$$
\operatorname {h e a d} _ {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 ]
$$
The experimental results are shown in Table 4. We find that the performance of MSTA w/gate vector  $\pmb{g}_k$ , w/o softmax and MSTA w/o gate vector  $\pmb{g}_k$  are almost the same, which means that using the gate
Table 4: Ablation study of the gate mechanism within the mixture of attention heads
Method	Pubmed	CoraFull	Computer	Photo	CS	Physics
STAGNN (origin)	90.46±0.22	72.65±0.36	91.72±0.30	95.64±0.27	95.77±0.16	97.09±0.18
w/ gate, w/o softmax	90.37±0.23	71.62±0.39	91.89±0.27	95.37±0.30	94.72±0.19	96.96±0.20
w/o gate	90.31±0.25	71.67±0.36	91.80±0.28	95.32±0.28	94.70±0.18	96.97±0.18
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 GNN-assisted 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.