Graph neural networks in TensorFlow

我们周围的世界中，对象及其关系无处不在，而关系对于理解一个对象的重要性，可能不亚于孤立地审视该对象自身的属性——例如交通网络、生产网络、知识图谱或社交网络。离散数学和计算机科学长期以来一直将这类网络形式化为[图](https://en.wikipedia.org/wiki/Graph_(discrete_mathematics))，图由以各种不规则方式通过边连接的节点组成。然而，大多数机器学习（ML）算法只允许输入对象之间存在规则且统一的关系，例如像素网格、单词序列，或者完全没有关系。

图神经网络（简称 GNN）已成为一种强大的技术，能够同时利用图的连通性（如较早的算法 DeepWalk 和 Node2Vec 所做的那样）以及各个节点和边上的输入特征。GNN 可以对整个图做出预测（这个分子是否会以某种方式发生反应？），对单个节点做出预测（鉴于这篇文档的引用情况，它的主题是什么？），或对潜在的边做出预测（这个产品是否可能与那个产品一起被购买？）。除了对图做出预测之外，GNN 还是一种强大的工具，用于弥合与更典型的神经网络用例之间的鸿沟。它们以连续的方式编码图的离散、关系信息，从而可以自然地将其纳入另一个深度学习系统之中。

我们很高兴宣布 TensorFlow GNN 1.0（TF-GNN）的发布，这是一个经过生产环境验证的库，用于大规模构建 GNN。它支持在 TensorFlow 中进行建模和训练，以及从海量数据存储中提取输入图。TF-GNN 从头开始为异构图而构建，其中对象和关系的类型由不同的节点和边集合表示。现实世界中的对象及其关系以不同的类型出现，TF-GNN 对异构图的关注使得表示它们变得自然。

在 TensorFlow 内部，此类图由 tfgnn.GraphTensor 类型的对象表示。这是一种复合张量类型（一个 Python 类中的张量集合），在 tf.data.Dataset、tf.function 等中被接受为一等公民。它既存储图结构，也存储附加到节点、边以及整个图上的特征。GraphTensor 的可训练变换可以定义为高层 Keras API 中的 Layers 对象，或直接使用 tfgnn.GraphTensor 原语。

GNN：为上下文中的对象进行预测

举例来说，让我们看一个 TF-GNN 的典型应用：预测由大型数据库的交叉引用表所定义的图中某一类型节点的属性。例如，一个计算机科学（CS）arXiv 论文的引用数据库，具有一对多的引用关系和多对一的被引用关系，我们希望预测每篇论文的主题领域。

与大多数神经网络一样，GNN 是在包含大量带标签样本（约数百万）的数据集上训练的，但每个训练步骤只包含一小批训练样本（比如数百个）。为了扩展到数百万规模，GNN 会在来自底层图的一连串规模合理的小子图上进行训练。每个子图包含足够的原始数据，以便为其中心带标签节点计算 GNN 结果并训练模型。这一过程——通常称为子图采样——对 GNN 训练极为关键。大多数现有工具以批处理方式完成采样，为训练生成静态子图。TF-GNN 提供了工具，通过动态且交互式采样来改进这一点。

TF-GNN 1.0 首次推出灵活的 Python API，可在所有相关规模上配置动态或批量子图采样：既可在 Colab notebook 中交互式进行（如这个），也可对存储在单台训练主机主内存中的小数据集进行高效采样，还可通过 Apache Beam 分布式处理存储在网络文件系统上的超大数据集（多达数亿个节点和数十亿条边）。详情请分别参阅我们的内存内和基于 beam 采样的用户指南。

在这些相同的采样子图上，GNN 的任务是计算根节点处的隐藏（或潜在）状态；该隐藏状态聚合并编码了根节点邻域的相关信息。一种经典方法是消息传递神经网络。在每一轮消息传递中，节点沿入边从邻居接收消息，并据此更新自身的隐藏状态。经过 n 轮后，根节点的隐藏状态反映了 n 条边范围内所有节点的聚合信息（下图以 n = 2 为例）。消息和新的隐藏状态由神经网络的隐藏层计算得出。在异构图（heterogeneous graph）中，通常有必要针对不同类型的节点和边分别使用单独训练的隐藏层。

训练设置的最后一步是将输出层置于 GNN 隐藏状态之上以处理有标签节点，计算损失（衡量预测误差），并通过反向传播更新模型权重，这与任何神经网络训练的常规做法一致。

除了监督训练（即最小化由标签定义的损失），GNN 也可以以无监督方式训练（即不使用标签）。这使我们能够计算节点及其特征的离散图结构的连续表示（或向量嵌入）。这些表示随后通常用于其他机器学习系统。通过这种方式，图所编码的离散关系信息可以被纳入更典型的神经网络用例中。TF-GNN 支持为异构图精细指定无监督目标。

构建 GNN 架构

TF-GNN 库支持在不同抽象层次上构建和训练 GNN。

在最高层次上，用户可以使用库中捆绑的、以 Keras 层表达的任意预定义模型。除了一小部分来自研究文献的模型外，TF-GNN 还附带一个高度可配置的模型模板，该模板提供了一套精选的建模选择，我们发现这些选择在许多内部问题上提供了强大的基线。模板实现了 GNN 层；用户只需初始化 Keras 层。

在最底层，用户可以基于在图上传递数据的原语从零编写 GNN 模型，例如将数据从节点广播到其所有出边，或将数据从所有入边池化到节点（例如计算入站消息之和）。TF-GNN 的图数据模型在特征或隐藏状态方面对节点、边和整个输入图一视同仁，因此不仅可以直观地表达上述以节点为中心的模型（如 MPNN），还能表达更一般的 GraphNets 形式。这可以但不一定需要在核心 TensorFlow 之上以 Keras 作为建模框架来完成。更多细节以及中间层次的建模方式，请参阅 TF-GNN 用户指南和模型集合。

训练编排

虽然高级用户可以自由进行自定义模型训练，但 TF-GNN Runner 也提供了一种简洁的方式来编排常见情况下的 Keras 模型训练。一个简单的调用可能如下所示：

Runner 为分布式训练以及在 Cloud TPU 上为固定形状进行 tfgnn.GraphTensor 填充等机器学习痛点提供了开箱即用的解决方案。除了像上面展示的那样在单个任务上进行训练之外，它还支持在多个（两个或更多）任务上协同进行联合训练。例如，可以将无监督任务与有监督任务混合，以应用特定的归纳偏置来为最终的连续表示（或向量嵌入）提供信息。调用者只需将 task 参数替换为任务映射即可：

此外，TF-GNN Runner 还包含integrated gradients的实现，用于模型归因。Integrated gradients 的输出是一个 GraphTensor，其连接结构与所观测的 GraphTensor 相同，但其特征被替换为梯度值，其中较大的值在 GNN 预测中的贡献大于较小的值。用户可以检查梯度值，以了解其 GNN 最常使用哪些特征。

结论

简而言之，我们希望 TF-GNN 能够有助于推动 GNN 在 TensorFlow 中的大规模应用，并促进该领域的进一步创新。如果您想了解更多，请尝试我们使用流行的 OGBN-MAG 基准测试的 Colab 演示（在浏览器中即可运行，无需安装），浏览其余的用户指南和 Colab，或查看我们的论文。

致谢

TF-GNN 1.0 版本由 Google Research：Sami Abu-El-Haija、Neslihan Bulut、Bahar Fatemi、Johannes Gasteiger、Pedro Gonnet、Jonathan Halcrow、Liangze Jiang、Silvio Lattanzi、Brandon Mayer、Vahab Mirrokni、Bryan Perozzi、Anton Tsitsulin、Dustin Zelle，Google Core ML：Arno Eigenwillig、Oleksandr Ferludin、Parth Kothari、Mihir Paradkar、Jan Pfeifer、Rachael Tamakloe，以及 Google DeepMind：Alvaro Sanchez-Gonzalez 和 Lisa Wang 合作开发。