论文笔记：AAAI 2021 Identity-aware Graph Neural Networks

前言

消息传递图神经网络(GNNs)为关系数据提供了强大的建模框架。然而，现有的 GNN 的表达能力受到1-Weisfeiler-Lehman (1-WL) 图同构检验的上界，即不能预测节点聚类系数和最短路径距离的 GNN。基于此本文提出了一类传递消息的 GNN，称为身份感知图神经网络 (ID-GNNs)，具有比 1-WL 测试更强的表达能力。ID-GNN 为现有 GNN 的局限性提供了一个新的解决方案。ID-GNN 通过在消息传递过程中归纳地考虑节点的身份来扩展现有的 GNN 体系结构。
为了嵌入一个给定的节点，ID-GNN 首先提取以该节点为中心的网络，然后进行轮次异构消息传递，中心节点与网络中其他周围节点应用不同的参数集。并且作者进一步提出了一个简单并且更快的 ID-GNN，它将节点标识信息作为增强节点特征注入。ID-GNN的两个版本都表示消息传递 GNN 的一般扩展。

论文链接：https://cs.stanford.edu/people/jure/pubs/idgnn-aaai21.pdf
code：http://snap.stanford.edu/idgnn/

1. Weisfeiler-Lehman(WL)算法和WL Test

内容参考于知乎：https://zhuanlan.zhihu.com/p/90645716

GCN模型可以看作图上非常有名的 Weisfeiler-Lehman 算法的一种变形

一维的Weisfeiler-Lehman如下所示：对于任意的节点 $v_i\in \mathcal{G}$ :

1. 获取节点 $v_i$ 的所有邻居节点 $j$ 的特征 $h_{v_j}$
2. 更新该节点的特征 h i ← h a s h { ∑ j h v j } h_i \leftarrow hash\{\sum_jh_{v_j}\} hi​←hash{∑j​hvj​​}，其中 $hash$ 是一个单射函数，也就是说对于不同的节点特征一定有不同的消息传递信息
3. 重复以上步骤 $K$ 次直到收敛

事实上，Weisfeiler-Lehman 算法在大多数图上会得到一个独一无二的特征集合，这意味着图上的每一个节点都有着独一无二的角色定位（例外在于网格，链式结构等等）。因此，对于大多数非规则的图结构，得到的特征可以作为图是否同构的判别依据，也就是WL Test（例如两个图是否是同质的，取决于节点的排列方式）。

对于 GCN 的信息传播策略，可以表示为 $h_{v_i}^{l+1}=\sigma ({\sum }_j\frac{1}{c_{ij}}{h}_{v_j}^l{W}^l)$，其中 $c_{ij}=\sqrt{d_j{d}_j}，j\in N_i$，$N_i$ 是节点 $i$ 的邻居节点。$d_i,d_j$ 分别是节点 $i,j$ 的度。本质上是在模型中使用了对称规范化之后的邻接矩阵 $D^{-\frac{1}{2}}A{D}^{-\frac{1}{2}}$，可以看出这个传播规则其实是加了参数 $W^l$ 后略微变化的 hash函数。如果我们采用合适的非线性函数 $\sigma$ 然后随机初始化参数矩阵那么这个矩阵就会是正交的，这个更新策略就会非常可靠。这样子我们就可以通过局部图结构中的距离得到非常有意义的平滑嵌入表示。

2. Introduction

现有 GNN 的一个基本限制是，两个具有不同邻域结构的节点可以有相同的计算图，从而出现不可区分的现象。计算图就是产生节点嵌入的邻域节点信息传播的过程。这样的失效案例非常多，以图1中的节点分类为例，A上的节点 v1 和图B上的节点 v2 的分类为例，它们的 GNN 聚合子树完全一样，在缺少有效的节点特征时，聚合之后的节点表示会完全一样，即：节点v1和节点v2的表示完全一样，两者无法区分。但是如果我们将 v1 和 v2 与其邻居区分开来(图中以着色的形式区分)，那么 v1 和 v2 的 GNN 聚合子树就不一样，进而两者的表示也不一样。在链路预测任务中，不能区分相邻结构相同但到源节点最短路径距离不同的候选节点。而在图分类任务中，传统的 GNNs 无法区别拓扑结构不同但是连接关系相似的多图关系。虽然特定任务的特征增强可以用来缓解这些失效模式，但为不同任务发现有意义的特征的过程并不是通用的，而且可能会影响 GNN 的传播效率。

