编码注意力机制

#2026/01/01 #ai

本章内容

• 探索在神经网络中使用注意力机制的原因
• 介绍一个基础的自注意力框架，并进一步探讨一种增强型自注意力机制
• 实现一个因果注意力模块，使大语言模型能够一次生成一个词元
• 使用 dropout 随机掩码部分注意力权重来降低过拟合风险
• 将多个因果注意力模块堆叠成一个多头注意力模块

将文本分割成单词和子词 → 并将这些词元编码成向量表示（嵌入） → 从而为大语言模型提供输入数据。

本章将探讨大语言模型架构中的一个核心部分—— 注意力 （attention）机制

如图 3-1 所示。我们将大篇幅单独探讨注意力机制，并专注于其工作原理。

然后，我们将编写自注意力机制的其他部分代码，观察其运行过程，并构建一个生成文本的模型。

图 3-1　构建大语言模型的 3 个主要阶段。本章重点讨论第一阶段中的第 (2) 步：

• 实现注意力机制，这是大语言模型架构中的关键部分

本章将实现 4 种注意力机制的变体，如图 3-2 所示。

这些变体中的每一个都是在前一个的基础上逐步建立的，最终目的是实现一种紧凑、高效的多头注意力机制，并在第 4 章中将其集成到大语言模型的架构中。

图 3-本章将要实现的不同注意力机制。我们将从一个简化版本的自注意力机制开始，然后逐步加入可训练的权重。

因果注意力机制在自注意力的基础上增加了额外掩码，使得大语言模型可以一次生成一个单词。最后，多头注意力将注意力机制划分成多个头，从而使模型能够并行捕获输入数据的各种特征

目录

• 长序列建模中的问题
• 3.2　使用注意力机制捕捉数据依赖关系
• 通过自注意力机制关注输入的不同部分
• 3.4　实现带可训练权重的自注意力机制
  • 3.4.1　逐步计算注意力权重
  • 3.4.2　实现一个简化的自注意 Python 类
• 3.5　利用因果注意力隐藏未来词汇
  • 3.5.1　因果注意力的掩码实现
  • 3.5.2　利用 dropout 掩码额外的注意力权重
  • 3.5.3　实现一个简化的因果注意力类
• 3.6　将单头注意力扩展到多头注意力
  • 3.6.1　叠加多个单头注意力层
  • 3.6.2　通过权重划分实现多头注意力
• 3.7　小结

长序列建模中的问题

重新讲述："3.1　长序列建模中的问题"  
要求：
1. 以更利于程序员理解的方式来表达（比如多用代码来表达）
2. 代码中多添加注释；如果注释是英文的，转成中文；代码执行结果尽量完整打印出来
3. 说人话，多举例
4. 不要丢失原文信息
5. 更结构化的表达（借助 markdown 中的标题层级、列表缩进、代码块等表达形式）
6. 多画流程图或者架构图（尽量使用 ASCII Art 来画，不行再用 Mermaid 来画）

详见

• 1. 长序列建模中的问题：传统模型处理序列数据的局限性
• 1. 长序列建模中的问题

3.2　使用注意力机制捕捉数据依赖关系

尽管 RNN 在翻译短句时表现良好，但在处理较长文本时效果不佳，因为它无法直接访问输入中靠前的单词。这种方法的一个主要缺点是，RNN 在将信息传递给解码器之前，必须将整个编码后的输入存储到一个单独隐藏状态中

因此，研究人员在 2014 年为 RNN 开发了 Bahdanau 注意力机制 ，该机制对编码器-解码器 RNN 进行了修改，使得解码器在每个解码步骤中可以选择性地访问输入序列的不同部分，如图 3-5 所示。

图 3-5　通过使用注意力机制，网络的生成文本解码器部分可以有选择地访问所有输入词元。这意味着对于生成一个特定的输出词元，某些输入词元比其他输入词元更重要。这种重要性由注意力权重决定，我们将在后面计算这些权重。需要注意的是，这里展示的是注意力机制的基本概念，并未描述 Bahdanau 机制（一种 RNN 方法，但其超出了本书的讨论范畴）的具体实现

有趣的是，仅仅 3 年后，研究人员发现 RNN 并不是构建自然语言处理深度神经网络的必需架构，并提出了最初的 Transformer 架构（在第 1 章中讨论过），其中包括受 Bahdanau 注意力机制启发的自注意力机制。

在计算序列表示时，自注意力机制允许输入序列中的每个位置关注同一序列中的所有位置。自注意力机制是基于 Transformer 架构的当代大语言模型（如 GPT 系列）的关键组成部分。

本章的重点是编码并理解类 GPT 模型中使用的自注意力机制，如图 3-6 所示。在第 4 章中，我们将对大语言模型的其余部分进行编码。

图 3-6　自注意力是 Transformer 模型中的一种机制，它通过允许一个序列中的每个位置与同一序列中的其他所有位置进行交互并权衡其重要性，来计算出更高效的输入表示。本章将从头开始编码这种自注意力机制，之后在第 4 章中我们再实现类 GPT 大语言模型的其余部分

通过自注意力机制关注输入的不同部分

