从零开始构建一个大语言模型-第三章第四节

第三章目录

3.1 长序列建模的问题
3.2 用注意力机制捕捉数据依赖关系
3.3 用自注意力关注输入的不同部分
3.4 用可训练的权重实现自注意力
3.5 用因果注意力隐藏未来的词
3.6 将单头注意力扩展到多头注意力‘

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

我们的下一步是实现最初在Transformer架构、GPT模型以及其他大多数流行大语言模型（LLM）中使用的自注意力机制。这种自注意力机制也被称为缩放点积注意力。图3.13展示了这种自注意力机制是如何融入到大语言模型实现的更广泛背景中的。

图 3.13：之前，我们编写了一个简化的注意力机制代码，以理解注意力机制背后的基本原理。现在，我们要为这个注意力机制添加可训练的权重。之后，我们会通过添加因果掩码和多头机制来扩展这种自注意力机制。

如图 3.13 所示，带有可训练权重的自注意力机制是建立在之前的概念之上的：我们希望将上下文向量计算为特定输入元素对应输入向量的加权和。正如你将看到的，与我们之前编写的基本自注意力机制相比，只有细微的差别。
最显著的差别在于引入了在模型训练期间会更新的权重矩阵。这些可训练的权重矩阵至关重要，这样模型（具体而言，是模型内部的注意力模块）才能学会生成 “优质” 的上下文向量。（我们将在第 5 章对大语言模型进行训练。）
我们将在两个小节中探讨这种自注意力机制。首先，我们会像之前一样逐步编写代码。其次，我们会将代码整理成一个简洁的 Python 类，以便可以导入到大语言模型架构中。

3.4.1 逐步计算注意力权重

我们将通过引入三个可训练的权重矩阵 W q W_q Wq​、 W k W_k Wk​ 和 W v W_v Wv​ 来逐步实现自注意力机制。如图 3.14 所示，这三个矩阵分别用于将嵌入后的输入标记 x ( i ) x^{(i)} x(i) 投影为查询（query）、键（key）和值（value）向量。

图3.14：在使用可训练权重矩阵的自注意力机制的第一步中，我们为输入元素 x x x 计算查询（ q q q）、键（ k k k）和值（ v v v）向量。和前面章节类似，我们指定第二个输入 x ( 2 ) x^{(2)} x(2) 作为查询输入。查询向量 q ( 2 ) q^{(2)} q(2) 是通过输入 x ( 2 ) x^{(2)} x(2) 与权重矩阵 W q W_q Wq​ 进行矩阵乘法得到的。同样地，我们通过涉及权重矩阵 W k W_k Wk​ 和 W v W_v Wv​ 的矩阵乘法得到键向量和值向量。
之前，在计算简化的注意力权重以得到上下文向量 z⁽²⁾时，我们把第二个输入元素 x⁽²⁾ 定义为查询。之后，我们对其进行了推广，为包含六个单词的输入句子 “Your journey starts with one step.” 计算了所有的上下文向量 z⁽¹⁾ … z^(T)。
同样地，为了便于说明，我们这里先只计算一个上下文向量 z⁽²⁾。然后我们会修改这段代码，使其能够计算所有的上下文向量。 让我们从定义几个变量开始：

x_2 = inputs[1] 
d_in = inputs.shape[1] 
d_out = 2

请注意，在类似 GPT 的模型中，输入和输出维度通常是相同的。但为了更便于理解计算过程，我们这里将使用不同的输入维度（ d i n = 3 d_{in}=3 din​=3）和输出维度（ d o u t = 2 d_{out}=2 dout​=2）。 接下来，我们要初始化图 3.14 中所示的三个权重矩阵 W q W_q Wq​、 W k W_k Wk​ 和 W v W_v Wv​：

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

权重参数与注意力权重

在权重矩阵 里，“权重” 是 “权重参数” 的简称，这些是神经网络在训练过程中需要优化的值。不要把它和注意力权重弄混。正如我们之前了解到的，注意力权重决定了上下文向量在多大程度上依赖于输入的不同部分（也就是网络对输入不同部分的关注程度）。
总之，权重参数是定义网络连接的基本且可学习的系数，而注意力权重则是动态的、与上下文相关的值。
尽管我们当前的临时目标只是计算一个上下文向量 z⁽²⁾，但我们仍然需要所有输入元素的键向量和值向量，因为在计算相对于查询 q⁽²⁾ 的注意力权重时会用到它们（见图 3.14）。 我们可以通过矩阵乘法得到所有的键向量和值向量：

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

从输出结果我们可以看出，我们成功地将六个输入标记从三维嵌入空间投影到了二维嵌入空间。

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

第二步是计算注意力分数，如图3.15所示。

图 3.15：注意力分数的计算是一种点积计算，与我们在 3.3 节简化的自注意力机制中使用的计算方式类似。这里的新特点在于，我们并非直接计算输入元素之间的点积，而是使用通过各自的权重矩阵对输入进行变换后得到的查询向量和键向量来进行计算。
首先，让我们来计算注意力分数 ω \omega ω ₂₂。

keys_2 = keys[1] 
attn_score_22 = query_2.dot(keys_2) 
print(attn_score_22)

未归一化注意力分数的计算结果如下。

tensor(1.8524)

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

attn_scores_2 = query_2 @ keys.T 
print(attn_scores_2)

正如我们所见，快速检查一下，输出中的第二个元素与我们之前计算的注意力分数 ω \omega ω₂₂是相符的。

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：在计算出注意力分数 ω \omega ω之后，下一步是使用 softmax 函数对这些分数进行归一化处理，以得到注意力权重 α \alpha α。
得到的注意力权重如下：

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

