SGFormer：简化并增强Transformer以应对大型图表示的挑战

人工智能咨询培训老师叶梓转载标明出处

大型图数据的表示学习面临的主要挑战是如何在有限的计算资源下，有效地捕捉节点间的依赖关系并生成有用的节点表示。现有的基于Transformer的方法通常采用多层多头注意力机制，这虽然能够捕获全局信息，但计算复杂度高，难以扩展到拥有数百万甚至更多节点的大型图。而且这些方法往往需要额外的位置编码、特征预处理或额外的损失函数，增加了模型的复杂性和计算负担。上海交通大学的研究者团队提出了一种新的简化图Transformer（SGFormer），旨在解决大型图表示学习中的可扩展性和效率问题。

SGFormer通过单层注意力机制实现对任意节点间信息的高效传播。这种设计不仅减少了计算开销，而且由于其线性复杂度，使得模型能够扩展到亿级别节点的超大规模图。SGFormer的主要优点为：

简化的全局注意力机制：SGFormer采用单层单头注意力模型，通过一次传播层和线性复杂度，有效地在节点间传播信息。
无需位置编码或预处理：与需要位置编码或其他预处理步骤的方法不同，SGFormer直接使用节点特征，简化了模型的复杂性。
高效的计算性能：SGFormer的计算复杂度与节点数成线性关系，这使得它在处理大型图时具有显著的效率优势。

方法

研究者首先使用一个神经网络层将输入特征映射到潜在空间中的节点嵌入，其中f_I可以是一个浅层（例如，单层）MLP。这些节点嵌入将用于后续的注意力计算和传播。

SGFormer的核心是其简单的全局注意力模型，该模型能够以线性复杂度捕捉节点间的隐式依赖关系。通过单层注意力机制，任意节点对之间的信息可以被计算出来。与传统的Transformer不同，SGFormer发现仅需单层全局注意力就足够，因为在密集连接的注意力图上的单层传播可以自适应地将每个节点的信息传播到批次中的任意节点。因此，尽管模型简单，但足以捕捉任意节点对之间的隐式依赖关系，同时显著减少了计算开销。

为了结合输入图G的先验信息，研究者采用了一种简单而有效的方案，将全局注意力的输出与GNN传播的嵌入在输出层组合。具体为：，其中α是一个权重超参数，GN模块可以是一个简单的GNN架构（例如，GCN），具有良好的可扩展性。输出函数f_O将最终表示Z_O映射到预测，具体取决于下游任务。

模型的总体计算复杂度为O(N+E)，其中E=∣E∣，因为GN模块需要O(E)。由于图的典型稀疏性（即2E≪N^2），模型可以线性地相对于图大小进行扩展。另外由于只有一层全局注意力和简单的GNN架构，模型相当轻量，使得训练和推理更加高效。

对于连GCN也无法在单个GPU上使用全批量处理训练的大型图，可以使用随机小批量划分方法。这种方法在训练过程中只增加了微不足道的额外成本，并且允许模型扩展到任意大型图。

由于全局注意力机制的线性复杂度，可以采用大批量大小，这有助于模型在每个小批量中捕获信息性的全局交互。模型还与高级技术（如邻居采样、图聚类和历史嵌入）兼容。

表1展示了图Transformer在架构、表达性和可扩展性方面的对比。大多数现有Transformer针对小图分类任务开发，而部分专注于节点分类任务，面临的挑战在于大型图的规模。

在架构方面，现有模型可能采用边缘/位置嵌入或增强训练损失来捕获图信息。然而，这些方法可能需要额外的预处理步骤，对于大型图可能耗时且占用内存，也可能使优化过程复杂化。相比之下，SGFormer无需位置嵌入、增强损失或预处理，仅使用单层单头全局注意力，提供了一种既高效又轻量的解决方案。

