A Mathematical Framework for Transformer Circuits

type

status

date

slug

summary

1 总结

1.1 逆向工程结论

Zero layer transformers model bigram statistics. The bigram table can be accessed directly from the weights.

One layer attention-only transformers are an ensemble of bigram and “skip-trigram” (sequences of the form "A… B C") models. The bigram and skip-trigram tables can be accessed directly from the weights, without running the model. These skip-trigrams can be surprisingly expressive. This includes implementing a kind of very simple in-context learning.

Two layer attention-only transformers can implement much more complex algorithms using compositions of attention heads. These compositional algorithms can also be detected directly from the weights. Notably, two layer models use attention head composition to create “induction heads”, a very general in-context learning algorithm.

One layer and two layer attention-only transformers use very different algorithms to perform in-context learning. Two layer attention heads use qualitatively more sophisticated inference-time algorithms — in particular, a special type of attention head we call an induction head — to perform in-context-learning, forming an important transition point that will be relevant for larger models.

1.2 概念要点

Attention heads can be understood as independent operations, each outputting a result which is added into the residual stream. Attention heads are often described in an alternate “concatenate and multiply” formulation for computational efficiency, but this is mathematically equivalent.

Attention-only models can be written as a sum of interpretable end-to-end functions mapping tokens to changes in logits. These functions correspond to “paths” through the model, and are linear if one freezes the attention patterns.

Transformers have an enormous amount of linear structure. One can learn a lot simply by breaking apart sums and multiplying together chains of matrices.

Attention heads can be understood as having two largely independent computations: a QK (“query-key”) circuit which computes the attention pattern, and an OV (“output-value”) circuit which computes how each token affects the output if attended to.

Key, query, and value vectors can be thought of as intermediate results in the computation of the low-rank matrices and . It can be useful to describe transformers without reference to them.

Composition of attention heads greatly increases the expressivity of transformers. There are three different ways attention heads can compose, corresponding to keys, queries, and values. Key and query composition are very different from value composition.

All components of a transformer (the token embedding, attention heads, MLP layers, and unembedding) communicate with each other by reading and writing to different subspaces of the residual stream. Rather than analyze the residual stream vectors, it can be helpful to decompose the residual stream into all these different communication channels, corresponding to paths through the model.

2 Transformer Overview

许多情况下我们发现等价重构 transformer 是有帮助的，虽然这种方法并不是标准的。这需要我们首先找到一种更易解释模型的表示。

但是当一种计算有多个等价表示方法的时候，最容易被人类理解的与最高效计算的表示往往是不同的。

2.1 模型简化

只关注 “toy transformers”

只关注 attention 且没有 MLP 这种简化部分动机是：带有注意力头的回路带来了 Distill 回路中没有的新挑战，并且单独考虑它们让我们能够给出更优雅的解法。对 MLP 的理解尚未取得很大成功，也是本文的主要弱点，下文有一些讨论

未考虑 biases，但是带有偏差的模型总是可以在没有偏差的情况下进行模拟，方法是将它们折叠到权重中并创建一个始终为一的维度。此外，attention-only 的 transformers 中的 biases 大多会和 logits 的 biases 相乘。

忽略 layer normalization，它增加了表示的复杂性，并且在可变缩放范围内，层范数可以合并到相邻权重中

我们还期望层范数可以代替批量归一化（它可以完全折叠到相邻参数中）

2.2 高层架构

只关注自回归的 decoder-only 的模型

一个 transformer 始于 token embedding，紧接着一系列的 “residual block”，最后接一个 token unembedding

每一个 residual block 有一个 attention layer 和一个 MLP layer. Attention 和 MLP 各自从残差流中“读”他们的输入（通过一个线性映射），然后将他们的结果“写”到残差流中（通过添加一个线性映射回来）

每个 attention layer 有多个并行操作的头

2.3 虚拟权重和作为通信通道的残差流

高层架构中的一个主要的特点是每一层将其结果加到所谓的“残差流”中

残差流只是所有先前层的输出与原始嵌入的总和

我们将其视为一个通信的通道，因为其本身没有任何操作且所有层都通过它进行交流

残差流具有深度线性结构

每一层都执行任意线性变换以在开始时从残差流中“读入”信息，然后在“写回”残差流之前执行另一个随机线性变换

残差流的这种线性、加性结构具有许多重要影响

一个最基本的结果是残差流没有 “Privileged Basis”；我们可以旋转所有与其交互的矩阵来旋转它，同时不改变模型行为

💡

