Tailoring Self-Attention for Graph via Rooted Subtrees

TL;DR

This paper introduces Subtree Attention (STA), a novel multi-hop graph attention mechanism that captures hierarchical and long-range information effectively, improving over existing methods in graph neural networks.

Contribution

The paper proposes STA, a new attention mechanism that bridges global and local attention in graphs, with a linear-time implementation and theoretical approximation guarantees.

Findings

STA outperforms existing graph transformers and GNNs on node classification tasks.

The linear-time STA implementation maintains high performance with reduced computational cost.

Theoretical proof shows STA approximates global attention in extreme cases.

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.