TokenGT:Transformer是强大的图学习器

论文标题:Pure Transformers are Powerful Graph Learners
论文链接:https://arxiv.org/abs/2207.02505
论文来源:NIPS 2022

一、概述

由于Transformer的完全注意力架构接收、处理和关联任意结构的输入和输出的能力,消除了将特定于数据和任务的归纳偏差融入网络架构的需要,这使得其在NLP、CV等诸多领域成为了通用架构。与大规模训练相结合,它为构建一个通用模型开辟了新的篇章,该模型可以解决涉及多种数据模式甚至混合模式的广泛问题。

在图学习领域,受这些突破的启发,许多研究尝试将自注意力机制结合到之前以消息传递为主的图神经网络(GNN)架构中。然而,由于跨节点的全局自注意力不能反映图结构,这些方法引入了特定于图的架构修改。这些方法包括
①限制自注意力只关注节点的局部邻域;
②使用全局自注意力与消息传递GNN相结合;
③通过attention的偏置将边信息注入全局自注意力机制中。

尽管性能不错,但就通用性而言,这样的修改可能是一个限制约束,特别是考虑到未来与多任务和多模态通用注意力架构的集成。此外,由于其偏离纯粹的自注意力机制,这些方法可能会继承消息传递的问题,如过平滑,并且与有用的工程技术不兼容,例如为标准自注意力开发的linear attention。

在本文中我们采用标准的Transformer架构,将其直接应用在图上。我们将所有的节点和边看做独立的token,并且使用相应的token-wise embedding来增强(augment)他们,最后将这些token作为标准Transformer的输入。模型命名为Tokenized Graph Transformer (TokenGT),该模型与用在语言和视觉上的Transformer相同,每个节点或边都被视为一个token,类似于句子中的word或图像中的patch。本文也证明了这种简单的方法在理论和实践中都能成为强大的图学习器。在这篇博客中只介绍方法部分。

本文的一个关键理论结果是我们证明了通过适当的token-wise embedding,节点和边token上的自注意力可以近似图上的任何置换等变线性算子(permutation equivariant linear operator)。另外,我们还发现一种非常简单的embedding就可以实现这种精确的近似。这提供了一个坚实的理论保证,通过embedding和足够的注意力头,Transformer至少具备二阶不变图网络(2-IGN)一样的表达能力,这已经超过了所有的消息传递GNN的能力。这赋予了模型至少与二维Weisfeiler-Lehman (WL)图同构检验相当的表达能力,这对于现实时间的图数据已经是足够的了。我们进一步将我们的理论结果扩展到具有k阶超边的超图上,表明具有k阶广义token embedding的Transformer至少与k-IGN和k-WL测试一样具有表达能力。

二、方法

TokenGT是一个纯Transformer的架构,使用token-wise embedding来表征图信息,包括node identifier和type identifier两部分。对于一个图\mathcal{G}=(\mathcal{V},\mathcal{E})\mathcal{V}=\left \{v_{1},\cdots ,v_{n}\right \}m条边\mathcal{E}=\left \{e_{1},\cdots ,e_{m}\right \}\subseteq \mathcal{V}^{2},以及节点和边的特征X^{\mathcal{V}}\in \mathbb{R}^{n\times C},X^{\mathcal{E}}\in \mathbb{R}^{m\times C}。在TokenGT中,我们将图的每个节点和边当做一个独立的token(也就是说一共有n+m个token),其特征也就是X=[X^{\mathcal{V}};X^{\mathcal{E}}]\in \mathbb{R}^{(n+m)\times C}

一种朴素的处理图的方法是直接将X输入给Transformer,然而这是不合适的,因为图的连通性被丢弃了。为了能够完整地表示图信息,TokenGT采用一种token-wise embedding来增强X,具体的,也就是使用标准正交的node identifier和可训练的type identifier来编码一个token是一个节点还是一条边。本文证明了采用这样一种简单表示方法的Transformer具有很强的Graph学习能力。

  1. Node Identifiers

TokenGT采用的token-wise embedding的第一个部分是标准正交的node Identifier,用来表示给定的输入图的连通结构。对于一个输入图\mathcal{G}=(\mathcal{V},\mathcal{E}),我们首先提供n个node-wise的标准正交向量P\in \mathbb{R}^{n\times d_{p}}作为node Identifier,然后按照以下方式来增强图的token特征向量X
①对于每个节点v\in \mathcal{V},增强X_v变成[X_{v},P_{v},P_{v}]
②对于每条边(u,v)\in \mathcal{E},增强X_{(u,v)}变成[X_{(u,v)},P_{u},P_{v}]

