DataWhale图网络学习（五）超大图上的节点表征学习

0 节点表征学习

首先我们回忆一下图网络的节点表征学习，对于给定图$G=(\mathcal{V},\mathcal{E},A)$，它由$N=|\mathcal{V}|$个节点和$|\mathcal{E}|$条边组成，其邻接矩阵记为$A$。节点的原始表征为$X\in {\mathbb{R}}^{N\times F}$的$F$维特征。对于有L层的图卷积层的图神经网络，每一层都通过聚合上一层的图中节点的邻居的嵌入来构建每个节点的嵌入：
$$Z^{(l+1)}=A^{\prime }{X}^{(l)}{W}^{(l)},X^{(l+1)}=\sigma \left(Z^{(l+1)}\right)\text{\hspace{1em}(1)}$$
其中$X^{(l)}\in {\mathbb{R}}^{N\times F_l}$是第$l$层$N$个节点的嵌入，$A^{\prime }$是归一化和规范化后的邻接矩阵，$W^{(l)}\in {\mathbb{R}}^{F_l\times F_{l+1}}$是特征转换矩阵，也就是要学习的参数。

当图神经网络应用于半监督节点分类任务时，训练的目标是通过最小化损失函数来学习公式（1）中的权重矩阵：
$$\mathcal{L}=\frac{1}{\left|{\mathcal{Y}}_L\right|}\sum _{i\in {\mathcal{Y}}_L}\mathrm{loss}\left(y_i,z_i^L\right)\text{\hspace{1em}(2)}$$
其中，$y_L$包含所有被标记节点的标签；$z_i^{(L)}$是$Z^{(L)}$的第$i$行，表示节点$i$的最终预测，并且其对应真实标签为$y_i$。

1 图网络面临的问题

图神经网络已经成功的应用于许多图节点预测和边预测任务中，然而这些图都不太大，对于在超大图上进行图神经网络学习依然具有较大的挑战：

• 要么面临随着图神经网络层数增加计算成本呈指数增加的问题；
• 要么面临保存整个图的信息和每一层每个节点的嵌入（embedding）到GPU内存而消耗巨大内存空间的问题；
• 要么面临需要计算所有节点的损失产生的梯度，这种训练方式需要非常巨大的计算开销和GPU内存开销。

针对这个问题论文Cluster-GCN: An Efficient Algorithm for Training Deep and Large Graph Convolutional Network提出了一种新的图神经网络，以及一种新的训练方法。Cluster-GCN论文提出：

• 利用图节点聚类算法将一个图的节点划分为$c$个簇，每一次选择几个组的节点和这些节点对应的边构成一个子图，然后对子图做训练。
• 由于是利用图节点聚类算法将节点划分为多个簇，所以簇内边的数量要比簇间边的数量多得多。
• 基于小图进行训练，不会消耗很多内存空间，于是我们可以训练更深的神经网络，从而得到更高的精度。
• 采用mini-batch SGD的方式训练可以提高图神经网络的训练速度并减少GPU内存需求。

2 Cluster-GCN分析

2.1 简单的Cluster-GCN

Cluster-GCN的技术是由这样的问题驱动的：在mini-batch SGD的参数更新中，我们能否设计一个将节点分成多个batch，对应地将图划分成多个计算子图，来最大化嵌入利用率？

对于一个图$G$，我们将其节点划分为$c$个簇：$\mathcal{V}=\left[{\mathcal{V}}_1,\cdots {\mathcal{V}}_c\right]$，其中${\mathcal{V}}_t$由第$t$个簇中的节点组成。因此我们有$c$个子图：
G ˉ = [ G 1 , ⋯   , G c ] = [ { V 1 , E 1 } , ⋯   , { V c , E c } ] \bar{G}=\left[G_{1}, \cdots, G_{c}\right]=\left[\left\{\mathcal{V}_{1}, \mathcal{E}_{1}\right\}, \cdots,\left\{\mathcal{V}_{c}, \mathcal{E}_{c}\right\}\right] Gˉ=[G1​,⋯,Gc​]=[{V1​,E1​},⋯,{Vc​,Ec​}]
其中每个${\mathcal{E}}_t$只由${\mathcal{V}}_t$中的节点之间的边组成。重组节点后，邻接矩阵被划分为大小为$c^2$的子矩阵，即
$$A=\bar{A}+\Delta =\left[\begin{array}{ccc}{A}_{11}& \cdots & A_{1c}\\ ⋮& \ddots & ⋮\\ A_{c1}& \cdots & A_{cc}\end{array}\right]\text{\hspace{1em}(4)}$$
其中
$$\bar{A}=\left[\begin{array}{ccc}{A}_{11}& \cdots & 0\\ ⋮& \ddots & ⋮\\ 0& \cdots & A_{cc}\end{array}\right],\Delta =\left[\begin{array}{ccc}0& \cdots & A_{1c}\\ ⋮& \ddots & ⋮\\ A_{c1}& \cdots & 0\end{array}\right]\text{\hspace{1em}(5)}$$
其中，对角线每个块$A_{tt}$都是大小为$\left|{\mathcal{V}}_t\right|\times \left|{\mathcal{V}}_t\right|$的邻接矩阵，它由$G_t$内的边构成。$\bar{A}$是图$\bar{G}$的邻接矩阵。$A_{st}$由两个簇${\mathcal{V}}_s$和${\mathcal{V}}_t$之间的边构成。

