Understanding Transformer Reasoning Capabilities via Graph Algorithms
https://arxiv.org/abs/2405.18512
我们对 Transformer 模型的图推理能力进行了全面的评估,并证明它们通常优于更多特定领域的图神经网络。
Transformer ()是Google 于 2017 年推出的通用顺序深度学习架构。最近,Transformer表现出了其擅长学习基于图的算法任务的惊人能力,例如连通性(例如,“从节点 A 到节点 B 有路径吗?”)、最短路径(“从节点 A 到节点 B 需要经过多少条边?”)和循环检查(“图中是否存在循环,例如,从 A → B、B → C、C → A 的边?”)。
为什么这令人惊讶?因为与消息传递神经网络(message passing neural networks,MPNN) 不同,消息传递神经网络的参数明确编码为输入图的结构,而 Transformer 是一种通用架构,最广为人知,用于语言任务。其骨干是一种计算成本高昂的自注意力机制(self-attention mechanism),用于编码输入单词/tokens之间的关联,并非专门为图而设计。因此,虽然基于 Transformer 的神经网络取得了巨大的经验进步,但我们显然仍然缺乏对其在现实场景中的算法推理能力的理论理解。了解哪种架构最适合哪些现实世界的任务意味着花费更少的时间通过精心挑选的任务来定义和评估新架构,并更好地理解架构决策权衡。
在“Understanding Transformer Reasoning Capabilities via Graph Algorithms”一文中,我们从理论和实证两个角度对 Transformer 的图推理问题进行了广泛的研究。我们发现 Transformer 可以以高度参数高效的方式解决可并行化任务(如连接性),并且在需要分析长距离依赖关系的任务上优于图神经网络(graph neural networks,GNN)。我们还引入了一种新颖的图推理任务表示层次结构,在几种实际的参数缩放方案中形式化了 Transformer 的推理能力。
实验
MPNN 是一种定制的神经架构,专门用于学习遵从图结构的函数,而 transformer 则更为通用。但 transformer 和 MPNN 都可以使用类似的 ML 框架进行训练,以解决图算法问题。
对于实验的经验部分,我们研究了 Transformer 是否可以像 MPNN 那样对图任务进行编码,如果可以,是否可以更有效地进行编码。
为了测试这一点,我们将图任务编码为 Transformer 输入。给定一个图实例(具有节点和边)和一个查询(确定节点𝑣 2 和𝑣 4 是否连接),我们将每个节点 ID 编码为离散标记(类似于 LLM 将单词映射到标记的方式),将每条边编码为一对标记。然后,我们将问题编码为节点标记(node tokens)列表、边标记(edge tokens)列表和查询的标记编码(token encoding of the query)。
如何将图推理任务实例(这里是图连接)嵌入为transformer输入。
由于 transformer 和 MPNN 并不是用于图结构分析的唯一 ML 方法,我们还比较了各种其他基于 GNN 和 transformer 的架构的分析能力。对于 GNN,我们将 transformer 和 MPNN 与图卷积网络(GCN) 和图同构网络(GIN) 等模型进行了比较。
此外,我们将我们的 Transformer 与更大的语言模型进行了比较。语言模型也是 Transformer,但参数要多出几个数量级。我们将 Transformer 与Talk Like a Graph(https://arxiv.org/abs/2310.04560)中描述的语言建模方法进行了比较,该方法将图编码为文本,使用自然语言来描述关系,而不是将输入图视为抽象标记的集合。
我们要求训练有素的语言模型使用多种提示方法解决各种检索任务:
零样本(Zero-shot),仅提供单一提示并要求提供解决方案,而不会提供进一步的提示。
Few-shot,在要求模型解决任务之前,提供几个已解决的提示-响应对的示例。
思维链(Chain-of-thought,CoT),它提供了一组示例(类似于少样本),每个示例在要求模型解决任务之前包含一个提示、一个响应和一个解释。
零样本 CoT,要求模型展示其工作,而不包括额外的计算示例作为背景。
CoT-bag,它要求 LLM 在提供相关信息之前构建一个图。
对于实验的理论部分,我们创建了一个任务难度层次结构,以评估 transformer 可以用小模型解决哪些任务。
我们只考虑适用于有界大小的无向和无加权图的图推理任务:节点数、边数、边存在性、节点度、连通性、节点连通性(对于无向图)、循环检查和最短路径。
第 3 节的理论层次总结,直观地展示了哪种类型的图推理任务可以在哪种转换器缩放机制(Depth1 (D1)、LogDepth (LD)、LogDepthWide (LDW) 和 LogDepthPause (LDP))下解决。
在这个层次结构中,我们根据深度(Transformer 中自注意力层的数量,按顺序计算)、宽度(每个图标记使用的向量的维度)、空白标记的数量和三种不同类型对图任务难度进行分类:
检索任务:简单、本地聚合任务。
可并行化任务:从并行操作中受益匪浅的任务。
搜索:并行操作带来的好处有限的任务。
结果
经验
对于在专门数据上训练的小型模型(约 60M 个参数),Transformer 的表现优于没有经过针对性训练的大型 LLM(如下面的条形图所示)。此外,只需少量样本和参数,Transformer 的算法推理能力就可以得到增强。
TFGraphs2-结果
所有检索任务(顶部)以及可并行化和搜索任务(底部)的准确率(按训练模式划分),对比提示的 PaLM2 LLM(左)和训练过的 transformers(右)。
如下所示,MPNN 在“本地”任务上表现更好,在这些任务中,我们无需进行多次推理即可聚合有关邻近节点的信息。无论是在 1K 还是 100K 样本上进行训练,MPNN 在诸如节点度(与节点相邻的最大边数)和循环检查等本地任务上都具有更高的准确率。
此外,Transformer 在“全局”任务上表现更好,因为信息需要在图实例之间传播。全局任务的一个例子是图连通性,因为图中的任何两个节点可能相距很远。
TFGraphs3-架构
按架构执行任务,将节点度与循环检查和连通性进行比较。
这些结果表明,对于 transformer 来说,小样本情况和大样本情况之间的差异要大得多,并且在 100K 样本上训练时,GNN 的性能无法接近 transformer 的性能。
训练神经网络时,有两个关键因素在起作用:容量和可学习性。容量表示神经网络架构的基本限制。也就是说,容量决定了神经架构可以表示哪些类型的数学函数。可学习性询问这些函数是否真的可以在具有相对较少训练样本的 ML 设置中学习。如果架构倾向于使用较少样本来选择某些类型的解决方案,我们就说它对它们具有正归纳偏差。
如果我们尝试在此框架下理解我们的 Transformer 和 MPNN 结果,我们会发现 GNN 确实在图任务上具有更好的归纳偏差,这从它们在数据较少的环境中表现出色就可以看出。另一方面,Transformer 在较大数据集上的成功表明 GNN 可能受到容量的限制,因为这些反映了无法通过更多数据克服的基本模型缺陷。
理论
我们用理论结果支持了实证结果,对比了两种架构的容量或表达能力。先前的研究表明,Transformer 可以模拟并行算法,而 MPNN 等其他模型则受到缺乏全局连接的限制。使用我们的任务难度层次结构,我们证明了 Transformer 在MPC 模型下对可并行任务具有更大的表达能力,并且可以使用对数深度和有界宽度来解决它们(即,模型的宽度最多是某个数量m,其增长速度比图标记的数量N慢得多,即m ≤ N0.1)。另一方面,除非这些任务在计算和参数数量方面都大得多,否则没有 MPNN 可以解决这些任务。
我们还发现,检索任务可以用单层 Transformer 模型解决,而并行任务则可以通过对数深度轻松完成(即 Transformer 的深度与图中节点数成对数关系)。搜索任务也可以用对数深度完成,但不是参数高效的方式。
TFGraphs4-层次结构
任务难度层次的可视化,其中检索、并行和搜索任务根据解决它们所需的转换器的复杂性进行组织。
这些结果解释了为什么 Transformer 在连接性方面表现更佳——与 GNN 不同,模型的大小(就深度和参数数量而言)不需要随着输入图的大小而快速增长。
结论
我们对 Transformer 模型的图推理能力进行了全面评估,阐明了它们在不同图推理任务中的有效性。通过引入一种新颖的表示层次结构,该研究区分了检索、可并行化和搜索推理任务,并深入了解了 Transformer 在不同粒度级别的性能。实证研究表明,Transformer 在基于图的推理问题中表现出色,通常可以匹敌或超越专门的图模型。此外,该研究还强调了 Transformer 有效捕捉全局图模式的卓越能力,展示了它们理解长距离依赖关系的能力,这是解决涉及全局图结构的任务的关键因素。总体而言,这项工作明确了反映 Transformer 基本推理能力的精确表示权衡,并证明了用于量化这些能力的任务确实可以以样本高效和参数高效的方式进行学习。
致谢
这项工作是与 Google 的同事合作完成的,包括 Ethan Hall、Anton Tsitsulin、Mehran Kazemi、Bryan Perozzi、Jonathan Halcrow 和 Vahab Mirrokni。