Exphormer：图结构数据的可扩展Transformer

概述

图（由节点和边表示对象及其关系）在计算和机器学习中无处不在。社交网络、道路网络和分子结构等领域的底层数据集都具有自然的图结构。机器学习可用于学习节点、边或整个图的属性。

图神经网络（GNN）是学习图的常见方法，通过对节点、边和全局属性应用可优化变换来操作图数据。最典型的GNN类通过消息传递框架操作，其中每一层聚合节点与其直接邻居的表示。

图Transformer的挑战

最近，图Transformer模型已成为消息传递GNN的流行替代方案。这些模型基于Transformer架构在自然语言处理（NLP）中的成功，将其适应于图结构数据。图Transformer中的注意力机制可以通过交互图建模，其中边表示相互关注的对节点。与消息传递架构不同，图Transformer具有与输入图分离的交互图。

典型的交互图是一个完全图，表示对所有节点对之间直接交互建模的完整注意力机制。然而，这产生了二次计算和内存瓶颈，限制了图Transformer在最多几千个节点的小图数据集上的适用性。使图Transformer可扩展被认为是该领域最重要的研究方向之一。

Exphormer的解决方案

在ICML 2023上提出的"Exphormer：图的稀疏Transformer"中，我们通过引入专门为图数据设计的稀疏注意力框架来解决可扩展性挑战。Exphormer框架利用扩展图（谱图论中的强大工具），并能够在各种数据集上实现强大的实证结果。我们的Exphormer实现现已可在GitHub上获得。

扩展图

Exphormer的核心思想是使用扩展图，这些图稀疏但连接良好，具有一些有用特性：

图的矩阵表示具有与完全图相似的线性代数特性
它们表现出随机游走的快速混合，即从任何起始节点的随机游走只需少量步骤即可确保收敛到节点上的"稳定"分布

常见的扩展图类是d-正则扩展图，其中每个节点有d条边（即每个节点的度为d）。扩展图的质量通过其谱隙（其邻接矩阵的代数特性）来衡量。那些最大化谱隙的图称为Ramanujan图。

Exphormer架构

Exphormer将标准Transformer的密集、完全连接的交互图替换为稀疏d-正则扩展图的边。直观上，扩展图的谱近似和混合特性允许远距离节点在图Transformer架构中堆叠多个注意力层后相互通信，即使这些节点可能不直接相互关注。

此外，通过确保d是常数（与节点数量无关），我们在结果交互图中获得线性数量的边。

构建稀疏交互图

Exphormer将扩展边与输入图和虚拟节点相结合。具体来说，Exphormer的稀疏注意力机制构建了一个由三种类型边组成的交互图：

来自输入图的边（局部注意力）
来自常数度扩展图的边（扩展注意力）
从每个节点到一小组虚拟节点的边（全局注意力）

每个组件都有特定目的：来自输入图的边保留了输入图结构的归纳偏差（通常在完全连接的注意力模块中丢失）；扩展边允许良好的全局连接性和随机游走混合特性（用更少的边谱近似完全图）；虚拟节点作为全局"内存汇"，可以直接与每个节点通信。

与序列稀疏Transformer的关系

将Exphormer与序列的稀疏注意力方法进行比较很有趣。在概念上最类似于我们方法的结构可能是BigBird，它通过组合不同组件构建交互图。BigBird也使用虚拟节点，但与Exphormer不同，它使用窗口注意力和来自Erdős-Rényi随机图模型的随机注意力。

BigBird中的窗口注意力查看序列中标记周围的标记——Exphormer中的局部邻域注意力可以看作是窗口注意力对图的泛化。

实验结果

早期工作展示了在最多5,000个节点的图数据集上使用基于完整图Transformer的模型。为了评估Exphormer的性能，我们基于著名的GraphGPS框架构建，该框架结合了消息传递和图Transformer，并在许多数据集上实现了最先进的性能。

我们表明，在GraphGPS框架中用Exphormer替换密集注意力用于图注意力组件，可以实现相当或更好的性能模型，通常具有更少的可训练参数。

此外，Exphormer显著允许图Transformer架构扩展到远超上述通常的图大小限制。Exphormer可以扩展到10,000+节点图的数据集，如Coauthor数据集，甚至更大的图，如著名的ogbn-arxiv数据集（一个包含170K节点和110万边的引文网络）。

最后，我们观察到Exphormer通过扩展器创建小直径覆盖图，表现出有效学习长程依赖关系的能力。长程图基准是一套五个图学习数据集，旨在衡量模型捕捉长程交互的能力。结果显示，基于Exphormer的模型优于标准GraphGPS模型（在发布时在五个数据集中的四个上先前是最先进的）。

结论

图Transformer已成为机器学习的重要架构，将NLP中高度成功的基于序列的Transformer适应于图结构数据。然而，可扩展性一直是在大型图数据集上使用图Transformer的主要挑战。

在本文中，我们提出了Exphormer，一个使用扩展图提高图Transformer可扩展性的稀疏注意力框架。Exphormer被证明具有重要的理论特性并表现出强大的实证性能，特别是在学习长程依赖关系至关重要的数据集上。