互联网资讯 / 人工智能 · 2023年12月30日

Transformer逆向工程为数学框架,研究团队发表25篇文章

Transformer是Google团队于2017年6月推出的自然语言处理(NLP)领域的经典模型,最初由Ashish Vaswani等人在论文《Attention Is All You Need》中提出。自Transformer问世以来,它在NLP、计算机视觉(CV)、语音处理、生物学和化学等多个领域都带来了显著的进展。

Transformer的应用日益广泛,诸如GPT-3、LaMDA和Codex等模型均基于该架构构建。然而,随着Transformer模型的不断扩展,其开放性和高容量也为潜在的、甚至有害的行为提供了更大的空间。即使在大型模型训练完成多年后,创建者和用户仍然会不断发现之前未曾遇到的问题。

解决这些问题的一种方法是机械可解释性(Mechanistic Interpretability),即对Transformer的计算过程进行逆向工程,类似于程序员将复杂的二进制文件转换为可读源代码的过程。

如果逆向工程能够成功实施,我们将能够更系统地解释当前模型的安全隐患,识别潜在问题,甚至可能预见未来尚未构建的模型安全风险。这一过程将Transformer的黑箱操作变得更加透明。尽管之前有研究者尝试对视觉模型进行逆向工程,例如Distill Circuits Thread项目,但迄今为止,尚未有针对Transformer或语言模型的类似研究。

本论文由25位研究者共同撰写,尝试从最基本的步骤开始逆向Transformer。论文由Chris Olah起草,他在Anthropic人工智能安全与研究公司工作,专注于神经网络逆向工程研究。随后,Neel Nanda对论文初稿进行了重要修改,Nanda目前是DeepMind的一名研究工程实习生。Nelson Elhage负责详细编辑论文,以提高各章节的清晰度,他曾在Stripe科技公司任职。

经逆向工程,Transformer「翻译」成数学框架 | 25位学者撰文

左:Neel Nanda;右:ChristopheR Olah

考虑到语言模型的复杂性和规模,该研究发现,从最简单的模型开始逆向Transformer是最有效的。研究的目标是识别简单的算法模式、主题(Motifs)或框架,然后将这些发现应用于更复杂的模型。具体而言,他们的研究仅包括具有两层或更少层的仅含注意力块的Transformer模型,这与层数多达96层的GPT-3形成鲜明对比。

经逆向工程,Transformer「翻译」成数学框架 | 25位学者撰文

论文地址:https://tRansfoRMeR-ciRcuITs.pub/2021/fRaMewoRk/index.htMl#acknowledgMents

该研究通过以一种新颖但数学上等效的方式对Transformer的操作进行概念化,能够深入理解这些小模型及其内部运作方式。值得注意的是,研究发现特定的注意力头,称为归纳头(Induction Heads),可以在这些小模型中解释上下文学习,并且这些注意力头只在至少有两个注意层的模型中发展。此外,研究还展示了这些注意力头在特定数据上的操作实例。

各章节内容概览

为了探讨逆向工程Transformer所面临的挑战,研究者对几个仅含注意力的玩具模型进行了逆向分析。

首先是零层Transformer模型的二元统计。研究者发现,二元表可以直接通过权重进行访问。

经逆向工程,Transformer「翻译」成数学框架 | 25位学者撰文

在讨论更复杂的模型之前,研究者认为考察零层(Zero Layer)Transformer非常有用。这种模型接收一个Token,进行嵌入,然后解除嵌入,以预测下一个Token的logits。

由于该模型无法从其他Tokens传递信息,它仅能基于当前Token预测下一个Token。这意味着,W_UW_E的最佳行为是近似二元对数似然。

经逆向工程,Transformer「翻译」成数学框架 | 25位学者撰文

零层注意力型Transformer模型。

接下来,单层注意力型Transformer模型可以看作是二元和skIP三元模型的集合。与零层Transformer类似,二元和skIP三元表也可以直接通过权重访问,而无需运行模型。这些skIP三元模型表现出惊人的表达能力,包括实现非常简单的上下文内学习。

对于单层注意力型Transformer模型,研究者提供了一些路径扩展(Path Expansion)技巧。

如下图所示,单层注意力型Transformer由一个Token嵌入组成,接着是一个注意力层(单独应用每个注意力头),最后是解除嵌入:

经逆向工程,Transformer「翻译」成数学框架 | 25位学者撰文

利用之前获得的张量标记(Tensor Notation)和注意力头的替代表征,研究者将Transformer的表示形式转化为三个项的乘积,具体如下图所示:

经逆向工程,Transformer「翻译」成数学框架 | 25位学者撰文

研究者采用的核心技巧是简单地扩展这个乘积,即将乘积(每个项对应一个层)转化为一个和,其中每个项对应一个端到端的路径。他们指出,每个端到端路径项都易于理解,可以独立推理,并能够组合叠加以形成模型行为。

经逆向工程,Transformer「翻译」成数学框架 | 25位学者撰文

最后,两层注意力型Transformer模型可以通过组合注意力头实现更复杂的算法。这些组合算法也能够直接通过权重检测到。需要强调的是,两层模型适应于注意力头的组合,以创建“归纳头”(Induction Heads),这是一种非常通用的上下文内学习算法。

具体而言,当注意力头有以下三种组合选择时:

Q – 组合:W_Q在一个受前面头影响的子空间中读取; K – 组合:W_K在一个受前面头影响的子空间中读取; V – 组合:W_V在一个受前面头影响的子空间中读取。

研究者表示,Q – 组合和K – 组合与V – 组合截然不同。前两者对注意力模式产生影响,使注意力头能够表达更复杂的模式。而V – 组合则影响一个注意力头专注于特定位置时所传递的信息。最终结果是,V – 组合头更像一个单一单元,并可能被用于创建额外的“虚拟注意力头”。

对于Transformer,有一个基本问题是“如何计算logits”?与单层模型使用的方法相似,研究者写出的乘积,其中每个项在模型中都是一个层,并扩展以形成和,其中每个项对应一个端到端路径。

经逆向工程,Transformer「翻译」成数学框架 | 25位学者撰文

其中,直接路径项和单个头项与单层模型中的相同。最后的“虚拟注意力头”项对应于V – 组合。尽管虚拟注意力头在概念上很有趣,研究者发现它们在小规模的两层模型中通常难以发挥重要作用。

此外,这些项中的每一个都对应于模型可以实现更复杂注意力模式的一种方式。虽然理论上难以进行推理,但在讨论归纳头时,它们很快会在具体实例中被应用。