查看原文
其他

​NeurIPS 2023 | SGFormer: 仅使用一层全局注意力的简化图Transformer

让你更懂AI PaperWeekly
2024-08-22


设计针对图结构数据的 Transformer 模型(通常简称 graph Transformer)目前已成为了一个备受关注的研究方向。有别于传统图神经网络(简称 GNN)每层更新只能聚合邻居节点的信息,Transformer 通过全局注意力机制在每层更新中可以聚合图中所有其他节点的信息。这种设计可以有效避免 GNN 的过度挤压问题、改善对异质图的建模、缓解噪声/残缺/冗余连边对性能的影响 [1]
然而,目前大部分 graph Transformer 只能处理小规模图(如分子图)[2,3,4]。对于普通规模的图数据,例如节点数目达到上千的量级,全局注意力的平方级复杂度会使得传统的模型难以扩展。
另一方面,目前的 Transformer 模型通常采用比较复杂的模型结构,例如堆叠多层的多头注意力机制。对于大规模图数据,一个常见的现象是图中带标签的节点比例较低,这种情况下复杂的模型结构会影响到模型的泛化性能。
为此,在这项研究工作中,我们试图探索一条有别于大众方法的技术路径,即对 Transformer 进行简化,以实现对大规模图的高效建模。

具体的,我们提出了 Simplified Graph Transformer(简称 SGFormer),在只使用一层线性复杂度的全局注意力网络的情况下,我们的方法在 12 个不同性质和规模的节点性质预测 benchmark 上都表现出极具竞争力的精度,首次实现了将 Transformer 方法扩展到上亿节点规模的超大图上,并在中等规模图上相比同类 graph Transformer 方法实现了大幅度的计算效率提升。

论文题目:

Simplifying and Empowering Transformers for Large-Graph Representations

论文链接:

https://arxiv.org/pdf/2306.10759.pdf

代码链接:
https://github.com/qitianwu/SGFormer


三分钟读论文

Powered by AI © PaperWeekly








模型介绍

本文提出的方法主要基于一个观察:由于全局注意力机制已经能够实现任意两两节点间的信息传递,因此原则上只需要一层注意力网络模型就能具备建模任意两两间的影响交互的表达能力基于此,我们试图将全局注意力网络从“堆叠多层”简化到“一层”。
模型结构:SGFormer 的模型结构非常简单,主体部分采用一层全局注意力网络和普通 GNN 的结合。假设图中含有  个节点,节点输入特征表示 首先使用一个浅层网络(如一层 MLP)将节点特征映射到隐空 这里  定义为节点的初始表征,而后被用于全局注意力的计算和传播更新。
SGFormer 的全局注意力层采用一种简化版的计算方式,可以在不丢失精度的情况下实现线性复杂度的加速。
上式 经过特征变换后的节点特征(这里 , 被实例化成线性特征变换网络), 是注意力权重的归一化矩阵。在注意力权重的计算部分,传统的 Softmax 注意力网络采用的 这种计算需要 的复杂度。
而上式中的注意力计算省去了 Softmax 操作,并且通过矩阵乘法结合律交换乘法次 这样只需要 的计算复杂度。值得强调的是,针对如何实现线性复杂度的加速,现有的 graph Transformer 也做了类似的探索,如 NodeFormer [1] 中采用随机特征对 Softmax 算子进行近似。不同的是,这里 SGFormer 采用的方法不需要引入任何近似方法,可以在不丢失精度的情况下实现线性复杂度的加速。

上述的全局注意力机制没有考虑观测图结构,为了将输入图信息融入节点表征,我们采用一种简单有效的思路,直接将全局注意力更新后的节点表征与 GNN 输出的节点表征进行线性加和:

这里的 GN 可以是任意的 GNN 网络, 是两项的结合权重, 是用于输出的浅层 MLP。下图展示了 SGFormer 前馈计算过程。

▲ SGFormer 的前馈计算示意图。假设输入图为 ,其中 表示节点特征, 表示邻接矩阵。当图的规模比较大时,需要先采用 mini-batch 方法将一个完整的图随机划分为固定大小的块 ,然后输入进模型;对于中等规模的图,则不需要进行节点划分。SGFormer 的主体结构采用一层全局注意力和普通 GNN 的结合,最终得到节点的表征和输出预测。
和同类方法的对比:为了更直观的展现 SGFormer 在设计上的不同之处,下表展示了在多个维度上 graph Transformer 方法的对比。可以看到,相比于其他 graph Transformer,我们提出的方法 SGFormer 不需要位置编码(positional encoding),连边表征(edge embedding),增广的损失函数(augmented loss)和任何预处理步骤
同时,在只使用线性  复杂度的情况下,仍然保留了建模任意两两节点影响交互的表达力(all-pair expressivity)。值得强调的是,相比之前的大图 Transformer(如 NodeFormer [1])最多只在百万级节点规模的图上进行了实验,SGFormer 成功扩展到了上亿节点规模的超大图。

▲ 与现有 graph Transformer 方法的对比,其中简称 PE 表示位置编码(positional encoding),MA 表示多头注意力(multi-head attention),AL 表征增广的损失函数(augmented loss),EE 表示连边表征(edge embedding),all-pair expressivity 表示模型是否能建模任意两两节点的影响交互,largest demo 表示实验使用的最大规模图的节点数目。


实验结果

我们在 12 个不同性质和规模的节点性质预测 benchmark 上通过和同类方法进行比较验证了提出方法的有效性。下表展示了在中等规模图上的测试分类效果,结果显示 SGFormer 虽然只使用一层注意力机制,仍然可以取得极具竞争力的性能。而其他 graph Transformer 方法在部分情况会遇到 out-of-memory(OOM)的问题。

