Research map / Graph learning / GNN architectures
ENAHPool: The Edge-Node Attention-based Hierarchical Pooling for Graph Neural Networks
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes ENAHPool, a cluster-based graph pooling operator with hard node assignment and node/edge attention, plus MD-MPNN to mitigate over-squashing, evaluated on eight graph classification benchmarks.
Problem
Cluster-based graph pooling relies on ambiguous soft node assignment and uniform summation of node features and edge connectivity strengths, while standard MPNNs suffer from over-squashing for long-range dependencies.
Approach
Assign each node exclusively to one cluster via hard assignment (with straight-through estimator), aggregate node features with node-level attention and inter-cluster edge strengths with edge-level attention, and use a Multi-Distance MPNN that propagates messages directly among nodes at each shortest distance weighted by path counts.
Key results
- ENAHPool with MD-MPNN achieves best accuracy on seven of eight graph classification datasets, e.g., 79.91% on D&D and 88.56% on REDDIT-B.
- ENAHPool without MD-MPNN (ENAHPool(-MD)) still beats most alternatives, e.g., 78.93% on NCI1.
- Hard assignment outperforms soft assignment on all datasets in ablation, e.g., 78.34% vs 77.33% on D&D.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use ENAHPool with MD-MPNN instead of the strongest reported baseline | Graph classification · D&D | accuracy: 78.45 ± 0.40 -> 79.91 ± 0.25 | Helped |
| Use ENAHPool with MD-MPNN instead of the strongest reported baseline | Graph classification · PROTEINS | accuracy: 76.42 ± 0.39 -> 77.12 ± 0.15 | Helped |
| Use ENAHPool with MD-MPNN instead of the strongest reported baseline | Graph classification · NCI1 | accuracy: 78.64 ± 1.53 -> 79.34 ± 0.31 | Helped |
| Use ENAHPool with MD-MPNN instead of the strongest reported baseline | Graph classification · FRANKENSTEIN | accuracy: 66.26 ± 0.47 -> 67.68 ± 0.13 | Helped |
| Use ENAHPool with MD-MPNN instead of the strongest reported baseline | Graph classification · IMDB-B | accuracy: 74.12 ± 0.56 -> 74.54 ± 0.42 | No clear effect |
| Use ENAHPool with MD-MPNN instead of the strongest reported baseline | Graph classification · IMDB-M | accuracy: 51.53 ± 0.65 -> 51.74 ± 0.16 | No clear effect |
| Use ENAHPool with MD-MPNN instead of the strongest reported baseline | Graph classification · COLLAB | accuracy: 81.28 ± 0.15 -> 81.09 ± 0.51 | No clear effect |
| Use ENAHPool with MD-MPNN instead of the strongest reported baseline | Graph classification · REDDIT-B | accuracy: 82.75 ± 0.82 -> 88.56 ± 0.25 | Helped |
| Add MD-MPNN instead of conventional MPNNs (ablation) | ENAHPool · D&D | accuracy: 79.21 ± 0.51 -> 79.91 ± 0.25 | Helped |
| Add MD-MPNN instead of conventional MPNNs (ablation) | ENAHPool · REDDIT-B | accuracy: 85.60 ± 1.06 -> 88.56 ± 0.25 | Helped |
| Use hard rather than soft node assignment (ablation) | ENAHPool without attention mechanisms · D&D | accuracy: 77.33 ± 0.43 -> 78.34 ± 0.62 | Helped |
| Use hard rather than soft node assignment (ablation) | ENAHPool without attention mechanisms · REDDIT-B | accuracy: 82.12 ± 1.06 -> 83.31 ± 1.01 | Helped |
Surprising findings
- Edge-based attention is intentionally omitted in the first pooling layer because it would make all connectivity strengths equal to 1 after pooling.
Limitations
- MD-MPNN complexity is O(N^3) due to adjacency matrix multiplications.
- Performance declines beyond two pooling layers and a threshold of MD-MPNN layers, attributed to overfitting.
- Evaluated only on graph classification with GCN backbone.
Experimental setup
Datasets: D&D PROTEINS NCI1 FRANKENSTEIN IMDB-B IMDB-M COLLAB REDDIT-B
Benchmarks: D&D PROTEINS NCI1 FRANKENSTEIN IMDB-B IMDB-M COLLAB REDDIT-B
Metrics: classification accuracy
Baselines: Set2Set SortPool SAGPool GMT DiffPool TopKPool ASAP StructPool MinCutPool SEP-G ABDPool
Builds on: DiffPool GCN ABDPool SEP-G straight-through estimator
Closest papers by meaning
- ASAP: Adaptive Structure Aware Pooling for Learning Hierarchical Graph Representations (AAAI 2020 · similarity 0.90)
- Accurate Learning of Graph Representations with Graph Multiset Pooling (ICLR 2021 · similarity 0.90)
- Geometry-Aware Edge Pooling for Graph Neural Networks (NeurIPS 2025 · similarity 0.89)
- SkipPool: Improved Sparse Hierarchical Graph Pooling with Differentiable Exploration (AAAI 2025 · similarity 0.88)
- MaxCutPool: differentiable feature-aware Maxcut for pooling in graph neural networks (ICLR 2025 · similarity 0.87)
- Edge Prompt Tuning for Graph Neural Networks (ICLR 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.