Privileged Basis vs Basis Free: A privileged basis occurs when some aspect of a model’s architecture encourages neural network features to align with basis dimensions, for example because of a sparse activation function such as ReLU. In a transformer, the only vectors with privileged bases are tokens, attention patterns and MLP activations. See the appendix table of variables for a full list of which activations have or don’t have privileged activations.

Some types of interpretability only make sense for activations with a privileged basis. For example, it doesn't really make sense to look at the "neurons" (ie. basis dimensions) of activations like the residual stream, keys, queries or values, which don't have a privileged basis. This isn't to say that there aren't ways to study them; a lot of interesting work is done on word embeddings without assuming a privileged basis. But having a privileged basis does open up helpful approaches.

2.3.1 虚拟权重

线性残差流的一个特别有用的结果是：我们可以认为隐式“虚拟权重”直接连接任何一对层（即使是那些被许多其他层分隔的层）

这些虚拟权重是一层的输出权重与另一层的输入权重的乘积，并且描述后一层读取前一层写入的信息的程度

2.3.2 子空间和残差流带宽

残差流是一个高维向量空间，在小模型中可能是几百维，在大模型中可能达到几万维

这意味着层可以通过将不同的信息存储在不同的子空间中来实现将不同的信息发送到不同的层。这对于注意力头来说尤其重要，因为每个单独的头都在相对较小的子空间（通常是 64 或 128 维）上运行，并且可以很容易地写入完全不相交的子空间并且不交互

添加后，信息将保留在子空间中，除非另一层主动删除它。从这个角度来看，残差流的维度就变成了“内存”或“带宽”之类的东西。原始令牌嵌入以及未嵌入大多与相对较小的维度交互。这使得大多数维度“自由”供其他层存储信息。

一般来说，"计算维度"（如神经元和注意力头结果维度）远远多于残差流的信息移动维度。仅一个 MLP 层的神经元数量通常就比残差流的维数多四倍。所以，例如在 50 层 transformer 的第 25 层，残留流的神经元数量是它之前维数的 100 倍，它试图与 100 倍于它之后维数的神经元进行交流，以某种方式进行叠加交流。我们把这样的张量称为 "Bottleneck Activations"，并希望它们在解释上具有非同寻常的挑战性。这也是我们为什么要尝试用虚拟权重来拆分通过残差流发生的不同通信流，而不是直接研究它的主要原因。

💡

Bottleneck Activations: We say that an activation is a bottleneck activation if it is a lower-dimensional intermediate between two higher dimensional activations. For example, the residual stream is a bottleneck activation because it is the only way to pass information between MLP activations, which are typically four times larger than it. (Additionally, in addition to it being a bottleneck for the communication between adjacent MLP layers, it’s also the only pathway by which arbitrary early MLP layers can communicate with arbitrary late MLP layers, so the stream may simultaneously be conveying different pieces of information between many different pairs of MLP layers, a much more extreme bottleneck than just 4x!) Similarly, a value vector is a bottleneck activation because it’s much lower dimensional than the residual stream, and it’s the only way to move information from the residual stream for one token position in the context to another (without the model dedicating an additional attention head to the same pattern, which it sometimes will). See above table for a list of which activations are or are not bottleneck activations.

也许正因为对残差流带宽的这种高要求，我们看到了一些暗示，即某些 MLP 神经元和注意头可能会扮演一种 "内存管理 "角色，通过读入信息并写出负版本来清除其他层设置的残差流维度

2.4 注意力头是独立的和可加的

将注意力层视为一些完全独立的并行操作的注意力头，每个头的结果加回其残差流中

Transformer 原论文中将注意力层的输入描述为堆叠各个头的结果向量 , , ，然后乘输出矩阵

这里我们将分成与各个头对应的大小相同的子块，结果为：

2.5 注意力头是一种信息移动

探究独立的注意力头在做什么

最基本的动作是移动信息：从一个 token 的残差流中读取信息，写入到另一个 token 的残差流中

本节得到的主要观察结果是：从哪些 token 移动信息与“读取”哪些要移动的信息以及如何“写入”目的地是完全分开的

要看到这一点，以非标准方式描述注意力会很有帮助

给定一个注意力模式，计算注意力头的输出通常分为三个步骤：

计算残差流中每一个 token 的 value 向量：

根据注意力模式线性组合 value 向量来计算“结果向量”：

计算每个 token 的头的输出向量：

每一步都可以写成一个矩阵乘法，为什么将他们拆成单个步骤？如果把视为一个 2d 矩阵（由每个 token 的向量组成），我们将它在不同的侧面相乘：和在 “vector per token” 这一侧乘，在 “position” 这一侧乘。张量可以为我们提供一种更自然的语言来描述矩阵之间的这种映射。可能有用的一个动机是注意到我们想要表达从矩阵到矩阵的线性映射：。数学家将这种线性映射称为“(2,2)-张量”（它们将两个输入维度映射到两个输出维度）。因此张量是表达这种变换的自然语言。

