[论文解读] Exphormer: Sparse Transformers for Graphs

Hamed Shirzad, Ameya Velingker+3 more|arXiv (Cornell University)|Mar 10, 2023

Advanced Graph Neural Networks被引用 7

一句话总结

Exphormer 通过引入基于虚拟全局节点和扩展图的稀疏注意力机制，实现了图变换器在图规模上的线性复杂度，从而在大规模图上实现可扩展且精确的建模。它在多个基准测试中达到最先进性能，包括在 ogbn-arxiv 上实现 72.44% 的准确率，并可扩展至最多 169K 个节点的图，而全注意力模型因内存限制无法运行。

ABSTRACT

Graph transformers have emerged as a promising architecture for a variety of graph learning and representation tasks. Despite their successes, though, it remains challenging to scale graph transformers to large graphs while maintaining accuracy competitive with message-passing networks. In this paper, we introduce Exphormer, a framework for building powerful and scalable graph transformers. Exphormer consists of a sparse attention mechanism based on two mechanisms: virtual global nodes and expander graphs, whose mathematical characteristics, such as spectral expansion, pseduorandomness, and sparsity, yield graph transformers with complexity only linear in the size of the graph, while allowing us to prove desirable theoretical properties of the resulting transformer models. We show that incorporating Exphormer into the recently-proposed GraphGPS framework produces models with competitive empirical results on a wide variety of graph datasets, including state-of-the-art results on three datasets. We also show that Exphormer can scale to datasets on larger graphs than shown in previous graph transformer architectures. Code can be found at \url{https://github.com/hamed1375/Exphormer}.

研究动机与目标

• 为解决标准图变换器的二次方内存和计算复杂度问题，后者限制了其在大规模图上的可扩展性。
• 在超越为序列优化设计（如 BigBird 和 Performer）的稀疏注意力机制方面，提升图神经网络的性能与效率。
• 开发一种理论基础坚实的稀疏注意力机制，在减少参数量和训练时间的同时保持表达能力。
• 实现对大规模数据集（如包含 169K 个节点的 ogbn-arxiv）的图变换器训练，这些数据集因内存限制而使全注意力模型无法处理。
• 将稀疏注意力机制集成到 GraphGPS 框架中，以在多样化图学习任务中实现具有竞争力的性能。

提出的方法

• 引入连接到所有其他节点的虚拟全局节点，通过增加 O(|V|) 条边，提供长距离归纳偏置。
• 采用具有 O(|V|) 条边的扩展图作为稀疏叠加网络，利用其谱扩展特性和小直径实现高效的消息传递。
• 将局部邻域、虚拟节点和扩展图边结合，形成一种混合注意力机制，从谱上近似全自注意力。
• 利用扩展图的数学性质（如伪随机性和谱近似），确保模型以最少的边捕捉全局依赖关系。
• 在 GraphGPS 框架内实现该架构，支持与消息传递 GNN 及多种位置编码的模块化集成。
• 通过将注意力限制在局部邻居、虚拟节点和扩展图边，实现线性复杂度注意力，将 O(|V|²) 的复杂度降低至 O(|V|)。

实验结果

研究问题

• RQ1 基于虚拟节点和扩展图的稀疏注意力机制是否能在保持与全注意力图变换器相当性能的同时，实现对更大图的可扩展性？
• RQ2 扩展图与虚拟节点的结合是否为图表示学习提供了优于现有稀疏注意力机制的归纳偏置？
• RQ3 Exphormer 是否能在全注意力模型因内存限制而无法运行的大图（如包含 169K 个节点的 ogbn-arxiv）上进行训练？
• RQ4 局部邻域、虚拟节点和扩展图边这三部分组件各自对性能的贡献如何？其中哪些组件可被移除而不造成性能损失？
• RQ5 Exphormer 是否在标准图基准测试中优于或至少匹配全注意力模型及其他稀疏变换器的性能？

主要发现

• Exphormer 在 ogbn-arxiv 上实现了 72.44% 的最先进准确率，该数据集为包含 169K 个节点和 110 万条边的大规模引用网络，而全注意力模型因内存限制无法运行。
• 在 Amazon-Computer 数据集上，Exphormer 达到 90.74% 的准确率，与最佳非变换器方法持平，并优于其他稀疏变换器。
• 由于内存使用减少，Exphormer 在 MalNet-Tiny 上可使用 256 的批量大小进行训练，而全注意力 GraphGPS 仅能支持 16 的批量大小。
• 同时包含虚拟节点和扩展图边的模型在 ogbn-arxiv 上实现 72.44% 的准确率，优于仅使用虚拟节点的版本（72.22%），并显著优于 MPNN 基线模型。
• Exphormer 可扩展至最多 35K 个节点的图（如 Coauthor CS），而标准 GraphGPS 使用全注意力时因 OOM 问题而失败。
• 消融研究显示，尽管局部邻域始终能提升性能，但虚拟节点与扩展图边的相对贡献因数据集而异，二者组合使用时效果最佳。