Cluster-wise Graph Transformer with Dual-granularity Kernelized Attention
About
In the realm of graph learning, there is a category of methods that conceptualize graphs as hierarchical structures, utilizing node clustering to capture broader structural information. While generally effective, these methods often rely on a fixed graph coarsening routine, leading to overly homogeneous cluster representations and loss of node-level information. In this paper, we envision the graph as a network of interconnected node sets without compressing each cluster into a single embedding. To enable effective information transfer among these node sets, we propose the Node-to-Cluster Attention (N2C-Attn) mechanism. N2C-Attn incorporates techniques from Multiple Kernel Learning into the kernelized attention framework, effectively capturing information at both node and cluster levels. We then devise an efficient form for N2C-Attn using the cluster-wise message-passing framework, achieving linear time complexity. We further analyze how N2C-Attn combines bi-level feature maps of queries and keys, demonstrating its capability to merge dual-granularity information. The resulting architecture, Cluster-wise Graph Transformer (Cluster-GT), which uses node clusters as tokens and employs our proposed N2C-Attn module, shows superior performance on various graph-level tasks. Code is available at https://github.com/LUMIA-Group/Cluster-wise-Graph-Transformer.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Graph Classification | PROTEINS | Accuracy76.48 | 742 | |
| Graph Classification | MUTAG | Accuracy87.11 | 697 | |
| Graph Classification | COLLAB | Accuracy80.43 | 329 | |
| Graph Classification | ogbg-molpcba (test) | AP29.61 | 206 | |
| Graph Regression | ZINC (test) | MAE0.071 | 204 | |
| Graph Regression | Peptides struct LRGB (test) | MAE0.2464 | 178 | |
| Graph Classification | Peptides-func LRGB (test) | AP0.6975 | 136 | |
| Graph Classification | D&D | Accuracy79.15 | 110 | |
| Graph Classification | IMDB MULTI | Accuracy52.13 | 109 | |
| Graph Regression | ZINC | MAE0.071 | 96 |