$\left[Y_1,\cdots ,Y_c\right]$，其中$X_t$和$Y_t$分别由$V_t$中节点的特征和标签组成。$\Delta$是由$A$的所有非对角线块组成的矩阵。同样，我们可以根据$\left[{\mathcal{V}}_1,\cdots ,{\mathcal{V}}_c\right]$划分特征矩阵$X$和训练标签$Y$，得到$\left[X_1,\cdots ,X_c\right]$和$\left[Y_1,\cdots ,Y_c\right]$，其中$X_t$和$Y_t$分别由$V_t$中节点的特征和标签组成。

用块对角线邻接矩阵$\bar{A}$去近似邻接矩阵$A$的好处是，图神经网络的目标函数变得可以分解为不同的batch。以${\bar{A}}^{\prime }$表示$\bar{A}$的归一化版本，最终的嵌入矩阵成为
$$\begin{array}{rl}{Z}^{(L)}& ={\bar{A}}^{\prime }\sigma \left({\bar{A}}^{\prime }\sigma \left(\cdots \sigma \left({\bar{A}}^{\prime }X{W}^{(0)}\right)W^{(1)}\right)\cdots \right)W^{(L-1)}\\ & =\left[\begin{array}{c}{\bar{A}}_{11}^{\prime }\sigma \left({\bar{A}}_{11}^{\prime }\sigma \left(\cdots \sigma \left({\bar{A}}_{11}^{\prime }{X}_1{W}^{(0)}\right)W^{(1)}\right)\cdots \right)W^{(L-1)}\\ ⋮\\ {\bar{A}}_{cc}^{\prime }\sigma \left({\bar{A}}_{cc}^{\prime }\sigma \left(\cdots \sigma \left({\bar{A}}_{cc}^{\prime }{X}_c{W}^{(0)}\right)W^{(1)}\right)\cdots \right)W^{(L-1)}\end{array}\right]\end{array}\text{\hspace{1em}(6)}$$
由于$\bar{A}$的块对角形式（注意${\bar{A}}_{tt}^{\prime }$是${\bar{A}}^{\prime }$的相应对角块）。损失函数也可以分解为
$${\mathcal{L}}_{{\bar{A}}^{\prime }}=\sum _t\frac{\left|{\mathcal{V}}_t\right|}{N}{\mathcal{L}}_{{\bar{A}}_{tt}^{\prime }}\text{ and }{\mathcal{L}}_{{\bar{A}}_{tt}^{\prime }}=\frac{1}{\left|{\mathcal{V}}_t\right|}\sum _{i\in {\mathcal{V}}_t}\mathrm{loss}\left(y_i,z_i^{(L)}\right)\text{\hspace{1em}(7)}$$
然后，Cluster-GCN基于公式(6)和公式(7)的分解形式。在每一步参数更新，我们采样一个簇${\mathcal{V}}_t$，然后根据${\mathcal{L}}_{{{\bar{A}}^{\prime }}_{tt}}$的梯度进行参数更新。这种训练方式，只需要用到子图$A_{tt}$，当前batch的$X_t,Y_t$以及神经网络参数${\left\{W^{(l)}\right\}}_{l=1}^L$。 实现起来只需要（公式(6)的一个区块的）矩阵乘积的前向和梯度反向传播。

我们使用图节点聚类算法来划分图。图节点聚类方法将图上节点分成多个簇，使簇内边远多于簇间边。这些正是我们所需要的，因为 1）如前所述，每个batch的嵌入的利用率相当于簇内边数量。也就是，每个节点和它的邻居大部分情况下都位于同一个簇中，因此多跳后的邻居节点大概率仍然在同一个簇中。2）由于我们用块对角线近似邻接矩阵$\bar{A}$代替邻接矩阵$A$，并且误差与簇间连接$\Delta$成正比，所以我们需要找到一个分区方式，在这种分区方式下簇间连接的数量最小。

我们可以看到，Cluster-GCN可以避免高代价的邻域搜索（图右），而专注于每个簇内的邻居。表1显示了两种不同的节点划分策略：随机划分与聚类划分。我们可以看到，在相同的epoches下，使用聚类分区可以达到更高的精度。这表明使用图聚类很重要，分区不应该随机形成。