目前已经有针对于图分类任务和链路预测的具体方法来克服上述的缺陷，虽然这些方法有丰富的理论基础，但它们通常是特定于任务的(图分类或连接预测)，并且在计算或实现中经常受到增加的复杂性的影响。相反，消息传递 GNN 具有跨节点、链路和图级任务的高预测性能，同时实现起来简单有效。因此，将通过 GNN 传递的消息扩展到 1-WL 测试之外，以克服当前 GNN 的局限性，是一个非常重要的问题。

基于此，作者提出了 Identity-aware Graph Neural Networks (ID-GNNs)，一种可以超越 1-WL 测试的消息传递 GNN 框架。ID-GNNs 提供了一个通用扩展，使任何传递 GNN 的现有消息更具表现力。ID-GNN通过在消息传递过程中归纳地考虑节点的身份来嵌入每个节点。这种方法不同于用 one-hot 编码标记每个节点，这种编码是转导的(不能推广到看不见的图)。如图1所示，ID-GNNs 使用归纳身份着色技术来区分节点本身(计算图中的根节点)和其各自计算图中的本地邻域内的其他节点。添加的身份信息允许ID-GNN区分节点、边和图级任务之间的相同计算图，这种方式克服了前面讨论的限制。

作者提出了 ID-GNN 的两个版本。作为一种通用方法，身份信息通过应用多轮异构消息传递来得到。具体来说，要嵌入一个给定的节点，ID-GNN 首先提取以该节点为中心的自我网络，然后应用消息传递，来自中心节点(图1中彩色节点)和其他节点的消息使用不同的参数集进行计算。这种方法自然适用于涉及节点或边缘特征的应用。同时作者还考虑了 ID-GNN 的简化版本，其中通过来自给定节点的循环计数作为扩展节点特征注入身份信息。这些循环计数通过计算 ID-GNN 计算图的每一层中的彩色节点来捕获节点身份信息，并可以通过图的邻接矩阵的幂有效地计算。

3. ID-GNN

图可以表示为 $\mathcal{G}=(\mathcal{V},\mathcal{E})$，其中 V = { v 1 , … , v n } \mathcal{V} = \{v_1,\dots,v_n\} V={v1​,…,vn​} 代表节点集，$\mathcal{E}$ 代表图中节点的连接关系，节点特征可以表示为 X = { x v ∣ ∀ v ∈ V } \mathcal{X =\{\mathbf{x}_v | \forall v\in\mathcal{V}\}} X={xv​∣∀v∈V}，同时边也可以存在特征，表示为 F = { f u v ∣ ∀ e u v ∈ E } \mathcal{F}=\{f_{uv}|\forall e_{uv}\in\mathcal{E}\} F={fuv​∣∀euv​∈E}。消息传递的 GNN 的目标是基于局部邻域网络信息聚合的迭代过程学习到具有实际含义的节点特征嵌入 $h_v$。因此迭代 $k$ 次的信息传递的迭代或者 $k$ 层的GNN 网络可以表示为：
m u ( k ) = M S G ( k ) ( h u ( k − 1 ) ) h v ( k ) = A G G ( k ) ( { m u ( k ) , u ∈ N ( v ) } , h v ( k − 1 ) ) (1) \mathbf{m}_u^{(k)}=\mathbf{MSG}^{(k)}(h_u^{(k-1)})\\h_v^{(k)}=\mathbf{AGG}^{(k)}(\{\mathbf{m}_u^{}(k),u\in\mathcal{N}(v)\},h_v^{(k-1)})\tag{1} mu(k)​=MSG(k)(hu(k−1)​)hv(k)​=AGG(k)({mu​(k),u∈N(v)},hv(k−1)​)(1)

