[Paper Review] Self-Attention Graph Pooling

Junhyun Lee, Inyeop Lee+1 more|arXiv (Cornell University)|Apr 17, 2019

Advanced Graph Neural Networks44 references375 citations

TL;DR

SAGPool introduces a self-attention based graph pooling layer for GNNs that accounts for node features and graph topology, enabling end-to-end hierarchical graph representations with competitive performance. It compares global and hierarchical pooling architectures across multiple datasets, showing strong results especially on larger graphs.

ABSTRACT

Advanced methods of applying deep learning to structured data such as graphs have been proposed in recent years. In particular, studies have focused on generalizing convolutional neural networks to graph data, which includes redefining the convolution and the downsampling (pooling) operations for graphs. The method of generalizing the convolution operation to graphs has been proven to improve performance and is widely used. However, the method of applying downsampling to graphs is still difficult to perform and has room for improvement. In this paper, we propose a graph pooling method based on self-attention. Self-attention using graph convolution allows our pooling method to consider both node features and graph topology. To ensure a fair comparison, the same training procedures and model architectures were used for the existing pooling methods and our method. The experimental results demonstrate that our method achieves superior graph classification performance on the benchmark datasets using a reasonable number of parameters.

Motivation & Objective

• Motivate improved graph pooling that leverages both node features and topology.
• Develop a differentiable, scalable pooling layer using self-attention on graphs.
• Evaluate SAGPool in global and hierarchical pooling architectures against strong baselines.
• Show that SAGPool achieves high accuracy with a parameter count that is relatively independent of input graph size.

Proposed method

• Define SAGPool as a self-attention based pooling layer using graph convolution to compute attention scores.
• Compute Z = sigma(à D^{-1/2} à D^{-1/2} X Theta_att) where à is A plus self-loops and D is the degree matrix.
• Select top-kN nodes by Z to form X_out and A_out, retaining a fixed fraction k of nodes.
• Allow variants using different GNNs (ChebConv, SAGE, GAT) and two-hop connectivity (augmentation, serial, parallel) for attention.
• Provide global and hierarchical pooling architectures for fair comparison with baselines (Set2Set, SortPool, DiffPool, gPool).
• Maintain a consistent parameter count independent of input graph size.

Experimental results

Research questions

• RQ1 Can self-attention based pooling that incorporates graph topology improve graph classification performance?
• RQ2 How does SAGPool compare to existing pooling methods under global versus hierarchical architectures?
• RQ3 What is the impact of using topology-aware attention and two-hop information on pooling quality and efficiency?
• RQ4 How do variants of SAGPool (different GNNs, augmentation/serial/parallel, multi-GNN averaging) affect performance across datasets?

Key findings

• SAGPool g achieves 76.19% ±0.94 on D&D, 70.04% ±1.47 on PROTEINS, 74.18% ±1.20 on NCI1, 74.06% ±0.78 on NCI109, and 62.57% ±0.60 on FRANKENSTEIN using global pooling architecture.
• SAGPool h achieves 76.45% ±0.97 on D&D, 71.86% ±0.97 on PROTEINS, 67.45% ±1.11 on NCI1, 67.86% ±1.41 on NCI109, and 61.73% ±0.76 on FRANKENSTEIN using hierarchical pooling architecture.
• SAGPool generally outperforms baselines Set2Set, SortPool, DiffPool, and gPool across datasets, with notable gains on D&D and PROTEINS.
• Incorporating graph topology via the normalized adjacency in attention scores (as in SAGPool) improves performance over topology-agnostic pooling like gPool.
• Variants that incorporate two-hop connectivity and multi-GNN averaging show dataset-dependent improvements, indicating SAGPool’s flexibility can be leveraged for further gains.
• Sparsity and end-to-end learning yield favorable memory and parameter characteristics, with SAGPool’s parameter count remaining stable as graph size grows.