WWW 2023 | 更少的内存，更快的速度，更高的准确度——GNN大图训练暨工具包发布

近年来，图神经网络（Graph Neural Network，GNN）在图数据学习上表现出了优越的性能。然而，当面临现实世界中百万乃至亿级节点数量的大图时，现有的GNN仍面临着诸多局限性，主要包括：

（1）消息传递模式中，邻居节点数量随着传播距离指数式增长，这将带来可拓展性问题和过平滑问题；

（2）嵌入向量表的参数数量随着节点数量线性增长，当节点数量较多时，巨大的嵌入向量表将难以优化；

（3）GNN卷积层和嵌入向量表有着不同的性质，同时优化它们可能导致次优的结果。

为了克服这些问题，我们提出了xGCN——一种采用全新训练方式的GNN模型。xGCN将节点嵌入向量表视为“静态特征”，而不是需要梯度训练的参数，使用无监督的传播过程和一个DNN来更新嵌入向量，能够显著减少模型参数数量、加速模型训练，同时达到更好的预测准确率。文章在4个大规模社交网络数据集上进行了实验，结果表明xGCN在准确率和训练速度方面均超过了基线模型。目前，xGCN模型已被应用在Xbox游戏社交平台的社交推荐业务中。

同时，我们推出了一个轻量的GNN模型工具包：XGCN，旨在帮助用户快速地构建和训练大规模网络上的GNN图嵌入模型。XGCN包含了xGCN以及GraphSAGE、LightGCN等常见GNN模型的实现。我们提供了详细的说明文档，方便用户快速上手进行图数据集处理、模型构建和运行。

该成果“xGCN: An Extreme Graph Convolutional Network for Large-scale Social Link Prediction”发表于2023年国际万维网会议（WWW’23）。WWW是互联网技术领域的顶级会议，是中国计算机学会（CCF）推荐的A类会议。

论文链接：https://dl.acm.org/doi/10.1145/3543507.3583340
项目地址：https://github.com/CGCL-codes/XGCN_librar

一、 问题背景

大规模图数据广泛地存在于现代社会中，例如微信、推特等线上社交平台中的社交网络，和淘宝、亚马逊等购物平台中的商品交易网络。图嵌入（Graph Embedding）作为一种图数据挖掘方法，近年来备受学术界和工业界关注，其中的图神经网络（Graph Neural Network，GNN）方法引入了显式编码图结构的卷积层，表现出了优越的性能。然而，当面临现实世界中的大规模网络时，现有的GNN仍面临着诸多局限性，主要包括：

（1）消息传递模式中，邻居节点数量随着传播距离指数式增长，这将带来可拓展性问题和过平滑问题。

（2）在经典的图嵌入任务中，每个节点拥有一个d维的可学习的嵌入向量e。设节点数量为N，则全体嵌入向量构成了一个N x d的矩阵E，即嵌入向量表。该嵌入向量表的参数数量随着节点数量线性增长，当节点数量较多时，嵌入向量表将变得十分庞大。例如，在一个有着一亿用户的社交网络中，假设嵌入向量的维度为32，那么整个嵌入向量表的参数数量为3.2x109。在GNN中，大量的嵌入向量参数将导致过高的训练开销。

（3）GNN卷积层和嵌入向量表有着不同的性质，同时优化这两部分参数可能导致次优的结果。嵌入向量表E是稀疏的，每一个batch的训练只涉及部分参数的更新。而卷积层参数是稠密的，每个训练batch都会被梯度信息全体更新。

二、模型设计

如何更好地在大图上训练GNN？为了解决这个问题，我们提出了xGCN——一种采用全新训练方式的GNN模型。下面首先介绍模型的设计思路，然后介绍模型的具体组成和训练方式。

如图1所示，在模型训练中，经典的随机梯度下降（Stochastic Gradient Descent，SGD）方法直接使用梯度信息学习整个嵌入向量表，从而优化损失函数，每个batch的训练中，其中的嵌入向量都会被动态地更新。在大规模网络上进行训练时，该庞大的嵌入向量表将导致极大的计算开销。

众所周知，大模型的训练是十分困难的，那么能否设计一种减少参数数量的策略，让参数在大部分时候保持“静态”？如图2(a)所示，我们注意到，如果把嵌入向量表视为静态特征，使用一个有着更少参数的深度神经网络（Deep Neural Network，DNN）来学习如何提升嵌入向量的质量，或许同样能达到优化损失函数的效果。此时只有DNN的参数被SGD动态更新，随着训练的进行，DNN输出的嵌入向量将产生更小的损失。DNN的可学习参数数量一般为O(cd^2)（其中c为常数），与节点数量无关，远小于大规模嵌入向量表的参数数量，更加容易训练。

当DNN训练收敛后，我们可以用它输出一个新的、有着更小损失的嵌入向量表。需要注意的是，只训练一趟很可能是不够的，因为开始训练时嵌入向量表是随机初始化的，仅仅训练一个有着少量参数的DNN难以产生较高质量的输出。因此，如图2(b)所示，我们考虑用DNN的输出替换旧的嵌入向量表，再次重复DNN的训练。如果不断地重复上述的“训练DNN”和“刷新嵌入向量表”操作，我们希望节点嵌入向量的质量能够不断提升。

单独的DNN无法像GNN一样显式地编码图结构信息，可能导致较弱的表达能力，所以我们考虑在DNN训练前再加入一个无监督的传播操作，直接汇聚邻居节点的嵌入向量，如图3所示。