缩放点积注意力背后的原理

通过嵌入维度大小进行归一化的原因是为了避免梯度过小，从而提升训练性能。例如，当增大嵌入维度时（对于类似 GPT 的大语言模型，嵌入维度通常大于 1000），较大的点积在经过 softmax 函数处理后，会在反向传播过程中导致梯度变得极小。随着点积的增大，softmax 函数的表现会更接近阶跃函数，导致梯度趋近于零。这些极小的梯度会显著减缓学习速度，甚至导致训练停滞。 通过嵌入维度的平方根进行缩放，这就是这种自注意力机制也被称为缩放点积注意力的原因。

现在，最后一步是计算上下文向量，如图 3.17 所示

图 3.17：在自注意力计算的最后一步，我们通过注意力权重组合所有的值向量来计算上下文向量。

类似于我们在计算上下文向量时将其作为输入向量的加权和（见 3.3 节），现在我们将上下文向量计算为值向量的加权和。在这里，注意力权重充当了一个加权因子，用于衡量每个值向量的相对重要性。同样和之前一样，我们可以使用矩阵乘法一步得出输出结果。

context_vec_2 = attn_weights_2 @ values 
print(context_vec_2)

得到的向量内容如下：

tensor([0.3061, 0.8210])

到目前为止，我们仅计算了一个上下文向量z⁽²⁾。接下来，我们将对代码进行改进，以计算输入序列中所有的上下文向量，即从z⁽¹⁾到z^(T)。

[!NOTE] 为什么是查询、键和值？
注意力机制中的 “键”“查询” 和 “值” 这些术语借鉴自信息检索和数据库领域，在这些领域中，类似的概念被用于存储、搜索和检索信息。
查询（Query）：类似于数据库中的搜索查询。它代表模型当前关注或试图理解的项目（例如，句子中的一个单词或标记）。查询用于探测输入序列的其他部分，以确定应该对它们给予多少关注。
键（Key）：类似于数据库中用于索引和搜索的键。在注意力机制中，输入序列中的每个项目（例如，句子中的每个单词）都有一个关联的键。这些键用于与查询进行匹配。
值（Value）：在这种情况下，它类似于数据库中键 - 值对中的值。它代表输入项目的实际内容或表示。一旦模型确定哪些键（以及输入的哪些部分）与查询（当前关注的项目）最相关，它就会检索相应的值。

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

至此，我们已经历经诸多步骤来计算自注意力机制的输出。我们这样做主要是为了便于讲解，这样我们可以一步一步地进行推导。实际上，考虑到下一章中对大语言模型（LLM）的实现，将这些代码组织成一个 Python 类会很有帮助，如下所示。

在这段 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))

由于输入包含六个嵌入向量，这会得到一个存储六个上下文向量的矩阵。

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

快速检查一下，你会注意到第二行（[0.3061, 0.8210]）与上一节中的上下文向量 的内容相匹配。图 3.18 总结了我们刚刚实现的自注意力机制。
自注意力机制涉及可训练的权重矩阵 、 和 。这些矩阵分别将输入数据转换为查询、键和值，它们是注意力机制的关键组成部分。正如我们将在后续章节中看到的，在训练过程中，随着模型接触到更多的数据，它会调整这些可训练的权重。
我们可以通过利用 PyTorch 的 nn.Linear 层来进一步改进 SelfAttention_v1 的实现。当禁用偏置单元时，nn.Linear 层能有效地执行矩阵乘法。此外，使用 nn.Linear 层的一个显著优势是…（你提供的内容此处未完整）

图 3.18：在自注意力机制中，我们使用三个权重矩阵W_q 、W_k和W_v对输入矩阵X中的输入向量进行变换。然后，根据得到的查询向量（Q）和键向量（K）计算注意力权重矩阵。接着，利用注意力权重和值向量（V）计算上下文向量（Z）。为了视觉上的清晰，我们聚焦于一个包含n个词元的单一输入文本，而非多个输入组成的批次。因此，在这种情况下，三维输入张量被简化为二维矩阵。这种方法能让我们更直观地可视化和理解相关过程。为了与后续图表保持一致，注意力矩阵中的值并未展示真实的注意力权重。（为减少视觉干扰，本图中的数字保留两位小数。每行的值相加应等于 1.0 或 100%。）
使用 nn.Linear 而非手动实现 nn.Parameter(torch.rand(...)) 的优势在于，nn.Linear 具备优化后的权重初始化方案，有助于实现更稳定、高效的模型训练。

你可以像使用 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=)

请注意，SelfAttention_v1 和 SelfAttention_v2 会产生不同的输出，这是因为它们在权重矩阵上采用了不同的初始权重。具体而言，nn.Linear 使用了更为复杂精细的权重初始化方案。

[!NOTE] 练习 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 以转置形式存储权重矩阵。）完成赋值后，你应该会观察到两个实例产生相同的输出。

接下来，我们将对自注意力机制进行改进，尤其专注于融入因果性和多头机制。因果性方面需要对注意力机制进行修改，以防止模型获取序列中的未来信息，这对于像语言建模这样的任务至关重要，在语言建模中，每个单词的预测应该仅依赖于之前的单词。

多头部分则是将注意力机制拆分为多个 “头”。每个头学习数据的不同方面，使模型能够同时关注来自不同位置的不同表示子空间的信息。这有助于提升模型在复杂任务中的性能。