表1：图的随机分区与聚类分区的对比（采用mini-batch SGD训练）。聚类分区得到更好的性能（就测试F1集得分而言），因为它删除了较少的分区间边。

Dataset	random partition	clustering partition
Cora	78.4	82.5
Pubmed	78.9	79.9
PPI	68.1	92.9

2.2 随机多分区

尽管简单Cluster-GCN实现了良好的计算和内存复杂性，但仍有两个潜在问题：

• 图被分割后，一些边（公式(4)中的$\Delta$部分）被移除。性能可能因此会受到影响。
• 图聚类算法倾向于将相似的节点聚集在一起。因此，单个簇中节点的分布可能与原始数据集不同，导致在进行SGD参数更新时对梯度的估计有偏差。

对于一个样本标签不均衡的例子，与随机划分相比，我们清楚地看到大多数簇的熵值较小，表明簇的标签分布偏向于某些特定的标签。这增加了不同batch的差异，并可能影响训练的收敛性。

为了解决上述问题，我们提出了一种随机多簇方法，来合并簇之间的边并减少各簇之间的熵值差异。我们首先将图划分为$p$簇，${\mathcal{V}}_1,\cdots ,{\mathcal{V}}_p$，$p$是一个较大的值。在构建用于SGD参数更新的batch时，我们不是只考虑一个簇，而是随机选择$q$个簇，表示为$t_1,\dots ,t_q$，得到的数据batch包含节点$\left\{{\mathcal{V}}_{t_1}\cup \cdots \cup {\mathcal{V}}_{t_q}\right\}$ 、簇内边 { A i i ∣ i ∈ t 1 , … , t q } \left\{A_{i i} \mid i \in t_{1}, \ldots, t_{q}\right\} {Aii​∣i∈t1​,…,tq​}和簇间边 { A i j ∣ i , j ∈ t 1 , … , t q } \left\{A_{i j} \mid i, j \in t_{1}, \ldots, t_{q}\right\} {Aij​∣i,j∈t1​,…,tq​}。数据batch中包含了簇间边，从而不同batch间的差异减小。

下图展示了此算法，在每个epoch中，不同的簇被选择来组成一个数据batch，并且我们可以观察到，使用多个簇来组成一个batch可以提高收敛性。最终的Cluster-GCN算法在算法1中呈现。

2.3 深层GCNs的训练问题

深度GCN模型的优化变得很困难，因为它可能会阻碍前几层的信息被传递。在[9]中，他们采用了一种类似于残差连接的技术，使模型能够将前一层的信息带到下一层。具体来说，他们修改了公式(1)，将第$l$层的隐藏表征添加到下一层，如下所示
$$X^{(l+1)}=\sigma \left(A^{\prime }{X}^{(l)}{W}^{(l)}\right)+X^{(l)}\text{\hspace{1em}(8)}$$
在原始的GCN设置中，每个节点都聚合了其邻居在上一层的表示。然而，在深层GCN的设置下，该策略可能并不适合，因为它没有考虑到层数的问题，也就是说，附近的邻居应该比远处的节点贡献更大。在这里，我们提出了另一种简单的技术来改善深度GCN的训练，放大每个GCN层中使用的邻接矩阵$A$的对角线部分。通过这种方式，我们在每个GCN层的聚合中对来自上一层的表示赋予更多的权重。一个例子是给$\bar{A}$添加一个标识，如下所示，
$$X^{(l+1)}=\sigma \left(\left(A^{\prime }+I\right)X^{(l)}{W}^{(l)}\right)\text{\hspace{1em}(9)}$$
虽然公式(9)似乎是合理的，但对所有节点使用相同的权重而不考虑其邻居的数量可能不合适。
此外，它可能会受到数值不稳定的影响，因为当使用更多的层时，数值会呈指数级增长。因此，我们提出了一个修改版的公式(9)，以更好地保持邻居信息和数值范围。我们首先给原始的$A$添加一个标识，并进行归一化处理
$$\tilde{A}=(D+I{)}^{-1}(A+I)\text{\hspace{1em}(10)}$$
然后考虑，
$$X^{(l+1)}=\sigma \left((\tilde{A}+\lambda \mathrm{diag}(\tilde{A}))X^{(l)}{W}^{(l)}\right)\text{\hspace{1em}(11)}$$

3 Cluster-GCN实践

PyG为Cluster-GCN提出的训练方式和神经网络的构建提供了良好的支持。我们无需在意图节点是如何被划分成多个簇的，PyG提供的接口允许我们像训练普通神经网络一样在超大图上训练图神经网络。

3.1 数据集准备

from torch_geometric.datasets import Reddit
from torch_geometric.data import ClusterData, ClusterLoader, NeighborSampler