其中 $h_v^{(k)}$ 代表经过 $k$ 次迭代的节点特征嵌入，因此存在关系 $h_v^{(0)}={\mathbf{x}}_v$，$m_v^{(k)}$ 代表消息嵌入，$\mathcal{N}(v)$ 代表 $v$ 的局部邻居。不同的 GNN 具有不同的 ${\mathbf{M}\mathbf{S}\mathbf{G}}^{(k)}(\cdot )$ 和 ${\mathbf{A}\mathbf{G}\mathbf{G}}^{(k)}(\cdot )$。

3.1 Identity-aware Graph Neural Networks

ID-GNN 的两个重要组成部分：

1. 归纳身份着色，将身份信息注入到每个节点。
2. 异构消息传递，在消息传递中使用身份信息。

Inductive identity coloring

ID-GNN 生成节点嵌入，首先要生成中心节点 $v$ 的自我 $K$ 跳网络 ${\mathcal{G}}_v^{(K)}$，然后 ID-GNN 给自我网络的中心节点分配一个独特的颜色。总的来说，${\mathcal{G}}_v^{(K)}$ 中的节点在整个嵌入过程中可以分为两种类型：有颜色的节点和没有颜色的节点。这种着色技术是归纳性的，因为即使对节点进行排列，自我网络的中心节点仍然可以与其他邻近节点区分开来。

Heterogeneous message passing

将 $K$ 轮消息传递应用到所有提取的自我网络中。为了嵌入节点 $u\in {\mathcal{G}}_v^{(K)}$，ID-GNN 定义异构消息传递框架如下：
m s ( k ) = M S G 1 [ s = v ] ( k ) ( h s ( k − 1 ) ) h u ( k ) = A G G ( k ) ( { m s ( k ) , s ∈ N ( u ) } , h u ( k − 1 ) ) (3) \mathbf{m}_s^{(k)}=\mathbf{MSG}_{\mathbb{1}[s=v]}^{(k)}(h_s^{(k-1)})\\h_u^{(k)}=\mathbf{AGG}^{(k)}(\{\mathbf{m}_s^{}(k),s\in\mathcal{N}(u)\},h_u^{(k-1)})\tag{3} ms(k)​=MSG1[s=v](k)​(hs(k−1)​)hu(k)​=AGG(k)({ms​(k),s∈N(u)},hu(k−1)​)(3)
其中，只有当经过 $K$ 轮消息传递后，$h_v^{(K)}$ 才被用于节点 $v$ 的嵌入。并且与式（1）的区别在于，应用了两种 ${\mathbf{M}\mathbf{S}\mathbf{G}}^{(k)}$，其中 ${\mathbf{M}\mathbf{S}\mathbf{G}}_1^{(k)}(\cdot )$ 被用于经过分配颜色的节点，相反 ${\mathbf{M}\mathbf{S}\mathbf{G}}_0^{(k)}(\cdot )$用于没有被分配过颜色的节点。因此 $1[s=v]$ 代表着如果 $s=v$ 则为 1，否则为 0。这样，归纳身份着色被编码到 ID-GNN 计算图中。

这种异构消息传递方法的一个优点是，它适用于任何传递 GNN 的消息。可以支持在消息传递过程中包含边的特征信息 $f_{su}$

m s u ( k ) = M S G 1 [ s = v ] ( k ) ( h s ( k − 1 ) , f s u ) h u ( k ) = A G G ( k ) ( { m s u ( k ) , s ∈ N ( u ) } , h u ( k − 1 ) ) (4) \mathbf{m}_{su}^{(k)}=\mathbf{MSG}_{\mathbb{1}[s=v]}^{(k)}(h_s^{(k-1)},f_{su})\\h_u^{(k)}=\mathbf{AGG}^{(k)}(\{\mathbf{m}_{su}^{}(k),s\in\mathcal{N}(u)\},h_u^{(k-1)})\tag{4} msu(k)​=MSG1[s=v](k)​(hs(k−1)​,fsu​)hu(k)​=AGG(k)({msu​(k),s∈N(u)},hu(k−1)​)(4)

