Hierarchical Graph Representation Learning with Differentiable Pooling
Hierarchical Graph Representation Learning with Differentiable Pooling
Rex Ying Christopher Morris William L. Hamilton Jiaxuan You Xiang Ren Jure Leskovec

Abstract
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.
Code Repositories
Benchmarks
| Benchmark | Methodology | Metrics |
|---|---|---|
| graph-classification-on-collab | GNN (DiffPool) | Accuracy: 75.48% |
| graph-classification-on-dd | S2V (with 2 DiffPool) | Accuracy: 82.07% |
| graph-classification-on-dd | GNN (DiffPool) | Accuracy: 80.64% |
| graph-classification-on-enzymes | S2V (with 2 DiffPool) | Accuracy: 63.33% |
| graph-classification-on-enzymes | GNN (DiffPool) | Accuracy: 62.53% |
| graph-classification-on-proteins | GNN (DiffPool) | Accuracy: 76.25% |
| graph-classification-on-reddit-multi-12k | GNN (DiffPool) | Accuracy: 47.08 |
| graph-property-prediction-on-ogbg-code2 | DiffPool w/ graphSAGE | Ext. data: No Number of params: 10095826 Test F1 score: 0.1401 ± 0.0012 Validation F1 score: 0.1405 ± 0.0012 |
Build AI with AI
From idea to launch — accelerate your AI development with free AI co-coding, out-of-the-box environment and best price of GPUs.