重新讲述："3.3　通过自注意力机制关注输入的不同部分"  
要求：
1. 以更利于程序员理解的方式来表达（比如多用代码来表达）
2. 原文中的代码一定要保留，代码中多添加注释；如果注释是英文的，转成中文；代码执行结果尽量完整打印出来
3. 说人话，多举例
4. 不要丢失原文信息
5. 更结构化的表达（借助 markdown 中的标题层级、列表缩进、代码块等表达形式）
6. 多画流程图或者架构图（尽量使用 ASCII Art 来画，不行再用 Mermaid 来画）

详见 3. 通过自注意力机制关注输入的不同部分

3.4　实现带可训练权重的自注意力机制

接下来，我们将实现在原始 Transformer 架构、GPT 模型和大多数其他流行的大语言模型中使用的自注意机制。这种自注意力机制也被称为 缩放点积注意力 （scaled dot-product attention）。图 3-13 展示了在实现整个大语言模型的过程中，自注意力机制是如何嵌入的。

图 3-13　之前，我们实现了一个简化的注意力机制，以理解注意力机制背后的基本原理。现在，我们将为此注意力机制添加可训练的权重。之后，我们将通过添加因果掩码和多头机制来扩展这种自注意力机制

如图 3-13 所示，带有可训练权重的自注意力机制是建立在先前概念之上的：我们希望将上下文向量计算为某个特定输入元素对于序列中所有输入向量的加权和。你会看到，带有可训练权重的自注意力机制与我们之前实现的基础自注意力机制只有些微的不同。

最显著的区别是这里引入了在模型训练期间更新的权重矩阵。这些可训练的权重矩阵至关重要，这样模型（特别是模型内部的注意力模块）才能学会产生“好的”上下文向量。（请注意，我们将在第 5 章中训练大语言模型。）

3.4.1 节和 3.4.2 节将讨论这种自注意力机制。首先，我们会像以前一样逐步编写代码。其次，我们会把代码组织成一个紧凑的 Python 类，以便能够导入到大语言模型架构中。

3.4.1　逐步计算注意力权重

本节将通过引入 3 个可训练的权重矩阵 、 和 ，一步一步地实现自注意力机制。这 3 个矩阵用于将嵌入的输入词元 分别映射为查询向量、键向量和值向量，如图 3-14 所示。

图 3-14　在实现具有可训练权重矩阵的自注意机制的第一步中，我们计算了输入元素 的查询向量（）、键向量（）和值向量（）。与之前类似，我们将第二个输入元素 指定为查询输入。查询向量 是通过第二个输入元素 与权重矩阵 之间的矩阵乘法得到的。同样，我们通过包含权重矩阵 和 的矩阵乘法得到键向量和值向量

之前，当我们通过注意力权重计算上下文向量 时，将第二个输入元素 定义为了查询。然后，我们将这一方法推广到了计算所有上下文向量 ，应用于 6 个词的输入句子“Your journey starts with one step.”

同样，为了便于说明，这里我们只计算一个上下文向量 。之后我们会修改这段代码来计算所有上下文向量。

首先，定义几个变量：

请注意，在类 GPT 模型中，输入和输出的维度通常是相同的，但为了便于理解计算过程，这里我们使用不同的输入维度（ d_in=3 ）和输出维度（ d_out=2 ）。

然后，初始化图 3-14 中的 3 个权重矩阵 、 和 ：

torch.manual_seed(123)
W_query = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)
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 ，以便在训练中更新这些矩阵。

接下来，计算查询向量、键向量和值向量：

query_2 = x_2 @ W_query
key_2 = x_2 @ W_key
value_2 = x_2 @ W_value
print(query_2)

因为我们通过 d_out 将对应的权重矩阵的列数设置为了 2，所以查询的输出结果是一个二维向量。

tensor([0.4306, 1.4551])

权重参数与注意力权重

在权重矩阵 中，“权重”是“权重参数”的简称，表示在训练过程中优化的神经网络参数。这与注意力权重是不同的。正如我们已经看到的，注意力权重决定了上下文向量对输入的不同部分的依赖程度（网络对输入的不同部分的关注程度）。

总之，权重参数是定义网络连接的基本学习系数，而注意力权重是动态且特定于上下文的值。

虽然目前我们的目标只是计算一个上下文向量 ，但仍然需要所有输入元素的键向量和值向量，因为它们参与了计算相对于查询 的注意力权重（参见图 3-14）。

可以通过矩阵乘法得到所有的键向量和值向量：

keys = inputs @ W_key
values = inputs @ W_value
print("keys.shape:", keys.shape)
print("values.shape:", values.shape)

从输出中可以看出，我们成功地将 6 个输入词元从三维空间映射到了二维嵌入空间。

keys.shape: torch.Size([6, 2])
values.shape: torch.Size([6, 2])

接下来是计算注意力分数，如图 3-15 所示。

图 3-15　注意力分数的计算是一种点积计算，与 3.3 节中使用的方法类似。不同之处在于，我们不是直接计算输入元素之间的点积，而是使用通过各自权重矩阵变换后的查询向量和键向量进行计算

首先，计算出注意力分数 ：

未归一化的注意力评分结果如下所示：

tensor(1.8524)

同样，可以通过矩阵乘法将这个计算推广到所有的注意力分数：

如你所见，经过快速检查，输出中的第二个元素与之前计算的 attn_score_22 一致。