基于上述思路，我们的xGCN具体由如下3个模块组成：

（1）嵌入向量传播，以显式编码图结构信息。该过程可表示为：E←PZ，其中Z为初始的嵌入向量表，P为图结构矩阵。矩阵P可以有多种选择，例如归一化的邻接矩阵或者PPR top-k邻居矩阵等。该操作是无监督的且高效的，在DNN训练之前执行。

（2）训练DNN，以提升嵌入向量的质量。如图4所示，我们的DNN由一个前馈网络（Feed-Forward Network，FFN）和一个放缩网络（Scaling Neural Network，SNN）构成。FFN主要负责进行非线性变换，输入和输出都是d维。SSN则输出一个标量，负责进行线性放缩以调整输出向量的尺度。DNN的计算可表示为：X=SNN(E)∙FFN(E)。我们将该DNN称为“RefNet”（Refinement Network）。

（3）嵌入向量刷新：用DNN的输出替换旧的嵌入向量表，从而开始新的一轮嵌入向量传播和DNN训练。

如图5所示，在xGCN的训练中，上述三个模块被依次执行：首先进行嵌入传播，编码图结构信息；接着进行DNN训练，提升嵌入向量的质量；然后是嵌入向量表的刷新，以进行新一轮的传播和提升。我们设计了一个简单的策略来控制嵌入向量刷新的时机和次数：当刷新次数小于K时，固定每T个epoch进行一次刷新，使得模型能够快速热身；当刷新次数达到K时，仅当验证集得分难以持续增长时才触发一次刷新，使得DNN能够充分地被训练。

三、 实验结果

我们在四个大规模社交网络数据集上进行了社交链路预测实验。其中包括三个百万级节点的网络（Pokec、LiveJournal和Xbox-3m）和一个有着一亿用户节点的Xbox游戏社交网络（Xbox-100m）。我们使用了Recall@k和NDCG@k评价指标（在图表中分别缩写为“R@k”和“N@k”）。

模型准确率上，表1展示了xGCN和基线模型在百万级数据集上的结果，我们的结论如下：

（1）xGCN在所有数据集和所有评价指标上都显著超过了最佳的基线模型。例如，在Recall@100指标上，xGCN在Pokec、LiveJournal和Xbox-3m数据集上分别达到了15.34%、16.20%和48.18%的性能提升。

（2）相对于浅层的图嵌入方法，GNN方法普遍有着更好的性能，验证了GNN显式编码邻居信息的有效性。

（3）简化的GNN（如LightGCN和PPRGo）在大多数情况下优于经典的GNN（如GraphSAGE、GAT和GIN），这和LightGCN等工作中得到的结论一致，即：卷积层中的线性变换与激活函数对嵌入的学习可能有着不利影响。

为了研究不同模型的训练效率，我们记录了训练过程中xGCN、LightGCN和PPRGo在验证集上的Recall@100得分。如图6所示，结果显示xGCN使用更少的时间达到了更高的得分。我们观察到xGCN的每个epoch拥有明显更少的耗时。

为了研究xGCN在超大规模数据集上的表现，我们选择若干具有代表性的基线模型在节点数量达到一亿的Xbox-100m数据集上进行了实验。表2展示了模型准确率、参数数量、以及训练时间（对于GraphSAGE和GBP，我们使用了node2vec训练得到的嵌入向量进行初始化）。xGCN在准确率上显著超过了基线模型，并且用时仅11个小时，而其它GNN的训练时长则普遍达到了3天。目前，xGCN模型已被应用在Xbox游戏社交平台的社交推荐业务中。

四、XGCN工具包

我们推出了一个轻量的GNN模型工具包：XGCN，旨在帮助用户快速地构建和训练大规模网络上的GNN图嵌入模型。XGCN包含了xGCN以及GraphSAGE、LightGCN等常见GNN模型的实现。如图7所示，XGCN有着详细的说明文档（见文后链接），方便用户快速上手进行图数据集处理、模型构建和运行。

XGCN的主要特性有：

（1）高效的大规模GNN图嵌入实现。许多现有模型的官方开源实现仅考虑了小数据集的情形，无法高效地处理大图。而XGCN充分利用了DGL和Pytorch的机制，进行了高效地实现。如表3所示，我们测试了在不同开源实现中，LightGCN和UltraGCN的单个epoch的运行时长，结果显示XGCN具有更好的效率和可拓展性。

（2）完整的图数据集处理流水线。针对链接预测任务（如user-item推荐），XGCN提供了从原始文本文件处理、正样本划分，到验证集/测试集生成的处理模块。

（3）友好的模型构建基础设施。XGCN提供了BaseEmbeddingModel等抽象类，用户可在XGCN的基础设施上快速定制新的DataLoader和新模型。

详细内容请参见：
Xiran Song, Jianxun Lian, Hong Huang, Zihan Luo, Wei Zhou, Xue Lin, Mingqi Wu, Chaozhuo Li, Xing Xie, and Hai Jin. 2023. xGCN: An Extreme Graph Convolutional Network for Large-scale Social Link Prediction. In Proceedings of the ACM Web Conference 2023 (WWW ’23), April 30–May 04, 2023, Austin, TX, USA. ACM, New York, NY, USA, 11 pages. https://doi.org/ 10.1145/3543507.3583340

XGCN项目地址：
https://github.com/CGCL-codes/XGCN_library

作者：宋熙然

来源：公众号【穿过丛林】

Illustration by IconScout Store from IconScout