▲ 中等规模图的测试分类精度,采用常用的数据划分和评测准则
下表展示了在大规模图上的测试分类效果,SGFormer 在不同数据集上都取得了最好的性能。特别地,在包含上亿节点规模的图数据集 ogbn-papers100M 上,现有的 graph Transformer 在有限计算资源下无法扩展,而 SGFormer 可以在单卡条件下只使用 23GB 显存在 3.5 小时内完成一轮训练。

▲ 大规模图的测试分类精度

下面是模型的计算效率和可扩展性的对比,结果显示 SGFormer 可以在中等规模图上实现了大幅度的效率提升,并且随着图规模的增加,时间和空间消耗都呈线性增长。当节点数目达 0.1M 时,SGFormer 只占用了近 3GB 的显存,这说明 SGFormer 具有很好的适配于大图的可扩展性。

▲ 计算效率对比:单个 epoch 训练时间(Tr),推理时间(Inf),空间消耗情况(Mem)

▲ 可扩展性对比:单个 epoch 训练时间与空间消耗情况
除此之外,本文也做了相关的消融实验。结果显示当全局注意力层数增加,SGFormer 的测试性能不一定会得到改善(甚至在部分情况下会恶化),而计算时间开销则会显著增加。
与此同时,对于其他模型(包括使用 Softmax 注意力网络,不使用自连接,NodeFormer [1]),使用一层注意力网络也能带来足够优越的性能。这也验证了我们最初的猜想,即当模型采用了全局注意力机制后,仅使用一层注意力网络就可以建模任意两两节点间的影响交互,带来足够的表达能力。

▲ 不同注意力网络层数对 SGFormer 性能和时间效率的影响

▲ 不同注意力网络层数对其他方法(SGFormer w/ Softmax 表示将 SGFormer 中的注意力网络替换为含 Softmax 的网络,SGFormer w/o self-loop 表示去掉模型中的自连接,NodeFormer [1])性能的影响


理论分析

除实验验证外,本文也从图信号处理的角度阐释了一层注意力机制为什么有效,并讨论了一层注意力网络和堆叠多层网络之间的关系。对这部分感兴趣的读者可以阅读我们论文的第 6 节。



结语

本文针对构建图数据的 Transformer 方法提出了一种简化版的模型设计(称作 SGFormer),其主体结构仅采用一层全局注意力机制,且不需要使用位置编码、增广的损失函数和任何预处理操作,这种简单直接的设计带来了线性复杂度的优势且大幅提升了计算效率。有别于现有大多数方法的设计理念,本文探索了一条新的技术路径,即利用更小的代价来实现更强大的模型,希望能为后续的相关工作提供启发。

此外,对 graph Transformer 这一新兴领域感兴趣的读者也欢迎关注我们的另外两项工作,分别在可扩展性和可解释性上作了相应的探索:



参考文献

[1] Qitian Wu, Wentao Zhao, Zenan Li, David Wipf, and Junchi Yan. Nodeformer: A scalable graph structure learning transformer for node classification. In NeurIPS 2022.
[2] Vijay Prakash Dwivedi and Xavier Bresson. A generalization of transformer networks to graphs. CoRR, abs/2012.09699, 2020.
[3] Zhanghao Wu, Paras Jain, Matthew A. Wright, Azalia Mirhoseini, Joseph E. Gonzalez, and Ion Stoica. Representing long-range context for graph neural networks with global attention. In NeurIPS 2021
[4] Dexiong Chen, Leslie O’Bray, and Karsten M. Borgwardt. Structure-aware transformer for graph representation learning. In ICML 2022.
[5] Qitian Wu, Chenxiao Yang, Wentao Zhao, Yixuan He, David Wipf, and Junchi Yan. DIFFormer: Scalable (graph) transformers induced by energy constrained diffusion. In ICLR 2023


更多阅读



#投 稿 通 道#

 让你的文字被更多人看到 



如何才能让更多的优质内容以更短路径到达读者群体,缩短读者寻找优质内容的成本呢?答案就是:你不认识的人。


总有一些你不认识的人,知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁,促使不同背景、不同方向的学者和学术灵感相互碰撞,迸发出更多的可能性。 


PaperWeekly 鼓励高校实验室或个人,在我们的平台上分享各类优质内容,可以是最新论文解读,也可以是学术热点剖析科研心得竞赛经验讲解等。我们的目的只有一个,让知识真正流动起来。


📝 稿件基本要求:

• 文章确系个人原创作品,未曾在公开渠道发表,如为其他平台已发表或待发表的文章,请明确标注 

• 稿件建议以 markdown 格式撰写,文中配图以附件形式发送,要求图片清晰,无版权问题

• PaperWeekly 尊重原作者署名权,并将为每篇被采纳的原创首发稿件,提供业内具有竞争力稿酬,具体依据文章阅读量和文章质量阶梯制结算


📬 投稿通道:

• 投稿邮箱:hr@paperweekly.site 

• 来稿请备注即时联系方式(微信),以便我们在稿件选用的第一时间联系作者

• 您也可以直接添加小编微信(pwbot02)快速投稿,备注:姓名-投稿


△长按添加PaperWeekly小编



🔍


现在,在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧


·
·

继续滑动看下一个
PaperWeekly
向上滑动看下一个

您可能也对以下帖子感兴趣

文章有问题?点此查看未经处理的缓存