[论文解读] TF-GNN: Graph Neural Networks in TensorFlow

Oleksandr Ferludin, Arno Eigenwillig+25 more|arXiv (Cornell University)|Jul 7, 2022

Advanced Graph Neural Networks被引用 15

一句话总结

TF-GNN 是一个开源的 TensorFlow 库，旨在实现可扩展的图神经网络（GNNs），通过四层 API 层次结构支持异构、多关系图数据。它使研究人员和开发者能够以不同抽象层次构建、训练和部署 GNN 模型，从原始 TensorFlow 操作到高级 Keras 模型均可支持，实现了在 OGBN-MAG 上的最先进性能，其中一个更小、经过优化的 MPNN 模型优于更大、更复杂的架构。

ABSTRACT

TensorFlow-GNN (TF-GNN) is a scalable library for Graph Neural Networks in TensorFlow. It is designed from the bottom up to support the kinds of rich heterogeneous graph data that occurs in today's information ecosystems. In addition to enabling machine learning researchers and advanced developers, TF-GNN offers low-code solutions to empower the broader developer community in graph learning. Many production models at Google use TF-GNN, and it has been recently released as an open source project. In this paper we describe the TF-GNN data model, its Keras message passing API, and relevant capabilities such as graph sampling and distributed training.

研究动机与目标

• 解决在复杂异构关系数据上训练图神经网络时缺乏可扩展、生产就绪框架的问题。
• 提供分层 API，支持从初学者到专家研究人员的各类开发者。
• 与 TensorFlow 生态系统无缝集成，支持分布式训练、TPU 加速以及跨模态预训练模型的访问。
• 实现 GNN 的高效模型开发、超参数调优和在真实应用场景中的部署。
• 证明简单、经过良好调优的 GNN 模型可在基准图学习任务上超越更复杂的架构。

提出的方法

• TF-GNN 实现了四层 API：数据层用于异构图表示，Keras 消息传递 API 用于构建 GNN 层，GraphUpdate 用于定义消息传递操作，编排器用于训练和推理。
• 支持以边为中心、以节点为中心和以图为中心的 GNN 构建模块，支持对复杂关系结构的灵活建模。
• 该库采用消息传递框架，通过在图邻域上应用可学习函数来更新节点和边的特征。
• 与 TensorFlow 生态系统集成，支持分布式训练、TPU 加速以及 SavedModel 导出以供推理。
• 提供模型库，包含预构建组件如 VanillaMPNN、GraphSAGE、GATv2 和 MultiHeadAttention，支持快速原型开发。
• 通过与 Vizier 集成，支持超参数调优，实现对模型容量、dropout、正则化以及消息维度的自动化优化。

实验结果

研究问题

• RQ1 能否在 TensorFlow 上构建一个统一的、可扩展的 GNN 框架，以支持异构、多关系图数据？
• RQ2 分层 API 设计在多大程度上能有效服务具有不同机器学习专业水平的用户？
• RQ3 在大规模图基准测试中，简单且经过良好优化的 MPNN 模型能否超越更复杂的基于 Transformer 的架构？
• RQ4 与 TensorFlow 生态系统的集成在多大程度上能提升 GNN 的训练效率和部署准备度？
• RQ5 自动化超参数调优在不增加架构复杂性的情况下，发现高性能 GNN 配置的有效性如何？

主要发现

• 一个仅 589 万参数的 MPNN 模型，通过 Vizier 进行超参数优化后，在 OGBN-MAG 上实现了 0.5149 的验证准确率，优于参数量达 2680 万的异构图 Transformer（HGT）模型。
• 经 Vizier 优化的 MPNN 模型在测试集上达到 0.5027 的准确率，超过 HGT 模型的 0.4982，证明了超参数调优在性能上的优势，而非架构复杂性。
• TF-GNN 框架已成功在谷歌多个团队中实现 GNN 模型的生产级部署，验证了其可扩展性和可用性。
• 该库对以边为中心和以图为中心的消息传递的支持，使得像图 Transformer 这类模型能够原生表达，而在以节点为中心的框架中则较为繁琐。
• 与 TensorFlow 生态系统的集成实现了 TPU 上的高效训练，并可无缝导出为 SavedModels 用于推理，显著降低了部署复杂度。
• 模型库和高级 API 显著缩短了开发时间，支持 GNN 的快速实验和部署。