Extension to edge-level tasks

在此部分继续扩展ID-GNN框架，以正确解决 edge-level 任务中现有的GNN限制(图1，中间)。假设我们想预测节点对 $u,v$ 的边标签。对于 ID-GNN，由一个嵌入 $h_{u|v}$ 的条件节点进行预测，通过在节点 $u$ 的计算图中分配节点 $v$ 而不是节点 $u$ 的身份着色来计算，如图1所示。当节点 $v$ 不在 $u$ 的 $K$ 跳自我网络内时，不使用身份着色，ID-GNNs仍然会遭遇已有的GNNs失效案例。因此在实践中使用更深层的 ID-GNN 来进行 edge-level 预测任务。

3.2 ID-GNNs Expressive Power: Theoretical Results

ID-GNNs are strictly more expressive than existing message passing GNNs

通过 1-WL 测试已经证明，通过 GNN 传递的现有消息具有表达能力上限，该上限可以通过图同构网络(GIN)实例化。ID-GNN 版本的 GIN 可以区分任何 GIN 可以区分的图，同时能够区分某些 GIN 无法区分的图。

ID-GNNs can count cycles
前面 ID-GNN 增加的表达能力。其中的一个具体方面在于 ID-GNN 计数周期的能力。计数周期的能力是直观的，它对于预测节点聚类系数等任务是至关重要的。
对于任意节点 $v$，存在一个 $K$ 层 ID-GNN 实例化，它可以学习嵌入 $h_v^{(K)}$，对于 $j=1，\dots ，K$，其中第 $j$ 维 $h_v^{(K)}[j]$ 等于从节点 $v$ 开始到节点 $v$ 结束的长度 $j$ 的周期数。

3.3 ID-GNNs Expressive Power: Case Studies

Node-level: Predicting clustering coefficient

现有的消息传递 GNN 不能纯粹从图结构归纳地预测聚类系数，而 ID-GNN 可以。聚类系数是一个广泛使用的度量，它表征了在一个节点的1跳邻居中闭合三角形的比例。图1中的节点分类失败案例表明， GNNs 无法预测聚类系数，因为GNNs无法区分聚类系数分别为 1 和 0 的节点 v1 和 v2 。通过使用 one-hot 节点特征，GNNs 可以解决这种问题。然而，在这种情况下，GNN 会记住每个节点的聚类系数，因为 one-hot 编码会防止泛化到看不见的图。

ID-GNNs可以学习节点嵌入$h_v^{(K)}$ ，其中 $h_v^{(K)}[j]$ 等于从节点 $v$ 开始到结束的长度为 $j$ 的周期数。有了这些循环次数，我们就可以计算节点 $v$ 的聚类系数 $c_v$：
c v = ∣ { e s u : s , u ∈ N ( v ) , e s u ∈ E } ∣ ( d u ) ( d v − 1 ) / 2 = h v ( K ) [ 3 ] h v ( K ) [ 2 ] ∗ ( h v ( K ) [ 2 ] − 1 ) (5) c_v=\frac{|\{e_{su}:s,u\in\mathcal{N}(v),e_{su}\in\mathcal{E}\}|}{(d_u)(d_v-1)/2}\\=\frac{h_v^{(K)}[3]}{h_v^{(K)}[2]*(h_v^{(K)}[2]-1)}\tag{5} cv​=(du​)(dv​−1)/2∣{esu​:s,u∈N(v),esu​∈E}∣​=hv(K)​[2]∗(hv(K)​[2]−1)hv(K)​[3]​(5)
其中 $d_v$ 代表节点 $v$ 的度。并且 $c_v$ 是 $h_v^{(K)}$ 的连续函数，因此在误差 $ϵ$ 的范围内，我们可以通过多层感知机来逼近它的值。