使用张量积，我们可以将应用注意力的过程描述为：

应用混合乘积性质和折叠等式即可得出：

2.5.1 关于注意力头的观察

以这种形式重写注意力头的一个主要好处是它显示了许多以前可能难以观察的结构：

注意力头将信息从一个 token 的残差流移动到另一个 token 的残差流中

由此推论，残差流向量空间（通常被解释为 "上下文词嵌入"）一般会有线性子空间，对应于从其他 token 复制而来的信息，而不是与当前 token 直接相关的信息。

注意力头实际上应用了两个线性运算和，他们独立操作于不同维度

决定了信息从哪个 token 转移到哪个 token
决定了从源 token 中读取哪些信息以及如何写入目标 token

是唯一的非线性部分（softmax），这意味着如果我们固定注意力模式，注意力头会执行线性操作；这也意味着如果不固定，注意力头某种意义上来说是“半线性”的，因为每个 token 的线性运算是恒定的

和总是一起运算，他们从不独立，和也同样如此

尽管他们被参数化为单独的矩阵，和始终可以被视为单独的低秩矩阵
从某种意义上说，键、查询和值向量都是肤浅的，它们是计算这些低阶矩阵的中间副产品。人们可以轻松地重新参数化低秩矩阵的两个因子以创建不同的向量，但功能仍然相同。
和总是一起运算，定义：，

注意力头乘积的行为与注意力头本身非常相似。根据分配律，有。乘积结果可以被视为在功能上等同于一个具有两个头组合而成的注意力模式的注意力头和一个输出值矩阵。我们称之为“虚拟注意力头”。

3 Zero-Layer Transformers

由于模型无法从其他 token 中移动信息，因此我们只是根据当前 token 来预测下一个 token

这意味着的最优行为是对二元对数似然进行估计

在每个 transformer 的扩展公式中都会出现，它们和“直接路径”对应，即一个 token embedding 在残差流中直接流向 unembedding 而不经过任何层。他唯一能影响的只有二元对数似然。由于模型的其他方面将预测二元对数似然的部分，因此它不会准确地代表较大模型中的二元统计，但它确实代表了一种“残差”。特别是似乎经常有助于表示更通用的语法规则无法描述的二元统计数据，例如“Barack”后面经常跟着“Obama”这一事实。

虽然是一种 “unembedding”，但他并不是的逆

4 One-Layer Attention-Only Transformers

我们认为 one-layer attention-only transformers 可以被理解为一个“二元模型”和几个 “skip-trigram” 模型的组合（“skip-trigram” 模型会影响序列 “A…BC” 的概率，“skip-trigram” 用来描述 “A…BC” 这种形式的序列）。直观地说，这是因为每个注意头可以有选择地从当前标记（"B"）注意到先前的标记（"A"），并复制信息以调整下一个可能标记（"C"）的概率。

本节的目标是严格展示这种对应关系，并演示如何将 transformer 的原始权重转换为可解释的 “skip-trigram” 概率调整表。

4.1 路径扩展技巧

使用张量表示方法和前文关于注意力头表示方法，上述过程可以写为三项乘积：

技巧是将乘积（每一项对应一层）转换为求和（每一项对应一条端到端路径）

在 zero-layer 中也存在，这一项不在两个位置之间移动信息，它唯一能影响的只有二元统计，它会填补其他项无法处理的 missing gap

4.2 将注意力头这一项拆分为 QK 和 OV 回路

对于每个注意力头，有

其中

这些项其中有两个的矩阵：

我们称之为 QK 回路，它为每个 query 和 key token 提供注意力分数，即每个条目描述一个给定的 query token 多大程度上关注给定的 key token

我们称之为 OV 回路，用来描述一个给定的 token 如何影响输出 logits

为了便于理解，将这些乘积视为通过模型的多个路径，从 token 起始，到 token 结束

QK 回路是通过跟踪注意力头中 query 和 key 向量的计算形成的；OV 回路跟踪从 value 向量计算直至变为 logits 的路径

注意力模式是一个从所有可能的源令牌到目标令牌的函数，但是一旦目标标记决定了多大程度上关注一个源标记，对输出的影响就是一个关于源标记的函数。也就是说如果多个目标标记同等程度地关注同一个源标记，那这个源标记对被预测的标记的 logits 的影响就是一样的

4.2.1 QK 和 OV 独立性