直观地来看,Transformer可以通过这些增强的token来完全识别图的连通结构,这是因为比较两个token之间的node Identifier可以揭示它们的关联信息(incidence information)。举例来说,对于一条边e=(u,v)是否与节点k相连,当且仅当k\in (u,v)时才有[P_{u},P_{v}][P_{k},P_{k}]^{T}=1否则为0。这允许Transformer识别和利用图的连通性结构,例如,当局部的信息很重要时,通过对相关联的pair对施加更多权重。

TokenGT对Node Identifier矩阵P的唯一要求是标准正交,因此在实际应用中有多种选择,本文列举了两种:
①正交随机特征(Orthogonal random features, ORFs),也就是对随机高斯矩阵G\in \mathbb{R}^{n\times n}进行QR分解得到Q矩阵Q\in \mathbb{R}^{n\times n},采用Q的行向量。
②对图的拉普拉斯矩阵进行特征分解\Delta =I-D^{-1/2}AD^{-1/2}=U^{T}\Lambda U,然后采用U的行向量。

ORFs并没有编码任何图结构信息,这是因为它们是完全随机的。这意味着基于ORFs的node Identifier只需要从node Identifier提供的关联信息中编译和识别图结构。虽然这是很有挑战性的,但是实验结果表明Transformer具备这样的能力。

与ORFs相比,拉普拉斯特征向量提供了一种图的positional embeddings,它描述了图上节点之间的距离。实验结果显示拉普拉斯特征向量比ORFs展示出更好的性能,这是因为其编码了位置信息。拉普拉斯特征向量的一个有趣的方面是,它们可以被看作是NLP Transformer对正余弦位置编码的推广,因为一维链图的特征向量是正弦和余弦函数。因此,通过选择拉普拉斯特征向量作为node Identifier,TokenGT可以被解释为涉及关系结构的输入的NLP Transformer的直接扩展。

  1. Type Identifiers

Token-wise embedding的第二个部分是可训练的type Identifier,用于标识一个token是节点还是边。对于一个输入图\mathcal{G}=(\mathcal{V},\mathcal{E}),首先准备一个可训练的参数矩阵E=[E^{\mathcal{V}};E^{\mathcal{E}}]\in \mathbb{R}^{2\times d_{e}},包含两个type IdentifierE^{\mathcal{V}}E^{\mathcal{E}},分别对应节点和边。然后对上面增强过的token继续进行增强:
①对于每个节点v\in \mathcal{V},增强[X_{v},P_{v},P_{v}]变成[X_{v},P_{v},P_{v},E^{\mathcal{V}}]
②对于每条边(u,v)\in \mathcal{E},增强[X_{(u,v)},P_{u},P_{v}]变成[X_{(u,v)},P_{u},P_{v},E^{\mathcal{E}}]

这些embedding提供了关于给定token是节点还是边的信息,这是至关重要的,例如,当注意力头试图专门关注节点token而忽略边token时。

  1. Main Transformer

有了node Identifier和type Identifier,现在我们可以得到输入给Transformer的特征向量X^{in}\in \mathbb{R}^{(n+m)\times (C+2d_{e}+d_{e})},将X^{in}通过一个参数矩阵w^{in}\in \mathbb{R}^{(C+2d_{e}+d_{e})\times d}映射到d维就可以输入给Transformer了。对于图级的预测任务,TokenGT会加入一个特殊token[graph]与其对应的可训练的embeddingX_{[graph]}\in \mathbb{R}^{d},这类似于BERT和ViT的做法。因此在面对图级任务时,模型的输入为Z^{(0)}=[X_{[graph]};X^{in}w^{in}]\in \mathbb{R}^{(1+n+m)\times d}

  1. 归纳偏差

类似于语言和视觉领域的Transformer模型,TokenGT将输入的节点和边视为独立的token,并对它们应用自注意力。与当前的图神经网络(GNN)相比,这种方法导致了更少的归纳偏差,其中稀疏图结构或者更基本的图的排列对称性被故意整合到每一层中。对于TokenGT,这些信息完全作为输入的一部分,通过token-wise embedding提供,模型必须学会如何从数据中解释和利用这些信息。尽管这种弱归纳偏差可能会引发关于模型表达能力的质疑,但本文的理论分析表明,得益于token-wise embedding和自注意力的表达能力,TokenGT是一个强大的图学习器。

三、实验

  1. 大规模图回归任务
大规模图
  1. 每层的Attention距离
每层的Attention距离
最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 206,214评论 6 481
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 88,307评论 2 382
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 152,543评论 0 341
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 55,221评论 1 279
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 64,224评论 5 371
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 49,007评论 1 284
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 38,313评论 3 399
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 36,956评论 0 259
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 43,441评论 1 300
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 35,925评论 2 323
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 38,018评论 1 333
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 33,685评论 4 322
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 39,234评论 3 307
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 30,240评论 0 19
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 31,464评论 1 261
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 45,467评论 2 352
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 42,762评论 2 345

推荐阅读更多精彩内容