tensor([1.2705, 1.8524, 1.8111, 1.0795, 0.5577, 1.5440])

现在，我们想要将注意力分数转换为注意力权重，如图 3-16 所示。我们通过缩放注意力分数并应用 softmax 函数来计算注意力权重。不过，此时是通过将注意力分数除以键向量的嵌入维度的平方根来进行缩放（取平方根在数学上等同于以 0.5 为指数进行幂运算）。

d_k = keys.shape[-1]
attn_weights_2 = torch.softmax(attn_scores_2 / d_k**0.5, dim=-1)
print(attn_weights_2)

图 3-16　在计算完注意力分数 后，下一步是使用 softmax 函数对这些分数进行归一化，以获得注意力权重

运行上述代码得到的注意力权重如下所示。

tensor([0.1500, 0.2264, 0.2199, 0.1311, 0.0906, 0.1820])

缩放点积注意力的原理

对嵌入维度进行归一化是为了避免梯度过小，从而提升训练性能。例如，在类 GPT 大语言模型中，嵌入维度通常大于 1000，这可能导致点积非常大，从而在反向传播时由于 softmax 函数的作用导致梯度非常小。当点积增大时，softmax 函数会表现得更像阶跃函数，导致梯度接近零。这些小梯度可能会显著减慢学习速度或使训练停滞。

因此，通过嵌入维度的平方根进行缩放解释了为什么这种自注意力机制也被称为缩放点积注意力机制。

现在，最后一步是计算上下文向量，如图 3-17 所示。

与计算上下文向量时对输入向量进行加权求和（参见 3.3 节）的方式类似，现在通过对值向量进行加权求和来计算上下文向量。在这里，注意力权重作为加权因子，用于权衡每个值向量的重要性。和之前一样，可以使用矩阵乘法一步获得输出结果：

context_vec_2 = attn_weights_2 @ values
print(context_vec_2)

图 3-17　在自注意力计算的最后一步，通过注意力权重将所有值向量进行加权求和，从而计算上下文向量

所生成的向量内容如下所示：

tensor([0.3061, 0.8210])

到目前为止，我们只计算了一个上下文向量 。在 3.4.2 节中，我们将扩展代码来计算输入序列中 到 的所有上下文向量。

为什么要用查询、键和值

在注意力机制中，“键”（key）、“查询”（query）和“值”（value）这些术语借用自信息检索和数据库领域，这些领域使用类似的概念来进行信息存储、搜索和检索。

查询 类似于数据库中的搜索查询。它代表了模型当前关注或试图理解的项（比如句子中的一个单词或词元）。查询用于探测输入序列中的其他部分，以确定对它们的关注程度。

键 类似于用于数据库索引和搜索的键。在注意力机制中，输入序列中的每个项（比如句子中的每个单词）都有一个对应的键。这些键用于与查询进行匹配。

在这种背景下， 值 类似于数据库中键-值对中的值。它表示输入项的实际内容或表示。一旦模型确定哪些键以及哪些输入部分与查询（当前关注的项）最相关，它就会检索相应的值。

3.4.2　实现一个简化的自注意 Python 类

到目前为止，我们已经完成了多个步骤来计算自注意力的输出。这些步骤主要是为了演示清晰，以便逐步了解每个环节。在实际操作中，为了实现第 4 章中的大语言模型，最好将这些代码组织成一个 Python 类，如代码清单 3-1 所示。

代码清单 3-1 　一个简化的自注意力类

import torch.nn as nn
class SelfAttention_v1(nn.Module):
    def __init__(self, d_in, d_out):
        super().__init__()
        self.W_query = nn.Parameter(torch.rand(d_in, d_out))
        self.W_key   = nn.Parameter(torch.rand(d_in, d_out))
        self.W_value = nn.Parameter(torch.rand(d_in, d_out))

    def forward(self, x):
        keys = x @ self.W_key
        queries = x @ self.W_query
        values = x @ self.W_value
        attn_scores = queries @ keys.T # omega
        attn_weights = torch.softmax(
            attn_scores / keys.shape[-1]**0.5, dim=-1
        )
        context_vec = attn_weights @ values
        return context_vec

在这段 PyTorch 代码中， SelfAttention_v1 是一个从 nn.Module 派生出来的类。 nn.Module 是 PyTorch 模型的一个基本构建块，它为模型层的创建和管理提供了必要的功能。

__init__ 方法初始化了可训练的权重矩阵（ W_query 、 W_key 和 W_value ），这些矩阵用于查询向量、键向量和值向量，每个矩阵将输入维度 d_in 转换为输出维度 d_out 。

在前向传播过程中，我们通过使用 forward 方法将查询向量和键向量相乘来计算注意力分数（ attn_scores ），然后使用 softmax 对这些分数进行归一化。最后，我们通过使用这些归一化的注意力分数对值向量进行加权来创建上下文向量。

可以通过以下方式来使用这个类：

torch.manual_seed(123)
sa_v1 = SelfAttention_v1(d_in, d_out)
print(sa_v1(inputs))

由于输入包含 6 个嵌入向量，因此我们会得到一个用于保存这 6 个上下文向量的矩阵：

