Source-linked AI summary
Hierarchical Graph Representation Learning with Differentiable Pooling
Rex Ying, Jiaxuan You, Christopher Morris, Xiang Ren, William L. Hamilton, Jure Leskovec
TL;DR
Existing GNNs are flat, which limits hierarchical representations for graph classification. DIFFPOOL learns soft node-to-cluster assignments and feeds coarsened graphs through stacked GNNs. It improves graph-classification performance across benchmarks and reaches state-of-the-art results on four of five benchmarks.
Problem
GNNs do not learn hierarchical graph representations, while global node pooling for graph classification ignores hierarchical structure.
Method
DIFFPOOL learns differentiable soft assignments that map nodes to clusters, using the resulting coarsened graph as input to the next GNN layer.
Results
DIFFPOOL achieves the highest average performance among GNN pooling approaches, improves GraphSAGE by an average of 6.27%, and reaches state-of-the-art results on 4 out of 5 benchmarks.
Takeaways & Limitations
DIFFPOOL provides a general hierarchical pooling strategy that benefits different GNN architectures and can learn interpretable hierarchical clusters.
Takeaways & Limitations
DIFFPOOL can assign structurally similar but distant or disconnected nodes to the same cluster, particularly in some molecular graphs.
Abstract
from arXiv · showhide
Recently, graph neural networks (GNNs) have revolutionized the field of graph representation learning through effectively learned node embeddings, and achieved state-of-the-art results in tasks such as node classification and link prediction. However, current GNN methods are inherently flat and do not learn hierarchical representations of graphs---a limitation that is especially problematic for the task of graph classification, where the goal is to predict the label associated with an entire graph. Here we propose DiffPool, a differentiable graph pooling module that can generate hierarchical representations of graphs and can be combined with various graph neural network architectures in an end-to-end fashion. DiffPool learns a differentiable soft cluster assignment for nodes at each layer of a deep GNN, mapping nodes to a set of clusters, which then form the coarsened input for the next GNN layer. Our experimental results show that combining existing GNN methods with DiffPool yields an average improvement of 5-10% accuracy on graph classification benchmarks, compared to all existing pooling approaches, achieving a new state-of-the-art on four out of five benchmark data sets.
1 Introduction
GNNs learn node embeddings but remain flat, limiting hierarchical graph understanding and graph classification. DIFFPOOL addresses this by learning soft node clusters and repeatedly applying GNNs to coarsened graphs.
- GNNs propagate information across graph edges but do not infer or aggregate hierarchical structure.
- At each layer, DIFFPOOL clusters learned node embeddings, applies another GNN to the coarsened graph, and repeats this process for L layers before graph classification.
- Global pooling for graph classification combines node embeddings while ignoring hierarchical structure in the graph.
- DIFFPOOL combines GNNs with differentiable graph pooling in an end-to-end hierarchical architecture.
- DIFFPOOL learns soft assignments that map nodes to clusters, enabling increasingly coarsened graph representations.
2 Related Work
Prior GNN research developed message-passing models and graph-classification pooling strategies, while hierarchical methods generally relied on deterministic clustering. DIFFPOOL instead targets end-to-end learning of hierarchical structure.
- General graph neural networks: GNNs generate node representations by iteratively aggregating neighboring-node features through differentiable message-passing functions.
- Graph classification with graph neural networks: Graph classification requires converting node embeddings into a representation of the entire graph, commonly through summation or averaging.
- Hierarchical graph representations: Existing hierarchical graph methods combine GNNs with deterministic clustering algorithms in a two-stage procedure.
- Hierarchical graph representations: DIFFPOOL seeks to learn hierarchical graph structure end-to-end rather than relying on a deterministic graph-clustering subroutine.
3 Proposed Method
DIFFPOOL stacks GNN modules with learned soft assignments to coarsen graphs hierarchically. Separate GNNs produce node embeddings and assignments, while pooling equations create cluster features and connectivity for subsequent layers.
- Proposed method: DIFFPOOL is a differentiable module for constructing deep, multilayer GNNs through hierarchical graph pooling.
- Graph neural networks: A GNN computes node embeddings from graph structure and initial node features using iterative message passing.
- Stacking GNNs and pooling layers: The pooling layer targets a coarsened graph with fewer nodes, allowing stacked GNN modules to operate on progressively coarser representations.
- Differentiable Pooling via Learned Assignments: The assignment matrix S^(l) softly maps each node or cluster at layer l to a cluster at layer l+1.
- Pooling with an assignment matrix: DIFFPOOL aggregates node embeddings according to assignments and transforms the input adjacency matrix into coarsened cluster connectivity.
- Learning the assignment matrix: Separate GNNs generate embedding and assignment matrices from the input cluster features and adjacency matrix.
- Auxiliary Link Prediction Objective and Entropy Regularization: An auxiliary link-prediction objective trains the pooling GNN by encouraging agreement between adjacency structure and reconstructed assignments.
4 Experiments
The experiments evaluate DIFFPOOL across graph-classification baselines, datasets, architectures, and cluster assignments. DIFFPOOL improves performance and learns hierarchical clusters, while its behavior depends on training objectives, graph structure, and the maximum cluster count.
- Experimental setup: The evaluation compares DIFFPOOL with GNN pooling methods, kernel-based approaches, and deterministic-clustering variants across five graph-classification benchmarks.The datasets include ENZYMES, PROTEINS, D&D, REDDIT-MULTI-12K, and COLLAB, using 10-fold cross-validation.
- Graph classification results: DIFFPOOL achieves the highest average performance among GNN pooling approaches, improves GRAPHSAGE by 6.27% on average, and reaches state-of-the-art results on 4 of 5 benchmarks.DIFFPOOL-DET reaches state-of-the-art performance on COLLAB, where many graphs exhibit single-layer community structure.
- Graph classification results: DIFFPOOL also significantly improves S2V on ENZYMES and D&D, with similar performance trends on the other datasets.These results support DIFFPOOL as a pooling strategy applicable to different GNN architectures.
- Efficiency: DIFFPOOL was observed to run 12× faster than GRAPHSAGE with SET2SET pooling while retaining competitive performance.Although DIFFPOOL computes an assignment matrix, graph coarsening reduces the size of later convolution inputs.
- Cluster assignment analysis: The learned clusters capture hierarchical community structure, and link-prediction auxiliary objectives improve membership quality and training stability.The assignment network also tends to collapse densely connected subgraphs while separately pooling meaningful structures in sparse regions.
- Cluster assignment analysis: Larger maximum cluster counts model more complex hierarchies but can add noise and reduce efficiency, while end-to-end training may leave some clusters unused.The assignment network can also pool distant nodes when they have similar features and neighborhood structure.
5 Conclusion
DIFFPOOL extracts hierarchical structure from real-world graphs and, combined with existing GNNs, achieves state-of-the-art results on several graph classification benchmarks.
- DIFFPOOL extracts complex hierarchical structure from real-world graphs.
- Combining DIFFPOOL with existing GNN models achieves new state-of-the-art results on several graph classification benchmarks.
- Future directions include differentiable hard cluster assignments and applying hierarchical pooling to downstream tasks requiring whole-graph modeling.