Hierarchical Graph Representation Learning with Differentiable Pooling

TL;DR

Addresses graph classification by enabling hierarchical representations in GNNs through a differentiable pooling layer, DiffPool. The method learns layer-wise soft cluster assignments to coarsen graphs and stacks GNN modules end-to-end, using embedding and pooling GNNs plus an auxiliary link-prediction objective and entropy regularization. Empirical results show DiffPool improves accuracy over existing pooling methods and achieves state-of-the-art on most benchmarks, with interpretable cluster structures. This approach broadens the applicability of deep GNNs to graph-level tasks and provides a scalable path to deeper architectures.

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.

Figures (2)

• Figure 1: High-level illustration of our proposed method DiffPool. At each hierarchical layer, we run a GNN model to obtain embeddings of nodes. We then use these learned embeddings to cluster nodes together and run another GNN layer on this coarsened graph. This whole process is repeated for $L$ layers and we use the final output representation to classify the graph.
• Figure 2: Visualization of hierarchical cluster assignment in DiffPool, using example graphs from Collab. The left figure (a) shows hierarchical clustering over two layers, where nodes in the second layer correspond to clusters in the first layer. (Colors are used to connect the nodes/clusters across the layers, and dotted lines are used to indicate clusters.) The right two plots (b and c) show two more examples first-layer clusters in different graphs. Note that although we globally set the number of clusters to be $25\%$ of the nodes, the assignment GNN automatically learns the appropriate number of meaningful clusters to assign for these different graphs.

Theorems & Definitions (2)

• Proposition 1
• proof