tensor([[0.2996, 0.8053],
        [0.3061, 0.8210],
        [0.3058, 0.8203],
        [0.2948, 0.7939],
        [0.2927, 0.7891],
        [0.2990, 0.8040]], grad_fn=<MmBackward0>)

通过快速检查，可以看到第 2 行（ [0.3061, 0.8210] ）与 3.4.1 节中的 context_vec_2 内容相符。

图 3-18 总结了我们刚刚实现的自注意力机制。

图 3-18　在自注意力机制中，我们用 3 个权重矩阵 、 和 来变换输入矩阵 中的输入向量。新方法根据所得查询矩阵（）和键矩阵（）计算注意力权重矩阵。然后，使用注意力权重矩阵和值矩阵（）计算上下文向量（）。为了视觉清晰，我们关注具有 个词元的单个输入文本，而不是一批多个输入。因此，在这种情况下，三维输入张量被简化为二维矩阵。这种方法允许更直观地可视化和理解所涉及的过程。为了与后面的图保持一致，注意力矩阵中的值不代表真正的注意力权重（该图中的数值被截断为小数点后两位，以减少视觉混乱。每行中的值加起来应为 1.0 或 100%）

自注意力机制包含了可训练的权重矩阵 、 和 。这些矩阵将输入数据转换为查询向量、键向量和值向量，这些组件在注意力机制中至关重要。随着模型在训练中接触更多数据，它会调整这些可训练的权重，后续章节会对此进行介绍。

可以通过使用 PyTorch 的 nn.Linear 层来进一步优化 SelfAttention_v1 的实现，当偏置单元被禁用时， nn.Linear 层可以有效地执行矩阵乘法。相比手动实现 nn.Parameter(torch.rand(...)) ，使用 nn.Linear 的一个重要优势是它提供了优化的权重初始化方案，从而有助于模型训练的稳定性和有效性，如代码清单 3-2 所示。

代码清单 3-2 　一个使用 PyTorch 线性层的自注意力类

class SelfAttention_v2(nn.Module):
    def __init__(self, d_in, d_out, qkv_bias=False):
        super().__init__()
        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)

    def forward(self, x):
        keys = self.W_key(x)
        queries = self.W_query(x)
        values = self.W_value(x)
        attn_scores = queries @ keys.T
        attn_weights = torch.softmax(
            attn_scores / keys.shape[-1]**0.5, dim=-1
        )
        context_vec = attn_weights @ values
        return context_vec

可以像使用 SelfAttention_v1 一样使用 SelfAttention_v2 ：

torch.manual_seed(789)
sa_v2 = SelfAttention_v2(d_in, d_out)
print(sa_v2(inputs))

输出结果如下所示：

tensor([[-0.0739,  0.0713],
        [-0.0748,  0.0703],
        [-0.0749,  0.0702],
        [-0.0760,  0.0685],
        [-0.0763,  0.0679],
        [-0.0754,  0.0693]], grad_fn=<MmBackward0>)

请注意， SelfAttention_v1 和 SelfAttention_v2 因为使用了不同的初始权重矩阵而给出了不同的输出，这是由 nn.Linear 使用了一个更复杂的权值初始化方案所导致的。

练习 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 对象中，这样两个对象就会产生相同的结果。

你的任务是将 SelfAttention_v2 实例中的权重正确地分配给 SelfAttention_v1 实例。为此，需要了解两个版本中的权重之间的关系。（提示： nn.Linear 以转置的形式存储权重矩阵。）完成分配后，应该能够观察到这两个实例产生了相同的输出。

接下来，我们将改进自注意力机制，重点是在机制中引入因果机制和多头机制。因果机制的作用是调整注意力机制，防止模型访问序列中未来的信息，这在语言建模等任务中尤为重要，因为每个词的预测只能依赖之前出现的词。

多头机制涉及将注意力机制分成多个“头”。每个头会学习数据的不同特征，使模型能够在不同的位置同时关注来自不同表示子空间的信息。这能够提升模型在复杂任务中的性能。

3.5　利用因果注意力隐藏未来词汇

对于许多大语言模型任务，你希望自注意力机制在预测序列中的下一个词元时仅考虑当前位置之前的词元。因果注意力（也称为 掩码注意力 ）是一种特殊的自注意力形式。它限制模型在处理任何给定词元时，只能基于序列中的先前和当前输入来计算注意力分数，而标准的自注意力机制可以一次性访问整个输入序列。

现在，我们将通过修改标准自注意力机制来创建 因果注意力 机制，这是在后续章节中开发大语言模型的关键步骤。要在类 GPT 模型中实现这一点，对于每个处理的词元，需要掩码当前词元之后的后续词元，如图 3-19 所示。我们会掩码对角线以上的注意力权重，并归一化未掩码的注意力权重，使得每一行的权重之和为 1。稍后，我们将通过代码来实现这一掩码和归一化过程。

图 3-19　在因果注意力机制中，我们掩码了对角线以上的注意力权重，以确保在计算上下文向量时，大语言模型无法访问未来的词元。例如，对于第 2 行的单词“journey”，仅保留当前词（“journey”）和之前词（“Your”）的注意力权重

3.5.1　因果注意力的掩码实现

接下来我们将在代码中实现因果注意力掩码。如图 3-20 所示，为了实现应用因果注意力掩码的步骤并得到掩码后的注意力权重，我们将使用前面介绍的注意力分数和权重来编写因果注意力机制代码。