QK 和 OV 回路我们认为是分开，独立的，但为什么可以这样考虑

假设我们运行一个模型两次，第一次收集每个头的注意力模式，这只依赖 QK 回路；第二次将注意力模式冻结，会得到一个 token 和 logits 之间的一个线性函数

4.3 解释 Skip-trigram

Skip-trigram 描述形如 [source]... [destination][out] 的序列，QK 回路决定 destination token 关注并从哪些 source token 拷贝信息，OV 回路描述对 out 的影响

4.3.1 复制/原始情境学习

单层模型中的大多数注意头都将其大部分能力用于复制

OV 电路的设定是，如果注意头注意到标记，就会增加该标记的概率，其次是类似标记的概率

QK 电路只关注有可能成为下一个标记的标记

因此，标记会被复制，但只会被复制到大语法统计数据使其看起来可信的地方

4.3.2 其他有趣的例子

略

4.3.4 PRIMARILY POSITIONAL ATTENTION HEADS

略

4.3.4 SKIP-TRIGRAM "BUGS"

例如：一个头增加了 keep… in mind and keep… at bay 的概率，那么 keep… in bay and keep… at mind 的概率也会增加

略

4.4 QK/OV 矩阵小结

QK、OV 矩阵轻易就有数十亿个元素，需要找到更好的方法帮助我们理解它们

QK 和 OV 矩阵是低秩的，只有阶（64 或 128）

查看各个条目通常会发现结构简单得多的迹象，比如一个头中人名对应最常见的查询是 “by”，而地名对应最常见的查询是 “from”。这暗示了矩阵中的簇结构

复制行为在 OV 矩阵中很普遍，并且可以说是最有趣的行为之一。看来应该可以将其形式化

我们还没有找到明确的正确答案，但我们乐观地认为，正确的矩阵分解或降维可能会提供大量信息。

4.4.1 检测复制行为

复制行为可以自动检测数来，复制的本质是将同样的向量映射到自身（拥有某个 token 会增加其自身的概率）。但这个概念不好确切表示，因为无法准确界定“复制矩阵”的边界。比如小模型中没有但是大模型中可以观察到一些注意力头它从附近的单词中“复制”一些性别、复数和时态的混合，帮助模型使用正确的代词和共轭动词。这些注意力头的矩阵并不完全复制单个标记，但它们似乎在某种非常有意义的意义上复制。所以复制行为要比看上去的复杂的多。

一个自然的方法是使用特征向量和特征值，我们知道对一个矩阵如果则是矩阵的特征向量是矩阵的特征值。

带入到 OV 回路中，如果是一个正实数，那么我们认为存在一个关于 token 的线性组合，该线性组合增加了同样是这些 token 它们的 logits 的线性组合。简单来说可以认为这一组能够相互增加自身概率的 token。

通常特征向量既有正元素也有负元素，它们可能对应两组意义不同的 token（例如代表雄性或雌性、单数或复数），它们增加自己所在集合中其他 token 的概率，降低其他集合中 token 的概率

对于随机矩阵，我们期望具有相同数量的正特征值和负特征值，并且许多特征值都是复数。但复制需要正特征值，事实上我们观察到许多注意力头具有正特征值，显然反映了复制结构：

将其进一步折叠，并获得有多少注意力头正在复制的直方图：

看来 12 个头中有 10 个在明显地复制！（这与扩展权重的定性检查一致。）

但是，虽然复制矩阵必须具有正特征值，但尚不清楚所有具有正特征值的矩阵都是我们必须考虑复制的矩阵

矩阵的特征向量不一定是正交的，这允许病态的例子：可能存在具有所有正特征值的矩阵，这些矩阵实际上将某些 token 映射到减少同一 token 的 logits。正特征值仍然意味着矩阵在某种意义上是“平均复制”，并且它们仍然是复制的有力证据，因为默认情况下它们似乎不太可能，并且从经验上看似乎与复制一致。但它们不应该被视为矩阵在所有合理意义上的复制的决定性证据。

尝试以其他方式形式化“复制矩阵”，一种可能性是查看矩阵的对角线，它描述了每个标记如何影响其自身的概率。不出所料，对角线上的条目非常正向。我们还可以问，一个随机标记比其他任何标记（或者说是 k 个增加最多的标记之一，以考虑大写字母不同或空格相同的标记）增加自身概率的频率。所有这些似乎都指向了这些注意头是复制矩阵的方向，但还不清楚其中是否有任何一个是 "此矩阵的主要行为是复制 "的完全稳健的形式化。

值得注意的是，所有这些潜在的复制概念都是通过特征值之和等于迹等于对角线之和这一事实联系在一起的。

