消息传递图神经网络

消息传递范式
是一种聚合邻接节点信息来更新中心节点信息的范式
(1)邻接节点信息变换、(2)邻接节点信息聚合到中心节点、(3)聚合信息变换
用${\mathbf{x}}_i^{(k-1)}\in {\mathbb{R}}^F$表示$(k-1)$层中节点$i$的节点特征，${\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),$$
其中$\square$表示可微分的、具有排列不变性（函数输出结果与输入参数的排列无关）的函数。具有排列不变性的函数有，和函数、均值函数和最大值函数。$\gamma$和$\varphi$表示可微分的函数，如MLPs（多层感知器）。此处内容来源于CREATING MESSAGE PASSING NETWORKS。
Pytorch Geometric中的MessagePassing基类
Pytorch Geometric(PyG)提供了MessagePassing基类，它实现了消息传播的自动处理，继承该基类可使我们方便地构造消息传递图神经网络，我们只需定义函数$\varphi$，即message()函数，和函数$\gamma$，即update()函数，以及使用的消息聚合方案，即aggr="add"、aggr="mean"或aggr="max"。这些是在以下方法的帮助下完成的：

• MessagePassing(aggr="add", flow="source_to_target", node_dim=-2)：

  • aggr：定义要使用的聚合方案（“add”、"mean "或 “max”）；
  • flow：定义消息传递的流向（"source_to_target "或 “target_to_source”）；
  • node_dim：定义沿着哪个轴线传播。

• MessagePassing.propagate(edge_index, size=None, **kwargs)：

  • 开始传播消息的起始调用。它以edge_index（边的端点的索引）和flow（消息的流向）以及一些额外的数据为参数。
  • 请注意，propagate()不仅限于在形状为[N, N]的对称邻接矩阵中交换消息，还可以通过传递size=(N, M)作为额外参数。例如，在二部图的形状为[N, M]的一般稀疏分配矩阵中交换消息。
  • 如果设置size=None，则假定邻接矩阵是对称的。
  • 对于有两个独立的节点集合和索引集合的二部图，并且每个集合都持有自己的信息，我们可以传递一个元组参数，即x=(x_N, x_M)，来标记信息的区分。

• MessagePassing.message(...)：

  • 首先确定要给节点$i$传递消息的边的集合，如果flow="source_to_target"，则是$(j,i)\in \mathcal{E}$的边的集合；
  • 如果flow="target_to_source"，则是$(i,j)\in \mathcal{E}$的边的集合。
  • 接着为各条边创建要传递给节点$i$的消息，即实现$\varphi$函数。
  • MessagePassing.message(...)函数接受最初传递给MessagePassing.propagate(edge_index, size=None, **kwargs)函数的所有参数。
  • 此外，传递给propagate()的张量可以被映射到各自的节点$i$和$j$上，只需在变量名后面加上_i或_j。我们把$i$称为消息传递的目标中心节点，把$j$称为邻接节点。

• MessagePassing.aggregate(...)：

  • 将从源节点传递过来的消息聚合在目标节点上，一般可选的聚合方式有sum, mean和max。

• MessagePassing.message_and_aggregate(...)：

  • 在一些场景里，邻接节点信息变换和邻接节点信息聚合这两项操作可以融合在一起，那么我们可以在此函数里定义这两项操作，从而让程序运行更加高效。

• MessagePassing.update(aggr_out, ...):

  • 为每个节点$i\in \mathcal{V}$更新节点表征，即实现$\gamma$函数。该函数以聚合函数的输出为第一个参数，并接收所有传递给propagate()函数的参数。

  继承MessagePassing类的GCNConv
  GCNConv的数学定义为
  x i ( k ) = ∑ j ∈ N ( i ) ∪ { i } 1 deg ⁡ ( i ) ⋅ deg ⁡ ( j ) ⋅ ( Θ ⋅ x j ( k − 1 ) ) , \mathbf{x}_i^{(k)} = \sum_{j \in \mathcal{N}(i) \cup \{ i \}} \frac{1}{\sqrt{\deg(i)} \cdot \sqrt{\deg(j)}} \cdot \left( \mathbf{\Theta} \cdot \mathbf{x}_j^{(k-1)} \right), xi(k)​=j∈N(i)∪{i}∑​deg(i) ​⋅deg(j) ​1​⋅(Θ⋅xj(k−1)​),
  其中，相邻节点的特征首先通过权重矩阵$\mathbf{\Theta }$进行转换，然后按端点的度进行归一化处理，最后进行加总。这个公式可以分为以下几个步骤：

1. 向邻接矩阵添加自环边。
2. 线性转换节点特征矩阵。
3. 计算归一化系数。
4. 归一化$j$中的节点特征。
5. 将相邻节点特征相加（"求和 "聚合）。

步骤1-3通常是在消息传递发生之前计算的。步骤4-5可以使用MessagePassing基类轻松处理。该层的全部实现如下所示

import torch
from torch_geometric.nn import MessagePassing
from torch_geometric.utils import add_self_loops, degree

class GCNConv(MessagePassing):
    def __init__(self, in_channels, out_channels):
        super(GCNConv, self).__init__(aggr='add', flow='source_to_target')
        # "Add" aggregation (Step 5).
        # flow='source_to_target' 表示消息从源节点传播到目标节点
        self.lin = torch.nn.Linear(in_channels, out_channels)

    def forward(self, x, edge_index):
        # x has shape [N, in_channels]
        # edge_index has shape [2, E]

        # Step 1: Add self-loops to the adjacency matrix.
        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))

        # Step 2: Linearly transform node feature matrix.
        x = self.lin(x)

        # Step 3: Compute normalization.
        row, col = edge_index
        deg = degree(col, x.size(0), dtype=x.dtype)
        deg_inv_sqrt = deg.pow(-0.5)
        norm = deg_inv_sqrt[row] * deg_inv_sqrt[col]

        # Step 4-5: Start propagating messages.
        return self.propagate(edge_index, x=x, norm=norm)

    def message(self, x_j, norm):
        # x_j has shape [E, out_channels]
        # Step 4: Normalize node features.
        return norm.view(-1, 1) * x_j

GCNConv继承了MessagePassing并以"求和"作为领域节点信息聚合方式。该层的所有逻辑都发生在其forward()方法中。在这里，我们首先使用torch_geometric.utils.add_self_loops()函数向我们的边索引添加自循环边（步骤1），以及通过调用torch.nn.Linear实例对节点特征进行线性变换（步骤2）。

归一化系数是由每个节点的节点度得出的，它被转换为每个边的节点度。结果被保存在形状[num_edges,]的张量norm中（步骤3）。

在message()函数中，我们需要通过norm对相邻节点特征x_j进行归一化处理。这里，x_j包含每条边的源节点特征，即每个中心节点的邻接。

这就是创建一个简单的x传递层的全部内容。我们可以把这个层作为深度架构的构建块。我们可以很方便地初始化和调用它：

conv = GCNConv(16, 32)
x = conv(x, edge_index)

propagate函数

def propagate(self, edge_index: Adj, size: Size = None, **kwargs):
    r"""开始消息传播的初始调用。
    Args:
        edge_index (Tensor or SparseTensor): 定义了消息传播流。
        	当flow="source_to_target"时，节点`edge_index[0]`的信息将被发送到节点`edge_index[1]`，
        	反之当flow="target_to_source"时，节点`edge_index[1]`的信息将被发送到节点`edge_index[0]`
        kwargs: 图其他属性或额外的数据。
    """

其他
Torch geometric GCNConv 源码分析