在表达性方面，一些图Transformers将注意力计算限制在节点的子集上，允许线性扩展，但牺牲了表达性。SGFormer则在每一层都保持对所有节点的注意力计算，同时实现O(N)复杂度，无需任何近似或随机组成部分，训练过程更稳定。

在可扩展性方面，现有模型通常因全局所有节点对注意力而具有O(N^2)复杂度，限制了对中等规模图的扩展。SGFormer则实现线性扩展，支持在具有多达0.1M个节点的大型图上进行全批量训练。对于更大的图，SGFormer与大批量大小的小批量训练兼容，允许模型在对性能影响极小的情况下捕获信息性的全局信息。值得关注的是SGFormer可以扩展到web规模的图，如ogbn-papers100M，展示了其强大的可扩展性。

想要掌握如何将大模型的力量发挥到极致吗？叶老师带您深入了解 Llama Factory —— 一款革命性的大模型微调工具。实战专家1小时讲解让您轻松上手，学习如何使用 Llama Factory 微调模型。

评论留言“参加”或扫描微信备注“参加”，即可参加线上直播分享，叶老师亲自指导，互动沟通，全面掌握Llama Factory。关注享粉丝福利，限时免费录播讲解。

LLaMA Factory 支持多种预训练模型和微调算法。它提供灵活的运算精度和优化算法选择，以及丰富的实验监控工具。开源特性和社区支持使其易于使用，适合各类用户快速提升模型性能。

实验

实证评估聚焦于SGFormer在不同规模图数据集上的表现，检验了其学习有效表示和在大型图上扩展的能力。评估涵盖了从中等规模（2K至30K节点）到大规模（0.1M至0.1B节点）图的数据集。

评估始于常用图数据集，包括三个高同质性比率的引文网络Cora、CiteSeer和PubMed，以及四个异质性图Actor、Squirrel、Chameleon和Deezer-europe。这些图的节点数在2K至30K之间。对于引文网络，数据划分遵循了半监督设置。Actor和Deezer-europe使用了随机划分，而Squirrel和Chameleon则依据评估论文进行划分，以过滤掉重叠节点。

SGFormer与多个GNN和Transformer模型进行了比较，包括标准的GCN、GAT和SGC，以及更先进的模型如JKNet、APPNP、SIGN、H2GCN、CPGNN和GloGNN。此外，还与专为节点分类设计的最先进模型NodeFormer，以及为图分类设计的Graphormer和GraphTrans进行了比较。

表2展示了所有模型的结果，显示SGFormer在所有数据集上显著优于标准的GNN模型（GCN、GAT和SGC），在Actor数据集上的最大改进达到25.9%。这一发现表明，尽管SGFormer模型简单，但其单层全局注意力机制效果显著。另外与NodeFormer等其他先进模型相比，SGFormer在大多数情况下也展现出更高的性能。

进一步在节点数从百万到十亿的大型图数据集上评估SGFormer，包括引文网络ogbn-arxiv、蛋白质相互作用网络ogbn-proteins、物品共现网络Amazon2M和社交网络pokec。这些评估使用了OGB的公共划分。

由于规模庞大，与MLP、GCN、SGC、SIGN以及使用邻居采样的GCN-NSampler和GAT-NSampler进行了比较。主要竞争对手为NodeFormer，这是一个最近提出的具有全对注意力的可扩展图Transformer。

表3展示了实验结果，SGFormer在所有数据集上均取得了优异的结果，相较于GNN竞争对手显示出显著的性能提升。这证实了SGFormer的全局注意力机制在捕捉大量节点间隐式相互依赖性方面的有效性。特别是在最大的数据集ogbn-papers100M上，SGFormer表现出了其在极大规模图上的强大性能。

表4报告了在Cora、PubMed和Amazon2M上的每个时期的训练时间、推理时间和GPU内存成本。结果显示，SGFormer的速度比其他竞争对手快了几个数量级，尤其在Cora数据集上比Graphormer快了38倍至141倍。