Source-linked AI summary
Self-Attention Graph Pooling
Junhyun Lee, Inyeop Lee, Jaewoo Kang
TL;DR
Graph pooling for graph neural networks remains difficult, particularly because existing methods face complexity or topology-awareness limitations. SAGPool uses graph-convolution-based self-attention to select nodes while incorporating node features and graph topology. It generally performs well across benchmark settings, outperforming hierarchical pooling methods on all evaluated datasets, while retaining a consistent parameter count across input graph sizes.
Problem
Graph pooling remains less developed, with existing methods facing quadratic storage or node-dependent parameters, or failing to account for graph topology.
Method
SAGPool uses graph-convolution-based self-attention to select nodes for hierarchical pooling from node features and graph topology.
Results
SAGPool generally performs well and outperformed hierarchical pooling methods on all evaluated datasets.
Takeaways & Limitations
SAGPool provides hierarchical, end-to-end graph representation learning with reasonable complexity and a consistent parameter count across input graph sizes.
Takeaways & Limitations
SAGPool cannot parameterize pooling ratios to optimize them for each graph, and binary node-preservation decisions do not completely solve this issue.
Abstract
from arXiv · showhide
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.
1. Introduction
Graph learning extends CNN-style convolution and pooling to non-Euclidean data, but graph pooling remains less developed than graph convolution. SAGPool addresses this gap with self-attention that uses graph convolution to consider node features and topology while learning hierarchical representations with relatively few parameters.
- Graphs represent social, biological, and molecular structures as nodes and edges in non-Euclidean domains.
- Graph convolution has achieved strong performance across applications, whereas graph pooling has fewer established methods.Prior work has developed spectral and non-spectral graph convolutions for recommender systems, chemical research, and natural language processing.
- Existing hierarchical pooling methods trade off scalability and topology: DiffPool has quadratic storage complexity and node-dependent parameters, while gPool and related methods omit graph topology.
- SAGPool uses self-attention with graph convolution to select nodes for retention or removal, incorporating both node features and graph topology.
- SAGPool learns hierarchical graph representations end-to-end using relatively few parameters.
2. Related Work
Graph pooling methods reduce graph representations through topology-based, global, or hierarchical strategies. Related work motivates SAGPool by combining hierarchical representation learning with feature- and topology-aware pooling while addressing complexity concerns.
- Graph pooling reduces representation size and is grouped into topology-based, global, and hierarchical approaches.
- Topology based pooling: Topology-based pooling uses graph coarsening algorithms, including Graclus, rather than learned neural pooling.Spectral clustering relies on eigendecomposition, while Graclus avoids eigenvectors through an equivalent weighted kernel k-means objective.
- Global pooling: Global pooling aggregates node representations across a graph, allowing graphs with different structures to be processed but not learning hierarchical representations.
- Hierarchical pooling: Hierarchical pooling learns node assignments or clusters to capture structural information across layers.DiffPool learns assignment matrices end-to-end from node features and adjacency information.
- gPool has O(|V| + |E|) storage complexity versus DiffPool's O(k|V|^2), but its projection scores do not use graph topology.
- SAGPool is proposed to combine node features and topology for hierarchical representations with reasonable time and space complexity.
3. Proposed Method
SAGPool uses graph-convolution-based self-attention to select nodes while considering both node features and graph topology. It supports hierarchical pooling with a small parameterization and several multi-hop or multi-GNN attention variants.
- Self-Attention Graph Pooling: SAGPool uses a GNN to generate self-attention scores for selecting nodes during graph pooling.The attention mechanism distinguishes nodes to retain from nodes to drop.
- Self-Attention Graph Pooling: Graph convolution makes SAGPool’s attention scores depend on both node features and graph topology.The method accepts node features and the adjacency matrix when computing attention.
- Graph Pooling: The pooling ratio k determines how many nodes remain: SAGPool selects the top ⌈kN⌉ nodes by attention score.The selected node features are masked, and the adjacency matrix is restricted to the retained nodes.
- Variation of SAGPool: SAGPool can incorporate two-hop connections through edge augmentation or stacked GNN layers, and can average attention scores from multiple GNNs.These variants are named SAGPoolaugmentation, SAGPoolserial, and SAGPoolparallel.
- Model Architecture: The evaluation uses shared model architectures and graph convolution settings across baselines and SAGPool for fair comparison.The global architecture uses three graph convolutional layers with concatenated outputs, while the hierarchical architecture uses three convolution-and-pooling blocks.
4. Experiments
The experiments evaluate global and hierarchical graph-pooling methods for graph classification on benchmark datasets using matched training procedures and architectures. SAGPool generally performs well, especially on D&D and PROTEINS, and outperforms hierarchical pooling methods across all datasets.
- Datasets and Evaluation: The study evaluates global and hierarchical pooling methods on graph classification using five benchmark datasets containing more than 1k graphs.The datasets include protein, molecular, and other graph-classification benchmarks.
- Training Procedure: The experiments use 20 random seeds with 10-fold cross-validation, producing 200 testing results per method and dataset.Ten percent of the training data is used for validation, with shared early stopping and hyperparameter-selection procedures.
- Baselines: Baselines include Set2Set, SortPool, DiffPool, and gPool, compared with SAGPool under matched global or hierarchical architectures.The same hyperparameter search strategy is used for all baselines and SAGPool.
- Parameter Analysis: gPool and SAGPool maintain a consistent parameter count regardless of input graph size and pooling ratio.Figure 3 uses 128 input features, hidden size 128, and two classes for the hierarchical pooling models.
- Results: SAGPool generally performs well against global pooling methods, especially on D&D and PROTEINS.The reported results are average accuracies and standard deviations.
- Results: SAGPool outperforms the hierarchical pooling methods on all evaluated datasets.The performance of SAGPool variants differs across datasets and GNN types, while the variants show potential for improvement.
5. Analysis
The analysis compares pooling architectures, topology-aware attention, sparse implementation, efficiency, and SAGPool variants. SAGPool generally performs well while retaining practical complexity advantages, but its pooling ratio remains a limitation.
- 5.1. Global and Hierarchical Pooling: Global pooling performs better on smaller-node datasets, whereas hierarchical pooling is more effective on larger graphs.Global pooling minimizes information loss; hierarchical pooling efficiently extracts useful information from large-scale graphs.
- 5.1. Global and Hierarchical Pooling: SAGPool tends to perform well with both global and hierarchical architectures.
- 5.2. Effect of Considering Graph Topology: The graph Laplacian term incorporates graph topology into SAGPool attention scores without requiring recalculation within the same block.The term was already used in a previous graph convolutional layer.
- 5.2. Effect of Considering Graph Topology: SAGPool achieves superior graph-classification performance to gPool with the same number of parameters.
- 5.3. Sparse Implementation: Sparse SAGPool reduces computational complexity to O(|E|) and storage to O(|V|+|E|), compared with dense pooling's O(|V|^2) storage.The comparison is made against dense matrix operations and dense pooling methods such as DiffPool.
- 5.4. Relation with the Number of Nodes: Unlike DiffPool, SAGPool's parameter count is independent of cluster size, and its cluster size can vary with the number of input nodes.DiffPool requires a predefined cluster size tied to the maximum number of nodes.
- 5.5. Comparison of the SAGPool Variants: Two-hop variants may improve performance, while selecting the appropriate number of averaged attention-score GNNs can stabilize results.
- 5.6. Limitations: SAGPool cannot learn optimal pooling ratios separately for each graph, and binary node-preservation classification does not completely solve this issue.
6. Conclusion
The paper proposes SAGPool, a self-attention graph pooling method with hierarchical pooling, topology and feature awareness, reasonable complexity, and end-to-end representation learning. Its parameter count remains consistent across input graph sizes.
- 6. Conclusion: SAGPool is a self-attention graph pooling method designed for hierarchical graph representation learning.
- 6. Conclusion: SAGPool considers both node features and graph topology while supporting end-to-end representation learning.
- 6. Conclusion: SAGPool uses a consistent number of parameters regardless of input graph size.