最近看到了一篇文章,从数学角度对Tranformer进行解读的文章《A Mathematical Framework for Transformer Circuits》,文章以结构较为简单的Transformer为例,通过对其进行逆向工程来实现数学层面的解析,并逐步推广至复杂结构的Transformer,很有启发性,特此记录。
主要是关于三个部分,【一种新的观点方式】
- Zero layer transformers
- One layer attention-only transformers
- Two layer attention-only transformers
主要为翻译,但不全是,有笔者自己的理解与解释。
在本文中,我们试图对逆向工程变压器采取非常简化的步骤。鉴于现代语言模型令人难以置信的复杂性和规模,我们发现从尽可能简单的模型开始并从那里开始工作是最有成效的。我们的目标是发现简单的算法模式、图案或框架,这些模型、图案或框架随后可以应用于更大、更复杂的模型。具体来说,在本文中,我们将研究只有注意力块的两层或两层以下的变压器——这与GPT-3这样的大型现代变压器形成鲜明对比,GPT-3有96层,并将注意力块与MLP块交替使用。
我们发现,通过以一种新的但数学等效的方式概念化变压器的运行,我们能够理解这些小模型,并深入了解它们在内部的运行方式。特别值得注意的是,我们发现,我们称之为“感应头(induction heads)”的特定注意力头可以解释这些小模型中的上下文学习,而这些头只在至少有两个注意力层的模型中发展。我们还浏览了一些这些负责人对特定数据采取行动的例子。
在第一篇论文中,我们不试图将我们的见解应用于更大的模型,但在即将发表的一篇论文中,我们将表明,我们理解变压器的数学框架和感应头的概念,仍然至少部分地与更大、更现实的模型相关——尽管我们离能够完全反向工程此类模型还有很长的路要走。
先来看本文主要结论
为了探索逆向工程 transformers 面临哪些挑战,研究者对几个 attention-only 的 toy 模型进行了逆向功能。
我们发现,变压器架构的许多微妙细节要求我们以与InceptionV1 Circuits work
Thread: Circuits (作者之前的工作) 截然不同的方式进行逆向工程。我们将在下面的部分中解开其中的每一点,但现在简要总结一下。一旦进入具体的每部分,作者还将扩展在这里介绍的许多术语。
在我们尝试逆向工程transfromer之前,简要回顾一下transfromer的高级结构并描述下对它们的看法。
在许多情况下,我们发现以等效但非标准的方式重构变压器是有帮助的。机械可解释性要求我们将模型分解为人类可解释的部分。重要的第一步是找到表示形式,这样更容易对模型进行推理。在现代深度学习中,有充分的理由!非常强调计算效率,我们对模型的数学描述往往反映了如何编写高效代码来运行模型的决定。但是,当有许多等效的方法来表示相同的计算时,最人性化可解释的表示和最高效的表示可能会有所不同。
回顾transformer还将让我们在术语上保持一致,这些术语有时会有所不同。在这个过程中,我们还将介绍一些符号,但由于这个符号在许多章节中使用,我们在符号附录中提供了所有符号的详细描述,作为读者的简明参考。
为了以最干净的形式展示本文中的想法,作者专注于“toy transformer”,并进行了一些简化。
toy 类似一个大楼和一个大楼的模型,很形象的描述
在本文的大部分内容中,作者将做出非常实质性的改变:只专注于没有MLP层的“attention-only”Transformer。这么做的一部分原因是具有attention heads的数据流程提出了蒸馏数据流程(Distill circuits work)从未面临的问题,而且单独考虑具有attention heads的数据流程能够使作者更为优雅地处理这些问题。
本文作者主要研究自回归的decoder-only transformer 语言模型。
Transformer模型可分为三部分:token embedding,一组residual blocks,token unembedding。其中,每个residual block由一个attention层和MLP层按序组成。attention层和MLP层均通过线形投影从residual stream中“读取”其输入,然后通过添加一个线形投影将其结果“写入”到residual stream。每个attention层由多个并行执行的heads组成。Transformer的高层架构图如图1. 所示。
在Transformer的高层架构图中可知,(上图很清晰了)每一层都将其输出的结果添加到“Residual Stream”,Residual Stream仅是其前面所有层的输出和原始embedding的加和。因为Residual Stream不会对其自身进行处理,并且所有层都通过其进行通信,所以作者将其视为一个通信通道。
Residual Stream是一个深度线形结构。在Residual Stream起始阶段,每一层都通过执行一个任意的线性转换来从residual stream中“读取”信息,此外,在执行加法操作前,每一层都通过执行另一个任意的线性转换来将其输出“写入”到residual stream中。Residual Stream这种线性、累加结构具有很多重要的意义。其中一个基础的意义就是residual stream没有“privileged basis”,可以通过轮换所有与residual stream交互的矩阵来轮换residual stream,并且模型的行为不会被改变。
residual stream就是指从embed—>unembed的直接过程,而不经过任何block。原文有个很可爱的小图。
作者在这里特别强调了线形结构,在神经网络结构中完全线性的Residual Stream是极其不寻常的,即使是ResNet这种被广泛使用的、与Residual Stream最为相似的网络结构,也在其Residual Stream结构上,或访问Residual Stream的数据时使用了非线性激活函数。
Residual Stream结构中忽略了每一层开始时的层归一化,但是由于层归一化是由一个常数标量决定的,所以层归一化是一个常数仿射变换,并且可以合并至线性变换。
Residual Stream是线性的一个特别有用的结果,即可以认为隐式的虚拟权重可以在residual stream中通过乘法操作得到两个层的交互,以此直接连接任意两个层,即使两个层之间隔着多个层。虚拟权重是一个层的输出权重与另一个层的输入权重的乘积,通过这种方式描述了后一层(被指向层)读入由前一层(指向层)写入的信息的程度。
residual stream是一个高维度的向量空间。在小模型中可能是数百个维度;在大模型中可以达到数万。这意味着Residual Stream结构中的层可以通过将信息存储在不同的子空间来将不同的信息发送至不同的层。这一点在attention heads中极其重要,因为每个独立的head在相对较小的子空间上运行(通常是64或128维),在heads之间不发生交互的情况下可以很容易地将信息写入完全不相交的子空间。
一但信息写入到了某个子空间中,信息就会一直存储在该子空间,直到另一个层主动删除该信息。从这个角度来看,redidual stream的维度就像“内存”或“带宽”。原始的token embeddings和unembedding主要和redidual stream维度中相对小的一部分交互。这就留下了大多数维度来存储其他层的信息。
Residual Stream带宽通常应该较大。因为Residual Stream带宽中有更多的“计算维度”,例如,神经元和attention head结果维度,远多于用于移动信息的维度。仅是一个单独的MLP层,其神经元的数目是residual stream维度的四倍。因此,例如在一个50层transformer中的第25层中,residual stream中神经元数目是其之前维度的100倍,并且尝试与数目是其之后维度100倍的神经元进行通信,可能采用叠加方式进行通信。作者将这类tensor称为“bottleneck activations”,并且认为对其进行解释非常具有挑战性。这就是作者尝试以虚拟权重的方式将residual stream中发生的不同通信streams分开而不是直接研究这些通信streams的主要原因。
也许是因为对residual stream带宽的高需求,作者观察到一些迹象表明,某些MLP神经元和attention heads可能会担任一种“记忆管理”(“memory management”)的角色,通过读入信息和写出负面信息来清除其它层设置的residual stream维度。
一些MLP神经元的输入权重和输出权重之间具有负的余弦相似度,这可能表明这些神经元将从residual stream中删除信息。类似地,一些attention heads在其 W O W V W_OW_V WOWV矩阵中具有较大的负特征值,并且主要关注当前token,这些attention heads可能作为删除信息的一种机制。值得注意的是,尽管这些可能是“记忆管理”(“memory management”)删除信息的通用机制,但是它们也可能是有条件地删除信息的机制,仅在某些条件下运行。
如上文所述,作者认为Transformer的attention层是一些完全独立的attention heads h ∈ H h\in H h∈H
组成(这点其实不难理解,跟一般的MSA认知一致),这些heads 完全并行的运行,并且每个head都将其输出添加会residual stream。但这并不是Transformer层的典型呈现方式,而且它们的等价性并不明显。
先来看普遍认知的结构,nlp和cv其实都一致
在Transforer的原始文献中[2],attention层的输出表示为堆叠的向量 r h 1 , r h 2 , . . . r^{h_1},r^{h_2},... rh1,rh2,... 和输出矩阵 W O H W_O^H WOH的乘积。如果为每个head将 W O H W_O^H WOH分割为尺寸相等的多个块,那么可以得到如下公式 W O H [ r h 1 r h 2 … ] = [ W O h 1 , W O h 2 , … ] ⋅ [ r h 1 r h 2 … ] = ∑ i W O h i r h i (1) W_{O}^{H}\left[\begin{array}{c} r^{h_{1}} \\ r^{h_{2}} \\ \ldots \end{array}\right]=\left[W_{O}^{h_{1}}, W_{O}^{h_{2}}, \ldots\right] \cdot\left[\begin{array}{c} r^{h_{1}} \\ r^{h_{2}} \\ \ldots \end{array}\right]=\sum_{i} W_{O}^{h_{i}} r^{h_{i}} \tag{1} WOH rh1rh2… =[WOh1,WOh2,…]⋅ rh1rh2… =i∑WOhirhi(1)公式(1)揭示了将 W O H W_O^H WOH拆分后等价于独立地执行各个heads,将每个head乘以各自的输出矩阵,并将结果添加至residual stream。在实现上述公式时,拼接定义通常是首选,因为能够生成一个更大、计算效率更高的矩阵乘法。但是为了从理论上理解Transformers,作者在此更倾向于将其视为独立的加法。[we prefer to think of them as independently additive.]
如果各个attention heads独立地运行,那么它们将会产生什么样的效果?attention heads的基本行为是移动信息。它们从一个token的residual stream中读取信息,然后将其写入到另一个token的residual stream。从这一章节中得到的主要观察结果是:从其中移动信息的那些tokens,与将被移动的“读取”到的信息以及如何“写入”到目的地是完全分开的。
为了验证这一点,需要以非标准的方式实现attention。给定一个注意力模式,计算attention head的输出通常分为如下三步:
如前文所述,输出矩阵的乘法通常实现为应用于所有heads结果的拼接的一个矩阵乘法,与上述步骤是等效的。
上述的每一步都可以通过矩阵乘法实现,但是作者并未将上述步骤合并为一个步骤,原因如下。
假设 x x x 是一个 2d 矩阵,由表示每个token的向量组成,并且将在不同的边相乘。 W V W_V WV和
W O W_O WO 乘以“表示每个token的向量”(“vector per token”)的边【该边的每个元素表示一个token】,A乘以“位置”(“position”)的边。Tensors能够提供更为自然的一种语言来描述矩阵之间的这种映射关系。将计算attention head输出分为三步的一个动机是可能有助于作者表达矩阵到矩阵之间的线性映射: [ n c o n t e x t , d m o d e l ] → [ n c o n t e x t , d m o d e l ] [n_{context}, d_{model}]\to [n_{context}, d_{model}] [ncontext,dmodel]→[ncontext,dmodel]。数学上称这种线性映射为“(2,2)-tensors”【其实也就和transformer block中的输入输出】。也因此,tensors自然是描述这种变换的语言。
后面会用到一些张量积的表示 Tensor Product / Kronecker Product Notation ⊗ \otimes ⊗
记住就好,定义而已
乘积符合混合乘积性质: ( A ⊗ B ) ( C ⊗ D ) = A C ⊗ B D (A\otimes B)(C\otimes D)=AC\otimes BD (A⊗B)(C⊗D)=AC⊗BD
Left-right multiplying:将 x 乘以一个张量积 A ⊗ W A⊗W A⊗W 等价于同时左乘、右乘: ( A ⊗ W ) x = A x W T (A⊗W)x=AxW^T (A⊗W)x=AxWT。当将它们相加时,等同于将乘积的结果进行相加: ( A 1 ⊗ W 1 + A 2 ⊗ W 2 ) x = A 1 x W 1 T + A 2 x W 2 T (A_1⊗W_1+A_2⊗W_2)x=A_1xW_1^T+A_2xW_2^T (A1⊗W1+A2⊗W2)x=A1xW1T+A2xW2T
What about the attention pattern?
注意力模式:通常计算keys k i = W K x i k_i=W_Kx_i ki=WKxi ,计算queries q i = W Q x i q_i=W_Qx_i qi=WQxi ,通过每个key向量和query向量的点积计算注意力模式 A = s o f t m a x ( q T k ) A=softmax(q^Tk) A=softmax(qTk)。也可以不引入key和query直接计算 A = s o f t m a x ( x T W Q T W K x ) A=softmax(x^TW_Q^TW_Kx) A=softmax(xTWQTWKx).
尽管这个公式在数学上是等价的,但是实际中如果采用这种方式实现注意力,例如,乘以 W O W V W_OW_V WOWV或 W Q T W K W_Q^TW_K WQTWK ,其计算效率极低。
作者以上述方式重写了attention head,这么做可以揭示出许多以前难以观察到的结构:
当移动信息时, W O V = W O W V W_{OV}=W_OW_V WOV=WOWV 管理attention heads从哪个residual stream子空间读写。通过奇异值分解 U S V = W O V USV=W_{OV} USV=WOV 可以很好地理解这一观点。由于 d h e a d < d m o d e l d_{head} < d_{model} dhead<dmodel, W O V W_{OV} WOV 是低秩的,并且S中只有对角线的一个子集是非零的。右边的奇异向量 V描述了被关注的residual stream中哪个子空间被“读入”(以某种方式存储为value向量),而左奇异向量 U描述了被“读入”的子空间被“写入”到目标residual stream的哪个子空间中。
至此前言结束,下面就是分析每种模式