Transformer Circuits 的数学框架（2021）

# Transformer 电路的数学框架 来源：https://transformer-circuits.pub/2021/framework/index.html ### 目录 - 模型简化 (https://transformer-circuits.pub/2021/framework/index.html#model-simplifications) - 高层次架构 (https://transformer-circuits.pub/2021/framework/index.html#high-level-architecture) - 虚拟权重与作为通信通道的残差流 (https://transformer-circuits.pub/2021/framework/index.html#residual-comms) - 注意力头是独立且可加的 (https://transformer-circuits.pub/2021/framework/index.html#architecture-attn-independent) - 注意力头作为信息移动 (https://transformer-circuits.pub/2021/framework/index.html#architecture-attn-as-movement) - 路径扩展技巧 (https://transformer-circuits.pub/2021/framework/index.html#onel-path-expansion) - 将注意力头项拆分为查询-键电路和输出-值电路 (https://transformer-circuits.pub/2021/framework/index.html#splitting-attention-head-terms-into-circuits) - 作为跳字三元组的解释 (https://transformer-circuits.pub/2021/framework/index.html#interpretation-as-skip-trigrams) - 总结OV/QK矩阵 (https://transformer-circuits.pub/2021/framework/index.html#summarizing-ovqk-matrices) - 我们是否“完全理解”单层模型？ (https://transformer-circuits.pub/2021/framework/index.html#do-we-fully-understand-one-layer-models) - 三种组合方式 (https://transformer-circuits.pub/2021/framework/index.html#three-kinds-of-composition) - Logits的路径扩展 (https://transformer-circuits.pub/2021/framework/index.html#path-expansion-of-logits) - 注意力分数QK电路的路径扩展 (https://transformer-circuits.pub/2021/framework/index.html#path-expansion-of-attention-scores-qk-circuit) - 分析一个两层模型 (https://transformer-circuits.pub/2021/framework/index.html#analyzing-a-two-layer-model) - 归纳头 (https://transformer-circuits.pub/2021/framework/index.html#induction-heads) - 项重要性分析 (https://transformer-circuits.pub/2021/framework/index.html#term-importance-analysis) - 虚拟注意力头 (https://transformer-circuits.pub/2021/framework/index.html#virtual-attention-heads) Transformer语言模型是一项正在兴起的技术，其现实世界应用正变得日益广泛，例如GPT-3、LaMDA、Codex、Meena、Gopher以及类似模型中的系统。然而，随着这些模型规模的扩大，其开放性和高容量也为出乎意料且有时有害的行为创造了更大的空间。即使在一个大型模型训练完成多年之后，其创建者和用户也常常会发现他们之前未曾察觉的模型能力——包括有问题的行为。 解决这些问题的一个途径是**机制可解释性**，试图逆向工程Transformer执行的详细计算过程，类似于程序员如何尝试将复杂的二进制文件逆向工程为人类可读的源代码。如果这成为可能，它可能为解释当前的安全问题、发现新问题，甚至可能预见尚未构建的强大未来模型的安全问题提供一个更系统的方法。 之前的一个项目，**Distill Circuits系列文章** (https://distill.pub/2020/circuits/)，曾尝试逆向工程视觉模型，但到目前为止，还没有针对Transformer或语言模型的类似项目。在本文中，我们试图迈出初步的、非常初步的逆向工程Transformer的步伐。考虑到现代语言模型令人难以置信的复杂性和规模，我们发现从最简单的模型入手逐步向上推进最为有效。我们的目标是发现简单的算法模式、模块或框架，这些可以随后应用于更大、更复杂的模型。 具体而言，在本文中，我们将研究**不超过两层且仅包含注意力块**的Transformer——这与GPT-3这样的大型现代Transformer形成对比，后者有96层，并且在注意力块与MLP块之间交替。我们发现，通过以一种新的但在数学上等价的方式来概念化Transformer的操作，我们能够理解这些小模型并获得对其内部运作的显著理解。特别值得注意的是，我们发现我们称之为“归纳头”的特定注意力头可以解释这些小模型中的**上下文学习**，并且这些头仅在具有至少两个注意力层的模型中才会发展。我们还通过一些示例展示了这些头在特定数据上的实际运作。 在这第一篇论文中，我们不打算将我们的洞察应用于更大的模型，但在即将发表的一篇论文中 (https://transformer-circuits.pub/2022/in-context-learning-and-induction-heads/index.html)，我们将展示我们理解Transformer的数学框架以及归纳头的概念，至少对于更大、更现实的模型来说仍然是部分相关的——尽管我们距离能够完全逆向工程此类模型还有很长的路要走。 --- ## 结果总结 #### 逆向工程结果 为了探索逆向工程Transformer的挑战，我们逆向工程了几个玩具性的、仅注意力的模型。在此过程中我们发现： - **零层Transformer建模二元组统计。** 二元组表可以直接从权重中获取。 - **单层仅注意力Transformer是二元组和“跳字三元组”（形如“A... B C”的序列）模型的集合。** 二元组和跳字三元组表可以直接从权重中获取，无需运行模型。这些跳字三元组的表达能力可能出乎意料地强。这包括实现一种非常简单的上下文学习。 - **两层仅注意力Transformer可以利用注意力头的组合实现更复杂的算法。** 这些组合算法也可以直接从权重中检测出来。值得注意的是，两层模型利用注意力头组合来创建“归纳头”，这是一种非常通用的上下文学习算法。我们将在即将发表的一篇论文中更详细地探讨归纳头。 - **单层和两层仅注意力Transformer使用截然不同的算法来执行上下文学习。** 两层注意力头使用定性上更复杂的推理时算法——特别是我们称为归纳头的一种特殊类型的注意力头——来进行上下文学习，形成了一个重要的过渡点，这对于更大的模型也具有相关性。 #### 概念要点 我们发现Transformer架构的许多细微之处要求我们以一种与InceptionV1 Circuits工作相当不同的方式来进行逆向工程。我们将在下文各节中详细阐述这些要点，但在此先做简要总结。我们还将在适当的章节中对这里引入的许多术语进行扩展。（需要澄清的是，我们并非意图声称这些观点中任何一点必然是新颖的；其中许多已在其他论文中隐式或显式地存在。） - **注意力头可以被理解为独立操作，每个头输出一个结果，该结果被加到残差流中。** 注意力头通常为了计算效率而被描述为另一种“拼接并相乘”的形式，但这在数学上是等价的。 - **仅注意力模型可以被写成一系列可解释的、从令牌到logits变化的端到端函数之和。** 这些函数对应于模型中的“路径”，并且如果固定注意力模式，它们是线性的。 - **Transformer具有巨量的线性结构。** 仅仅通过拆分求和以及将矩阵链相乘，就可以学到很多。 - **注意力头可以被理解为具有两个基本独立的计算：一个QK（“查询-键”）电路，用于计算注意力模式；一个OV（“输出-值”）电路，用于计算如果某个令牌被关注，它如何影响输出。** - **键、查询和值向量可以被视为低秩矩阵 $W_Q^TW_K$ 和 $W_OW_V$ 计算的中间结果。** 不提及它们来描述Transformer可能是有用的。 - **注意力头的组合极大地增加了Transformer的表达能力。** 注意力头有三种不同的组合方式，分别对应于键、查询和值。键和查询的组合与值的组合非常不同。 - **Transformer的所有组件（令牌嵌入、注意力头、MLP层和解嵌层）通过读写残差流的不同子空间彼此通信。** 与其分析残差流向量，不如将残差流分解为所有这些不同的通信通道（对应于模型中的路径），这可能更有帮助。 --- ## Transformer概述 在我们尝试逆向工程Transformer之前，简要回顾一下Transformer的高层次结构并描述我们如何看待它们是很有帮助的。在许多情况下，我们发现以等价但非标准的方式重新构建Transformer是有帮助的。 机制可解释性要求我们把模型分解成人类可解释的部分。一个重要的第一步是找到最便于推理模型的表示。在现代深度学习中，有充分的理由！人们非常重视计算效率，我们对模型的数学描述往往反映了如何编写高效代码来运行模型的决策。但当存在多种等价的计算表示方式时，最便于人类解释的表示和计算效率最高的表示很可能是不同的。 回顾Transformer也将使我们统一术语，因为术语有时会有所不同。我们还将在此过程中引入一些符号，但由于此符号用于多个章节，我们在符号附录中提供了所有符号的详细描述 (https://transformer-circuits.pub/2021/framework/index.html#notation)，作为读者的简明参考。 ### 模型简化 为了以最清晰的形式展示本文的思想，我们专注于进行了一些简化的“玩具Transformer”。在本文的大部分内容中，我们将做出一个非常实质性的改变：我们专注于“仅注意力”Transformer，它们没有MLP层。这是对Transformer架构的一个非常重大的简化。 我们部分动机在于，包含注意力头的电路带来了Distill Circuits工作未曾面临的新挑战，单独考虑它们使我们能够对这些问题进行特别优雅的处理。但我们也确实在理解MLP层方面进展甚微；在同时包含注意力和MLP层的普通Transformer中，有许多主要由注意力头介导的电路我们可以研究，其中一些似乎非常重要，但MLP部分则更难取得进展。这是我们工作的一个主要弱点，我们计划在未来重点解决。 尽管如此，我们将在后面的章节中讨论包含MLP层的Transformer。我们还做出了几个我们认为更表面化的更改，主要是为了清晰和简单。我们不考虑偏置，但有偏置的模型总是可以通过将偏置合并到权重中并创建一个始终为1的维度来模拟，从而无需偏置。此外，在仅注意力Transformer中，偏置大多相乘后功能上成为logits上的偏置。我们也忽略层归一化。显式考虑它会增加相当多的复杂性，并且在变量缩放范围内，层归一化可以合并到相邻的权重中。我们还预期，除了一些实现上的麻烦，层归一化可以用批归一化替代（批归一化可以完全合并到相邻参数中）。 ### 高层次架构 有几种变体的Transformer语言模型。我们专注于自回归、仅解码器的Transformer语言模型，例如GPT-3。（原始Transformer论文有一个特殊的编码器-解码器结构来支持翻译，但许多现代语言模型不包括这一点。） Transformer以令牌嵌入开始，随后是一系列“残差块”，最后是令牌解嵌层。每个残差块由一个注意力层，后跟一个MLP层组成。注意力层和MLP层都从残差流中“读取”它们的输入（通过执行线性投影），然后通过添加一个线性投影回残差流中来“写入”它们的结果。每个注意力层由多个头组成，它们并行操作。 ### 虚拟权重与作为通信通道的残差流 Transformer高层次架构的主要特征之一是，每一层都将其结果添加到我们称之为“残差流”的东西中。构建带有残差流的模型可以追溯到Schmidhuber小组的早期工作，例如高速公路网络和LSTM，这些在最近流行的残差网络架构中取得了重大成功。在Transformer中，残差流向量通常被称为“嵌入”。我们更喜欢残差流这个术语，不仅因为它强调了残差性质（我们认为这很重要），而且因为我们认为残差流通常为当前令牌以外的令牌保留了子空间，打破了嵌入术语所暗示的直觉。 残差流仅仅是所有前一层的输出以及原始嵌入的总和。我们将残差流视为一个通信通道，因为它本身不进行任何处理，所有层都通过它进行通信。残差流具有深度的线性结构。值得注意的是，完全线性的残差流在神经网络架构中是非常不寻常的：即使是ResNets，在广泛使用的最相似架构中，其残差流上或在访问残差流时也具有非线性激活函数！ 每一层在开始时执行任意线性变换以从残差流中“读入”信息（这忽略了每一层开始时的层归一化，但在恒定标量范围内，层归一化是恒定的仿射变换，可以合并到线性变换中。关于如何处理层归一化，请参见附录中的讨论），并在添加到残差流以“写入”其输出之前执行另一个任意线性变换。 残差流的这种线性、可加的结构具有许多重要含义。一个基本的推论是，残差流没有“特权基”；我们可以通过旋转与之交互的所有矩阵来旋转它，而不会改变模型行为。 #### 虚拟权重 残差流是线性的一个特别有用的推论是，人们可以思考直接连接任何一对层（甚至被许多其他层隔开的层）的隐式“虚拟权重”，通过将其相互作用通过残差流相乘。这些虚拟权重是一层输出权重与另一层输入权重的乘积（即 $W_{I}^2 W_{O}^1$）。注意对于注意力层，有三种不同的输入权重：$W_Q$、$W_K$ 和 $W_V$。为了简单和通用，我们在此认为层只有输入和输出权重。