DiffPool

Paper : Hierarchical Graph Representation Learning with Differentiable Pooling
Code :

摘要

对于Graph Classification 任务来说，GNN的最后一步需要关注图数据的全局信息，既包含图的结构信息，也包含各个节点的属性信息。如果采用一次性全局池化的方法聚合点上的特征表示，并没有将图的结构考虑进去，而且可能在节点重编号后聚合的结果会发生改变，也就是说聚合过程中依赖节点顺序。在CNN模型中，通常采用的做法是通过层次化池化机制来逐渐提取全局信息，作者仿照CNN提出了一种层次化图聚类的方式，支持添加到GNN模型中进行端到端的学习。

每层将该层的节点映射到若干个cluster中，作为下一层输入的节点。

定义

定义图 $G$ 表示为 $(A,F)$ ，其中 A ∈ { 0 , 1 } n × n A\in\{0,1\}^{n\times n} A∈{0,1}n×n 表示邻接矩阵，$F\in {\mathbb{R}}^{n\times d}$ 表示节点的特征矩阵。图分类任务定义为学习一个映射 $f:\mathcal{G}\to \mathcal{Y}$，这需要首先学习图的整体特征表示，即 $f^{\prime }:\mathcal{G}\to {\mathbb{R}}^D$

图神经网络的数学表示为

$$H^{(k)}=M(A,H^{(k-1)};{\theta }^{(k)})$$

其中，$H^{(k)}\in {\mathbb{R}}^{n\times d}$，$H^{(0)}=F$，一种常见的GNN模型表示为

$$H^{(k)}=\text{ReLU}({\tilde{D}}^{-\frac{1}{2}}\tilde{A}{\tilde{D}}^{-\frac{1}{2}}{H}^{(k-1)}{W}^{(k-1)})$$

其中，$\tilde{A}=A+I$，$\tilde{D}=\text{diag}({\sum }_j{\tilde{A}}_{ij})$，为简化起见，使用 $Z=\text{GNN}(A,X)$ 表示上式，其中 $Z,X\in {\mathbb{R}}^{n\times d}$，$Z$ 表示该层输出的点特征矩阵，$X$ 表示输入的点特征矩阵。对于池化层来说，可以表示为
$$(Z^{\prime },A^{\prime })=\text{Pool}(A,X)$$
其中 $A^{\prime }\in {\mathbb{R}}^{m\times m}$，$Z^{\prime }\in {\mathbb{R}}^{m\times d}$ ，且 $m<n$

池化

假定 $n_l$ 表示经过了 $l-1$ 层池化之后的节点个数，那么第 $l$ 层的簇分配矩阵 $S^{(l)}\in {\mathbb{R}}^{n_l\times n_{l+1}}$， $S^{(l)}$ 给出的是簇分配的概率形式，是一种软分配矩阵。如果已知簇分配矩阵，那么有

$$\begin{gathered} {X}^{(l+1)}=(S^{(l)}{)}^{\text{T}}{Z}^{(l)}\in {\mathbb{R}}^{n_{l+1}\times d} \\ {A}^{(l+1)}=(S^{(l)}{)}^{\text{T}}{A}^{(l)}{S}^{(l)}\in {\mathbb{R}}^{n_{l+1}\times n_{l+1}} \end{gathered}$$

使用两个独立的GNN网络层，一个用来学习点的特征矩阵，一个用来学习簇分配矩阵

$$\begin{gathered} {Z}^{(l)}={\text{GNN}}_{\text{embed}}^{(l)}(A^{(l)},X^{(l)}) \\ {S}^{(l)}=\text{softmax}({\text{GNN}}_{\text{pool}}^{(l)}(A^{(l)},X^{(l)})) \end{gathered}$$

对于最后一层池化层，强制 $n_{l+1}=1$，$S^{(l)}$ 为全1矩阵。

容易证明只要GNN具有排列不变性(Permutation invariance)，那么DiffPool具有排列不变性。

辅助损失函数

对于第 $l$ 层，定义损失函数

$$L_{\text{LP}}^{(l)}=||A^{(l)},S^{(l)}(S^{(l)}{)}^{\text{T}}|{|}_F$$

矩阵A的Frobenius范数定义为矩阵A各项元素的绝对值平方的总和
$$||A|{|}_F=\sqrt{\sum _i\sum _j|a_{i,j}{|}^2}$$

另一个辅助损失函数是为了使簇分配函数尽量接近one-hot向量，即定义正则项

$$L_{\text{E}}^{(l)}=\frac{1}{n_l}\sum _{i=1}^{n_l}H(S_i^{(l)})$$

其中，$S_i^{(l)}$ 表示矩阵 $S^{(l)}$ 的第 $i$ 行，$H$ 表示熵函数

实验

作者在实验中观察到，尽管使用DiffPool性能显着提高，但可能无法稳定训练，并且即使使用相同的超参数设置，不同次运行之间的精度也存在显着差异。 可以看出，添加辅助损失函数可以使训练更加稳定，并减少不同次运行之间准确性的标准偏差。

作者在GraphSAGE和Structure2Vec两种算法上进行了实验，证明了DiffPool的通用性。

作者将不同层的聚类结果可视化显示出来，用来证明算法学到了有意义的聚类方式。

作者观察到DIFFPOOL学会了以非均匀的方式将节点折叠为簇，并且倾向于将密集连接的子图折叠为簇。由于GNN可以有效地对密集的，类似团的子图执行消息传递(直径较小)，因此在这样一个密集的子图中将节点集中在一起不太可能导致结构信息的丢失。这直观地解释了为什么折叠密集子图是DIFFPOOL的有用合并策略。

由于分配网络基于输入节点及其邻居的特征来计算簇分配矩阵，因此具有相似输入特征和邻域结构的节点将具有相似的集群分配。 尽管某两个节点相距较远，但它们对于自己和所有邻居都具有完全相同的邻域结构和特征。 在这种情况下，DIFFPOOL网络将它们分配到同一群集中，这与其他体系结构中的池化概念不同。

DIFFPOOL对于超参数最大集群数量是敏感的，作者发现簇分配根据网络深度和最大集群数而变化较大。使用较大的最大集群数，合并的GNN可以对更复杂的层次结构进行建模，但是很大的参数会导致更大的噪声和更低的效率。 特别的，簇分配矩阵可能不会使用某些群集。 对应于未使用集群的列所有节点的值都较低。

总结

作者提出的DIFFPOOL给出了一种有效的层次化聚类的方法，并给出了两个未来的研究方向

• 学习硬簇分配矩阵(one-hot)
• 将DIFFPOOL应用到其他领域中

个人认为，作者在实验中观察到的三个现象都可以继续往下研究，尤其是第三点需要改进为一个对超参数不敏感的聚类方法。