GNN的花式研究越来越多了~ 本来读了这篇后想写一下, 发现AI in Graph的小伙伴已经写的挺好了~
今天给大家介绍德克萨斯大学奥斯汀分校电子与计算机工程系 Atlas 教授团队发表在 ICML2021 的文章“A Unified Lottery Ticket Hypothesis for Graph Neural Networks”。本文将 LTH(彩票假设) 拓展到了图神经网络,实现了有效的图神经网络模型压缩。
随着图规模的迅速增长和深度图神经网络(GNN)的发展,GNN 的训练和推理变得越来越昂贵。现有的网络剪枝算法无法解决 GNNs 中由图的大小和连通性导致的空间和计算瓶颈问题。为此,本文首先提出了一个 unified GNN sparsification(UGS)框架,同时对图的邻接矩阵和模型权重进行剪枝,有效地加速了大规模图的 GNN 推理。利用这一新工具,首次将最近流行的彩票假说(LTH:lottery ticket hypothesis)推广到 GNNs,将图彩票(GLT:graph lottery ticket)定义为一对子数据集和稀疏子网络,通过迭代应用 UGS 可以从原始 GNN 和稠密图中联合识别。
与卷积神经网络中的 LT 一样,GLT 可以单独训练,使性能与完整的模型和图数据相当,并且可以从随机初始化和自监督的预训练神经网络中提取。GLT 在各种 GNN 架构和不同的任务中得到了实验验证,包括小规模的图数据集(Cora、Citeseer 和 PubMed)和来自 OGB 的大型数据集。具体来说,对于节点分类,GLT 在相同准确率下,在小图上节省了20%∼ 98%的 MACs(multiply–accumulate operations),在大图上节省了25%∼ 85%的 MACs。对于链路预测,GLT 在不影响预测性能的情况下,分别在小型和大型图数据集上节省了48%~97%和70%的 MACs。
GNNs 已经在各种图任务上取得了最优性能,但是 GNNs 的训练和推理效率低下,这给GNNs扩展到实际的大规模图应用带来了障碍,这一障碍来自于算法和硬件两个层面。
在算法层面上,GNN模型可以被认为是由传统的图组成,这些图的顶点特征配备了深度神经网络算法。GNN 推理的执行分为三类,它们具有独特的计算特性:图遍历、DNN 计算和聚合。GNN 广泛遵循递归的邻域聚合(或消息传递)方案,其中每个节点聚合其多跳邻居的特征向量以计算其新特征向量。当图很大并且具有密集/复杂的邻居连接时,聚合阶段需要大量计算。
在硬件层面上,GNN 的计算依赖于矩阵的稀疏和不规则结构,导致产生许多随机内存访问和有限的数据重用,但也需要相对较少的计算。因此,GNN 比其他神经网络具有更高的参考延迟,限制了其在离线计算推理中的应用。
本文的目的是从算法层面上降低爆炸性的 GNN 复杂度。主要有两种方案:简化图、简化模型。第一种方案主要是增加图的稀疏性,降低消息聚集的计算量。第二种方案主要是减少模型参数。
以两层的 GCN 模型为例:
我们将模型记为
,图
包括
,
。定义两个遮罩
和
,分别指示了图中不重要的连接和模型中不重要的参数。
的形状与图的邻接矩阵相同,
的形状与模型参数相同。给定
,
通过以下目标共同优化。
给定一个 GNN
和一个图
, GNN相关子网络和子图可以定义为
和
,
和
是二部的遮罩。如果一个子网络
训练在一个稀疏图
上有和在原图上训练的完整GNN模型有相当或更好的性能,则定义
) 为同一的图彩票 unified graph lottery tickets (GLTs),其中
是原GNN模型的初始化参数。然后使用如下算法找到GTL。
本文使用GCN、GIN、GAT作为图模型,在Cora、Citeseer、PubMed三个数据集上进行了节点分类和链路预测实验,部分结果如下:
文中使用GraphCL进行自监督预训练,从中获得GLT。下图展示了其与随机初始化的对比结果。
本文对稀疏化的子图做了可视化,结果如下。