[Paper Review] Hierarchical Graph Representation Learning with Differentiable Pooling

Rex Ying, Jiaxuan You+4 more|arXiv (Cornell University)|Jun 22, 2018

Advanced Graph Neural Networks34 references833 citations

TL;DR

DiffPool introduces a differentiable graph pooling module to learn hierarchical graph representations, yielding 5–10% average accuracy gains on graph classification benchmarks and state-of-the-art results on four of five datasets.

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.

Motivation & Objective

• Motivate the need for hierarchical graph representations for graph classification.
• Propose a differentiable pooling mechanism to coarsen graphs across multiple GNN layers.
• Enable end-to-end training of deep hierarchical GNNs with learned cluster assignments.
• Show that DiffPool improves performance across diverse datasets and GNN backbones.

Proposed method

• Introduce DiffPool, a differentiable pooling module that learns soft cluster assignments S^(l) to form coarser graphs.
• Compute cluster embeddings X^(l+1) = S^(l)ᵀ Z^(l) and coarsened adjacency A^(l+1) = S^(l)ᵀ A^(l) S^(l).
• Use two GNNs per layer: one for node embeddings (embedding GNN) and one for producing the pooling assignments (pooling GNN).
• Train with an auxiliary link prediction objective L_LP = ||A^(l) , S^(l) S^(l)ᵀ||_F and an entropy regularization L_E to encourage near one-hot assignments.
• Ensure permutation invariance of DiffPool under node permutations given permutation-invariant GNNs.
• Support end-to-end training with a final graph representation used for classification.

Experimental results

Research questions

• RQ1 Does DiffPool outperform existing pooling methods for GNNs on graph classification benchmarks?
• RQ2 Can DiffPool provide a scalable, interpretable hierarchical graph representation across different GNN architectures?
• RQ3 Does incorporating auxiliary objectives stabilize training and improve clustering interpretability?
• RQ4 How does hierarchical pooling affect performance relative to flat GNN architectures across datasets?
• RQ5 Is the approach generalizable to architectures beyond GraphSAGE (e.g., Structure2Vec)?

Key findings

• DiffPool yields an average improvement of 5–10% accuracy over all existing pooling approaches in graph classification tasks.
• Achieves state-of-the-art results on four of five benchmark datasets.
• DiffPool combined with GraphSAGE shows about a 6.27% average gain over baselines.
• DiffPool also improves Structure2Vec when used in conjunction with its embedding, indicating general applicability.
• Auxiliary link prediction objective improves training stability and reduces accuracy variance across runs.
• DiffPool learns interpretable hierarchical clusters that align with community structure in graphs.