（论文速读）CLUSTER-GCN：一种训练深度和大图卷积网络的有效算法

论文题目Cluster-GCN: An Efficient Algorithm for Training Deep and Large Graph Convolutional NetworksCLUSTER-GCN一种训练深度和大图卷积网络的有效算法会议KDD 2019摘要图卷积网络(GCN)已经成功地应用于许多基于图的应用中然而训练大规模的GCN仍然是具有挑战性的。当前基于SGD的算法要么存在随着GCN层数呈指数级增长的高计算代价要么存在保存整个图并将每个节点嵌入到内存中所需的大量空间。本文利用图的聚类结构提出了一种新的适用于基于SGD的训练的GCN算法--ClusterGCN。CLUSTER-GCN的工作原理如下在每一步它都会对与图聚类算法识别的密集子图相关联的节点块进行采样并将邻域搜索限制在该子图内。这种简单但有效的策略显著提高了内存和计算效率同时能够实现与以前的算法相当的测试精度。为了测试算法的可扩展性我们创建了一个新的Amazon2M数据它有200万个节点和6100万条边比之前最大的公开可用的数据集(Reddit)大了5倍以上。对于基于这些数据的3层GCN的训练Cluster-GCN比以前最先进的VR-GCN更快(1523秒比1961秒)并且使用的内存要少得多(2.2 GB比11.2 GB)。此外对于4层GCN的训练我们的算法可以在36分钟左右完成而现有的GCN训练算法由于内存不足的问题而无法训练。此外Club-GCN允许我们在不需要太多时间和内存开销的情况下训练更深的GCN这导致了预测精度的提高-使用5层Club-GCN我们在PPI数据集上获得了最先进的测试F1分数99.36而之前的最好结果是98.71分。代码在https://github.com/google-research/google-research/tree/master/cluster_gcn上公开提供。Cluster-GCN 详解用图聚类让大规模 GCN 训练又快又省内存图卷积网络GCN已经在半监督节点分类、链接预测、推荐系统等任务上大放异彩。但当图规模从几千节点扩展到百万级时训练 GCN 立刻变成一场内存和时间的噩梦。KDD 2019 上的 Cluster-GCN 给出了一个非常优雅的解决方案用图聚类的结构来构造 mini-batch。本文带你完整拆解这篇论文。1. 为什么训练大规模 GCN 这么难GCN 的每一层做的事情可以写成问题就在于 A 把邻居的嵌入耦合了进来——一个节点的损失依赖它 L-hop 内的所有邻居。GCN 越深依赖爆炸得越厉害。论文从内存 / 每 epoch 时间 / 收敛速度三个维度总结了已有算法算法内存每 epoch 时间收敛Full-batch GD (Kipf 2017)❌ O(NFL)✅❌ 每 epoch 只更新一次Vanilla SGD / GraphSAGE✅❌ O(d^L) 邻居爆炸✅VR-GCN❌ 要存全部历史嵌入✅✅论文 Table 1GCN 训练算法的时间与空间复杂度对比一图胜千言邻居扩展的可怕之处见下论文 Figure 1传统 GCN 邻居扩展 vs Cluster-GCN 的簇内扩展。比较了传统的图卷积方法和本文提出的聚类法的邻域扩展差异。红色节点是邻域节点扩展的起始节点。传统的图卷积方法存在邻域指数扩张的问题而我们的方法可以避免昂贵的邻域扩张。左边是传统的邻居扩展——从一个红色起点出发几层之后几乎整张图都被拉进来。右边则是 Cluster-GCN 的思路把邻居搜索限制在一个稠密子图内。2. 关键概念嵌入利用率Embedding Utilization作者提出了一个很关键的度量如果第 l 层某节点的嵌入 z_i^(l) 被上一层重用 u 次那么它的嵌入利用率就是 u。全批量 GD每个嵌入被平均重用 d 次d 是平均度利用率最大。随机 mini-batch SGD因为图稀疏采出来的节点几乎没重叠利用率接近 0。如何让 mini-batch 的利用率变高答案是让一个 batch 内的节点之间边尽可能多也就是让 ‖A_{B,B}‖₀ 最大化。这恰好就是一个图聚类问题。3. Vanilla Cluster-GCN把图 G 用 METIS 划分成 c 个子图 {G₁, …, G_c}邻接矩阵重排后其中 Ā 是块对角部分簇内边Δ 是簇间边。作者直接用 Ā 近似 A于是 GCN 的前向传播天然沿着块对角分解开每个 batch 就是一个 cluster只加载 Aₜₜ、Xₜ、Yₜ做纯粹的稠密矩阵乘法不用邻居采样。时间复杂度O(‖A‖₀F NF²) ——对 L 线性内存复杂度O(bLF LF²) ——只存当前 batch 的嵌入那随机划分行不行作者做了实验论文 Table 2随机划分 vs METIS 聚类划分在 Cora / Pubmed / PPI 上的对比。图的随机分区与集群分区(在小批量SGD上训练)。集群分区会带来更好的性能(就测试F1分数而言)因为它删除了较少的分区间链接。这三个数据集都是GCN公共数据集。我们将在实验部分解释PPI数据。CORA有2,708个节点和13,264条边Pubmed有19,717个节点和108,365条边。结论PPI 上从 68.1 → 92.9聚类划分几乎是碾压式胜利。保留簇内边太重要了。4. Vanilla Cluster-GCN 的两个坑与随机多聚类改进Vanilla 版本有两个问题簇间边 Δ 被完全丢掉损失了信息聚类倾向于把相似节点聚在一起导致每个 cluster 的标签分布偏斜batch 间方差大SGD 估计有偏。作者在 Reddit 上算了每个 cluster 的标签熵论文 Figure 2随机划分 vs 聚类划分下 batch 内标签熵的直方图。基于标签分布的熵值直方图。在这里我们在每个批中使用随机分区而不是集群分区。大多数聚类划分的批次具有较低的标签熵这表明每个批次内的标签分布是倾斜的。相比之下随机划分将导致批处理中的标签熵更大尽管它的效率较低。在本例中我们使用300个聚类对Reddit数据集进行划分。大部分聚类 batch 的熵都非常低——标签严重偏斜。解决方案Stochastic Multiple Partitions先划分成 p 个 clusterp 取得偏大每个 batch 随机抽 q 个 cluster 合并把这 q 个 cluster 之间的簇间边加回来。论文 Figure 3随机多聚类方案示意。提出了一种随机多分区方案。在每个时期我们随机抽样集群(本例中使用Q2)和它们的集群间链接以形成新的批次。相同的色块在同一批次中。这一招同时解决两个问题簇间边被找回多 cluster 的组合稀释了标签偏斜。论文 Figure 4Reddit 上单 cluster vs 多 cluster 组合的收敛对比多 cluster 组合明显更稳、收敛更好。完整算法见论文 Algorithm 1。5. 训练更深的 GCN对角增强技巧之前的工作普遍认为GCN 加深没用甚至变差但作者认为那是因为数据集太小 优化困难。他们提出一个新的归一化再加一个对角增强项直觉是近邻的贡献应比远邻更大同时避免深层的数值不稳定。效果非常戏剧化论文 Table 11不同对角增强策略在 PPI 上 2–8 层的表现对比。使用不同对角增强技术的比较。对于所有方法我们给出了在200个历元内达到的最佳验证精度。本实验采用PPI丢失率为0.1。其他设置与第4.1节中的相同。标有红色的数字表示收敛较差。不加增强的方案在 7、8 层直接崩到 43 左右无法收敛而 λ1 的对角增强版本能稳稳跑到 8 层还保持 96.2。论文 Figure 58 层 GCN 的收敛曲线对比6. 实验结果数据集见论文 Table 3。PPI、Reddit 是公开常用集Amazon2M 是作者构造的新集——244.9 万节点6185.9 万边比 Reddit 大 5 倍以上。6.1 训练速度与精度论文 Figure 6Cluster-GCN vs VRGCN vs GraphSAGE 在 PPI / Reddit / Amazon 上 2–4 层 GCN 的训练时间 vs F1 曲线在 PPI 和 Reddit 上Cluster-GCN 全面领先。Amazon 上因为没节点特征、计算被稀疏矩阵操作主导2 层和 4 层略慢于 VRGCN——作者指出这是 PyTorch 稀疏张量支持不如 TensorFlow 成熟造成的Table 6 里做了 benchmark。6.2 内存占用论文 Table 5不同方法在 PPI / Reddit / Amazon 上 2/3/4 层的内存对比举一个直观的例子Reddit 上 4 层 GCN、512 隐层单元VRGCN 要2064 MBCluster-GCN 只要308 MB——差了近 7 倍。原因是 VRGCN 要存所有层的历史嵌入而 Cluster-GCN 只存当前 batch。6.3 Amazon2M 上的对决论文 Table 8Amazon2M 上时间、内存、F1 三项对比亮点3 层VRGCN 1961 秒 / 11218 MBCluster-GCN1523 秒 / 2235 MB4 层VRGCN直接 OOMCluster-GCN 2289 秒 / 2241 MB还拿到最高 F1 90.41。6.4 深度 GCN 的时间增长论文 Table 9PPI 上 2–6 层的训练时间对比VRGCN 从 2 层的 103.6s 一路飙到 6 层的 1956s指数级而 Cluster-GCN 从 52.9s 到 157.3s线性——完美印证了理论复杂度。6.5 SOTA 结果论文 Table 10PPI 和 Reddit 上的 SOTA 对比PPICluster-GCN 5 层 2048 隐层单元 → F1 99.36刷新此前 GaAN 的 98.71RedditCluster-GCN 4 层 128 隐层单元 → F1 96.60同样是最好成绩。7. 总结Cluster-GCN 的核心贡献可以浓缩成一句话通过图聚类构造 mini-batch最大化嵌入利用率让 GCN 训练同时拥有全梯度下降的时间效率和 SGD 的内存效率。三大记忆点概念嵌入利用率 一个 batch 内的边数聚类构造 batch 让它最大化算法METIS 划分 随机多聚类组合加回簇间边深度对角增强归一化让 GCN 能真正跑到 7–8 层不崩。在 Amazon2M 这个 200 万节点的图上Cluster-GCN 只要不到 1 小时、2GB 显存就能训练 4 层 GCN 达到 F1 90.41。这在当时是 GCN 可扩展性的一次巨大突破也为后续大规模图神经网络训练框架如 PyG、DGL 中的 ClusterLoader铺平了道路。