本文首发于Michael Bronstein教授的博客。此版本可能与首发版本有所不同。Michael Bronstein是牛津大学DeepMind人工智能教授、Twitter图学习研究负责人。本文由Michael Bronstein教授和来自Graphcore的Emanuele Rossi以及Daniel Justus合著。
用于图神经网络（GNN）的标准硬件的适用性是图机器学习（Graph ML）社区中经常被忽视的问题。本文探讨了在Graphcore开发的新硬件架构上时间图神经网络的实现，该架构专为图结构化的工作负载定制。
图结构化数据出现在许多处理交互实体的复杂系统的问题中。近年来，将机器学习方法应用于图结构数据中的方法，特别是图神经网络（GNN），越来越受欢迎。
大多数的GNN架构假设图是固定的。然而，这种假设往往过于简单：因为在许多应用中，底层系统是动态的，图会随着时间的推移而变化。例如，在社交网络或推荐系统中，描述用户与内容互动的图可以实时变化。最近，一些能够处理动态图的GNN架构被开发出来，包括我们自己的时间图网络（TGN）[1]。
在本文中，我们所探讨的是TGN在不同尺寸的动态图中的应用，并研究这一类模型的计算复杂程度。我们使用Graphcore的Bow IPU（智能处理器）来训练TGN，并证明了为什么IPU的架构非常适合解决这些复杂问题，单个IPU处理器与一个NVIDIA A100 GPU相比，其速度最多可提升一个数量级。
在我们以前的文章中曾经详细描述过TGN架构，它包括两个主要部分：首先，节点embeddings是通过经典的图神经网络架构生成的，在这里作为一个单层图注意力网络加以实现[2]。此外，TGN保留了一个总结每个节点过去所有交互的存储。该存储可通过稀疏的读/写操作进行访问，并使用门控循环网络（GRU）更新新的交互[3]。
我们关注的是随时间推移获得新的边从而产生变化的图。在这种情况下，一个给定节点的存储包含了所有针对该节点的边以及它们各自的目标节点的信息。通过间接作用，一个给定节点的存储也可以包含更远的节点的信息，从而使图注意力网络中的附加层变得可有可无。
应用于小型图
我们首先在JODIE Wikipedia数据集[4]上进行了TGN的实验，这是一个由Wikipedia文章和用户组成的二分图，用户和文章之间的每条边都代表用户对该文章的编辑。该图由9227个节点（8227个用户和1000篇文章）和157474条带有时间戳的边组成，由描述编辑情况的172维的LIWC特征向量[5]加以注释。
在训练过程中，边被逐batch插入到最初断开的节点集中，同时使用真实边和随机抽样的负边的对比损失来训练模型。验证结果作为识别随机抽样的负边中的真实边的概率加以报告。
直观地讲，大batch size对训练和推理都有不利结果：节点存储和图的连接性都是在处理完一个完整的batch后才更新的。因此，一个batch中的后期事件可能依赖于过时的信息，因为它们不了解batch中的早期事件。事实上，我们观察到大的batch size对任务性能的不利影响。事实上，我们观察到大batch size对任务性能的不利影响，如下图所示：
然而，使用小的batch强调在训练和推理过程中快速访问存储对于实现高吞吐量的重要性。因此，如下图所示，具有大型处理器内存储的IPU与具有较小batch size的GPU相比，表现出越来越大的吞吐量优势。特别是当使用10的batch size时，TGN在IPU上的训练速度大约是11倍，即使是200的大batch size，在IPU上的训练速度仍然是3倍左右。
为了更好地理解Graphcore IPU上TGN训练吞吐量的提高，我们研究了不同硬件平台在TGN的关键操作上所花费的时间。我们发现，GPU上所花费的时间是由注意力模块和GRU主导的，而这两个操作在IPU上的执行效率更高。此外，在所有操作中，IPU处理小batch size操作的效率要高得多。
尤其是我们观察到，IPU的优势随着较小和较零散的存储操作而增长。更广泛地说，我们可以得出结论：当计算和存储访问十分异质化时，相较于GPU，IPU架构表现出非常明显的优势。
扩展到大型图
虽然TGN模型在其默认配置下是相对轻量级的，有大约26万个参数，当把模型应用于大型图时，IPU处理器内存储的大部分被节点存储占用。然而，由于它是稀疏访问的，这个张量可以移动到片外存储，在这种情况下，处理器内存储的利用率便不再与图的尺寸相关。
为了测试大型图上的TGN架构，我们将其应用于一个匿名图，其中包含了1550万Twitter用户间的2.61亿个关注[6]。这些边被分配了728个不同的时间戳。这些时间戳按照日期排序，但不提供关于关注发生的实际日期的任何信息。由于这个数据集中没有节点或边的特征，该模型完全依赖于图的拓扑结构和时间演变来预测新的链接。
与单一的负样本相比，大量的数据使识别正边的任务过于简单，因此我们使用了1000个随机采样的负边中的真实边的平均倒数排名（MRR）作为验证指标。此外，我们发现，当数据集规模增加时，更大的hidden size更有利于模型的性能。对于给定的数据，我们发现256的潜伏尺寸是准确性和吞吐量之间的最佳点。
使用片外存储作为节点存储可将吞吐量降低大约两倍。然而，从使用不同尺寸的诱导子图以及具有10倍于Twitter图的节点数和随机连接的合成数据集看来，我们证明吞吐量几乎与图的尺寸无关（见下表）。在IPU上使用该技术，TGN可以应用于几乎任意尺寸的图，仅受可用主机存储的限制，同时在训练和推理期间保持可观的吞吐量。
正如我们之前反复指出的那样，为实现图机器学习模型选择硬件是一个关键却经常被忽视的问题。特别是在研究界，抽象出底层硬件的云计算服务的可用性导致了这方面的某些“懒惰”。然而，当涉及到实现在有实时时延要求的大规模数据集上工作的系统时，硬件因素的考量就不能那么轻率了。我们希望我们的研究能引起人们对这个重要话题的更多关注，并为未来更高效的图机器学习应用算法和硬件架构铺平道路。
[1] E. Rossi et al., Temporal graph networks for deep learning on dynamic graphs(2020) arXiv:2006.10637. See accompanying blog post.
[2] P. Veličković, et al., Graph attention networks (2018) ICLR.
[3] K. Cho et al., On the properties of neural machine translation: Encoder-Decoder approaches (2014), arXiv:1409.1259.
[4] S. Kumar et al., Predicting dynamic embedding trajectory in temporal interaction networks (2019) KDD.
[5] J. W. Pennebaker et al., Linguistic inquiry and word count: LIWC 2001 (2001). Mahway: Lawrence Erlbaum Associates 71.
[6] A. El-Kishky et al., kNN-Embed: Locally smoothed embedding mixtures for multi-interest candidate retrieval (2022) arXiv:2205.06205.