图是一种由边连接的潜在广泛节点网络，可用于表达和询问数据之间的关系，如社交联系、金融交易、交通、能源网格和分子相互作用。 随着研究人员收集更多数据并构建这些图形图片，研究人员将需要更快、更高效的方法以及更强的计算能力，以图神经网络 (GNN) 的方式对它们进行深度学习。
现在，麻省理工学院和 IBM 研究院的研究人员开发了一种称为 SALIENT（SAmpling、sLIcing 和数据移动）的新方法，它通过解决计算中的三个关键瓶颈来提高训练和推理性能。 这极大地减少了 GNN 在大型数据集上的运行时间，例如，大型数据集包含 100 亿个节点和 1 亿条边。 此外，该团队发现，当计算能力从 16 个增加到 XNUMX 个图形处理单元 (GPU) 时，该技术可以很好地扩展。 这项工作在第五届机器学习和系统会议上进行了介绍。
“我们开始研究当前系统在将最先进的图形机器学习技术扩展到真正的大数据集时遇到的挑战。 事实证明还有很多工作要做，因为许多现有系统主要在适合 GPU 内存的较小数据集上取得了良好的性能，”主要作者兼麻省理工学院计算机科学博士后蒂姆卡勒说和人工智能实验室 (CSAIL)。
专家所说的庞大数据集是指像整个比特币网络这样的规模，其中某些模式和数据关系可以阐明趋势或犯规行为。 “区块链上有近十亿笔比特币交易，如果我们想在这样一个联合网络中识别非法活动，那么我们将面临如此规模的图表，”共同作者、高级研究科学家兼经理陈杰说。 IBM Research 和 MIT-IBM Watson 人工智能实验室。 “我们希望构建一个能够处理这种图形并使处理尽可能高效的系统，因为我们每天都希望跟上生成的新数据的步伐。”
Kaler 和 Chen 的合著者包括 Jump Trading 的 Nickolas Stathas MEng '21，他在毕业论文中开发了 SALIENT； 前 MIT-IBM Watson AI Lab 实习生和 MIT 研究生 Anne Ouyang； MIT CSAIL 博士后 Alexandros-Stavros Iliopoulos； MIT CSAIL 研究科学家 Tao B. Schardl； 麻省理工学院电气工程 Edwin Sibley Webster 教授兼麻省理工学院-IBM 沃森人工智能实验室研究员 Charles E. Leiserson。
对于这个问题，该团队在开发他们的方法时采用了面向系统的方法：SALIENT，Kaler 说。 为此，研究人员实施了他们认为重要的、适合现有机器学习框架的组件的基本优化，例如 PyTorch Geometric 和深度图形库 (DGL)，它们是构建机器学习模型的接口。 Stathas 说这个过程就像更换引擎来制造一辆更快的汽车。 他们的方法旨在适应现有的 GNN 架构，以便领域专家可以轻松地将这项工作应用到他们指定的领域，以加快模型训练并在推理过程中更快地梳理出见解。 该团队确定的诀窍是让所有硬件（CPU、数据链路和 GPU）始终保持忙碌：CPU 对图形进行采样并准备小批量数据，然后通过数据链路传输这些数据，更关键的 GPU 正在努力训练机器学习模型或进行推理。
研究人员首先分析了一个常用的 GNN 机器学习库 (PyTorch Geometric) 的性能，该库显示可用 GPU 资源的利用率低得惊人。 应用简单的优化，研究人员将 GPU 利用率从 10% 提高到 30%，相对于公共基准代码，性能提高了 1.4 到两倍。 这种快速基线代码可以在 50.4 秒内通过算法（一个纪元）对大型训练数据集执行一次完整传递。
为了进一步提高性能，研究人员着手研究数据管道开始时出现的瓶颈：图形采样和小批量准备的算法。 与其他神经网络不同，GNN 执行邻域聚合操作，该操作使用图中其他附近节点中存在的信息来计算有关节点的信息——例如，在社交网络图中，来自用户朋友的朋友的信息。 随着 GNN 中层数的增加，网络必须接触的节点数量可能会激增，超过计算机的极限。 邻域采样算法通过选择较小的随机节点子集来提供帮助； 然而，研究人员发现，当前的实现速度太慢，无法跟上现代 GPU 的处理速度。 作为回应，他们确定了数据结构、算法优化等的组合，提高了采样速度，最终仅采样操作就提高了大约三倍，使每个时期的运行时间从 50.4 秒增加到 34.6 秒。 该团队指出，他们还发现可以在推理过程中以适当的速率进行采样，从而提高整体能源效率和性能，这一点在文献中被忽视了。
在以前的系统中，此采样步骤是一种多进程方法，会在进程之间创建额外的数据和不必要的数据移动。 研究人员通过创建一个具有轻量级线程的进程，将 CPU 上的数据保存在共享内存中，从而使他们的 SALIENT 方法更加灵活。 此外，SALIENT 利用了现代处理器的缓存，Stathas 说，并行化特征切片，在 CPU 核心缓存的共享内存中从感兴趣的节点及其周围的邻居和边缘中提取相关信息。 这再次将每个时期的总运行时间从 34.6 秒减少到 27.8 秒。
研究人员解决的最后一个瓶颈是使用预取步骤在 CPU 和 GPU 之间流水线传输小批量数据，这将在需要之前准备数据。 该团队计算出，这将最大限度地利用数据链路中的带宽，并使该方法达到完美利用； 然而，他们只看到了大约 90%。 他们发现并修复了一个流行的 PyTorch 库中的性能错误，该错误导致 CPU 和 GPU 之间不必要的往返通信。 修复此错误后，该团队使用 SALIENT 实现了每轮 16.5 秒的运行时间。
“我认为，我们的工作表明细节决定成败，”卡勒说。 “当你在训练图神经网络时密切关注影响性能的细节时，你可以解决大量的性能问题。 使用我们的解决方案，我们最终完全被 GPU 计算瓶颈，这是此类系统的理想目标。”
SALIENT 的速度是在三个标准数据集 ogbn-arxiv、ogbn-products 和 ogbn-papers100M 以及多机设置和不同级别的扇出（CPU 为 GPU 准备的数据量）和跨多种架构，包括最新的最先进的 GraphSAGE-RI。 在每个设置中，SALIENT 都优于 PyTorch Geometric，最显着的是在大型 ogbn-papers100M 数据集上，该数据集包含 100 亿个节点和超过 16 亿条边在这里，它在一个 GPU 上运行时比最初创建的优化基线快三倍这项工作； 使用 XNUMX 个 GPU，SALIENT 的速度提高了八倍。
虽然其他系统的硬件和实验设置略有不同，因此并不总是直接比较，但 SALIENT 仍然优于它们。 在达到类似精度的系统中，代表性的性能数字包括使用一个 GPU 和 99 个 CPU 的 32 秒，以及使用 13 个 CPU 的 1,536 秒。 相比之下，SALIENT 使用一个 GPU 和 20 个 CPU 的运行时间为 16.5 秒，而使用 16 个 GPU 和 320 个 CPU 的运行时间仅为两秒。 “如果您查看之前工作报告的底线数字，我们的 16 GPU 运行时间（两秒）比之前在此数据集上报告的其他数字快一个数量级，”Kaler 说。 研究人员将他们的性能改进部分归因于他们在转移到分布式设置之前为单台机器优化代码的方法。 Stathas 说，这里的教训是，为了你的钱，“在你开始扩展到多台计算机之前，更有效地使用你拥有的硬件更有意义，”这可以显着节省成本和碳排放可以通过模型训练来实现。
这种新能力现在将使研究人员能够处理和挖掘越来越大的图表。 例如，前面提到的比特币网络包含 100,000 个节点； SALIENT 系统可以处理大 1,000 倍（或三个数量级）的图形。
“在未来，我们不仅会考虑在我们为分类或预测每个节点的属性而实施的现有算法上运行这个图神经网络训练系统，而且我们还想做更深入的任务，比如识别图中的常见模式（子图模式），[这]对于指示金融犯罪实际上可能很有趣，”Chen 说。 “我们还想识别图中相似的节点，从某种意义上说，它们可能对应于金融犯罪中的同一个不良行为者。 这些任务将需要开发额外的算法，可能还需要开发神经网络架构。”
这项研究得到了麻省理工学院-IBM 沃森人工智能实验室的支持，并得到了美国空军研究实验室和美国空军人工智能加速器的部分支持。
