DiffPool: Hierarchical Graph Representation Learning with Differentiable Pooling

本文提出了一种分层的图池化机制DiffPool，通过将原始图的部分节点映射成粗粒度图上的一个节点模型能够学习到图的层级结构，在池化多次后将图映射成一个节点，进而得到整张图的表示。

本文被NIPS2018接收，地址：Hierarchical Graph Representation Learning with Differentiable Pooling

Motivation

在有机分子中，人们希望能够编码局部的分子结构（原子间的化学键）以及图的粗粒度结构（代表分子功能的结构单元），但当前的GNN无法以分层的方式聚合信息。为了解决这个问题，本文提出了一种图上的层级池化方法DiffPool，通过将多个节点映射成一个簇的方式得到粗粒度的图，在池化多次后，图上的节点个数逐渐减少，最终只剩下一个节点，这样在得到不同层级结构的基础上最终得到整张图的特征表示向量。

Method

Differentiable Pooling via Learned Assignments

pooling with an assignment matrix

通过$S^{(l)}\in {\mathbb{R}}^{n_l\times n_{l+1}}$可以实现从第$l$层向第$l+1$层节点的映射过程，$n_l$表示第$l$层节点的个数，$S^{(l)}[i,j]$是$l$层上节点$i$映射到$l+1$层上节点$j$的概率，第$l+1$层的节点特征矩阵$X^{(l+1)}$和邻接矩阵$A^{(l+1)}$通过$S^{(l)}$可以得到$$\begin{array}{l}{X}^{(l+1)}=S^{(l{)}^T}{Z}^{(l)}\in {\mathbb{R}}^{n_{l+1}\times d}\\ A^{(l+1)}=S^{(l{)}^T}{A}^{(l)}{S}^{(l)}\in {\mathbb{R}}^{n_{l+1}\times n_{l+1}}\end{array}$$注意到$A^{(l+1)}$可能所有的值都不为0，图上的节点可能全部两两相连，$A^{(l+1)}[i,j]$可以看做节点$i$到节点$j$的连接强度（connectivity strength）。

learning the assignment matrix

assignment matrix$S^{(l)}$ 和embedding matrix$Z^{(l)}$可通过两个GNN得到
$$Z^{(l)}={\mathrm{G}\mathrm{N}\mathrm{N}}_{l,\mathrm{e}\mathrm{m}\mathrm{b}\mathrm{e}\mathrm{d}}\left(A^{(l)},X^{(l)}\right)$$$$S^{(l)}=\mathrm{softmax}\left({\mathrm{G}\mathrm{N}\mathrm{N}}_{l,\mathrm{p}\mathrm{o}\mathrm{o}\mathrm{l}}\left(A^{(l)},X^{(l)}\right)\right)$$在${\mathrm{G}\mathrm{N}\mathrm{N}}_{l,\mathrm{p}\mathrm{o}\mathrm{o}\mathrm{l}}$的结果上按行做$softmax$，可以得到图上每一个节点对应下一层图上的每一个节点的可能性。

The embedding GNN generates new embeddings for the input nodes at this layer, while the pooling GNN generates a probabilistic assignment of the input nodes to $n_{l+1}$clusters.

Auxiliary Link Prediction Objective and Entropy Regularization

一个直观的想法是在原图上距离更近的节点们（更加连通）更容易被映射为粗粒度图上的一个节点，所以模型在训练时要minimize $L_{\mathrm{L}\mathrm{P}}={\Vert A^{(l)},S^{(l)}{S}^{(l{)}^T}\Vert }_F$，$||\cdot |{|}_F$是Frobenius norm，也就是说$A^{(l)}$要尽量地和$S^{(l)}{S}^{(l{)}^T}$接近。$A^{(l)}[i,j]$代表的是图上i,j之间的连通强度，又因为$S^{(l)}{S}^{(l{)}^T}[i,j]={\sum }_k{S}^{(l)}[i,k]S^{(l{)}^T}[k,j]$，当i,j同属于某一簇k的可能性高时$S^{(l)}{S}^{(l{)}^T}[i,j]$会比较大，所以通过优化$L_{\mathrm{L}\mathrm{P}}$可以让相邻的节点更容易被映射到同一个簇（下一层的节点）上。

另外，$S$的每一行$S_i$代表原图上的第$i$个节点向下一层的各个节点映射的概率，所以$S_i$应该接近一个one-hot的向量，即第$i$个节点仅向下一层的一个节点映射。所以还需要minimize $L_{\mathrm{E}}=\frac{1}{n}{\sum }_{i=1}^{n}H\left(S_i\right)$，通过减少熵的方式减少映射分布的不确定性。