Differentiable Pooling

Definition

Differentiable pooling (DiffPool) is a hierarchical graph pooling method that reduces the size of graph representations. The model is designed for graph-level tasks, such as graph classification, regression, and matching.

Architecture

DiffPool learns a soft assignment matrix $S^{(l)}$ that maps nodes in the input graph to a set of clusters and uses the assignment to generate a coarsened graph with fewer nodes.

The DiffPool layer is constructed by the two GNNs: embedding and pooling GNNs.

The embedding GNN is used to generate an embedding matrix $Z^{(l)}={\mathrm{GNN}}_{l,\text{embed}}(X^{(l)},A^{(l)})$ The pooling GNN is used to generate an assignment matrix representing a soft clustering of nodes. $S^{(l)}=\mathrm{softmax}\left({\mathrm{GNN}}_{l,\text{pool}}(X^{(l)},A^{(l)})\right)$ where $X^{(l)}$ and $A^{(l)}$ are the input cluster features and adjacency matrices at layer $l$ respectively, and the number of columns of $S^{(l)}$ is a hyperparameter and smaller than the number of input nodes.

DiffPool layer generates a new coarsened adjacency matrix and a new matrix embedding using the embedding and pooling GNNs. $(X^{(l+1)},A^{(l+1)}):=({S^{(l)}}^{⊺}{Z}^{(l)},{S^{(l)}}^{⊺}{A}^{(l)}{S}^{(l)})=\mathrm{DiffPool}(Z^{(l)},A^{(l)})$