大规模图神经网络高效训练新方法

大规模图神经网络高效训练方法

研究背景

图神经网络通过创建图中节点和边的向量表示（嵌入）来实现新的分析任务，如链接预测。节点的嵌入通常不仅考虑该节点本身的信息，还包括其直接邻居甚至邻居的邻居信息。在许多实际应用场景中（如推特用户关系图），单个节点可能拥有数千甚至数百万个连接，此时考虑所有邻居节点是不现实的。

混合CPU-GPU训练挑战

GPU是执行神经网络张量操作最高效的方式，但内存有限。对于无法完全放入GPU内存的大型图神经网络训练，通常使用CPU创建随机选择的图节点和边的小批量数据，然后与节点特征数据一起发送到GPU。递归邻居采样生成的小批量需要在CPU和GPU之间传输大量数据，这成为模型训练中最耗时的环节。

全局邻居采样方法

我们提出的全局邻居采样方法显著减少了从CPU到GPU的数据传输量。具体策略包括：

缓存机制

• 在创建小批量之前，从整个图中采样一组节点并将其特征加载到GPU内存中
• 创建小批量时，通过直接检索缓存中已有的邻居节点来采样
• 仅当缓存中邻居节点不足时，才从CPU获取额外节点

采样优化

• 优先采样高度节点（连接数多的节点），增加相关邻居在缓存中的概率
• 采样概率与节点度数成正比，确保缓存仍包含一定数量的低度节点
• 在嵌入过程中，对缓存节点进行加权，确保能良好近似考虑所有邻居时的嵌入结果

实验效果

在五个数据集上与三种其他方法比较，我们的方法在混合CPU-GPU设置下：

• 每个数据集上至少比第二优策略快两倍
• 两种采样策略在达到相当准确度时，训练速度比我们的方法慢一个数量级

技术优势

该方法被证明能够有效收敛到最优模型性能，既不会因偏向高度节点而影响性能，也不会因对多个小批量重复使用相同缓存而损害模型表现。

未来展望

当前研究局限于单个CPU和单个GPU，未来工作将探索如何将该方法推广到多GPU和分布式训练环境，例如在不同GPU上缓存不同节点集，并有效将每个小批量定向到缓存匹配度最高的GPU。
更多精彩内容 请关注我的个人公众号 公众号（办公AI智能小助手）或者 我的个人博客 https://blog.qife122.com/
对网络安全、黑客技术感兴趣的朋友可以关注我的安全公众号（网络安全技术点滴分享）

公众号二维码

公众号二维码