PyG教程(5)：剖析GNN中的消息传播机制

一.前言

众所周知，图神经网络可以从空域或谱域来对其进行研究。其中，空域角度主要借助消息传播机制来构建GNN。本文主要介绍的是消息传递机制，为下篇文章具体介绍PyG中是如何实现消息传播机制做好铺垫。

二.消息传递框架概述

消息传递是Gilmer等在Neural Message Passing for Quantum Chemistry中提出来的从空域角度定义GNN的范式。假设${\mathbf{x}}_i^{(k-1)}\in {\mathbb{R}}^F$表示节点$i$在第$k-1$层的特征，${\mathbf{e}}_{j,i}\in {\mathbb{R}}^D$表示节点$j$到节点$i$的边上的特征，则消息传播机制可以用如下公式来描述：
$${\mathbf{x}}_i^{(k)}={\gamma }^{(k)}\left({\mathbf{x}}_i^{(k-1)},{\square }_{j\in \mathcal{N}(i)}{\varphi }^{(k)}\left({\mathbf{x}}_i^{(k-1)},{\mathbf{x}}_j^{(k-1)},{\mathbf{e}}_{j,i}\right)\right)\text{\hspace{1em}(1)}$$
在消息传播机制中，主要分为三大步骤：消息生成、消息聚合、消息更新。

2.1 消息生成与传播

在本阶段中，每个节点将生成自己的消息，然后向自己的邻居节点”传播“自己的消息，也就是公式(1)中的：
$${\varphi }^{(k)}\left({\mathbf{x}}_i^{(k-1)},{\mathbf{x}}_j^{(k-1)},{\mathbf{e}}_{j,i}\right)$$
其中，${\varphi }^{(k)}$表示可微函数，例如MLP。在生成消息的过程中，可能会用到：

• 节点自己当前的特征（${\mathbf{x}}_i^{(k-1)}$）
• 节点邻居当前的特征（${\mathbf{x}}_j^{(k-1)}$）
• 节点与其邻居间边的特征（${\mathbf{e}}_{j,i}$）

当然上述三者并不都是必须的，具体使用什么来生成节点的消息取决于GNN的构建者。

2.2 消息聚合

在本阶段，每个节点会聚合来自邻居的消息，也就是公式(1)中的：
$${\square }_{j\in \mathcal{N}(i)}(\text{Message})$$
其中$\text{Message}$指代2.1节中每个节点的消息，$\mathcal{N}(i)$表示节点$i$的邻域，$\square$表示可微的、转置不变（permutation invariant）函数。转置不变指聚合邻居的消息的结果与邻居的聚合顺序无关，常见的包括sum,max,mean。

2.3 消息更新

在本阶段，每个节点利用聚合自邻居节点的消息来生成自己的消息，也就是公式(1)中的：
$${\gamma }^{(k)}\left({\mathbf{x}}_i^{(k-1)},\text{NeighborMsg}\right)$$
其中$NeighborMsg$指代2.2节中每个节点聚合自邻居的消息，${\gamma }^{(k)}$也表示可微函数，例如MLP。

2.4 消息传递机制小结

上述的消息传播机制可以用下图概括：

其中图中的方框便是的便是聚合邻居的神经网络。

经过前面的介绍可知：空域角度定义的GNN间的不同之处便在于它们关于消息生成、消息聚合和消息更新的实现的不同。

三.消息传播机制的示例

为了方便理解上面的消息传递机制，本节主要展示GCN、GraphSAGE中的消息传播机制。

3.1 GCN中的消息传播机制

GCN虽然是从谱域角度定义的，但同样从空域角度来对其进行解释，其所对应的消息传播机制如下：
$$h_i^{(l+1)}=\sigma (\frac{1}{\sqrt{|\mathcal{N}(i)|}}\sum _{j\in \mathcal{N}(i)}\frac{1}{\sqrt{|\mathcal{N}(j)|}}{h}_j^{(l)}{W}^{(l)})$$
生成消息：
$$\frac{1}{\sqrt{|\mathcal{N}(j)|}}{h}_j^{(l)}{W}^{(l)}$$
消息聚合：
$$\sigma (\frac{1}{\sqrt{|\mathcal{N}(i)|}}\sum _{j\in \mathcal{N}(i)}\text{Message})$$
其中$\sigma$指非线性激活，通常为ReLU。

消息更新：消息聚合来的消息。

3.2 GraphSAGE中的消息传播机制

GraphSAGE中的消息传播机制如下所示：
$$h_i^{(l+1)}=\sigma \left(W^{(l)}\cdot \mathrm{CONCAT}\left(h_i^{(l)},\mathrm{A}\mathrm{G}\mathrm{G}\left(\left\{h_j^{(l)},\forall j\in N(i)\right\}\right)\right)\right)$$
其中$\text{norm}$表示L2正则化。

消息生成：
$$h_j$$
两阶段聚合（首先聚合邻居的消息，然后聚合自身的消息）：
$$\begin{gathered} h_{N(i)}^{(l+1)}\leftarrow \mathrm{A}\mathrm{G}\mathrm{G}\left(\left\{h_j^{(l)},\forall j\in N(i)\right\}\right) \\ \sigma \left(W^{(l)}\cdot \mathrm{CONCAT}\left(h_i^{(l)},h_{N(i)}^{(l+1)}\right)\right) \end{gathered}$$
其中$\text{AGG}$在GraphSAGE中有三种实现：Mean、Pool或LSTM。

消息更新：消息聚合来的消息。

四.结语

参考资料：

• CREATING MESSAGE PASSING NETWORKS

GNN中的消息传播机制是借助PyG之类的图神经网络框架来编写自己的消息传播GNN的基础，只有对其了解较为深刻，你才能更好的设计自己的GNN模型。