图 3-20　在因果注意力中，获得掩码后的注意力权重矩阵的一种方法是对注意力分数应用 softmax 函数，将对角线以上的元素清零，并对所得矩阵进行归一化

在第 (1) 步中，按照之前的方法，通过 softmax 函数计算注意力权重：

结果得到如下注意力权重：

tensor([[0.1921, 0.1646, 0.1652, 0.1550, 0.1721, 0.1510],
        [0.2041, 0.1659, 0.1662, 0.1496, 0.1665, 0.1477],
        [0.2036, 0.1659, 0.1662, 0.1498, 0.1664, 0.1480],
        [0.1869, 0.1667, 0.1668, 0.1571, 0.1661, 0.1564],
        [0.1830, 0.1669, 0.1670, 0.1588, 0.1658, 0.1585],
        [0.1935, 0.1663, 0.1666, 0.1542, 0.1666, 0.1529]],
       grad_fn=<SoftmaxBackward0>)

可以使用 PyTorch 的 tril 函数来实现第 (2) 步，该函数可以创建一个对角线以上元素为 0 的掩码：

context_length = attn_scores.shape[0]
mask_simple = torch.tril(torch.ones(context_length, context_length))
print(mask_simple)

生成的掩码矩阵如下所示：

tensor([[1., 0., 0., 0., 0., 0.],
        [1., 1., 0., 0., 0., 0.],
        [1., 1., 1., 0., 0., 0.],
        [1., 1., 1., 1., 0., 0.],
        [1., 1., 1., 1., 1., 0.],
        [1., 1., 1., 1., 1., 1.]])

现在，可以把这个掩码矩阵和注意力权重矩阵相乘，使对角线上方的值变为 0：

masked_simple = attn_weights*mask_simple
print(masked_simple)

如你所见，对角线上方的元素已被成功地归 0：

tensor([[0.1921, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.2041, 0.1659, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.2036, 0.1659, 0.1662, 0.0000, 0.0000, 0.0000],
        [0.1869, 0.1667, 0.1668, 0.1571, 0.0000, 0.0000],
        [0.1830, 0.1669, 0.1670, 0.1588, 0.1658, 0.0000],
        [0.1935, 0.1663, 0.1666, 0.1542, 0.1666, 0.1529]],
       grad_fn=<MulBackward0>)

第 (3) 步是重新归一化注意力权重，使每一行的总和再次为 1。可以通过将每行中的每个元素除以每行中的和来实现这一点：

row_sums = masked_simple.sum(dim=-1, keepdim=True)
masked_simple_norm = masked_simple / row_sums
print(masked_simple_norm)

结果是一个注意力权重矩阵，其中对角线以上的注意力权重已被归 0，每一行之和为 1。

tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.5517, 0.4483, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.3800, 0.3097, 0.3103, 0.0000, 0.0000, 0.0000],
        [0.2758, 0.2460, 0.2462, 0.2319, 0.0000, 0.0000],
        [0.2175, 0.1983, 0.1984, 0.1888, 0.1971, 0.0000],
        [0.1935, 0.1663, 0.1666, 0.1542, 0.1666, 0.1529]],
       grad_fn=<DivBackward0>)

信息泄露

当我们应用掩码并重新归一化注意力权重时，初看起来，未来的词元（打算掩码的）可能仍然会影响当前的词元，因为它们的值会参与 softmax 计算。然而，关键的见解是，在掩码后重新归一化时，我们实际上是在对一个较小的子集重新计算 softmax（因为被掩码的位置不参与 softmax 计算）。

softmax 函数在数学上的优雅之处在于，尽管最初所有位置都在分母中，但掩码和重新归一化之后，被掩码的位置的效果被消除——它们不会以任何实际的方式影响 softmax 分数。

简而言之，掩码和重新归一化之后，注意力权重的分布就像最初仅在未掩码的位置计算一样。这保证了不会有来自未来或其他被掩码的词元的信息泄露。

尽管可以在此时完成对因果注意力的实现，但我们仍然可以进行改进。让我们利用 softmax 函数的数学特性，以更少的步骤更高效地计算掩码后的注意力权重，如图 3-21 所示。

图 3-21　在因果注意力中，获得掩码后的注意力权重矩阵的一种更有效的方法是在应用 softmax 函数之前将注意力分数用负无穷大值进行掩码

softmax 函数会将其输入转换为一个概率分布。当输入中出现负无穷大值（）时，softmax 函数会将这些值视为零概率。（从数学角度来看，这是因为 无限接近于 0。）

可以通过创建一个对角线以上是 1 的掩码，并将这些 1 替换为负无穷大（ -inf ）值，来实现这种更高效的掩码“方法”：

mask = torch.triu(torch.ones(context_length, context_length), diagonal=1)
masked = attn_scores.masked_fill(mask.bool(), -torch.inf)
print(masked)

这将产生以下掩码矩阵：

tensor([[0.2899,   -inf,   -inf,   -inf,   -inf,   -inf],
        [0.4656, 0.1723,   -inf,   -inf,   -inf,   -inf],
        [0.4594, 0.1703, 0.1731,   -inf,   -inf,   -inf],
        [0.2642, 0.1024, 0.1036, 0.0186,   -inf,   -inf],
        [0.2183, 0.0874, 0.0882, 0.0177, 0.0786,   -inf],
        [0.3408, 0.1270, 0.1290, 0.0198, 0.1290, 0.0078]],
       grad_fn=<MaskedFillBackward0>)

