【GNN】Diff Pool：网络图的层次化表达

阿泽 Crz

发布于 2020-07-21 11:33:22

1.1K0

发布于 2020-07-21 11:33:22

今天学习的是斯坦福大学的同学 2018 年的工作《Hierarchical Graph Representation Learning with Differentiable Pooling》，目前共有 140 多次引用。

目前，GNN 在图分类任务中的处理方法本质上是平面的（Flat），无法学习图形的层次化表达。对于一个包含多个标签的图来说，传统的方法都是为图中每个节点生成一个 Embedding 向量，然后利用这些 Embedding 向量来做全局池化或者输入到 MLP 中来预测图标签，但这种方法忽视了图的层次结构。

为此，作者提出了一个可微分的图池化模块——Diff Pool，用于完成图的层次化表达，并可以端到端的方式与目前多种模型相结合。

Diff Pool 为每一层的节点都学习了一个可微的软聚类，聚类后的集群作为下一层的输入，最后通过连接 MLP 完成图分类任务。

与现有方法相比，Diff Pool 在五个基准数据集中取得了 SOTA 的成绩。

1.Introduction

目前通用的做法是采用 GNN 来提取用于图分类的特征，基于消息传递的 GNN 架构为：

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

其中，

H^{(k)} \in \mathbb{R}^{n\times d}

为第 k 层节点的 Embedding 特征；M 为消息传递函数，与邻接矩阵和参数

\theta^{(k)}

2.Diff Pool

本文作者提出了可微的池化模块（Diff Pool），可应用于不同 GNN 模型中。Diff Pool 与 CNN 中的池化不同的是，前者不包含空间局部的概念，且每次 pooling 所包含的节点数和边数都不相同。

Diff Pool 在 GNN 的每一层上都会基于节点的 Embedding 向量进行软聚类，通过反复堆叠（Stacking）建立深度 GNN。因此，Diff Pool 的每一层都能使得图越来越粗糙，从而可以产生输入图的层级表征。具体过程如下图所示：

作者使用分配矩阵来进行池化。具体来说，设 l 层的聚类分配矩阵为

S^{(l)}\in \mathbb{R}^{n_l\times n_{l+1}}

，矩阵行表示 l 层的中

n_l

个节点，列表示 l+1 层中的

n_{l+1}

个节点，

S^{(l)}

表示从 l 层到 l+1 层的图节点的软聚类结果。

考虑 l 层的邻接矩阵

A^{(l)}

和节点特征矩阵

Z^{(l)}

，Diff Pool 每一层的粗化图为：

(A^{(l+1)}, X^{(l+1)} = \mathbf{DIFF\;POOL}(A^{(l)}, Z^{(l)}) \\

其中，

X^{l+1}

为下一层粗化图的节点输入特征。

我们将上式分成两部：

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})} \\

第一个式子是特征聚类，第二个式子是转换邻接矩阵。

接下来，我们讨论下如何完成上述目标。

首先，我们会使用两个独立的 GNN 分别学习不同的任务，一个是 Embedding GNN 用于学习 l 层的节点特征：

Z^{(l)} = \mathbf{GNN}_{l,embed}(A^{(l)}, X^{(l)}) \\

另一个是池化 GNN 用于学习 l 层到 l+1 层的聚类分配矩阵：

S^{(l)} = softmax(\mathbf{GNN}_{l,pool}(A^{(l)}, X^{(l)}) ) \\

softmax 每个输出矩阵都会应用到，用于学习聚类分配矩阵。

注意，这两个 GNN 输入都相同，但是参数不同，并且任务也不同。

其次，置换不变形对于图分类任务来说非常重要，池化层应该满足置换不变形。对于 Diff Pool 来说，作者表明只要 GNN 组件满足置换不变形，那么整体就会满足，即，若：

\mathbf{GNN} (A,X) = \mathbf{GNN} (PAP^T,X) \\

那么：

\mathbf{DIFF\; POOL}(A,Z) = \mathbf{DIFF\; POOL}(PAP^T,PX) \\

再者，在训练过程中，如果仅从图分类的角度来训练池化 GNN 是非常困难的，因为这是一个非凸优化问题。为了缓解这个问题，作者在训练过程中加上了一个链接预测的目标，从而促使邻近节点一起池化。

每一层中，我们最小化：

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

其中

||\cdot||_F

表示 Frobenius 范式。

最后，预测簇时的 softmax 结果应该更接近与 one-hot 向量，这样才能使得聚类结果更加清晰。所以我们还需要正则化聚类簇的熵：

L_E = \frac{1}{n} \sum_{i=1}^n H(S_i) \\

其中，H 为熵函数。

训练时，将

L_{LP}

和

L_E

加入到分类损失中，收敛速度会变慢，但是效果会变好，并且聚类的结果也更有解释性。

3.Experiment

简单看下实验。

下表为不同数据集下不同模型的实验结果：

在 Struct2Vec 的基础上应用 Diff Pool：

可视化 Diff Pool 的聚类结果：

4.Conclusion

作者引入了一种可微的池化方法，该方法能够基于网络图自适应的学习提取复杂的层次结构。Diff Pool 可以现有的 GNN 模型结合使用，并在多个基准数据集上取得了不错的成绩。

时间复杂度为

O(n^2)

5.Reference

Ying Z, You J, Morris C, et al. Hierarchical graph representation learning with differentiable pooling[C]//Advances in neural information processing systems. 2018: 4800-4810.

本文参与腾讯云自媒体分享计划，分享自微信公众号。

原始发表：2020-05-16，如有侵权请联系 cloudcommunity@tencent.com 删除

diff