从零构建 LLM 系列:Attention 机制的设计局限
深度教程剖析 Attention Head 的工作原理和性能瓶颈,帮助开发者理解 Transformer 核心设计。
深度教程剖析 Attention Head 的工作原理和性能瓶颈,帮助开发者理解 Transformer 核心设计。
完成了 Sebastian Raschka 著作《从零开始构建大语言模型》第 3 章的学习后——在上一篇文章中我已经详细讨论了多头注意力机制——在进入第 4 章之前,我觉得有必要暂停一下,进行一番总结和思考。
我想讨论两个问题:自注意力的"为什么"以及对上下文长度的一些思考。这篇文章关于"为什么"——也就是说,为什么书中所述的特定矩阵乘法集合能做到我们想要它们做的事情?
一如既往,我主要是为了理清自己的思路而做这些工作——当然还有可能对其他人有帮助。我当然会让多个 LLM 审阅,确保我没有发布完全的胡言乱语,但读者需谨慎!
让我们开始吧。正如我在本系列第 8 部分所写的:
我还认为值得指出的是,[书中的内容]非常是"机制性"的解释——它说明了我们如何进行这些计算,但没有说明为什么。我认为"为什么"实际上超出了本书的范围,但这是让我着迷的事情,我很快会在博客中讨论它。
那个"很快"现在就是了 :-)
我认为我在理解为什么这些方程有效时的核心问题,是我高估了单个注意力头能做什么。在第 6 部分,我谈到了句子"the fat cat sat on the mat"(肥猫坐在垫子上):
所以虽然"cat"的输入嵌入只是意味着"位置 3 的 cat",但这个句子中"cat"的上下文向量也有某种关于它是一只在坐着的猫的含义,也许不如它是一只特定的猫那么强("the"而不是"a"),以及关于它坐在垫子上的暗示。
我之前没有理解的是,这个说法是对的,但仅限于整个注意力机制的输出——而不是单个注意力头的输出。
每个单独的注意力头都很笨,它在做的事情比那简单得多!
让整个机制变得聪明的两个因素是多头注意力和分层。这本书已经详细介绍了多头注意力,所以让我们深入讨论第二部分。
早在开始的第 1 部分,我就写过:
Raschka 提到的另一件让我有点困惑的事情是,原始的 transformer 架构有六个编码器和六个解码器块,而 GPT-3 有 96 个 transformer 层。这与我对整个过程如何工作的模型不太相符。编码器和解码器似乎都是独立的东西,接受输入(tokens/embeddings)并产生输出(embeddings/tokens)。你会用多层来做什么?
现在我们已经介绍了注意力的工作原理,这变得更清楚了一些。多头注意力块获取一组输入嵌入(输入序列中每个 token 一个),并生成相同数量的上下文向量。没有什么阻止我们将这些上下文向量视为另一个注意力块的输入嵌入,并再做一次同样的事情。
(这也解释了为什么 Raschka 提到上下文向量中的维度数通常与输入嵌入中的维度数相匹配;这样使得对每一层使用相同"形状"的多头注意力计算更容易。)
在我的思维中,这类似于图像处理网络——比如说 CNN——的工作方式。在那些网络中,第一层可能检测边缘,第二层可能检测特定方向的线,下一层检测特定的形状,然后在某个地方,第 n 层可能识别狗的脸。
所以我上面描述的 token"cat"的表示不会是一个注意力头的输出的一部分,甚至注意力机制的第一层也可能没有那么丰富的东西。但它可能在多头注意力的第三层或第四层的输出中,或其他类似的地方。
到了 GPT-3 中的第 96 层,上下文向量中所代表的内容将被超级丰富化,在不同的 token 中分散了大量的信息。认识到这一点对我来说也是一个小的顿悟时刻。
如果你回想一下第 5 部分,没有注意力机制的编码器/解码器 RNN 的一个大问题是固定长度的瓶颈。你会把你的输入序列输入到一个编码器 RNN,它会试图将其含义表示在其隐藏状态中——一个特定固定长度的向量——为了准备将其传递给解码器。对于短输入很容易,但随着长度增加而变得越来越困难,最终随着它变得更长而变得不可能,因为你试图将越来越多的信息压缩到同样的"空间"中。
但有了注意力,从最后一个注意力层输出的这种超级丰富和组合的输入序列表示的长度与输入中的 token 数量成正比!当然,你仍然受到可用内存的限制(以及其他因素——参见下一篇文章),但你拥有的 token 越多,这个上下文向量的"隐藏状态"就越大。
所以,使用多头注意力加上分层允许我们即使在每个单独的注意力头很笨的时候也能建立复杂的表示。但是,回到这篇文章的核心,为什么这些笨注意力头使用它们所做的特定计算?
让我们用一个例子。
首先提个提醒/警告:注意力头正在学习它们自己的表示和模式来匹配作为深度学习梯度下降的一部分——所以无论它们学到什么可能都很奇怪和陌生,与我们理解的语法毫无关系。但对于这个例子,让我们假设情况并非如此,我们有一个注意力头已经学会了如何将冠词(如"a"、"an"和"the")与它们相关的名词相匹配。
那会怎样呢?让我们取"the fat cat sat on the mat",忽略除了两个"the"和名词"cat"与"mat"之外的所有东西。我们说我们的注意力头想为"cat"产生一个上下文向量,它将其与第一个"the"结合(意味着它会包含我们在谈论一个特定的猫而不仅仅是"一个"猫的概念),类似地它想将第二个"the"融合到"mat"中。
现在,请记住我们的输入序列是一系列输入嵌入,这些是 token 嵌入(这些是空间中的向量,指向 token 的某种抽象的"意义")和位置嵌入(代表它们在序列中的位置)的组合。
以"mat"为例,我们将其输入嵌入投影,意味着"token 'mat' 在位置 7"¹ 进入查询空间。对我的突破是查询空间是另一个嵌入空间,就像原始的输入嵌入空间一样,但对值有不同的表示。
比如说在这个新的嵌入空间中,表示要简单得多——它们的细节没有原始空间那么多。它只表示"这是一个冠词"或"这不是一个冠词",以及一些关于位置的信息——也就是说,位置 1 的冠词的嵌入接近位置 2 的嵌入,但与位置 69,536 的嵌入不太接近。而不是冠词的其他东西会在更远的地方。
在这个例子中,也许我们的注意力头学到的投影会将"位置 7 的 'mat'"映射到一个嵌入,指向"某个冠词——the 或 a——在位置 6 或更低,可能相当接近"的方向。换句话说,投影到查询空间会将一个 token 的输入嵌入转变成当处理该 token 时这个注意力头在寻找的东西。同样,"位置 2 的 'cat'"会被投影成意味着"某个冠词在位置 1 或更低,可能相当接近"的嵌入向量。
现在,除了将输入嵌入投影到查询空间,我们也在将它们投影到键空间。在这种情况下,我们想象的冠词匹配头会创建一个投影,将第一个"the"转变成意味着"一个冠词在位置 1"的东西,将第二个转变成意味着"一个冠词在位置 6"的。
所以,查询权重已将我们的输入嵌入投影到这个"低分辨率"的嵌入空间,指向意味着"这是我感兴趣的"的方向,而键权重已将输入嵌入投影到同一个嵌入空间,指向意味着"这是我"的方向。
这意味着当我们做点积时,"mat"的查询向量会指向与第二个"the"的键向量非常相似的方向,所以点积会很高——记住,只要向量的长度大致相同,点积就是它们相似度的指标。
这里重要的是查询和键向量使用的共享嵌入空间相比输入嵌入使用的丰富空间实际上可能相当贫瘠。在我们的案例中,这个头只关心的是 token 是名词、冠词还是其他什么,以及它们的位置。
让我们举个例子。这是我在第 6 部分想象注意力机制可能得出的注意力分数的虚构集合(改进为因果的,这样 token 不会注意它们"未来"的 token):
每一行是,对于第一列中的 token,所有其他单词的注意力分数。它基于我自己对单词重要性的直觉,这是你可能想象一个聪明的注意力头可能得出的那种东西。(记住 ω 是我们用来代表注意力分数的变量。)
但我们更接近现实世界的冠词-名词匹配头的例子真的很笨,所以它可能会想出更像这样的东西:
它所做的一切就是决定在考虑名词时注意"the"——而且当考虑"mat"时,它甚至付出一点注意力给第一个"the",因为它不知道它必须是它匹配的最近的"the"。²
(注意我们甚至失去了从左上到右下的对角线的 1——意味着 token 注意自己的分数!想想我们在做的投影和叉积——那里没有什么会创建它们。显然某个地方的某个头会做它(一个当试图找出该 token 在上下文中意味着什么时忽略一个 token 的 LLM 不太可能工作得很好),但在这个简化的例子中,没有什么会做它。)
现在,正如我之前所说的,真实的注意力头,经过数十亿 token 的梯度下降训练后,可能学到了一些奇怪的、抽象的、与我们对语言、语法和词性思考方式无关的东西。
但单独来看,它们会真的很笨,因为这个方程在做一些真的简单的事情:当考虑一个特定类型的东西时,寻找这种其他类型的东西。每个 token 被查询权重投影到共享嵌入空间中("我在寻找什么")并被键权重投影到同一空间中("我是什么"),而点积进行比较以找到匹配。
当然,这并不意味着我们失去任何信息。这个贫瘠的嵌入空间只用于进行匹配以计算我们的注意力分数。当我们计算上下文向量时,我们使用投影到值空间,这可以和我们喜欢的一样丰富。
值得注意的是,虽然 Raschka 在书中使用的例子对查询和键向量的共享空间以及值向量的空间有相同的维度,但实际上没有必要这样。我见过 LLM 的规范,其中 QK 空间有更少的维度——至少对这个琐碎的例子来说这是有意义的。
还值得注意的是,这个键/查询空间在这个例子中是贫瘠的,但在一个真实的"陌生"学到的例子中,它实际上可能相当复杂和丰富——但比这个例子要难得多理解。最终,该嵌入空间的性质将像其他一切一样被学习,并将匹配无论这个头学到做什么。
所以,那是(现在)我对缩放点积注意力工作原理的理解。我们只是在做简单的模式匹配,每个 token 的输入嵌入由查询权重投影到一个(学到的)嵌入空间,该空间能够代表它在某种意义上"寻找"的东西。它也由键权重投影到同一空间,但这次以一种使其指向它在同样意义上"是什么"的方式。然后点积匹配这些,以便我们可以将输入嵌入相互关联以计算我们的注意力分数。
这一切在我的脑子里都有意义,我希望它至少在其他几个人的脑子里也有意义 :-)
我会在这里总结一下;下次我会发布我现在对我们在书中所经历的一切对上下文长度意味着什么的理解。我们已经看到了随着输入序列增长而增长的隐藏状态的优势——缺点是什么?
这是本系列下一篇文章的链接。
我用一索引。↩
我用一索引。↩
我想不出有什么方式单个头可以做,老实说。它同时考虑所有其他 token,所以当它看第一个"the"时,它不知道还有另一个更接近的。↩
我想不出有什么方式单个头可以做,老实说。它同时考虑所有其他 token,所以当它看第一个"the"时,它不知道还有另一个更接近的。↩