dataset = Reddit('../dataset/Reddit')
data = dataset[0]
print(dataset.num_classes)
print(data.num_nodes)
print(data.num_edges)
print(data.num_features)

41
232965
114615873
602

该数据集包含41个分类任务，232,965个节点，114,615,873条边，节点维度为602维。

3.2 图节点聚类与数据加载器生成

cluster_data = ClusterData(data, num_parts=1500, recursive=False, save_dir=dataset.processed_dir)
train_loader = ClusterLoader(cluster_data, batch_size=20, shuffle=True, num_workers=0)
subgraph_loader = NeighborSampler(data.edge_index, sizes=[-1], batch_size=1024, shuffle=False, num_workers=0)

• train_loader，此数据加载器遵循Cluster-GCN提出的方案，图节点首先被聚类，返回的一个batch由多个簇组成。
• subgraph_loader，使用此数据加载器不对图节点聚类，计算一个batch中的节点的嵌入需要计算该batch中的所有节点的距离从$0$到$L$的邻居节点。

3.3 构造图神经网络的规范

class Net(torch.nn.Module):
    def __init__(self, in_channels, out_channels):
        super(Net, self).__init__()
        self.convs = ModuleList(
            [SAGEConv(in_channels, 128),
             SAGEConv(128, out_channels)])

    def forward(self, x, edge_index):
        for i, conv in enumerate(self.convs):
            x = conv(x, edge_index)
            if i != len(self.convs) - 1:
                x = F.relu(x)
                x = F.dropout(x, p=0.5, training=self.training)
        return F.log_softmax(x, dim=-1)

    def inference(self, x_all):
        pbar = tqdm(total=x_all.size(0) * len(self.convs))
        pbar.set_description('Evaluating')

        # Compute representations of nodes layer by layer, using *all*
        # available edges. This leads to faster computation in contrast to
        # immediately computing the final representations of each batch.
        for i, conv in enumerate(self.convs):
            xs = []
            for batch_size, n_id, adj in subgraph_loader:
                edge_index, _, size = adj.to(device)
                x = x_all[n_id].to(device)
                x_target = x[:size[1]]
                x = conv((x, x_target), edge_index)
                if i != len(self.convs) - 1:
                    x = F.relu(x)
                xs.append(x.cpu())

                pbar.update(batch_size)

            x_all = torch.cat(xs, dim=0)

        pbar.close()

        return x_all

此神经网络拥有forward和inference两个方法，forward函数的定义与普通的图神经网络无区别，inference方法应用于推理阶段，为了获取更高的预测精度，所有使用subgraph_loader。

3.4 训练、验证与测试

在训练过程中，我们从train_loader获取batch，每次只根据一个batch进行参数训练。但在测试节点，我们不做图节点聚类，因此不会有边丢失的情况。

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Net(dataset.num_features, dataset.num_classes).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.005)

def train():
    model.train()

    total_loss = total_nodes = 0
    for batch in train_loader:
        batch = batch.to(device)
        optimizer.zero_grad()
        out = model(batch.x, batch.edge_index)
        loss = F.nll_loss(out[batch.train_mask], batch.y[batch.train_mask])
        loss.backward()
        optimizer.step()

        nodes = batch.train_mask.sum().item()
        total_loss += loss.item() * nodes
        total_nodes += nodes

    return total_loss / total_nodes


@torch.no_grad()
def test():  # Inference should be performed on the full graph.
    model.eval()

    out = model.inference(data.x)
    y_pred = out.argmax(dim=-1)

    accs = []
    for mask in [data.train_mask, data.val_mask, data.test_mask]:
        correct = y_pred[mask].eq(data.y[mask]).sum().item()
        accs.append(correct / mask.sum().item())
    return accs


for epoch in range(1, 31):
    loss = train()
    if epoch % 5 == 0:
        train_acc, val_acc, test_acc = test()
        print(f'Epoch: {epoch:02d}, Loss: {loss:.4f}, Train: {train_acc:.4f}, '
              f'Val: {val_acc:.4f}, test: {test_acc:.4f}')
    else:
        print(f'Epoch: {epoch:02d}, Loss: {loss:.4f}')

Epoch: 01, Loss: 1.0916
Epoch: 02, Loss: 0.4641
Epoch: 03, Loss: 0.3921
Epoch: 04, Loss: 0.3615
Epoch: 05, Loss: 0.3500, Train: 0.9523, Val: 0.9504, test: 0.9482

Epoch: 26, Loss: 0.2411
Epoch: 27, Loss: 0.2390
Epoch: 28, Loss: 0.2366
Epoch: 29, Loss: 0.2562
Epoch: 30, Loss: 0.2400, Train: 0.9712, Val: 0.9524, test: 0.9506

可以看到经过学习，训练损失下降，训练准确度、评估准确度和测试准确度都有提升。

参考：

• DataWhale开源内容