TensorFlow中的图神经网络

2024年2月6日
发布者：Dustin Zelle（谷歌研究软件工程师）和 Arno Eigenwillig（CoreML软件工程师）

图神经网络的核心价值

现实世界中的对象及其关系无处不在，而关系对于理解对象的重要性往往不亚于对象自身的属性——例如交通网络、生产网络、知识图谱或社交网络。离散数学和计算机科学长期以来将此类网络形式化为图（graphs），即通过各种不规则方式由边连接的节点集合。然而，大多数机器学习（ML）算法仅支持输入对象之间的规则且统一的关系，例如像素网格、单词序列或完全没有关系。

图神经网络（GNN）作为一种强大技术，既能利用图的连通性（如DeepWalk和Node2Vec等传统算法），又能利用各个节点和边上的输入特征。GNN可以对整个图（例如：这个分子是否以某种方式反应？）、单个节点（例如：给定其引用，这篇文档的主题是什么？）或潜在边（例如：这个产品是否可能与该产品一起购买？）进行预测。除了对图进行预测外，GNN还是弥合与更典型神经网络用例之间差距的强大工具。它们以连续的方式编码图的离散关系信息，以便可以自然地包含在另一个深度学习系统中。

TensorFlow GNN 1.0发布

我们兴奋地宣布发布TensorFlow GNN 1.0（TF-GNN），这是一个经过生产测试的库，用于大规模构建GNN。它支持在TensorFlow中进行建模和训练，以及从大型数据存储中提取输入图。TF-GNN从头开始构建，专注于异构图（heterogeneous graphs），其中对象和关系的类型由不同的节点和边集合表示。现实世界的对象及其关系以不同的类型出现，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结果并训练模型。这个过程——通常称为子图采样（subgraph sampling）——对GNN训练极其重要。大多数现有工具以批处理方式完成采样，产生用于训练的静态子图。TF-GNN通过动态和交互式采样提供了改进此过程的工具。

图示：子图采样过程，其中从较大的图中采样小型、易处理的子图，以创建GNN训练的输入示例。

TF-GNN 1.0首次推出了灵活的Python API，用于在所有相关尺度上配置动态或批处理子图采样：在Colab笔记本中交互式（如此示例），用于高效采样存储在单个训练主机主内存中的小数据集，或通过Apache Beam分布式处理存储在网络文件系统上的巨大数据集（最多数亿节点和数十亿边）。有关详细信息，请分别参阅我们的内存和基于Beam的采样用户指南。

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

图示：一个简单的消息传递神经网络，其中在每个步骤中，节点状态从外部节点传播到内部节点，在那里汇集以计算新的节点状态。一旦到达根节点，就可以进行最终预测。

训练设置通过将输出层放置在标记节点的GNN隐藏状态之上来完成，计算损失（以测量预测误差），并通过反向传播更新模型权重，与任何神经网络训练一样。

超越监督训练

除了监督训练（即最小化由标签定义的损失）之外，GNN还可以以无监督方式（即没有标签）进行训练。这使我们能够计算节点及其特征的离散图结构的连续表示（或嵌入）。这些表示随后通常用于其他ML系统。通过这种方式，由图编码的离散关系信息可以包含在更典型的神经网络用例中。TF-GNN支持对异构图的无监督目标进行细粒度规范。

构建GNN架构

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

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

在最低级别，用户可以根据在图周围传递数据的原语（例如，将数据从节点广播到所有出边，或将数据从所有入边汇集到节点（例如，计算传入消息的总和））从头开始编写GNN模型。TF-GNN的图数据模型在特征或隐藏状态方面平等对待节点、边和整个输入图，使得不仅可以表达像上述MPNN这样的以节点为中心的模型，还可以表达更一般形式的GraphNets。这可以但不必在核心TensorFlow之上使用Keras作为建模框架来完成。有关更多详细信息和中间级别的建模，请参阅TF-GNN用户指南和模型集合。

训练编排

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

1
2
3
4

runner = tfgnn.KerasRunner()
model = tfgnn.keras.Model(...)
task = tfgnn.keras.SupervisedNodeClassification(...)
runner.run(model, task, ...)

Runner为ML痛点提供了即用型解决方案，例如分布式训练和用于Cloud TPU上固定形状的tfgnn.GraphTensor填充。除了在单个任务上训练（如上所示）之外，它还支持协同多个（两个或更多）任务的联合训练。例如，无监督任务可以与监督任务混合，以通过应用特定的归纳偏差通知最终的连续表示（或嵌入）。调用者只需用任务映射替换任务参数：

1
2
3
4

task = {
    "primary": tfgnn.keras.SupervisedNodeClassification(...),
    "auxiliary": tfgnn.keras.NodeEmbedding(...),
}

此外，TF-GNN Runner还包括用于模型归因的积分梯度（integrated gradients）实现。积分梯度输出是一个GraphTensor，其连接性与观察到的GraphTensor相同，但其特征被梯度值替换，其中较大值在GNN预测中的贡献比较小值更大。用户可以检查梯度值以查看他们的GNN最使用哪些特征。

结论

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

致谢

TF-GNN 1.0版本是由谷歌研究：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，谷歌Core ML：Arno Eigenwillig、Oleksandr Ferludin、Parth Kothari、Mihir Paradkar、Jan Pfeifer、Rachael Tamakloe，以及谷歌DeepMind：Alvaro Sanchez-Gonzalez和Lisa Wang之间的合作开发的。

标签： 算法与理论、机器智能、产品