腾讯Angel Graph团队刷新GNN最强榜单OGB世界纪录!

导读 / Introduction

近日，在国际顶级图学习标准OGB（Open Graph Benchmark）挑战赛中，腾讯大数据Angel Graph团队联合北京大学-腾讯协同创新实验室，以较大优势在三个最大的OGB分类数据集：ogbn-papers100M、ogbn-products和ogbn-mag三项任务榜单第一！

OGB是目前公认最权威的图学习通用性能评价基准数据集，由斯坦福大学Jure Leskovec教授团队建立并开源，并吸引了斯坦福大学、康奈尔大学、Facebook、NVIDIA、百度、阿里巴巴和字节跳动等国际顶尖高校与科技巨头参与。该数据集来源广泛——涵盖了生物网络、分子图、学术网络和知识图谱等领域，且囊括了基本的节点预测、边预测、图预测等图学习任务，数据真实、极具挑战性，素有图神经网络领域“ImageNet”之称，已成为全球图神经网络研究者检验自身功力的“试剑石”。

一、问题背景

图结构数据广泛存在于现实生活中。如图1所示，在推荐系统中可以用图来表示用户与商品之间的交互、在社交网络和知识图谱中图可以来表示各个用户或者实体之间的关系，图还可以用来建模各种药物和新材料。随着深度学习的兴起，研究者开始将深度学习应用到图数据，推动了图领域的相关研究蓬勃发展。作为图领域的重要的技术手段，图神经网络（GNN）已经成为解决图问题的重要算法，并受到了学术界和工业界的广泛关注。

图1 常见的图数据类型

二、技术难点

尽管GNN在多个应用场景中取得了巨大的成功，在工业界的大规模数据集上使用GNN还是存在着以下两个问题：

1. 可扩展性低

1.1 单机的存储问题 (High Memory Cost)

传统的GNN层中包含两个操作，特征传播和非线性变换。其中，特征传播操作涉及一个较大的稀疏邻接矩阵和特征矩阵的乘法，需要花费较多的时间；而且，如果在训练时使用GPU加速，那么就必须把稀疏邻接矩阵放置到GPU的显存上，但这一需求对于大图来说是无法被满足的。比如，OGB节点分类数据集中最大的ogbn-papers100M数据集的稀疏邻接矩阵大小超过50GB，目前仅有80GB显存版本的Tesla A100才能存下，而这还不包括存储特征和模型的空间。因此，如果要想在超大规模的图上应用传统的GNN模型，如GCN[1]，目前唯一可行的方法是使用分布式训练。

1.2 分布式训练的通信问题 (High Communication Cost)

图2 增加worker数量时两层GraphSAGE的加速比和训练瓶颈

GNN的每一次特征传播都需要拉取邻居特征，对于k层的GNN来说，每个节点需要拉取的k跳以内邻居节点特征随着层数增加以指数增加。对于稠密的连通图，每个节点在每次训练的时候几乎需要拉取全图的节点信息，造成海量的通信开销。尽管有一些采样算法能缓解这个问题，但并不能很好地解决。我们以两层的GraphSAGE[2]为例，图2(a) 显示当增加workers数量之后加速比和理想情况下相差甚远，这是因为当图节点被切分到多个workers上后需要去不同workers节点上拉取邻居的频率也显著增加。如图2(b) 所示，随着workers数量增加，通信开销会明显大于计算开销。因此，如何支持超大规模图数据的分布式训练一直是一个研究热点。

2. 灵活性差 (Low Flexibility)

2.1 约束非线性变换深度恒等于特征传播深度

如图3所示，对于常见的耦合的图神经网络，它在每一层中都强制要求非线性变换紧接着特征传播，去掉特征传播后GCN就退化为MLP。我们之前对深层GNN的评测工作[3]表明，图神经网络存在着两种深度：特征传播的深度和非线性变换的深度，他们分别导致了GNN里的过平滑以及模型退化问题。另外，这个工作还指出我们在图稀疏（标签、特征或者边稀疏）的时候需要增大，在图规模较大的时候需要增大。假设图稀疏且规模较小，这个时候我们需要大的和小的,强行约束会造成次优效果。  

图3 GCN和MLP的关系示意

2.2 没有考虑节点特异性

图4 感受野扩展速度的示例

图5 不同节点的预测准确性和特征传播步数的关系

我们通过一个简