diff --git a/Image/chapter3/figure3.23.png b/Image/chapter3/figure3.23.png new file mode 100644 index 0000000..0feb681 Binary files /dev/null and b/Image/chapter3/figure3.23.png differ diff --git a/Image/chapter3/figure3.24.png b/Image/chapter3/figure3.24.png new file mode 100644 index 0000000..13a9855 Binary files /dev/null and b/Image/chapter3/figure3.24.png differ diff --git a/Image/chapter3/figure3.25.png b/Image/chapter3/figure3.25.png new file mode 100644 index 0000000..ef31098 Binary files /dev/null and b/Image/chapter3/figure3.25.png differ diff --git a/Image/chapter3/figure3.26.png b/Image/chapter3/figure3.26.png new file mode 100644 index 0000000..a34d89e Binary files /dev/null and b/Image/chapter3/figure3.26.png differ diff --git a/Image/image3.22.png b/Image/image3.22.png deleted file mode 100644 index fcb8db7..0000000 Binary files a/Image/image3.22.png and /dev/null differ diff --git a/cn-Book/3.实现注意力机制.md b/cn-Book/3.实现注意力机制.md index 0b55334..9d5c59d 100644 --- a/cn-Book/3.实现注意力机制.md +++ b/cn-Book/3.实现注意力机制.md @@ -3,34 +3,53 @@ 本章涵盖以下内容: + **探讨在神经网络中使用注意力机制的原因** - + **介绍一个基本的自注意力框架,并逐步深入到改进的自注意力机制** - + **实现一个因果注意力模块,使 LLM 能够一次生成一个token** - + **使用 dropout 随机掩盖部分注意力权重,以减少过拟合** - +----- -在上一章中,你学习了如何准备输入文本以训练 LLM。这包括将文本拆分为单个单词和子词token,这些token可以被编码为向量表示,即所谓的嵌入,以供 LLM 使用。 +- [3.1 长序列建模的问题](#31-长序列建模的问题) +- [3.2 通过注意力机制捕捉数据依赖关系](#32-通过注意力机制捕捉数据依赖关系) +- [3.3 通过自注意力机制关注输入的不同部分](#33-通过自注意力机制关注输入的不同部分) + - [3.3.1 一种不含可训练权重的简化自注意力机制。](#331-一种不含可训练权重的简化自注意力机制) + - [3.3.2 为所有输入的 token 计算注意力权重](#332-为所有输入的-token-计算注意力权重) +- [3.4 实现带有可训练权重的自注意力机制](#34-实现带有可训练权重的自注意力机制) + - [3.4.1 逐步计算注意力权重](#341-逐步计算注意力权重) + - [3.4.2 实现一个简洁的自注意力机制 Python 类](#342-实现一个简洁的自注意力机制-python-类) +- [3.5 使用因果注意力机制来屏蔽后续词](#35-使用因果注意力机制来屏蔽后续词) + - [3.5.1 应用因果注意力掩码](#351-应用因果注意力掩码) + - [3.5.2 使用 dropout 遮掩额外的注意力权重](#352-使用-dropout-遮掩额外的注意力权重) + - [3.5.3 实现一个简洁的因果注意力类](#353-实现一个简洁的因果注意力类) +- [3.6 从单头注意力扩展到多头注意力](#36-从单头注意力扩展到多头注意力) + - [3.6.1 堆叠多层单头注意力层](#361-堆叠多层单头注意力层) + - [3.6.2 通过权重分割实现多头注意力机制](#362-通过权重分割实现多头注意力机制) +- [3.7 总结](#37-总结) + + +----- + + + +在上一章中,我们学习了如何准备输入文本以训练 LLM。这包括将文本拆分为单个单词和子词token,这些token可以被编码为向量,即所谓的嵌入,以供 LLM 使用。 在本章中,我们将关注 LLM 架构中的重要组成部分,即注意力机制,如图 3.1 所示。 -注意力机制是一个复杂的话题,因此我们将专门用一整章来讨论它。我们主要会将注意力机制独立来研究,关注其内部的工作原理。在下一章中,我们将编写环绕自注意力机制的 LLM 其他部分的代码,以便观察它的实际应用,并创建一个能够生成文本的模型。 +注意力机制是一个复杂的话题,因此我们将专门用一整章来讨论它。我们将注意力机制作为独立模块来研究,重点关注其内部的工作原理。在下一章中,我们将编写与自注意力机制相关的 LLM 的其他部分,以观察其实际运作并创建一个生成文本的模型。 本章中,我们将实现四种不同的注意力机制变体,如图 3.2 所示。 -图 3.2 中展示的这些不同的注意力变体是逐步构建的,其目标是在本章末尾实现一个紧凑且高效的多头注意力机制实现,以便在下一章中可以将其整合到我们将编写的 LLM 架构中 +图 3.2 中展示的这些不同的注意力变体是逐步构建的,其目标是在本章末尾实现一个简单且高效的多头注意力机制,以便在下一章中可以将其整合到我们将编写的 LLM 架构中。 + - ## 3.1 长序列建模的问题 -在深入了解自注意力机制之前(这是大语言模型的核心),让我们先探讨一下缺乏注意力机制的架构存在哪些问题(这些架构在大语言模型之前已经存在)。假设我们想要开发一个将一种语言翻译成另一种语言的翻译模型。如图 3.3 所示,我们无法简单地逐词翻译文本,因为源语言和目标语言的语法结构存在差异。 +在深入了解自注意力机制之前(这是大语言模型的核心),让我们先探讨一下缺乏注意力机制的架构存在哪些问题(这些架构在大语言模型之前已经存在)。假设我们想要开发一个将一种语言翻译成另一种语言的翻译模型。如图 3.3 所示,我们无法简单地逐词翻译文本,因为源语言和目标语言的语法结构往往存在差异。 @@ -38,15 +57,15 @@ 在第 1 章(1.4 节,使用 LLM 进行不同任务)介绍 Transformer 架构时,我们已经简要讨论过编码器-解码器网络。在 Transformer 出现之前,循环神经网络(RNN)是最流行的用于语言翻译的编码器-解码器架构。 -**循环神经网络(RNN)**是一种神经网络类型,其中前一步的输出会作为当前步骤的输入,使其非常适合处理像文本这样的序列数据。如果您不熟悉 RNN 的工作原理,不必担心,您无需了解 RNN 的详细机制也可以参与这里的讨论;我们在这里的重点更多是编码器-解码器结构的总体概念。 +**循环神经网络(RNN)**是一种神经网络类型,其中前一步的输出会作为当前步骤的输入,使其非常适合处理像文本这样的序列数据。如果您不熟悉 RNN 的工作原理,不必担心,您无需了解 RNN 的详细机制也可以参与这里的讨论;这一节学习的重点更多是编码器-解码器架构的总体概念。 -在编码器-解码器 RNN 中,输入文本被输入到编码器中,编码器按顺序处理文本内容。在每个步骤中,编码器会更新其隐状态(即隐藏层的内部值),试图在最终的隐状态中捕捉整个输入句子的含义,如图 3.4 所示。然后,解码器使用该最终隐状态来开始逐词生成翻译句子。解码器在每一步也会更新其隐状态,用于携带生成下一个词所需的上下文信息。 +在编码器-解码器架构的 RNN 网络中,输入文本被输入到编码器中,编码器按顺序处理文本内容。在每个步骤中,编码器会更新其隐状态(即隐藏层的内部值),试图在最终的隐状态中捕捉整个输入句子的含义,如图 3.4 所示。随后,解码器使用该最终隐状态来开始逐词生成翻译句子。解码器在每一步也会更新其隐状态,用于携带生成下一个词所需的上下文信息。 -尽管我们不需要深入了解这些编码器-解码器结构 RNN 的内部工作原理,但关键思想在于,编码器部分将整个输入文本处理为一个隐藏状态(记忆单元)。解码器随后使用该隐藏状态生成输出。您可以将这个隐藏状态视为一个嵌入向量,这是我们在第 2 章中讨论的概念。 +尽管我们不需要深入了解这些编码器-解码器架构的 RNN 的内部工作原理,但这里的关键思想在于,编码器部分将整个输入文本处理为一个隐藏状态(记忆单元)。解码器随后使用该隐藏状态生成输出。您可以将这个隐藏状态视为一个嵌入向量,这是我们在第 2 章中已讨论过的概念。 -编码器-解码器结构的 RNN 的一个重大问题和限制在于,**在解码阶段 RNN 无法直接访问编码器的早期隐藏状态**。因此,它只能依赖当前隐藏状态来封装所有相关信息。这种设计可能导致上下文信息的丢失,特别是在依赖关系较长的复杂句子中,这一问题尤为突出。 +编码器-解码器架构的 RNN 的一个重大问题和限制在于,**在解码阶段 RNN 无法直接访问编码器的早期隐藏状态**。因此,它只能依赖当前隐藏状态来封装所有相关信息。这种设计可能导致上下文信息的丢失,特别是在依赖关系较长的复杂句子中,这一问题尤为突出。 对于不熟悉 RNN 的读者,不必深入理解或学习这种架构,因为本书中不会使用它。本节的重点是,编码器-解码器 RNN 存在一个缺点,这一缺点促使了注意力机制的设计。 @@ -83,7 +102,7 @@ > > 在这个注意力矩阵中,可以看到**“it”对“the cat”有较高的关注权重(0.3),而对其他词的关注权重较低**。这种直接的关注能力让模型能够高效捕捉长距离依赖关系,理解“it”与“the cat”的语义关联。 - + ## 3.2 通过注意力机制捕捉数据依赖关系 @@ -91,11 +110,11 @@ 这一方法的一个主要缺陷在于,RNN 必须将整个编码后的输入信息存储在一个隐藏状态中,然后再将其传递给解码器,如上一节的图 3.4 所示。 -因此,研究人员在 2014 年为 RNN 开发了所谓的 Bahdanau 注意力机制(该机制以论文的第一作者命名)。该机制对编码器-解码器 RNN 进行了改进,使得解码器在每个解码步骤可以选择性地访问输入序列的不同部分,如图 3.5 所示。 +因此,研究人员在 2014 年为 RNN 开发了所谓的 Bahdanau 注意力机制(该机制以论文的第一作者命名)。该机制对编码器-解码器架构的 RNN 进行了改进,使得解码器在每个解码步骤可以选择性地访问输入序列的不同部分,如图 3.5 所示。 -有趣的是,仅仅三年后,研究人员发现构建用于自然语言处理的深度神经网络并不需要 RNN 结构,并提出了基于自注意力机制的原始 Transformer 架构(在第 1 章中讨论),其灵感来自 Bahdanau 提出的注意力机制。 +有趣的是,仅仅三年后,研究人员发现构建用于自然语言处理的深度神经网络并不需要 RNN 结构,随后提出了基于自注意力机制的原始 Transformer 架构(在第 1 章中讨论),其灵感来自 Bahdanau 提出的注意力机制。 自注意力机制是一种允许输入序列中的每个位置在计算序列表示时关注同一序列中所有位置的机制。自注意力机制是基于Transformer架构的当代大语言模型(如GPT系列模型)的关键组成部分。 @@ -103,11 +122,11 @@ - + ## 3.3 通过自注意力机制关注输入的不同部分 -现在我们将深入了解自注意力机制的内部工作原理,并从零开始学习如何实现它。自注意力机制是基于 Transformer 架构的每一个大语言模型的核心。需要注意的是,这一部分内容可能需要大量的专注与投入(无双关含义),但一旦掌握了它的基本原理,你就攻克了本书及大语言模型实现中最困难的部分之一。 +现在我们将深入了解自注意力机制的内部工作原理,并从零开始学习如何实现它。自注意力机制是基于 Transformer 架构的所有大语言模型的核心。需要注意的是,这一部分内容可能需要大量的专注与投入(无双关含义),但一旦掌握了它的基本原理,你就攻克了本书及大语言模型实现中最困难的部分之一。 > [!NOTE] > @@ -115,29 +134,29 @@ > > 在自注意力机制中,“self”指的是该机制通过关联同一输入序列中的不同位置来计算注意力权重的能力。它评估并学习输入内部各部分之间的关系和依赖性,例如句子中的单词或图像中的像素。这与传统注意力机制不同,传统机制关注的是两个不同序列间的关系,例如序列到序列模型中,注意力可能存在于输入序列和输出序列之间,这一点在图 3.5 中有示例说明。 -由于自注意力机制可能显得较为复杂,尤其是对于初次接触的读者,我们将在下一小节中首先介绍一个简化版的自注意力机制。随后,在第 3.4 节中,我们将实现带有可训练权重的自注意力机制,这种机制被用于大语言模型(LLM)中。 +由于自注意力机制对于初次接触的读者可能显得较为复杂,我们将在下一小节中首先介绍一个简化版的自注意力机制。随后,在第 3.4 节中,我们将实现带有可训练权重的自注意力机制,这种机制被用于大语言模型(LLM)中。 - -### 3.3.1 一种不含可训练权重的简单自注意力机制。 + +### 3.3.1 一种不含可训练权重的简化自注意力机制。 在本节中,我们实现了一个简化的自注意力机制版本,没有包含任何可训练的权重,如图 3.7 所示。本节的目标是先介绍自注意力机制中的一些关键概念,然后在 3.4 节引入可训练的权重。 -图 3.7 显示了一个输入序列,记作 x,由 T 个元素组成,表示为 x(1) 到 x(T)。该序列通常代表文本,例如一个句子,并且该文本已被转换为 token 嵌入,如第 2 章所述。 +图 3.7 显示了一个输入序列,记作 x,由 T 个元素组成,表示为 x(1) 到 x(T)。该序列通常代表文本,例如一个句子,并且该文本已被转换为 token 嵌入(不记得嵌入概念的请回顾第 2 章)。 -举例来说,假设输入文本为 “Your journey starts with one step”。在这个例子中,序列中的每个元素(如 `x(1)`)对应一个 `d` 维的嵌入向量,用于表示特定的 token,例如 “Your”。在图 3.7 中,这些输入向量显示为 3 维的嵌入向量。 +举例来说,假设输入文本为 “Your journey starts with one step”。在这个例子中,序列中的每个元素(如 x(1))对应一个 `d` 维的嵌入向量,用于表示特定的 token,例如 “Your”。在图 3.7 中,这些输入向量显示为 3 维的嵌入向量。 -在自注意力机制中,我们的目标是为输入序列中的每个元素 x(i) 计算其对应的上下文向量 z(i) 。上下文向量可以被解释为一种增强的嵌入向量。 +在自注意力机制中,我们的目标是为输入序列中的每个元素 x(i) 计算其对应的上下文向量 z(i) 。上下文向量可以被解释为一种增强的嵌入向量(`别着急,后文会解释`)。 -为了说明这个概念,我们聚焦于第二个输入元素 x(2) 的嵌入向量(对应于词 "journey")以及相应的上下文向量 z(2),如图 3.7 底部所示。这个增强的上下文向量 z(2) 是一个嵌入向量,包含了关于 x(2) 以及所有其他输入元素 x(1) 到 x(T) 的信息。 +为了说明这个概念,我们聚焦于第二个输入元素 x(2) 的嵌入向量(对应于词 "journey")以及相应的上下文向量 z(2),如图 3.7 底部所示。这个增强的上下文向量 z(2) 也是一个嵌入向量,包含了关于 x(2) 以及序列中所有其他输入元素 x(1) 到 x(T) 的语义信息。 -在自注意力机制中,上下文向量起着关键作用。它们的目的是通过整合序列中所有其他元素的信息(如同一个句子中的其他词),为输入序列中的每个元素创建丰富的表示,正如图 3.7 所示。这对大语言模型至关重要,因为模型需要理解句子中各个词之间的关系和关联性。之后,我们将添加可训练的权重,以帮助大语言模型学习构建这些上下文向量,使其与生成下一个词的任务相关。 +在自注意力机制中,上下文向量起着关键作用。它们的目的是通过整合序列中所有其他元素的信息(如同一个句子中的其他词),为输入序列中的每个元素创建丰富的表示,正如图 3.7 所示。这对大语言模型至关重要,因为模型需要理解句子中各个词之间的关系和关联性。之后的章节中,我们将添加可训练的权重,以帮助大语言模型学习构建这些上下文向量,用于执行生成下一个词的任务。 -在本节中,我们实现了一个简化的自注意力机制,以逐步计算这些权重和生成的上下文向量。 +在本节中,我们将实现一个简化的自注意力机制,以逐步计算注意力权重和由此生成的上下文向量。 -请考虑以下输入句子,该句子已经根据第 2 章的讨论转换为三维向量表示。为了便于说明和展示,我们选择了较小的嵌入维度,以确保句子在页面上可以无换行地展示。 +请考虑以下输入句子,该句子已经根据第 2 章的讨论转换为三维向量。为了便于说明和展示,我们选择了较小的嵌入维度,以确保句子在页面上可以完整地展示。 ```python import torch @@ -151,11 +170,11 @@ inputs = torch.tensor( ``` -实现自注意力机制的第一步是计算中间值 **ω**,即注意力得分,如图 3.8 所示。请注意,图 3.8 中前一层输入张量的值为截断版本,例如 0.87 因空间限制被截断为 0.8。在这个截断版本中,单词 "journey" 和 "starts" 的嵌入向量可能会由于随机因素而看起来相似。 +实现自注意力机制的第一步是计算中间值 **ω**,即注意力得分,如图 3.8 所示。(请注意,图 3.8 中展示的输入张量值是截断版的,例如,由于空间限制,0.87 被截断为 0.8。在此截断版中,单词 "journey" 和 "starts" 的嵌入向量可能会由于随机因素而看起来相似)。 -图 3.8 说明了如何计算查询 token 与每个输入 token 之间的中间注意力得分。我们通过计算查询 token x(2)x(2)x(2) 与每个其他输入 token 的点积来确定这些得分。 +图 3.8 展示了如何计算查询 token 与每个输入 token 之间的中间注意力得分。我们通过计算查询 x(2) 与每个其他输入 token 的点积来确定这些得分: ```python query = inputs[1] #A @@ -209,7 +228,7 @@ tensor([0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865]) > > **理解点积** > -> 点积运算本质上是一种将两个向量按元素相乘后再求和的简洁方式,我们可以如下演示: +> 点积运算本质上是一种将两个向量按元素相乘后再求和的简单方式,我们可以如下演示: > > ```python > res = 0. @@ -289,7 +308,7 @@ Sum: tensor(1.) > > Softmax 函数的公式如下: > -> $$ \text{softmax}\left(z_{i}\right)=\frac{e^{z_{i}}}{\sum_{j} e^{z_{j}}} $$ +> $$ \text{softmax}\left(z_{i}\right)=\frac{e^{z_{i}}}{\sum_{j} e^{z_{j}}} $$ > > 其中zi是输入的每个分数(即未激活的原始值),e 是自然对数的底。这个公式的作用是将输入向量中的每个元素转换为一个概率值,且所有值的和为 1。 > @@ -307,11 +326,11 @@ Sum: tensor(1.) > + **概率解释**:Softmax 输出可以解释为“模型对每个类别的信心”,使得输出直观可理解。 > + **与交叉熵的结合**:Softmax 与交叉熵损失函数结合效果特别好,可以直接将模型预测的概率分布与真实标签比较,从而更快收敛,效果更好。 -现在我们已经计算出了归一化的注意力权重,接下来可以执行图 3.10 所示的最后一步:通过将嵌入后的输入 token x(i) 与相应的注意力权重相乘,再将所得向量求和,来计算上下文向量 z(2)。 +现在我们已经计算出了归一化的注意力权重,接下来可以执行图 3.10 所示的最后一步:通过将嵌入后的输入 token x(i) 与相应的注意力权重相乘,再将所得向量求和来计算上下文向量 z(2) -如图 3.10 所示,上下文向量 z(2)  是通过所有输入向量的加权和计算得到的。这一过程涉及将每个输入向量与其对应的注意力权重相乘: +如图 3.10 所示,上下文向量 z(2) 是所有输入向量的加权和。其计算方法为将每个输入向量与对应的注意力权重相乘后相加。 ```python query = inputs[1] # 2nd input token is the query @@ -327,13 +346,13 @@ print(context_vec_2) tensor([0.4419, 0.6515, 0.5683]) ``` -在接下来的部分,我们将把串行计算上下文向量的过程优化为并行计算所有输入token的上下文向量。 +在接下来的章节,我们将把串行计算上下文向量的过程优化为并行计算所有输入token的上下文向量。 + - ### 3.3.2 为所有输入的 token 计算注意力权重 -在前一节中,我们计算了第二个输入元素的注意力权重和上下文向量,如图 3.11 中的高亮行所示。现在,我们将扩展该计算,以对所有输入计算注意力权重和上下文向量。 +在前一节中,我们计算了第二个输入元素的注意力权重和上下文向量(如图 3.11 中的高亮行所示)。现在,我们将扩展该计算,以对所有输入计算注意力权重和上下文向量。 @@ -362,7 +381,7 @@ tensor([[0.9995, 0.9544, 0.9422, 0.4753, 0.4576, 0.6310], [0.6310, 1.0865, 1.0605, 0.6565, 0.2935, 0.9450]]) ``` -以上张量中的每个元素都表示每对输入之间的注意力得分,正如图 3.11 中所示。请注意,图 3.11 中的值已进行了归一化,因此它们与以上张量中的未经归一化的注意力得分不同。我们稍后会处理归一化。 +以上张量中的每个元素都表示每对输入之间的注意力得分(正如图 3.11 中所示)。请注意,图 3.11 中的值已进行了归一化,因此它们与以上张量中的未经归一化的注意力得分不同。我们稍后会处理归一化。 在上述代码中,我们使用了 Python 中的 for 循环来计算所有输入对的注意力得分。然而,for 循环通常较慢,我们可以通过矩阵乘法实现相同的结果。 @@ -371,7 +390,7 @@ attn_scores = inputs @ inputs.T print(attn_scores) ``` -我们可以直观地确认结果与之前一致: +可以看到,结果与之前一致: ``` tensor([[0.9995, 0.9544, 0.9422, 0.4753, 0.4576, 0.6310], @@ -417,7 +436,7 @@ Row 2 sum: 1.0 All row sums: tensor([1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000]) ``` -在第三步也是最后一步中,我们使用这些注意力权重通过矩阵乘法的方式来并行计算所有的上下文向量: +在第 3 步也是最后一步中,我们使用这些注意力权重通过矩阵乘法的方式来并行计算所有的上下文向量: ```python all_context_vecs = attn_weights @ inputs @@ -447,31 +466,31 @@ tensor([[0.4421, 0.5931, 0.5790], Previous 2nd context vector: tensor([0.4419, 0.6515, 0.5683]) ``` -以上内容完成了对简单自注意力机制的代码演示。在接下来的部分,我们将添加可训练的权重,使大语言模型能够从数据中学习并提升其在特定任务上的性能。 +以上内容完成了对简化自注意力机制的代码演示。在接下来的部分,我们将添加可训练的权重,使大语言模型能够从数据中学习并提升其在特定任务上的性能。 + - ## 3.4 实现带有可训练权重的自注意力机制 -在本节中,我们正在实现一种在原始 Transformer 架构、GPT 模型以及大多数其他流行的大语言模型中使用的自注意力机制。这种自注意力机制也被称为缩放点积注意力。图 3.13 提供了一个思维模型,展示了这种自注意力机制是如何应用在在大语言模型的架构设计中。 +在本节中,我们将实现一种在原始 Transformer 架构、GPT 模型以及大多数其他流行的大语言模型中使用的自注意力机制。这种自注意力机制也被称为缩放点积注意力。图 3.13 提供了一个概念框架,展示了这种自注意力机制如何应用在在大语言模型的架构设计中。 -如图 3.13 所示,带有可训练权重的自注意力机制是基于之前简化自注意力机制的改进:我们希望计算某个特定输入元素的嵌入向量的加权和来作为上下文向量。您将看到,与我们在 3.3 节中编码的基本自注意力机制相比,只有细微的差别。 +如图 3.13 所示,带有可训练权重的自注意力机制是基于之前简化自注意力机制的改进:我们希望计算某个特定输入元素的嵌入向量的加权和来作为上下文向量。您将看到,与我们在 3.3 节中编码的简化自注意力机制相比,只有细微的差别。 -最显著的区别在于引入了在模型训练过程中不断更新的权重矩阵。这些可训练的权重矩阵至关重要,它们使模型(特别是模型内部的注意力模块)能够学习生成“优质”的上下文向量。(请注意,我们将在第五章训练大语言模型。) +最显著的区别在于引入了在模型训练过程中不断更新的权重矩阵。这些可训练的权重矩阵至关重要,它们使模型(特别是模型内部的注意力模块)能够学习生成“优质”的上下文向量。(请注意,我们将在第 5 章训练大语言模型。) + +我们将通过两个小节来深入讲解自注意力机制。首先,我们会像之前一样,逐步编写该机制的代码。然后,我们会将代码整理成一个紧凑的 Python 类,以便在之后第 4 章编写的大语言模型(LLM)架构中使用。 -我们将通过两个小节来深入讲解自注意力机制。首先,我们会像之前一样,逐步编写该机制的代码。然后,我们会将代码整理成一个紧凑的 Python 类,以便在第 4 章编写的大型语言模型(LLM)架构中使用。 - ### 3.4.1 逐步计算注意力权重 -我们通过引入三个可训练的权重矩阵:Wq、Wk 和 Wv 来逐步实现自注意力机制。这三个矩阵用于将嵌入后的输入 token x(i) 映射为查询向量、键向量和值向量,如图 3.14 所示。 +我们通过引入三个可训练的权重矩阵:Wq、Wk 和 Wv 来逐步实现自注意力机制。这三个矩阵用于将嵌入后的输入 token x(i) 映射为查询向量、键向量和值向量(如图 3.14 所示)。 -在第 3.3.1 节中,我们将第二个输入元素 x(2) 定义为查询(query),通过计算简化的注意力权重来得到上下文向量 z(2)。随后,在第 3.3.2 节中,我们将这一过程推广到整个输入句子 "Your journey starts with one step",为这六个词的输入句子计算所有的上下文向量 z(1) 到 z(T)。 +在 3.3.1 节中,我们将第二个输入元素 x(2) 定义为查询(query),通过计算简化的注意力权重来得到上下文向量 z(2)。随后,在第 3.3.2 节中,我们将这一过程推广到整个输入句子 "Your journey starts with one step",为这六个词的输入句子计算所有的上下文向量 z(1) 到 z(T)。 同样地,为了便于说明,我们将先计算一个上下文向量 z(2)。接下来,我们将修改代码以计算所有的上下文向量。让我们从定义一些变量开始: ```python @@ -496,7 +515,7 @@ W_key = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False) W_value = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False) ``` -请注意,这里我们将 `requires_grad` 设置为 `False`,以便在输出结果中减少不必要的信息,使演示更加清晰。但如果要将这些权重矩阵用于模型训练,则需要将 `requires_grad` 设置为 `True`,以便在模型训练过程中更新这些矩阵。 +请注意,这里我们将 `requires_grad` 设置为 `False`,以便在输出结果中减少不必要的信息,从而使演示更加清晰。但如果要将这些权重矩阵用于模型训练,则需要将 `requires_grad` 设置为 `True`,以便在模型训练过程中更新这些矩阵。 接下来,我们计算之前在图 3.14 中展示的 query、key 和 value 向量: @@ -521,7 +540,7 @@ print(query_2) > > 总之,权重参数是神经网络的基本学习系数,用于定义网络层之间的连接关系,而注意力权重则是根据上下文动态生成的特定值,用于衡量不同词语或位置在当前上下文中的重要性。 -尽管我们当前的目标仅仅是计算一个上下文向量 z(2),但仍然需要获取所有输入元素的 key 和 value 向量,因为它们参与了与查询向量 q(2) 一起计算注意力权重的过程,如图 3.14 所示。 +尽管我们当前的目标仅仅是计算一个上下文向量 z(2),但仍然需要获取所有输入元素的 key 和 value 向量,因为它们将参与与查询向量 q(2) 一起计算注意力权重的过程,如图 3.14 所示。 我们可以通过矩阵乘法获取所有元素的key和value向量: @@ -539,11 +558,11 @@ keys.shape: torch.Size([6, 2]) values.shape: torch.Size([6, 2]) ``` -接下来的第二步是计算注意力得分,如图 3.15 所示。 +接下来的第二步是计算注意力得分(如图 3.15 所示)。 -首先,我们计算注意力得分ω22 : +首先,我们计算注意力得分ω22 : ```python keys_2 = keys[1] #A @@ -568,7 +587,7 @@ print(attn_score_22) > > 很明显,在这两个句子中,`it`的指代不同,第一个句子中,`it`指代`cat`,而在第二个句子中,`it`指代`milk`。 > -> 根据一下注意力得分的公式(Qcat和Kit分别为`cat`和`it`的查询向量和键向量)可知,句子1中`score_cat_it`是要大于句子2中的`score_cat_it`,因为句子1中,`it`和`cat`的相关度更高,但是从公式中如何推断出实现呢? +> 根据以下注意力得分的公式(Qcat和Kit分别为`cat`和`it`的查询向量和键向量)可知,句子1中`score_cat_it`是要大于句子2中的`score_cat_it`,因为句子1中,`it`和`cat`的相关度更高,但是从公式中如何推断出实现呢? > > **score_cat_it = Qcat · Kit** > @@ -578,7 +597,7 @@ print(attn_score_22) > > **Kit = Wk * (Eit + Posit)** > -> 其中 **Ecat**和**Eit**是这两个词的嵌入向量,表示该词的基本语义信息,在不同的上下文中是固定的,根据公式可知,要使最终算出的**score_cat_it**与上下文语义相关,最重要的是Wq 和Wk这两个权重参数应该能反映出不同上下文语义的相关性。在标准的自注意力机制中,W、K、V向量都是固定的,然而,由于 GPT 模型是由多层自注意力模块堆叠而成,每一层都会根据当前输入和上下文信息,动态调整查询、键和值向量的权重矩阵。因此,即使初始的词嵌入和权重矩阵是固定的,经过多层处理后,模型能够生成与当前上下文相关的 Q、K、V 向量权重矩阵,最终计算出的Q、K、V 向量也就能反映出上下文的语义了。GPT多层的实现的细节后文会详述。 +> 其中 **Ecat**和**Eit**是这两个词的嵌入向量,表示该词的基本语义信息,在不同的上下文中是固定的,根据公式可知,要使最终算出的**score_cat_it**与上下文语义相关,最重要的是**Wq** 和**Wk**这两个权重参数应该能反映出不同上下文语义的相关性。在标准的自注意力机制中,W、K、V向量都是固定的,然而,由于 GPT 模型是由多层自注意力模块堆叠而成,每一层都会根据当前输入和上下文信息,动态调整查询、键和值向量的权重矩阵。因此,即使初始的词嵌入和权重矩阵是固定的,经过多层处理后,模型能够生成与当前上下文相关的 Q、K、V 向量权重矩阵,最终计算出的Q、K、V 向量也就能反映出上下文的语义了。GPT多层的实现的细节后文会详述。 我们可以再次通过矩阵乘法将其应用到所有注意力得分的计算: @@ -597,7 +616,7 @@ print(attn_scores_2) -接下来,如图 3.16 所示,我们通过缩放注意力分数并使用前面提到的 softmax 函数来计算注意力权重。与之前的不同之处在于,现在我们通过将注意力分数除以`keys`嵌入维度的平方根来进行缩放(注意,取平方根在数学上等同于指数为 0.5 的运算)。 +接下来,如图 3.16 所示,我们通过缩放注意力得分并使用前面提到的 softmax 函数来计算注意力权重。与之前的不同之处在于,现在我们通过将注意力得分除以`keys`嵌入维度的平方根来进行缩放(注意,取平方根在数学上等同于指数为 0.5 的运算)。 ```python d_k = keys.shape[-1] @@ -613,13 +632,13 @@ print(attn_weights_2) > [!NOTE] > -> **缩放点积注意力机制的原理** +> **缩放点积注意力机制的原理** > > 对嵌入维度大小进行归一化的原因是为了避免出现小梯度,从而提高训练性能。例如,当嵌入维度增大时(在 GPT 类大型语言模型中通常超过一千),较大的点积在反向传播中应用 softmax 函数后,可能会导致非常小的梯度。随着点积的增大,softmax 函数的行为会更加类似于阶跃函数,导致梯度接近于零。这些小梯度可能会显著减慢学习速度,甚至导致训练停滞。 > > 通过嵌入维度的平方根进行缩放,正是自注意力机制被称为‘缩放点积注意力’的原因。 -> [!TIP] +> [!TIP] > > **个人思考:** 这里再稍微解释一下上述关于缩放点积注意力的机制。在自注意力机制中,查询向量(Query)与键向量(Key)之间的点积用于计算注意力权重。然而,当嵌入维度(embedding dimension)较大时,点积的结果可能会非常大。那么大的点积对接下来的计算有哪些具体影响呢? > @@ -659,7 +678,7 @@ tensor([0.3061, 0.8210]) > > **值**(value)类似于数据库中的键值对中的“值”。它表示输入项的实际内容或表示。当模型确定哪些键(即输入中的哪些部分)与查询(当前的关注项)最相关时,就会检索出对应的值。 - + ### 3.4.2 实现一个简洁的自注意力机制 Python 类 @@ -718,7 +737,7 @@ tensor([[0.2996, 0.8053], -如图3.18所示,自注意力机制涉及可训练的权重矩阵 Wq、Wk 和 Wv。这些矩阵将输入数据转换为查询、键和值,它们是注意力机制的重要组成部分。随着训练过程中数据量的增加,模型会不断调整这些可训练的权重,在后续章节中我们会学习相关细节。 +如图3.18所示,自注意力机制涉及可训练的权重矩阵 Wq、Wk 和 Wv。这些矩阵将输入数据转换为查询、键和值,它们是自注意力机制的重要组成部分。随着训练过程中数据量的增加,模型会不断调整这些可训练的权重,在后续章节中我们会学习相关细节。 我们可以通过使用 PyTorch 的 `nn.Linear` 层来进一步改进 SelfAttention_v1 的实现。当禁用偏置单元时,`nn.Linear` 层可以有效地执行矩阵乘法。此外,使用 `nn.Linear` 替代手动实现的 `nn.Parameter(torch.rand(...))` 的一个显著优势在于,`nn.Linear` 具有优化的权重初始化方案,从而有助于实现更稳定和更高效的模型训练。 @@ -761,11 +780,11 @@ tensor([[-0.0739, 0.0713], [-0.0754, 0.0693]], grad_fn=) ``` -`SelfAttention_v1` 和` SelfAttention_v2` 的输出不同,因为它们的权重矩阵使用了不同的初始权重,这是由于 `nn.Linear` 层采用了一种更复杂的权重初始化方案。 +`SelfAttention_v1` 和` SelfAttention_v2` 的输出不同,因为它们的权重矩阵使用了不同的初始权重,根本原因在于 `nn.Linear` 层采用了一种更复杂的权重初始化方案。 > [!NOTE] > -> **练习 3.1:比较`SelfAttention_v1`和 `SelfAttention_v2`** +> **练习 3.1:比较`SelfAttention_v1`和 `SelfAttention_v2`** > > 请注意,`SelfAttention_v2` 中的 `nn.Linear` 层使用了一种不同的权重初始化方式,而 `SelfAttention_v1` 则使用 `nn.Parameter(torch.rand(d_in, d_out))` 进行初始化。这导致两种机制生成的结果有所不同。为了验证 `SelfAttention_v1` 和 `SelfAttention_v2` 的其他部分是否相似,我们可以将 `SelfAttention_v2` 对象中的权重矩阵转移到 `SelfAttention_v1` 中,从而使两者生成相同的结果。 > @@ -775,7 +794,7 @@ tensor([[-0.0739, 0.0713], 多头组件将注意力机制分解为多个‘头’。每个头能够学习数据的不同方面,使模型能够同时关注来自不同表示子空间的不同位置的信息。这提高了模型在复杂任务中的性能。 - + ## 3.5 使用因果注意力机制来屏蔽后续词 @@ -783,15 +802,15 @@ tensor([[-0.0739, 0.0713], 因果注意力(也称为掩蔽注意力)是一种特殊的自注意力形式。它限制模型在处理任何给定的 token 时,只能考虑序列中的前一个和当前输入,而不能看到后续的内容。这与标准的自注意力机制形成对比,后者允许模型同时访问整个输入序列。 -因此,在计算注意力分数时,因果注意力机制确保模型只考虑当前 token 之前或之前的 token。 +因此,在计算注意力分数时,因果注意力机制确保模型只考虑当前 token 或之前的 token。 -在 GPT 类大语言模型中,为了实现这一点,我们会对每个处理的 token 屏蔽其后续 token,即在输入文本中当前词之后的所有词,如图 3.19 所示。 +在 GPT 类大语言模型中,要实现这一点,我们需要对每个处理的 token 屏蔽其后续 token,即在输入文本中当前词之后的所有词,如图 3.19 所示。 如图 3.19 所示,我们对注意力权重的对角线上方部分进行了掩码操作,并对未掩码的注意力权重进行归一化,使得每一行的注意力权重之和为 1。在下一节中,我们将用代码实现这个掩码和归一化过程。 - + ### 3.5.1 应用因果注意力掩码 @@ -940,7 +959,7 @@ tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000], 现在,我们可以使用修改后的注意力权重,通过 `context_vec = attn_weights @ values` 来计算上下文向量,这在第 3.4 节中介绍过。不过,在下一节中,我们将首先介绍一个对因果注意力机制的细微调整,这一调整在训练大语言模型时有助于减少过拟合现象。 - + ### 3.5.2 使用 dropout 遮掩额外的注意力权重 @@ -1037,3 +1056,380 @@ tensor([[2.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000], 在理解了因果注意力和 dropout 掩码的基础上,接下来的部分中我们将开发一个简洁的 Python 类,以便高效应用这两种技术。 + + +### 3.5.3 实现一个简洁的因果注意力类 + +在本节中,我们将把因果注意力和 dropout 的修改整合到在 3.4 节开发的 `SelfAttention` Python 类中。该类将作为模板,用于接下来一节中开发多头注意力(多头注意力将是我们在本章实现的最后一个注意力类)。 + +但在开始之前,还需确保代码能够处理由多个输入组成的批次,以便 `CausalAttention` 类能够支持我们在第 2 章中实现的数据加载器所生成的批次输出。 + +为了简单起见,我们复制输入文本示例以模拟批量输入: + +```python +batch = torch.stack((inputs, inputs), dim=0) +print(batch.shape) #A + +#A 2个输入,每个输入有6个token,每个token的嵌入维度为3 +``` + +以上代码生成一个三维张量,包含 2 个输入文本,每个文本包含 6 个 token,每个 token 表示为一个 3 维嵌入向量: + +``` + torch.Size([2, 6, 3]) +``` + +以下的 CausalAttention 类与我们之前实现的 SelfAttention 类类似,不同之处在于我们现在添加了dropout和因果掩码组件,如以下代码所示: + +```python +# Listing 3.3 A compact causal attention class +class CausalAttention(nn.Module): + def __init__(self, d_in, d_out, context_length, dropout, qkv_bias=False): + super().__init__() + self.d_out = d_out + self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias) + self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias) + self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias) + self.dropout = nn.Dropout(dropout) #A + self.register_buffer( + 'mask', + torch.triu(torch.ones(context_length, context_length), + diagonal=1) + ) #B + + def forward(self, x): + b, num_tokens, d_in = x.shape #C + keys = self.W_key(x) + queries = self.W_query(x) + values = self.W_value(x) + attn_scores = queries @ keys.transpose(1, 2) #C + attn_scores.masked_fill_( #D + self.mask.bool()[:num_tokens, :num_tokens], -torch.inf) + attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1) + attn_weights = self.dropout(attn_weights) + context_vec = attn_weights @ values + return context_vec + + +#A 与之前的 SelfAttention_v1 类相比,我们添加了一个 dropout 层 +#B register_buffer 调用也是新添加的内容(后续内容会提供更多相关信息) +#C 我们交换第 1 和第 2 个维度,同时保持批次维度在第1个位置(索引0) +#D 在 PyTorch 中,带有下划线后缀的操作会在原有内存空间执行,直接修改变量本身,从而避免不必要的内存拷贝 +``` + +虽然新增的代码行与之前章节介绍的内容基本一致,但我们现在在 `__init__` 方法中添加了 `self.register_buffer()` 的调用。`register_buffer` 在 PyTorch 中并非所有情况下都必须使用,但在这里有其独特的优势。例如,当我们在大语言模型(LLM)中使用 `CausalAttention` 类时,buffer 会自动随模型迁移到合适的设备(CPU 或 GPU)。这意味着我们无需手动确保这些张量与模型参数在同一设备上,从而避免设备不匹配错误。 + +我们可以按如下方式使用 `CausalAttention` 类(类似于之前的 `SelfAttention`): + +```python +torch.manual_seed(123) +context_length = batch.shape[1] +ca = CausalAttention(d_in, d_out, context_length, 0.0) +context_vecs = ca(batch) +print("context_vecs.shape:", context_vecs.shape) +``` + +生成的上下文向量是一个三维张量,其中每个 token 现在都表示为一个二维嵌入: + +``` +context_vecs.shape: torch.Size([2, 6, 2]) +``` + +图 3.23 提供了一个概念框架,总结了我们迄今为止完成的内容。 + + + +如图 3.23 所示,本节我们重点介绍了神经网络中的因果注意力的概念和实现。在下一节中,我们将进一步扩展这一概念,实现一个多头注意力模块,该模块可以并行实现多个因果注意力机制。 + + + +## 3.6 从单头注意力扩展到多头注意力 + +在本章的最后一部分中,我们将之前实现的因果注意力类扩展为多头形式,这也称为多头注意力。 + +多头’一词指的是将注意力机制划分为多个‘头’,每个头独立运作。在这种情况下,单个因果注意力模块可以视为单头注意力,即只有一组注意力权重用于按顺序处理输入。 + +在接下来的小节中,我们将讨论从因果注意力扩展到多头注意力的过程。第一小节将通过堆叠多个因果注意力模块,直观地构建一个多头注意力模块以作说明。第 2 小节将以一种更复杂但计算效率更高的方式实现相同的多头注意力模块。 + + + +### 3.6.1 堆叠多层单头注意力层 + +在实际应用中,实现多头注意力需要创建多个自注意力机制的实例(在 3.4.1 节的图 3.18 中已有展示),每个实例都具有独立的权重,然后将它们的输出合并。多个自注意力机制实例的应用属于计算密集型(CPU密集型)操作,但它对于识别复杂模式至关重要,这是基于 Transformer 的大语言模型所擅长的能力之一。 + +图3.24展示了多头注意力模块的结构,该模块由多个单头注意力模块组成,如图3.18所示,彼此堆叠在一起。 + + + +如前所述,多头注意力机制的核心思想是在并行运行多个注意力机制的过程中,对输入数据(如注意力机制中的 query、key 和 value 向量)使用不同的、可学习的线性投影。具体来说,就是将这些输入数据与权重矩阵相乘,得到不同的投影结果。 + +在代码中,我们可以通过实现一个简单的 `MultiHeadAttentionWrapper` 类来实现这一点,该类会堆叠多个我们之前实现的 `CausalAttention` 模块的实例: + +```python +# Listing 3.4 A wrapper class to implement multi-head attention +class MultiHeadAttentionWrapper(nn.Module): + def __init__(self, d_in, d_out, context_length, + dropout, num_heads, qkv_bias=False): + super().__init__() + self.heads = nn.ModuleList( + [CausalAttention(d_in, d_out, context_length, dropout, qkv_bias) + for _ in range(num_heads)] + ) + def forward(self, x): + return torch.cat([head(x) for head in self.heads], dim=-1) +``` + +例如,如果我们使用这个 MultiHeadAttentionWrapper 类,并通过设置 num_heads=2 使用两个注意力头,同时将 CausalAttention 的输出维度 d_out 设置为 2,那么生成的上下文向量将是 4 维的(d_out*num_heads=4),如图 3.25 所示。 + + + +为了通过一个具体的例子进一步说明图 3.25,我们可以按如下方式使用 MultiHeadAttentionWrapper 类(使用方式类似于之前的 CausalAttention 类): + +```python +torch.manual_seed(123) +context_length = batch.shape[1] # This is the number of tokens +d_in, d_out = 3, 2 +mha = MultiHeadAttentionWrapper(d_in, d_out, context_length, 0.0, num_heads=2) +context_vecs = mha(batch) +print(context_vecs) +print("context_vecs.shape:", context_vecs.shape) +``` + +以上代码输出的上下文向量如下所示: + +``` +tensor([[[-0.4519, 0.2216, 0.4772, 0.1063], + [-0.5874, 0.0058, 0.5891, 0.3257], + [-0.6300, -0.0632, 0.6202, 0.3860], + [-0.5675, -0.0843, 0.5478, 0.3589], + [-0.5526, -0.0981, 0.5321, 0.3428], + [-0.5299, -0.1081, 0.5077, 0.3493]], + [[-0.4519, 0.2216, 0.4772, 0.1063], + [-0.5874, 0.0058, 0.5891, 0.3257], + [-0.6300, -0.0632, 0.6202, 0.3860], + [-0.5675, -0.0843, 0.5478, 0.3589], + [-0.5526, -0.0981, 0.5321, 0.3428], + [-0.5299, -0.1081, 0.5077, 0.3493]]], grad_fn=) +context_vecs.shape: torch.Size([2, 6, 4]) +``` + +由此生成的 context_vecs 张量的第一个维度是 2,因为我们有两个输入文本(输入文本被复制,因此它们的上下文向量完全相同)。第二个维度对应每个输入中的 6 个 token。第三个维度对应每个 token 的 4 维嵌入向量。 + +> [!NOTE] +> +> **练习 3.2:返回二维嵌入向量** +> +> 更改 `MultiHeadAttentionWrapper(..., num_heads=2)` 调用中的输入参数,使输出的上下文向量为 2 维而不是 4 维,同时保持 `num_heads=2` 的设置。提示:无需修改类的实现,只需更改其中一个输入参数即可。 + +在本节中,我们实现了一个 `MultiHeadAttentionWrapper`,用于组合多个单头注意力模块。不过需要注意的是,在 `forward` 方法中,这些模块是通过 `[head(x) for head in self.heads]` 串行处理的。我们可以通过并行处理各个注意力头来优化该实现。实现这一目标的一种方法是,通过矩阵乘法同时计算所有注意力头的输出,我们将在下一节中详细探讨。 + + + +### 3.6.2 通过权重分割实现多头注意力机制 + +在前一节中,我们创建了一个 MultiHeadAttentionWrapper,通过堆叠多个单头注意力模块来实现多头注意力。这是通过实例化并组合多个 CausalAttention 对象实现的。 + +与其维护两个独立的类 MultiHeadAttentionWrapper 和 CausalAttention,我们可以将这两个概念合并为一个 MultiHeadAttention 类。此外,除了简单地合并 MultiHeadAttentionWrapper 和 CausalAttention 的代码外,我们还会进行一些额外的修改,以更高效地实现多头注意力机制。 + +在 `MultiHeadAttentionWrapper` 中,多头机制是通过创建一个包含多个 `CausalAttention` 对象的列表(`self.heads`)来实现的,每个对象代表一个独立的注意力头。`CausalAttention` 类独立执行注意力机制,每个头的结果最终被拼接起来。相比之下,接下来的 `MultiHeadAttention` 类则将多头功能集成在一个单一的类中。它通过对变换后的query、key和value张量进行重塑,将输入分割成多个头,并在计算注意力后将这些头的结果组合在一起。 + +在进一步讨论之前,让我们先看一下 MultiHeadAttention 类的实现: + +```python +# Listing 3.5 An efficient multi-head attention class +class MultiHeadAttention(nn.Module): + def __init__(self, d_in, d_out, + context_length, dropout, num_heads, qkv_bias=False): + super().__init__() + assert d_out % num_heads == 0, "d_out must be divisible by num_heads" + self.d_out = d_out + self.num_heads = num_heads + self.head_dim = d_out // num_heads #A + self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias) + self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias) + self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias) + self.out_proj = nn.Linear(d_out, d_out) #B + self.dropout = nn.Dropout(dropout) + self.register_buffer( + 'mask', + torch.triu(torch.ones(context_length, context_length), diagonal=1) + ) + + + def forward(self, x): + b, num_tokens, d_in = x.shape + keys = self.W_key(x) #C + queries = self.W_query(x) #C + values = self.W_value(x) #C + + keys = keys.view(b, num_tokens, self.num_heads, self.head_dim) #D + values = values.view(b, num_tokens, self.num_heads, self.head_dim) #D + queries = queries.view(b, num_tokens, self.num_heads, self.head_dim) #D + + keys = keys.transpose(1, 2) #E + queries = queries.transpose(1, 2) #E + values = values.transpose(1, 2) #E + + attn_scores = queries @ keys.transpose(2, 3) #F + mask_bool = self.mask.bool()[:num_tokens, :num_tokens] #G + + attn_scores.masked_fill_(mask_bool, -torch.inf) #H + + attn_weights = torch.softmax( + attn_scores / keys.shape[-1]**0.5, dim=-1) + attn_weights = self.dropout(attn_weights) + + context_vec = (attn_weights @ values).transpose(1, 2) #I + + context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out) #J + context_vec = self.out_proj(context_vec) #K + return context_vec + + +#A 将投影维度缩小,以匹配期望的输出维度 +#B 使用线性层组合头部输出 +#C 张量形状:(b, num_tokens, d_out) +#D 我们通过添加 num_heads 维度来隐式地拆分矩阵。然后展开最后一个维度,使其形状从 (b, num_tokens, d_out) 转换为 (b, num_tokens, num_heads, head_dim) +#E 将张量的形状从 (b, num_tokens, num_heads, head_dim) 转置为 (b, num_heads, num_tokens, head_dim) +#F 对每个注意力头进行点积运算 +#G 掩码被截断到 token 的数量 +#H 使用掩码填充注意力分数 +#I 张量形状:(b, num_tokens, n_heads, head_dim) +#J 将多个注意力头的输出结果合并,其中输出维度 self.d_out 等于注意力头数 self.num_heads 与每个头的维度 self.head_dim 的乘积 +#K 添加一个可选的线性投影层 +``` + +尽管 `MultiHeadAttention` 类中张量的重塑(.view)和转置(.transpose)操作看起来非常复杂,但从数学角度来看,`MultiHeadAttention` 类与之前的 `MultiHeadAttentionWrapper` 类实现的概念是相同的。 + +从宏观层面上看,在之前的 MultiHeadAttentionWrapper 中,我们通过堆叠多个单头注意力层的方式来组合成一个多头注意力层。而 MultiHeadAttention 类采用了一种集成的方法:它从一个多头注意力层开始,并在内部将该层分解为各个独立的注意力头,如图 3.26 所示。 + + + +如图 3.26 所示,query、key 和 value 张量的拆分是通过张量的重塑和转置操作实现的,这些操作分别使用了 PyTorch 的 `.view` 和 `.transpose` 方法。首先,通过线性层对输入进行投影(分别生成 query、key 和 value),然后将其重塑为多个注意力头的形式。 + +关键操作是将 `d_out` 维度拆分成 `num_heads` 和 `head_dim`,其中 `head_dim = d_out / num_heads`。这种拆分通过 `.view` 方法实现:将形状为 `(b, num_tokens, d_out)` 的张量重塑为 `(b, num_tokens, num_heads, head_dim)`。 + +接下来对张量进行转置操作,将 `num_heads` 维度移动到 `num_tokens` 维度之前,使其形状变为 `(b, num_heads, num_tokens, head_dim)`。这种转置对于在不同注意力头之间正确对齐查询(queries)、键(keys)和值(values),并高效执行批量矩阵乘法至关重要。 + +为了说明这种批量矩阵乘法,假设我们有如下示例张量: + +```python +a = torch.tensor([[[[0.2745, 0.6584, 0.2775, 0.8573], #A + [0.8993, 0.0390, 0.9268, 0.7388], + [0.7179, 0.7058, 0.9156, 0.4340]], + [[0.0772, 0.3565, 0.1479, 0.5331], + [0.4066, 0.2318, 0.4545, 0.9737], + [0.4606, 0.5159, 0.4220, 0.5786]]]]) + +#A 该张量的形状为 (b, num_heads, num_tokens, head_dim) = (1, 2, 3, 4) +``` + +接下来,我们在张量本身与张量的一个视图之间执行批量矩阵乘法操作,其中张量的视图将最后两个维度(num_tokens 和 head_dim)进行了转置: + +```python +print(a @ a.transpose(2, 3)) +``` + +结果如下: + +``` +tensor([[[[1.3208, 1.1631, 1.2879], + [1.1631, 2.2150, 1.8424], + [1.2879, 1.8424, 2.0402]], + [[0.4391, 0.7003, 0.5903], + [0.7003, 1.3737, 1.0620], + [0.5903, 1.0620, 0.9912]]]]) +``` + +在这种情况下,PyTorch 中的矩阵乘法实现能够处理四维输入张量,因此矩阵乘法会在输入张量的最后两个维度(即 `num_tokens` 和 `head_dim`)之间执行,并对每个注意力头重复该操作。 + +上述方法成为了一种更简洁的方式,可以单独计算每个头的矩阵乘法: + +```python +first_head = a[0, 0, :, :] +first_res = first_head @ first_head.T +print("First head:\n", first_res) +second_head = a[0, 1, :, :] +second_res = second_head @ second_head.T +print("\nSecond head:\n", second_res) +``` + +该结果与我们之前使用批量矩阵乘法 `print(a @ a.transpose(2, 3))` 时获得的结果完全相同: + +``` +First head: + tensor([[1.3208, 1.1631, 1.2879], + [1.1631, 2.2150, 1.8424], + [1.2879, 1.8424, 2.0402]]) +Second head: + tensor([[0.4391, 0.7003, 0.5903], + [0.7003, 1.3737, 1.0620], + [0.5903, 1.0620, 0.9912]]) +``` + +在多头注意力机制中,计算完注意力权重和上下文向量之后,将所有头的上下文向量转置回形状 `(b, num_tokens, num_heads, head_dim)`。然后将这些向量重新塑形(展平)为 `(b, num_tokens, d_out)` 的形状,从而有效地将所有头的输出组合在一起。 + +此外,我们在多头注意力机制中添加了一个称为输出投影层(self.out_proj)的模块,用于在组合多个头的输出后进行投影。而在因果注意力类中并没有这个投影层。这个输出投影层并非绝对必要(详见附录 B 的参考部分),但由于它在许多 LLM 架构中被广泛使用,因此我们在这里加上以保持完整性。 + +尽管 `MultiHeadAttention` 类由于额外的张量重塑和转置操作看起来比 `MultiHeadAttentionWrapper` 更复杂,但它更加高效。原因在于,我们只需执行一次矩阵乘法即可计算键(keys),对于查询(queries)和值(values)也是如此。而在 `MultiHeadAttentionWrapper` 中,我们需要对每个注意力头重复执行这一矩阵乘法操作,这种计算方式的开销非常大。 + +MultiHeadAttention 类的用法与我们之前实现的 SelfAttention 和 CausalAttention 类类似: + +```python +torch.manual_seed(123) +batch_size, context_length, d_in = batch.shape +d_out = 2 +mha = MultiHeadAttention(d_in, d_out, context_length, 0.0, num_heads=2) +context_vecs = mha(batch) +print(context_vecs) +print("context_vecs.shape:", context_vecs.shape) +``` + +从结果可以看出,输出维度是由`d_out`参数直接控制的: + +``` +tensor([[[0.3190, 0.4858], + [0.2943, 0.3897], + [0.2856, 0.3593], + [0.2693, 0.3873], + [0.2639, 0.3928], + [0.2575, 0.4028]], + [[0.3190, 0.4858], + [0.2943, 0.3897], + [0.2856, 0.3593], + [0.2693, 0.3873], + [0.2639, 0.3928], + [0.2575, 0.4028]]], grad_fn=) +context_vecs.shape: torch.Size([2, 6, 2]) +``` + +在本节中,我们实现了 MultiHeadAttention 类,这将在后续章节实现和训练 LLM 时使用。请注意,虽然代码功能齐全,但我们使用了较小的嵌入维度和注意力头数,以便让输出结果更易于阅读。 + +作为对比,最小的 GPT-2 模型(1.17 亿参数)具有 12 个注意力头和 768 的上下文向量嵌入大小。而最大的 GPT-2 模型(15 亿参数)则具有 25 个注意力头和 1600 的上下文向量嵌入大小。请注意,在 GPT 模型中,token 输入的嵌入大小与上下文嵌入大小是相同的(`d_in = d_out`)。 + +> [!NOTE] +> +> **练习 3.3:初始化 GPT-2 规模的注意力模块** +> +> 使用 MultiHeadAttention 类初始化一个多头注意力模块,该模块的注意力头数量与最小的 GPT-2 模型相同(12 个注意力头)。同时确保输入和输出的嵌入大小与 GPT-2 相似(768 维)。请注意,最小的 GPT-2 模型支持的上下文长度为 1024 个 tokens。 + + + +## 3.7 总结 + ++ 注意力机制将输入元素转换为增强的上下文向量表示,其中包含了所有输入的信息。 ++ 自注意力机制通过对输入的加权求和来计算上下文向量表示。 ++ 在简化的注意力机制中,注意力权重是通过点积计算的。 ++ 点积仅仅是对两个向量逐元素相乘后求和的一种简洁方式。 ++ 矩阵乘法虽然并不是绝对必要的,但通过替换嵌套的 for 循环,它帮助我们更高效和简洁地计算。 ++ 在 LLM 中使用的自注意力机制,也称为缩放点积注意力,我们引入可训练的权重矩阵,以计算输入的中间转换:查询、值和键。 ++ 在使用从左到右读取和生成文本的 LLM 时,我们添加一个因果注意力掩码,以防止模型访问未来的 token。 ++ 除了通过因果注意力掩码将注意力权重置为零之外,我们还可以添加 dropout 掩码,以减少 LLM 中的过拟合现象。 ++ 基于Transformer的 LLM 中的注意力模块包含多个因果注意力实例,这被称为多头注意力。 ++ 我们可以通过堆叠多个因果注意力模块的实例来创建一个多头注意力模块。 ++ 创建多头注意力模块的一种更高效的方法是采用批量矩阵乘法。 + +> [!TIP] +> +> **个人思考:** 毫无疑问,本章的注意力机制是整本书中最重要的内容,也是最难的内容 ,这里强烈建议读者能多读几遍,按照文中的示例代码完整地实现一遍,对于不理解的地方多去查阅相关资料及深入思考,力求真正理解和掌握每个细节点。 +