另一篇 Jure Leskovec 实验室的最新成果,对 GNN 进行了改进,使得 GNN 能应用于大规模知识图谱推理。

论文:GNNAutoScale: Scalable and Expressive Graph Neural Networks via Historical Embeddings

代码:PyGAS: Auto-Scaling GNN in PyG

0. 摘要

本文提出了 GNNAutoScale (GAS),一种显著提高 GNN 可扩展性,同时保持 GNN 的表达能力和精确性,同时不需要对边缘进行采样的框架。

GAS 通过使用先前训练迭代得到的历史嵌入,并在每次 GNN 计算中仅使用一小部分 (mini-batch) 的实体和其对应的单跳邻域实体进行计算,同时用历史嵌入作为补充信息,实现修剪计算图的整个子树。

GAS 的 GPU 内存消耗是固定的,而不会丢弃任何数据。虽然现有的解决方案由于边缘的子采样或不可训练的传播而削弱了消息传递的表达能力,但 GAS 可以证明能够保持原始 GNN 的表达能力。通过提供历史嵌入的近似误差界限来实现这一点,并展示了如何在实践中进行收敛。

本文开源了相关代码,通过实验证明了 GAS 框架快速、显存占用率低、表达能力和 full-batch GNN 相当,并在大规模知识图谱推理上达到最先进的性能。

1. 方法

GNN 的一个主要问题是无法推广到大规模知识图谱中,本文的方法是提出了一种 GNNAutoScale (GAS) 框架,从 GNN 的底层消息传递的实现中解决了 GNN 的可伸缩性问题。GAS 通过历史嵌入方法,将过去的训练迭代中得到的节点嵌入保存下来。GAS 构建 GNN 计算图时,只保留当前批次的节点(mini-batch)和直接的1跳邻居,同时将历史嵌入作为补充信息。通过这种方法,GAS 可以保证训练时的内存消耗是基本固定的,可以将 GNN 推广到大规模知识图谱中,同时仍然考虑所有的邻域信息。

image-20211104175751447

image-20211104175759738

  • 图:历史嵌入信息。

image-20211104175818787

  • 图:带有历史嵌入的 Mini-batch GNN 处理过程。与普通的 GNN 相比,GAS 通过采用历史嵌入,将显存占用转移到了内存中。在 Mini-batch 的节点将信息输入到内存中,同时 Mini-batch 的单跳邻域节点从内存中获取它们最新的嵌入值,用于后面的计算。

GAS 的优势:

  • GAS 在全体数据中更新。每一轮迭代中,每条边都被计算且仅被计算一次,进而得到历史嵌入,从而保证了线性复杂度。
  • GAS 的推理时间复杂度是固定的。
  • GAS 易于实现。
  • GAS 有理论保证。

随后,文章重点介绍了为什么要采用历史嵌入,证明了近似误差比较低,表达能力很强。

其中,文章认为历史嵌入主要有两点主要误差:估计值和精确值的误差,过时的误差。推导得到了误差界限,并证明能让框架在训练中收敛。

image-20211104181636354

image-20211104181643304

image-20211104181651897

image-20211104181658281

image-20211104181704608

经过上述分析和证明,本文认为采用小批次训练后,需要最小化批次之间的互联性质。通过采用图聚类算法,使得邻居节点更接近。结果是可以增加节点估计值的精确性和降低历史嵌入的过时性。此外,模型还应用了辅助损失(auxiliary loss)、梯度裁剪(gradient clipping)等技术。

image-20211104181810989

  • 图:通过对训练方法进行调整,本文采用的训练方法能够达到连续 mini-batch 训练效率的两倍,同时比 full-batch 训练使用极少的显存。

2. 实验

本文进行的实验包括:

  • 消融实验(两种)。
  • 基础知识图谱对比。
  • GPU 消耗和数据利用率。
  • 计算开销和 IO 开销分析。
  • 计算效率对比。
  • 大规模知识图谱对比。

image-20211104183218217

  • 图:GAS,仅采用历史嵌入的 mini-batch 模型(没有采用其他的 GAS 中的技术)和 full-batch 模型对比。GAS 效果最好。

image-20211104202602891

  • 表:不同配置的 GAS 和 Full-batch 模型的对比。

image-20211104204928086

  • 表:消融实验。

image-20211104204937094

  • 表:GPU 消耗和数据利用率。

image-20211104204946498

  • 图:计算开销和 IO 开销。本文提出的 GAS 框架使得计算开销和 IO 开销差不多,提高了效率。

image-20211104204952833

  • 表:效率对比。

image-20211104205000359

  • 表:大规模知识图谱表现。