[论文解读] SGFormer: Simplifying and Empowering Transformers for Large-Graph Representations

Qitian Wu, Wentao Zhao+6 more|arXiv (Cornell University)|Jun 19, 2023

Advanced Graph Neural Networks被引用 11

一句话总结

SGFormer 使用简单的一层、线性时间全局注意力来处理大规模图，实现具竞争力的节点表征，并扩展到网页尺度的图（高达 0.1B 节点），同时获得显著的效率提升。

ABSTRACT

Learning representations on large-sized graphs is a long-standing challenge due to the inter-dependence nature involved in massive data points. Transformers, as an emerging class of foundation encoders for graph-structured data, have shown promising performance on small graphs due to its global attention capable of capturing all-pair influence beyond neighboring nodes. Even so, existing approaches tend to inherit the spirit of Transformers in language and vision tasks, and embrace complicated models by stacking deep multi-head attentions. In this paper, we critically demonstrate that even using a one-layer attention can bring up surprisingly competitive performance across node property prediction benchmarks where node numbers range from thousand-level to billion-level. This encourages us to rethink the design philosophy for Transformers on large graphs, where the global attention is a computation overhead hindering the scalability. We frame the proposed scheme as Simplified Graph Transformers (SGFormer), which is empowered by a simple attention model that can efficiently propagate information among arbitrary nodes in one layer. SGFormer requires none of positional encodings, feature/graph pre-processing or augmented loss. Empirically, SGFormer successfully scales to the web-scale graph ogbn-papers100M and yields up to 141x inference acceleration over SOTA Transformers on medium-sized graphs. Beyond current results, we believe the proposed methodology alone enlightens a new technical path of independent interest for building Transformers on large graphs.

研究动机与目标

• 重新评估在大规模图上深度注意力的必要性并探索可扩展的 Transformer 设计。
• 开发一个简单但具有表达力的模型，能够高效地捕捉全对节点交互。
• 在数千到十亿节点的图上实现可扩展学习，而无需繁重的预处理。
• 展示相对于最先进的图 Transformer 的实证性能与效率优势。

提出的方法

• 提出 Simplified Graph Transformer (SGFormer)，采用对 N 个节点在 O(N) 时间内工作的单层全局注意力。
• 使用浅层输入映射 f_I 获取注意传播的节点嵌入 Z^(0)。
• 定义带有 Q、K、V 投影的线性注意力机制，以及将全局注意力与自环信息结合的残差传播步骤 Z（等式 2–3）。
• 避免使用位置编码、边嵌入、预处理和额外损失；不使用随机近似。
• 可选地通过输出 Z_O = (1-α)Z + α GN(Z^(0), A) 来融入图结构，并通过线性头进行预测（等式 4）。
• 通过小批量划分实现大规模训练，并与邻居采样、聚类和历史嵌入兼容。

Figure 1: Illustration of the proposed model SGFormer and its data flow. The input graph data entails node features $\mathbf{X}$ and graph adjacency $\mathbf{A}$ . For large graphs, we need to use mini-batch sampling that randomly partitions the input graph into mini-batches with smaller sizes. Each

实验结果

研究问题

• RQ1 单层、线性时间全局注意力 Transformer 是否能够在大规模图上达到甚至超过多层 Transformer 的性能？
• RQ2 与 GNN 和图 Transformer 相比，SGFormer 在中等规模和网页规模的图上表现如何？
• RQ3 紧凑架构在有限监督条件下对泛化能力的影响如何？
• RQ4 线性时间注意力是否足以在没有近似的情况下捕获全对交互？

主要发现

• SGFormer 在从数千到十亿节点的十二个节点属性基准测试上取得竞争性或更优的性能，且仅使用单层注意力。
• 在中等规模的图上，SGFormer 的表现优于标准 GNN，最高提升可达 25.9%（如 actor 数据集），并且在 Graphormer 与 GraphTrans 面前仍具高度竞争力。
• 在大规模图上，SGFormer 在五个数据集上优于 NodeFormer，在 ogbn-papers100M 上达到 66.0 的准确率，训练约 3.5 小时，单个 GPU 内存约 23.0 GB。
• SGFormer 能扩展到网页尺度的图（ogbn-papers100M，节点数达到 0.1B），具线性复杂度，显著减少了训练/推理时间（例如在中等图上比 SOTA Transformer 的推理快上 141 倍）。
• 相比二次注意力基线，SGFormer 实现了显著的效率提升（例如在 Cora 上比 Graphormer 的训练快 38 倍、推理快 141 倍）。
• 更深的多层注意力并未始终带来性能提升，且成本更高，突显了一层设计在大规模图上的有效性。
• 理论分析将一层注意力与去噪优化视角联系起来，表明在合适设置下其可达到多层效果。

Figure 2: Scalability test of training time per epoch and GPU memory usage w.r.t. graph sizes (a.k.a. node numbers). NodeFormer suffers out-of-memory when # nodes reaches more than 30K.