Tailoring Self-Attention for Graph via Rooted Subtrees
Siyuan Huang Yunchong Song Jiayue Zhou Zhouhan Lin
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.