现在只需要对这些掩码结果应用 softmax 函数，就可以完成整个过程：

attn_weights = torch.softmax(masked / keys.shape[-1]**0.5, dim=1)
print(attn_weights)

从输出结果来看，每行中的值总和为 1，因此不需要再进行额外的归一化处理：

tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.5517, 0.4483, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.3800, 0.3097, 0.3103, 0.0000, 0.0000, 0.0000],
        [0.2758, 0.2460, 0.2462, 0.2319, 0.0000, 0.0000],
        [0.2175, 0.1983, 0.1984, 0.1888, 0.1971, 0.0000],
        [0.1935, 0.1663, 0.1666, 0.1542, 0.1666, 0.1529]],
       grad_fn=<SoftmaxBackward0>)

现在可以利用修改后的注意力权重通过 context_vec = attn_weights @ values 计算上下文向量了，就像 3.4 节中所描述的那样。然而，在此之前，我们将先讨论一种对因果注意力机制的小调整，这在训练大语言模型时可以有效减少过拟合。

3.5.2　利用 dropout 掩码额外的注意力权重

dropout 是深度学习中的一种技术，通过在训练过程中随机忽略一些隐藏层单元来有效地“丢弃”它们。这种方法有助于减少模型对特定隐藏层单元的依赖，从而避免过拟合。需要强调的是，dropout 仅在训练期间使用，训练结束后会被取消。

在 Transformer 架构中，一些包括 GPT 在内的模型通常会在两个特定时间点使用注意力机制中的 dropout：一是计算注意力权重之后，二是将这些权重应用于值向量之后。如图 3-22 所示，我们将在计算注意力权重之后应用 dropout 掩码，因为这是实践中更常见的做法。

图 3-22　利用因果注意力掩码（左上），我们应用一个额外的 dropout 掩码（右上）来将额外的注意力权重置 0，以减少训练期间的过拟合

下面的代码示例中使用了 50% 的 dropout 率，这意味着掩码一半的注意力权重。（当我们在接下来的章节中训练 GPT 模型时，将使用较低的 dropout 率，比如 10% 或 20%。）为了便于操作，我们首先将 PyTorch 的 dropout 实现应用于一个由 1 组成的 6 × 6 张量：

如你所见，大约有一半的值被置 0 了：

tensor([[2., 2., 0., 2., 2., 0.],
        [0., 0., 0., 2., 0., 2.],
        [2., 2., 2., 2., 0., 2.],
        [0., 2., 2., 0., 0., 2.],
        [0., 2., 0., 2., 0., 2.],
        [0., 2., 2., 2., 2., 0.]])

在对注意力权重矩阵应用 50% 的 dropout 率时，矩阵中有一半的元素会随机被置为 0。为了补偿减少的活跃元素，矩阵中剩余元素的值会按 1/0.5 = 2 的比例进行放大。这种放大对于维持注意力权重的整体平衡非常重要，可以确保在训练和推理过程中，注意力机制的平均影响保持一致。

现在，对注意力权重矩阵进行 dropout 操作：

torch.manual_seed(123)
print(dropout(attn_weights))

在处理后的注意力权重矩阵中，部分元素已被置为 0，其余元素则被重新缩放：

tensor([[2.0000, 0.0000, 0 .0000, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.7599, 0.6194, 0.6206, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.4921, 0.4925, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.3966, 0.0000, 0.3775, 0.0000, 0.0000],
        [0.0000, 0.3327, 0.3331, 0.3084, 0.3331, 0.0000]],
       grad_fn=<MulBackward0>

请注意，由于操作系统的差异，最终的 dropout 输出可能会有所不同。有关这种不一致性的详细信息，请查看 PyTorch 的 issue tracker。

理解了因果注意力和 dropout 掩码之后，现在我们可以开发一个简洁的 Python 类。这个类的目的是高效地应用这两种技术。

3.5.3　实现一个简化的因果注意力类

现在我们将把因果注意力和 dropout 修改应用到 3.4 节中开发的 SelfAttention Python 类中。这个类将成为开发 多头注意力 的基础，而多头注意力是我们最终实现的注意力类。

但在开始之前，需要确保代码可以处理包含多个样本的批次，以便 CausalAttention 类能够支持第 2 章中实现的数据加载器产生的批量输出。

为简单起见，可以通过复制输入文本示例来模拟批量输入：

这将生成一个三维张量，其中包含两个输入文本，每个文本有 6 个词元，每个词元是一个三维的嵌入向量。

torch.Size([2, 6, 3])

代码清单 3-3 中的 CausalAttention 类与先前实现的 SelfAttention 类类似，唯一的不同是增加了 dropout 和因果掩码组件。

代码清单 3-3 　一个简化的因果注意力类

虽然此时所有新增的代码行都应该是熟悉的，但我们在 __init__ 方法中增加了一个 self.register_buffer() 调用。虽然在 PyTorch 中使用 register_buffer 并非所有情况下都是必需的，但在这里具有一些优势。例如，当我们在大语言模型中使用 CausalAttention 类时，缓冲区会与模型一起自动移动到适当的设备（CPU 或 GPU），这在训练大语言模型时非常重要。这意味着我们无须手动确保这些张量与模型参数在同一设备上，从而避免了设备不匹配的错误。

可以按照之前使用 SelfAttention 类的方式来使用 CausalAttention 类：

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)

