[Paper Review] Edge Contraction Pooling for Graph Neural Networks

Frederik Diehl|arXiv (Cornell University)|May 27, 2019

Advanced Graph Neural Networks18 references98 citations

TL;DR

EdgePool is a hard pooling layer based on edge contractions that outperforms several learned pooling methods, integrates into many GNN architectures, and improves both node and graph classification without changing training.

ABSTRACT

Graph Neural Network (GNN) research has concentrated on improving convolutional layers, with little attention paid to developing graph pooling layers. Yet pooling layers can enable GNNs to reason over abstracted groups of nodes instead of single nodes. To close this gap, we propose a graph pooling layer relying on the notion of edge contraction: EdgePool learns a localized and sparse hard pooling transform. We show that EdgePool outperforms alternative pooling methods, can be easily integrated into most GNN models, and improves performance on both node and graph classification.

Motivation & Objective

• Motivate the need for graph pooling layers to enable hierarchical reasoning over node groups.
• Introduce EdgePool, a pooling layer based on edge contractions that preserves graph structure.
• Demonstrate that EdgePool outperforms existing pooling methods and is easy to integrate into common GNN architectures.
• Show that EdgePool improves both graph-level and node-level classification tasks.

Proposed method

• Compute edge scores from concatenated node features via a learnable linear transformation (r(eij) = W * (ni || nj) + b).
• Apply a local softmax normalization per target node to obtain edge scores (sij = 0.5 + softmax_r*(rij)).
• Iteratively contract the highest-scoring edges while ensuring no incident node is already merged, yielding roughly 50% pooling per layer.
• Compute new node features after contraction as hij = sij * (ni + nj) to enable gradient flow (gating).
• Unpooling maps nodes back through successive EdgePool layers using inverse mappings and scales features by the edge scores during unpooling.
• EdgePool can optionally incorporate edge features by concatenating them into the raw score (r(eij) = W*(ni || nj || fij) + b).

Experimental results

Research questions

• RQ1 Q1: Does EdgePool outperform alternative pooling methods on graph tasks?
• RQ2 Q2: Can EdgePool be easily integrated into existing GNN architectures?
• RQ3 Q3: Can EdgePool be beneficial for node classification tasks?

Key findings

• EdgePool consistently improves performance over non-pooling baselines and TopKPool across graph classification datasets (proteins, rdt-b, rdt-12k, collab) and outperforms all pooling methods except DiffPool on some tasks.
• EdgePool integrated into existing architectures yields average improvements around 2 percentage points, with varying gains across models (largest for GraphSAGE; smaller for GIN/GIN0).
• EdgePool provides notable gains for node classification across multiple datasets and models, averaging around 3.5 percentage points overall, with GIN and GIN0 showing the strongest improvements.
• On graph classification benchmarks, EdgePool achieves 72.5% on proteins, 87.3% on rdt-b, 45.6% on rdt-12k, and 67.1% on collab (means ± std as reported).
• EdgePool supports sparse, linear-time operation with respect to edges, enabling scalability to large graphs.