Edge-level: Predicting reachability or shortest path distance
Vanilla GNNs 使用两个节点的嵌入来完成节点间的连接预测任务，这种方法不能预测节点对之间的可达性或最短路径距离(SPD)。例如，两个节点可以嵌入相同的 GNN 节点，而不管它们是否位于相同的连接组件中如图1。尽管适当的节点特征初始化可以预测可达性和SPD，但 ID-GNN 通过使用条件节点嵌入为这一限制提供了一个通用解决方案。正如前面说明的，ID-GNN 将 edge-level 预测重新表述为条件节点级预测。因此，一个 $K$ 层 ID-GNN 可以通过条件节点嵌入 $h_{u|v}^{(K)}$ 来预测节点 $u\in \mathcal{G}$ 是否可以从 $v\in \mathcal{G}$ 在 $K$ 跳内到达:

m s ∣ v ( k ) = { 1 , i f 1 [ s = v ] = 1 h u ∣ v ( K ) , e l s e h u ∣ v ( k ) = M A X ( { m s ∣ v ( k ) , s ∈ N ( u ) } ) (6) m_{s|v}^{(k)}= \begin{cases} 1,\quad if \mathbb{1}[s=v]=1\\ h^{(K)}_{u|v}, \quad else \end{cases}\\h_{u|v}^{(k)}=\mathbf{MAX}(\{\mathbf{m}_{s|v}^{(k)},s\in\mathcal{N}(u)\}) \tag{6} ms∣v(k)​={1,if1[s=v]=1hu∣v(K)​,else​hu∣v(k)​=MAX({ms∣v(k)​,s∈N(u)})(6)
其中 $h_{u|v}^{(0)}=0,\forall u\in \mathcal{G}$，如果 ID-GNN 认为 $u$ 和 $v$ 之间存在连接关系，则输出 $h_{u|v}^{(K)}=1$。

Graph-level: Differentiating random d-regular graphs
如图1所示，现有的消息传递 GNN 不能单纯从图结构中区分随机 d维正规图，因为无论层数多少，每个节点的计算图都是相同的。ID-GNNs 可以区分随机 d维正规图。具体来说，生成了100个非同构随机 d维正规图，考虑3种不同图大小( $n$ )和节点度( $d$ )的设置。使用 $K$ 个周期计数，$K$ 层 ID-GNN 可以成功表示，来计算这些 d维正规图可以被区分的百分比。表1的结果证实，身份信息的添加可以极大地帮助区分d维正规图。

3.4 ID-GNN-Fast: Injecting Identity via Augmented Node Features

考虑到:

• （1）GNN 的 minibatch 实现在提取自我网络时具有计算开销，这是异构消息传递的 ID-GNN 所需要的。
• （2）循环计数信息解释了 ID-GNN 增加的表达能力的一个重要方面。相对于现有的 GNN，作者提出了 ID-GNN-Fast，其中通过使用循环计数作为增强节点特征来注入身份信息。

类似的循环计数信息在图核研究的背景下也被证明是有用的。具体来说，ID-GNN 使用长度为 1…K 的周期数作为增广节点特征 ${\mathbf{x}}_v^{+}\in {\mathbb{R}}^K$，起始和结束于节点 $v$。这些额外的特征 ${\mathbf{x}}_v^{+}$ 可以通过 ${\mathbf{x}}_v^{+}[k]=Diag(A^k)[v]$ 的稀疏矩阵乘法有效地计算出来，其中 $A$ 是邻接矩阵。然后通过连接这个增强特性 ${\mathbf{x}}_v=\mathbf{C}\mathbf{O}\mathbf{N}\mathbf{C}\mathbf{A}\mathbf{T}({\mathbf{x}}_v,{\mathbf{x}}_v^{+})$ 来更新所有节点的输入特征。

4. Experiments