最终生成的上下文向量是一个三维张量，其中每个词元现在用二维嵌入来表示。

context_vecs.shape: torch.Size([2, 6, 2])

图 3-23 总结了目前我们所取得的进展。我们集中讨论了神经网络中因果注意力的概念和实现。接下来，我们将基于这一概念，开发一个并行实现了多个因果注意力机制的多头注意力模块。

图 3-23　到目前为止，我们完成了以下步骤：从一个简化的注意力机制开始，添加了可训练的权重，然后引入了因果注意力掩码。接下来，我们将扩展因果注意力机制并编写多头注意力模块，以在我们的大语言模型中使用

3.6　将单头注意力扩展到多头注意力

在本节中，我们将进行最后一步操作，即把先前实现的因果注意力类扩展到多个头上。这也被称为 多头注意力 。

“多头”这一术语指的是将注意力机制分成多个“头”，每个“头”独立工作。在这种情况下，单个因果注意力模块可以被看作单头注意力，因为它只有一组注意力权重按顺序处理输入。

我们将从因果注意力扩展到多头注意力。首先，我们将直观地通过堆叠多个 CausalAttention 模块来构建多头注意力模块。然后，我们将用一种更复杂但计算上更高效的方式来实现这个多头注意力模块。

3.6.1　叠加多个单头注意力层

在实际操作中，实现多头注意力需要构建多个自注意力机制的实例（参见图 3-18），每个实例都有其独立的权重，然后将这些输出进行合成。虽然这种方法的计算量可能会非常大，但它对诸如基于 Transformer 的大语言模型之类的模型的复杂模式识别是非常重要的。

图 3-24 展示了多头注意力模块的结构，它是由图 3-18 所示的多个单头注意力模块依次叠加在一起组成的。

图 3-24　多头注意力模块包含两个堆叠在一起的单头注意力模块。因此，我们不是使用一个单一的矩阵 来计算值矩阵，而是在一个有两个头的多头注意模块中，现在有两个值权重矩阵： 和 。这同样适用于其他的权重矩阵，比如 和 。我们得到了两组上下文向量 和 ，最终可以将它们合并成一个单一的上下文向量矩阵

正如前面提到的，多头注意力的主要思想是多次（并行）运行注意力机制，每次使用学到的不同的线性投影——这些投影是通过将输入数据（比如注意力机制中的查询向量、键向量和值向量）乘以权重矩阵得到的。在代码中，可以通过实现一个简单的 MultiHeadAttentionWrapper 类来达到这一目标， MultiHeadAttentionWrapper 类堆叠了多个之前实现的 CausalAttention 模块实例，如代码清单 3-4 所示。

代码清单 3-4 　一个实现多头注意力的封装类

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)

如果采用这个具有两个注意力头（ num_heads=2 ）以及 CausalAttention 输出维度为 d_out=2 的 MultiHeadAttentionWrapper 类，那么我们就会得到一个四维的上下文向量（ d_out*num_heads=4 ），如图 3-25 所示。

图 3-25　使用 MultiHeadAttentionWrapper ，我们指定了注意力头的数量（ num_heads ）。如果设置 num_heads=2 ，那么我们就会得到一个具有两组上下文向量矩阵的张量。在每个上下文向量矩阵中，行表示对应于词元的上下文向量，列则对应于通过 d_out=4 指定的嵌入维度。我们沿着列维度连接这些上下文向量矩阵。由于我们有两个注意力头并且嵌入维度为 2，因此最终的嵌入维度是 2 × 2 = 4

为了更详细地说明这一点，可以像之前使用 CausalAttention 类一样使用 MultiHeadAttentionWrapper 类：

torch.manual_seed(123)
context_length = batch.shape[1] # 这是词元的数量
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=<CatBackward0>)
context_vecs.shape: torch.Size([2, 6, 4])

结果中的 context_vecs 张量的第一维是 2，因为我们有两个输入文本（输入文本是重复的，所以这些上下文向量完全相同）。第二维表示每个输入中的 6 个词元。第三维表示每个词元的四维嵌入。

练习 3.2　返回二维嵌入向量

更改 MultiHeadAttentionWrapper(..., num_heads=2) 调用的输入参数，使输出上下文向量是二维而不是四维，同时保持设置 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 类会将多头功能整合到一个类内。它通过重新调整投影后的查询张量、键张量和值张量的形状，将输入分为多个头，然后在计算注意力后合并这些头的结果。

在进一步讨论之前，先来看一下 MultiHeadAttention 类，如代码清单 3-5 所示。

代码清单 3-5 　一个高效的多头注意力类

尽管 MultiHeadAttention 类中的张量重塑（.view ）和转置（.transpose ）在数学上看起来很复杂，但 MultiHeadAttention 类实现的概念仍与之前的 MultiHeadAttentionWrapper 类相同。

