Source-linked AI summary
ASAP: Adaptive Structure Aware Pooling for Learning Hierarchical Graph Representations
Ekagra Ranjan, Soumya Sanyal, Partha Pratim Talukdar
TL;DR
Graph pooling methods must summarize graphs hierarchically without losing substructure or scalability. ASAP combines sparse differentiable clustering with Master2Token attention, LEConv scoring, and sparse connectivity computation. It achieves state-of-the-art results across five graph-classification datasets, consistently outperforming baselines with a reported average improvement of 4% over the previous sparse hierarchical state of the art.
Problem
Existing pooling methods either fail to capture graph substructure effectively or do not scale easily to large graphs, limiting hierarchical graph representation learning.
Method
ASAP hierarchically clusters local subgraphs, uses Master2Token attention for node membership, scores clusters with LEConv, and computes sparse edge weights for selected clusters.
Results
ASAP consistently outperforms all baselines on five graph-classification datasets, with an average improvement of 4% over the previous sparse hierarchical state-of-the-art method.
Takeaways & Limitations
ASAP combines hierarchical subgraph learning, sparse graph operations, and structure-aware connectivity in a pooling method that performs strongly across graph-classification benchmarks.
Takeaways & Limitations
The scoring analysis assumes a graph with a positive adjacency matrix and considers functions based on differences between transformed nodes and their neighbors.
Abstract
from arXiv · showhide
Graph Neural Networks (GNN) have been shown to work effectively for modeling graph structured data to solve tasks such as node classification, link prediction and graph classification. There has been some recent progress in defining the notion of pooling in graphs whereby the model tries to generate a graph level representation by downsampling and summarizing the information present in the nodes. Existing pooling methods either fail to effectively capture the graph substructure or do not easily scale to large graphs. In this work, we propose ASAP (Adaptive Structure Aware Pooling), a sparse and differentiable pooling method that addresses the limitations of previous graph pooling architectures. ASAP utilizes a novel self-attention network along with a modified GNN formulation to capture the importance of each node in a given graph. It also learns a sparse soft cluster assignment for nodes at each layer to effectively pool the subgraphs to form the pooled graph. Through extensive experiments on multiple datasets and theoretical analysis, we motivate our choice of the components used in ASAP. Our experimental results show that combining existing GNN architectures with ASAP leads to state-of-the-art results on multiple graph classification benchmarks. ASAP has an average improvement of 4%, compared to current sparse hierarchical state-of-the-art method.
1 Introduction
Graph classification needs hierarchical pooling to capture local substructures, but existing methods trade off structural fidelity and scalability. ASAP addresses this gap with sparse, differentiable clustering, attention, and structure-aware scoring.
- Motivation: Hierarchical graph representations capture local substructures that flat GNN aggregation can miss.In molecular graphs, groups of atoms can form functional groups relevant to graph labels.
- Limitations of prior pooling: Existing pooling methods face complementary limitations: DiffPool is dense and poorly scalable, while TopK and SAGPool do not fully capture graph structure.SAGPool uses local structure for scoring but does not use it effectively to determine pooled-graph connectivity.
- ASAP: ASAP is a sparse pooling operator that hierarchically captures local subgraph information and learns global features with improved pooled-graph edge connectivity.It is designed to retain sparsity while aggregating subgraphs into pooled nodes.
- ASAP: Master2Token provides self-attention for cluster membership, while LEConv adaptively learns functions of local extrema in graph substructures.These components are introduced as mechanisms for pooling-specific attention and cluster scoring.
2 Related Work
Graph pooling methods range from spectral and non-spectral approaches to global and hierarchical architectures, with trade-offs between structure preservation, differentiability, and scalability. ASAP is proposed to retain desirable hierarchical-pooling properties while preserving sparse graph operations.
- Spectral GNN methods use Fourier transformation and graph Laplacians, whereas non-spectral methods define convolution locally and generalize more easily across graphs.
- Differentiable and data-driven pooling methods followed earlier deterministic clustering approaches, expanding graph pooling beyond fixed algorithms.
- Global pooling summarizes an entire graph in one step, while hierarchical pooling captures graph topological information through successive reductions.
- DiffPool computes a dense soft assignment matrix that is infeasible for large graphs, whereas TopK and SAGPool preserve sparsity but do not effectively preserve node and edge information.
- ASAP is designed to provide desirable hierarchical-pooling properties without compromising sparsity in graph operations.
3 Preliminaries
The preliminaries define graph classification inputs and introduce the GCN and self-attention components used to represent graph structure and node importance. They also define receptive fields for pooling operators in terms of neighborhood hops affecting pooled nodes and edges.
- A graph is represented by nodes, edges, node features, an adjacency matrix, and a graph label; graph classification learns a mapping from input graphs to labels.
- ASAP uses a Graph Convolution Network to extract discriminative features for graph classification, using the initial node features wherever they are provided.
- Self-attention computes alignment scores to represent the importance of candidate entities for target queries, with target and candidate representations drawn from the input entities.
- In the global-task setting, candidate importance is computed after dropping the target query term from the alignment calculation.
- The receptive-field node and edge of a pooling operator count the hops needed to cover neighborhood nodes and edges influencing pooled representations.
4 ASAP: Proposed Method
ASAP hierarchically clusters local subgraphs, scores them, selects top clusters, and constructs a sparse pooled graph with learned connectivity.
- Local clustering: ASAP considers fixed-radius local clusters centered at each node to capture graph substructure.Each cluster contains nodes within h hops of its medoid.
- Cluster assignment: Its sparse assignment matrix S records node membership in local clusters while maintaining O(|E|) space complexity.The construction preserves sparsity comparable to the original adjacency matrix.
- Cluster assignment: Master2Token attention uses intra-cluster information to learn membership strengths and cluster representations.Attention scores α_i,j define S_i,j, while a separate GCN captures structural information within each cluster.
- Cluster selection: LEConv computes cluster fitness scores, after which the top ⌈kN⌉ clusters are selected for the pooled graph.Fitness scores are ranked by TOPk to determine the retained clusters.
- Pooled connectivity: The pooled adjacency combines cluster membership and original edge weights, connecting clusters sharing nodes or neighboring constituent nodes.Because the pruned assignment matrix is sparse, this operation can be implemented efficiently.
5 Theoretical Analysis
The analysis motivates LEConv for cluster scoring and characterizes ASAP’s receptive field, connectivity, and permutation equivariance.
- Fitness scoring: Standard GCN scoring can bias selection toward neighboring clusters when a node has a high pre-score.The analysis motivates selecting local extrema to obtain representative clusters from different graph regions.
- Fitness scoring: LEConv can learn neighbor-difference fitness functions that GCN cannot learn under the stated graph and function assumptions.The theorem compares GCN and LEConv for functions depending on transformed node-neighbor differences.
- Fitness scoring: LEConv scores clusters using both global and local importance through self-loops and local-extrema functions.This motivates its use as ASAP’s cluster fitness function.
- Connectivity: ASAP and DiffPool have edge receptive field 2h + 1, compared with h for TopK and SAGPool.A larger edge receptive field is described as important for information flow in subsequent GCN layers.
- Connectivity: For trees as N →∞, the sampling ratio needed to guarantee at least one pooled edge approaches 0.5 for ASAP and 1 for TopK or SAGPool.The analysis therefore indicates that ASAP can retain connectivity while sampling fewer clusters on such graphs.
- Equivariance: ASAP is a graph permutation equivariant pooling operator.The proposition states this property without further qualification in the supplied passage.
6 Experimental Setup
The experiments evaluate ASAP on five graph-classification datasets against hierarchical and global pooling baselines using a shared training setup.
- Datasets: The evaluation uses five graph-classification datasets spanning proteins, anticancer activity, and molecular mutagen classification.The datasets are D&D, PROTEINS, NCI1, NCI109, and FRANKENSTEIN.
- Datasets: The dataset-statistics table reports average graph counts, class counts, node counts, and edge counts using G_avg, C_avg, V_avg, and E_avg.These symbols are defined in the table caption.
- Baselines: The benchmark includes hierarchical baselines DiffPool, TopK, and SAGPool, plus global pooling methods Set2Set, Global-Attention, and SortPool.The comparison covers both hierarchical and global pooling approaches.
- Training and evaluation: ASAP uses pooling ratio k = 0.5 and receptive-field parameter h = 1 for consistency with the baselines.Training uses 10-fold cross-validation and reports average accuracy over 20 random seeds.
7 Results
The experiments evaluate ASAP against alternative pooling methods and ablate its node aggregation, attention, cluster-scoring, and edge-weighting components. ASAP outperforms all baselines across five graph-classification datasets, while its components improve performance and edge information is necessary.
- 7.2 Effect of Node Aggregation: Node aggregation improves ASAP performance across the tested aggregation configurations.The study compares None, Only cluster, and Both aggregation methods on FRANKENSTEIN and NCI1 validation data.
- 7.3 Attention Framework: M2T attention performs better than S2T and T2T on NCI1 and comparably on FRANKENSTEIN.The comparison evaluates the three attention frameworks on validation data.
- 7.4 Effect of LEConv: LEConv and Basic-LEConv perform significantly better than GCN, while LEConv outperforms Basic-LEConv.The results are attributed to modeling local extrema and, for LEConv, using three linear transformations instead of one.
- 7.5 Soft Edge Weights: Removing soft edge-weight computation causes a significant performance drop.This supports capturing edge information during graph pooling.
8 Discussion
ASAP preserves hierarchical clustering while restricting clusters to local neighborhoods, producing a sparse assignment matrix and accounting for input graph size. Its M2T attention uses all entities in a cluster to determine node membership.
- Local Hierarchical Clustering: ASAP clusters nodes within h-hop neighborhoods, whereas DiffPool considers the entire graph.This local restriction produces a sparse assignment matrix while retaining hierarchical clustering properties.
- Local Hierarchical Clustering: ASAP selects the top k fraction of current nodes, so the selected cluster count reflects input graph size.DiffPool instead predetermines a fixed number of clusters for each pooling layer.
- Comparison with Sparse Pooling: TopK and SAGPool maintain sparsity by dropping nodes rather than forming clusters or computing soft edge weights.The passage states that this prevents them from preserving node and edge information effectively.
- Source2Token & Token2Token: M2T generates a query from all entities in a cluster, unlike S2T and T2T formulations.The passage reports that M2T better utilizes intra-cluster information and performs better than S2T and T2T in the reported comparison.
9 Conclusion
ASAP is presented as a sparse, differentiable pooling method that hierarchically clusters local subgraphs, scores clusters with LEConv, and computes sparse edge weights. The paper reports theoretical and empirical validation and state-of-the-art performance on multiple graph-classification datasets.
- 9 Conclusion: ASAP hierarchically clusters local subgraphs to learn global graph features while preserving sparse operations.Its design includes sparse and differentiable pooling for graph-structured data.
- 9 Conclusion: M2T captures node membership in clusters, while LEConv scores clusters using local and global importance.ASAP samples clusters according to LEConv fitness scores and computes sparse edge weights for selected clusters.
- 9 Conclusion: The paper validates ASAP components theoretically and empirically and reports state-of-the-art performance on multiple graph-classification datasets.Experiments use hierarchical and global pooling architectures with GCN-based model configurations.
- CNN Analogy: In the CNN analogy, ASAP’s RF node and k correspond to kernel size and stride, respectively.RF node determines the cluster neighborhood radius, while k controls the sampling ratio.
E Ablation on pooling ratio k
Increasing the pooling ratio k retains more information and empirically increases validation performance on NCI1, but it also enlarges the pooled graph. Thus, choosing k trades performance against computational requirements.
- Pooling Ratio: Higher k retains more information and is empirically associated with increased validation performance on NCI1.The relationship is shown in Figure 2.
- Pooling Ratio: Increasing k raises computational requirements because a larger pooled graph is propagated to later layers.The resulting choice of k involves a trade-off between performance and computation.
F Proof of Theorem 1
The proof distinguishes standard GCN fitness from LEConv by whether the fitness function can represent transformed node–neighbor differences. Because nonnegative adjacency aggregation cannot express such difference terms, GCN cannot learn the target function, whereas LEConv can.
- Theorem statement: GCN fitness cannot learn functions depending on transformed differences between a node and its neighbors.The theorem considers functions f(X, A) built from these node–neighbor differences after a linear transformation.
- LEConv expressiveness: LEConv fitness can learn the target difference-based function.The proof establishes this by showing that an LEConv formulation represents the required function.
- GCN limitation: Nonnegative adjacency aggregation prevents GCN fitness from containing terms of the form β_i,j(x_iW − x_jW).The GCN expression uses nonnegative adjacency coefficients, so it cannot produce the required difference term.
- LEConv formulation: LEConv is introduced through a formulation with learnable matrices W1, W2, and W3.Setting W3 = W2 = W1, α1 = 1, and β_i,j = A_i,j makes the formulation equal to the stated target equation.
Proof of Theorem 2
The proof identifies balanced starlike trees as extremal structures for selecting nodes separated by at least h hops, then uses this result to derive sampling guarantees for pooled-tree connectivity. It concludes that ASAP requires a smaller asymptotic sampling ratio than TopK or SAGPool under the stated settings.
- Inductive step: Removing a leaf preserves the relevant optimum-node property with the least reduction in shortest-path separations.Removing non-leaf nodes can shorten paths between multiple node pairs, whereas removing a leaf affects only paths involving that leaf.
- Extremal tree structure: A balanced starlike tree with height h/2 maximizes the number of nodes that can be selected at least h hops apart among trees with N vertices.The result is established by induction over tree size, with path graphs as the base case.
- Connectivity guarantee: A pooling method retains an edge irrespective of selected-cluster positions when its sampled clusters exceed the maximum number of mutually separated nodes for its receptive field.The argument applies the pigeonhole principle to the extremal tree bound.
- Arbitrary trees: As N →∞, the required sampling ratio for arbitrary trees approaches 1 for TopK or SAGPool and 0.5 for ASAP.The theorem uses 1-hop receptive fields for TopK and SAGPool and a 2h + 1 receptive field for ASAP.
- Path graphs: For path graphs, the asymptotic lower bound approaches 0.5 for TopK or SAGPool and 0.25 for ASAP.The path-graph bound follows from the inequality involving receptive-field edge distance and the sampling ratio.
- Connectivity comparison: ASAP can connect nodes up to k + 2h hops apart when TopK uses the kth graph power, giving ASAP better connectivity in this setting.The comparison assumes h ≥ 1.