图神经网络的统一彩票假设
机器学习
2021-06-08 v2 人工智能
机器学习
摘要
随着图规模快速增长以及更深层的图神经网络（GNNs）不断涌现，GNN 的训练与推理变得愈发昂贵。现有的网络权重剪枝算法无法解决 GNN 中由图的规模和连通性所导致的主要空间与计算瓶颈。为此，本文首先提出一个统一的 GNN 稀疏化（UGS）框架，同时剪枝图邻接矩阵和模型权重，以有效加速大规模图上的 GNN 推理。借助这一新工具，我们进一步首次将近期流行的彩票假设推广到 GNN，将图彩票（GLT）定义为一对核心子数据集与稀疏子网络，它们可通过迭代应用 UGS 从原始 GNN 和完整稠密图中联合识别出来。如同其在卷积神经网络中的对应概念，GLT 可独立训练以匹配使用完整模型和图进行训练的性能，并且可从随机初始化和自监督预训练的 GNN 中抽取。我们的方案已在多种 GNN 架构和不同任务上、在小型图数据集（Cora、Citeseer 和 PubMed）以及来自具有挑战性的开放图基准（OGB）的大规模数据集上得到实验验证。具体而言，对于节点分类，我们发现的 GLT 在小型图上以 20%~98% 的 MACs 节省、在大型图上以 25%~85% 的 MACs 节省取得相同准确率。对于链接预测，GLT 分别在小型和大型图数据集上带来 48%~97% 和 70% 的 MACs 节省，且不损害预测性能。代码见 https://github.com/VITA-Group/Unified-LTH-GNN。
引用
@article{arxiv.2102.06790,
title = {A Unified Lottery Ticket Hypothesis for Graph Neural Networks},
author = {Tianlong Chen and Yongduo Sui and Xuxi Chen and Aston Zhang and Zhangyang Wang},
journal= {arXiv preprint arXiv:2102.06790},
year = {2021}
}