在整体层面上，在之前的 MultiHeadAttentionWrapper 中，我们堆叠了多个单头注意力层，并将其合并成一个多头注意力层。 MultiHeadAttention 类采用了一种综合的方法。它从一个多头注意力层开始，然后在内部将这个层分割成单独的注意力头，如图 3-26 所示。

图 3-26　在具有两个注意力头的 MultiHeadAttentionWrapper 类中，我们初始化了两个权重矩阵 和 ，并计算了两个查询矩阵 和 （上）。在 MultiHeadAttention 类中，我们初始化了一个更大的权重矩阵 ，并只与输入矩阵进行一次矩阵乘法操作，得到一个查询矩阵 ，然后将查询矩阵分割成了 和 （下）。对键矩阵和值矩阵的操作与之类似，为了减少视觉混乱，这里没有展示

在 PyTorch 中，通过使用 .view 方法进行张量重塑以及使用 .transpose 方法进行张量转置，我们实现了对查询张量、键张量和值张量的分割。输入首先经过线性层进行变换（针对查询矩阵、键矩阵和值矩阵），然后被重塑为多个头。

关键操作是将 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) 的形状。这种转置对于正确对齐不同头的查询矩阵、键矩阵和值矩阵，以及有效地执行批处理矩阵乘法至关重要。

为了说明这个批处理矩阵乘法，假设有下面所列的张量的例子：

现在，在原始的张量和转置后的张量之间执行批处理矩阵乘法，其中我们转置了最后两个维度，即 num_tokens 和 head_dim ：

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 ）之间进行，并对每个头重复这一操作。

例如，上述方法提供了一种更简化的方式来单独计算每个头的矩阵乘法：

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]])

在 MultiHeadAttention 中，计算完注意力权重和上下文向量后，将所有头的上下文向量转置为 (b, num_tokens, num_heads, head_dim) 的形状。这些向量接着会被重塑（展平）为 (b, num_tokens, d_out) 的形状，从而有效地整合所有头的输出。

此外，我们在 MultiHeadAttention 中添加了一个输出投影层（ self.out_proj ），这是在合并多个头之后的步骤，而 CausalAttention 类中并不存在这个层。这个输出投影层并不是必需的（更多细节参见附录 B），但它在许多大语言模型架构中被广泛使用，这就是我们在这里添加它以保持完整性的原因。

尽管 MultiHeadAttention 类因额外的张量重塑和转置显得比 MultiHeadAttentionWrapper 更复杂，但它的效率更高。原因是我们只需进行一次矩阵乘法来计算键矩阵，例如， keys = self.W_key(x) （查询矩阵和值矩阵也是如此）。在 MultiHeadAttentionWrapper 中，我们需要对每个注意力头重复进行这种矩阵乘法，而矩阵乘法是计算资源消耗较大的操作之一。

MultiHeadAttention 类的用法与我们之前实现的 SelfAttention 类和 CausalAttention 类类似：

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=<ViewBackward0>)
context_vecs.shape: torch.Size([2, 6, 2])

现在，我们已经实现了将在实现和训练大语言模型时使用的 MultiHeadAttention 类。需要注意的是，尽管代码功能齐全，但为了保持输出的可读性，我们使用了相对较小的嵌入维度和注意力头数量。

相比之下，最小的 GPT-2 模型（参数量为 1.17 亿）有 12 个注意力头，上下文向量嵌入维度为 768，而最大的 GPT-2 模型（参数量为 15 亿）有 25 个注意力头，上下文向量嵌入维度为 1600。请注意，在 GPT 模型中，词元输入和上下文嵌入的嵌入维度是相同的（ d_in = d_out ）。

练习 3.3　初始化 GPT-2 大小的注意力模块

使用 MultiHeadAttention 类初始化一个多头注意力模块，该模块应具有与最小的 GPT-2 模型相同数量的注意力头（12 个）。同时，确保使用与 GPT-2 相似的输入和输出嵌入维度（768）。请注意，最小的 GPT-2 模型支持 1024 个词元的上下文长度。

3.7　小结

• 注意力机制可以将输入元素转换为增强的上下文向量表示，这些表示涵盖了关于所有输入的信息。
• 自注意力机制通过对输入进行加权求和来计算上下文向量表示。
• 在简化的注意力机制中，注意力权重是通过点积计算得出的。
• 点积是两个向量的元素逐个相乘并将这些乘积相加的一种简洁计算方法。
• 尽管矩阵乘法不是必需的，但它可以通过替代嵌套的 for 循环使计算更高效、更紧凑。
• 在用于大语言模型的自注意力机制（也被称为“缩放点积注意力”）中，我们引入了可训练的权重矩阵来计算输入的中间变换：查询矩阵、值矩阵和键矩阵。
• 在处理从左到右读取和生成文本的大语言模型时，我们会添加一个因果注意力掩码，以防止模型访问未来的词元。
• 除了使用因果注意力掩码将注意力权重置 0，还可以添加 dropout 掩码来减少大语言模型中的过拟合。
• 基于 Transformer 的大语言模型中的注意力模块涉及多个因果注意力实例，这被称为“多头注意力”。
• 可以通过堆叠多个因果注意力模块实例来创建多头注意力模块。
• 创建多头注意力模块的一种更高效的方法是使用批量矩阵乘法。