超大图上的节点表征学习

一、图神经网络的挑战

• 随着网络层数的增加, 计算成本呈指数增长
• 保存整个图的信息和每一层每个节点的嵌入(embedding)需要消耗巨大的内存空间

二、 Cluster-GCN的出现

• 可能损失预测精度或者对提高内存的利用率并不显著
• 无需保存整个图的信息和每一层每个节点的嵌入（embedding）

2.1 概览

• 1.利用图节点聚类算法将一个图的节点划分为c个簇, 每一次选择几个组的节点和边构造一个子图, 对子图进行训练
• 2.由于利用图节点聚类算法划分多个簇, 所以簇内边的数量要比簇间边的数量要多得多
• 3.基于子图进行训练, 不会消耗很多内存空间, 所以可以训练更深的神经网络来得到更优的结果

2.2 节点表征学习背景

• 给定一个图$G=(\mathcal{V},\mathcal{E},A)$，它由$N=|\mathcal{V}|$个节点和$|\mathcal{E}|$条边组成，其邻接矩阵记为$A$。每一个节点都关联一个$F$维的属性向量。$X\in {\mathbb{R}}^{N\times F}$表示$N$的特征。一个$L$层的图神经网络由$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$个节点的嵌入（embedding），并且有$X^{(0)}=X$。$A^{\prime }$是归一化和规范化后的邻接矩阵，$W^{(l)}\in {\mathbb{R}}^{F_l\times F_{l+1}}$是特征转换矩阵，也就是要学习的参数。为了简单起见，我们假设所有层的特征维度都是一样的，即 $\left(F_1=\cdots =F_L=F\right)$。激活函数$\sigma (\cdot )$通常被设定为ReLU。

• 当图神经网络应用于半监督节点分类任务时，训练的目标是通过最小化损失函数来学习公式（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$。

3.3 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]$$
  $$\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]$$

• $\Delta$是由 $A$ 的所有非对角线块组成的矩阵

• 用块对角线邻接矩阵$\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}$$

• 损失函数分解为:
  $${\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)$$

• 在每一步参数更新，我们采样一个簇${\mathcal{V}}_t$，然后根据${\mathcal{L}}_{{{\bar{A}}^{\prime }}_{tt}}$的梯度进行参数更新

• 图节点聚类方法将图上节点分成多个簇，使簇内边远多于簇间边

• Cluster-GCN可以避免高代价的邻域搜索，而专注于每个簇内的邻居

3.4 Cluster-GCN的问题

• 图被分割时一些边的 $\Delta$ 部分被移除, 性能可能会受到影响
• 图聚类算法倾向于将相似的节点聚集在一起。因此，单个簇中节点的分布可能与原始数据集不同，导致在进行SGD参数更新时对梯度的估计有偏差。
• 为了解决上述问题, 提出了一种随机多簇的算法

3.5 随机多分区

• 将图划分为$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间的差异减小。如图所示
  
• 对比随机多分区与聚类划分的效果
  

三、 Cluster-GCN实践

3.1 数据集分析

import torch
import torch.nn.functional as F
from torch.nn import ModuleList
from tqdm import tqdm
from torch_geometric.datasets import Reddit, Reddit2
from torch_geometric.data import ClusterData, ClusterLoader, NeighborSampler
from torch_geometric.nn import SAGEConv

dataset = Reddit('input/Reddit')
data = dataset[0]
print(dataset.num_classes)
# 41
print(data.num_nodes)
# 232965
print(data.num_edges)
# 114615892
print(data.num_features)
# 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=12)
subgraph_loader = NeighborSampler(data.edge_index, sizes=[-1], batch_size=1024, shuffle=False, num_workers=12)

3.3 构造神经网络

class Net(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super(Net, self).__init__()
        self.convs = ModuleList(
            [SAGEConv(in_channels, hidden_channels), SAGEConv(hidden_channels, 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):
                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):
        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())
            
            x_all = torch.cat(xs, dim=0)
        
        return x_all

3.4 训练, 测试

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Net(dataset.num_features, 128, 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():
    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}, Val: {val_acc:.4f}, test: {test_acc:.4f}')
    else:
        print(f'Epoch: {epoch:02d}, Loss: {loss:.4f}')

• 这边由于算力有限, 所以输出5轮的结果
  

四、作业

• 尝试将数据集切分成不同数量的簇进行实验，然后观察结果并进行比较。

cluster_data = ClusterData(data, num_parts=1000, recursive=False, save_dir=dataset.processed_dir)

cluster_data = ClusterData(data, num_parts=2000, recursive=False, save_dir=dataset.processed_dir)

• 对比发现1000个簇初始loss最大, 但是收敛最快, 2000个簇的初始loss最小, 但是收敛比较慢, 1500个簇促使loss和收敛速度都居中, 准确率方面1500个簇效果最好, 不过因为跑的轮数比较少, 不足以说明最终结果。

五、参考资料

• Cluster-GCN: An Efficient Algorithm for Training Deep and Large Graph Convolutional Network

• Datawhale GNN组队学习开源教程