4.5 我们“完全理解”单层模型吗？

人们常常怀疑，真正逆向工程神经网络是否可能或值得一试。既然如此，我们很容易就会指着单层只受关注的 transformer 说："看，如果我们用最简化的、玩具版的 transformer 来表示，至少这个最小版本是可以被完全理解的。

但这种说法实际上取决于人们对 "完全理解 "的理解。在我们看来，我们现在对这个简化模型的理解，就像人们看到一个巨大线性回归的权重并理解它，或者看到一个大型数据库并理解查询它意味着什么一样。这是一种理解。算法不再神秘。神经网络参数的上下文问题已经被剥离。但是，如果不对其进行进一步的总结，就会有太多的东西让人无法将模型记在脑子里。

鉴于普通的单层神经网络只是广义线性模型，并且可以如此解释，或许单注意层也大多是单层神经网络也就不足为奇了。

5 Two-Layer Attention-Only Transformers

注意力头的组合是单层与双层的关键区别，如果没有组合，那么双层模型只是用更多的头实现 skip-trigram，有了组合才能实现更强大的上下文学习机制

5.1 三种组合方式

我们将残差流视为一种通信通道，决定每个注意力头从哪些残差流子空间读取信息，决定写到哪些子空间中。由于注意力头向量远小于残差流的尺寸（的值范围在 1/10 到 1/100 之间），注意力头在较小的子空间上运行可以轻松避免显着的交互。

注意力头的组合有三种选择：

Q 组合：从受一个过去 token 影响的子空间读入

K 组合：从受一个过去 token 影响的子空间读入

V 组合：从受一个过去 token 影响的子空间读入

Q 组合和 K 组合不同于 V 组合，Q、K 组合影响着注意力模式，V 组合影响着给定一个关注位置注意力头移动什么信息。V 组合形成的头更像是一个独立的单元，可以被认为是创建了一个额外的“虚拟注意力头”：将信息移动与信息移动组合得出信息移动。

5.2 Logits 的路径扩展

Logits 如何被计算出来？

其中两个项（直接路径项和单个头部项）与单层模型相同。最后一个 "虚拟注意头 "项与 V-Composition 相对应。

5.3 Path Expansion of Attention Scores QK Circuit

只看 logit 展开可能会忽略了 two-layer attention-only transformer 最根本的不同特性： Q 组合和 K 组合使它们的第二层注意力模式更具表现力。

为了看到这一点，我们需要观察 QK 回路计算注意力模式。

一个头的注意力模式是，其中是 QK 回路，他将 token 映射到注意力分数。

对于第一层的头，QK 回路和 one-layer 模型矩阵一样：

对于第二层的 QK 回路，Q-composition 和 K-composition 都发挥作用，前一层注意力头可能会影响键和查询的构造。

第一层的规约是在 token embedding 上操作的：

但是第二层是在第一层上操作的：

我们试图表达以下形式的多线性函数：

用 (4,2)-张量表示，每一项形如

其中，描述 token 之间 query 侧的信息移动，描述 token 之间 key 侧的信息移动，描述它们如何乘在一起形成注意力分数

5.4 分析双层模型

下图显示了第一层和第二层之间的 QKV 组合

第一层的头的注意力模式很简单：最关注之前的 token，其次是当前 token 和前两个

第二层出现了归约头

图片展示了第一层和第二层之间的 QKV 组合，下面是第一层，上面是第二层。QKV 组合描述一个第二层的头的 Q、K、V 向量从一个给定的第一层的头中读了多少信息。这通过相关矩阵乘积的 F 范数除以单独矩阵的范数来衡量。

对于 Q 组合：

对于 K 组合：

对于 V 组合：

5.5 归纳头

在双层仅注意力模型中，组合看上去似乎只有一个目的：创建归约头

我们之前看到，单层模型将其大量能力用于复制头，这是实现情境学习的一种粗略方法

诱导头则是实现上下文学习的一种更为强大的机制

5.5.1 归纳头的作用

归纳头会在上下文中搜索当前 token 的之前的示例。如果没有找到，他们就会关注第一个标记，然后什么也不做。但如果找到了，他们就会查看示例的下一个标记并复制它。这样，他们就能准确或近似准确地重复之前的 token 序列

归纳头与 K 组合有关

5.5.2 归纳头如何工作

归纳头的核心技巧是 key 是由向后偏移一个的 token 计算出来的，query 在搜索相似 key 向量时，由于 key 是偏移的，所以搜索到了下一个 token

构造归纳头最简单的方法就是使用 K 组合，在一个 token 的头上让 key 向量向前偏移一个 token