从零开始构建一个大语言模型-第三章第三节（勘误重发）

第三章目录

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

3.3 用自注意力关注输入的不同部分

我们现在将介绍自注意力机制的内部工作原理，并学习如何从头开始编写其代码。自注意力是基于转换器架构的每个大语言模型的基石。这个主题可能需要大量的专注和注意力（并非双关），但一旦你掌握了它的基本原理，你就攻克了本书以及一般大语言模型实现中最困难的部分之一。

[!NOTE] 自注意力中的 “自”
在自注意力中，“自” 指的是该机制通过关联单个输入序列内不同位置来计算注意力权重的能力。它评估并学习输入自身各个部分之间的关系和依赖，例如句子中的单词或图像中的像素。

这与传统的注意力机制形成对比，传统注意力机制的重点在于两个不同序列元素之间的关系，比如在序列到序列的模型中，注意力可能处于输入序列和输出序列之间，如图 3.5 所示的示例。

由于自注意力机制可能看起来比较复杂，特别是当你第一次遇到它时，我们将首先研究它的一个简化版本。然后，我们将实现大语言模型中使用的带有可训练权重的自注意力机制。

3.3.1 一种不带可训练权重的简单自注意力机制

让我们首先实现一种简化的自注意力变体，它不包含任何可训练的权重，如图 3.7 所总结的那样。目标是在添加可训练权重之前说明自注意力中的几个关键概念。

图 3.7 自注意力的目标是为每个输入元素计算一个上下文向量，该向量结合了来自所有其他输入元素的信息。在这个例子中，我们计算上下文向量 z ⁽²⁾。用于计算 z ⁽²⁾ 的每个输入元素的重要性或贡献度是由注意力权重α²¹ 至 α^2T 决定的。在计算 z ⁽²⁾ 时，注意力权重是根据输入元素 x ⁽²⁾ 和所有其他输入元素来计算的。

图 3.7 展示了一个输入序列，记为 x，它由 T 个元素组成，这些元素表示为 x ⁽¹⁾ 到 x ^(T)。这个序列通常代表文本，例如一个句子，并且已经被转换为词元嵌入。
例如，考虑一个输入文本，如 “Your journey starts with one step.” 在这种情况下，序列的每个元素，例如 x ⁽¹⁾，对应于一个表示特定词元（如 “Your”）的 d 维嵌入向量。图 3.7 将这些输入向量显示为三维嵌入。
在自注意力机制中，我们的目标是为输入序列中的每个元素 x ⁽ⁱ⁾ 计算上下文向量 z ⁽ⁱ⁾。上下文向量可以被解释为一个增强的嵌入向量。
为了阐释这个概念，让我们关注第二个输入元素的嵌入向量 x ⁽²⁾（它对应于词元 “journey”）以及图 3.7 底部所示的相应上下文向量 z ⁽²⁾。这个增强的上下文向量 z ⁽²⁾ 是一个包含关于 x ⁽²⁾ 和所有其他输入元素 x ⁽¹⁾ 到 x ^(T) 信息的嵌入。
上下文向量在自注意力机制中起着至关重要的作用。它们的目的是通过整合序列中所有其他元素的信息，为输入序列（如一个句子）中的每个元素创建增强的表示（图 3.7）。这在大语言模型中是至关重要的，因为大语言模型需要理解句子中单词之间的关系和相关性。之后，我们将添加可训练的权重，这些权重有助于大语言模型学会构建这些上下文向量，以便它们对大语言模型生成下一个词元有用。但首先，让我们实现一个简化的自注意力机制，一步一步地计算这些权重和最终的上下文向量
考虑以下输入句子，它已经被嵌入到三维向量中（见第 2 章）。我选择了一个较小的嵌入维度，以确保它能在页面上显示而无需换行：

import torch 
inputs = torch.tensor( 
	[[0.43, 0.15, 0.89], # Your (x^1) 
	[0.55, 0.87, 0.66], # journey (x^2) 
	[0.57, 0.85, 0.64], # starts (x^3) 
	[0.22, 0.58, 0.33], # with (x^4) 
	[0.77, 0.25, 0.10], # one (x^5) 
	[0.05, 0.80, 0.55]] # step (x^6) 
)

实现自注意力机制的第一步是计算中间值ω，即所谓的注意力分数，如图3.8所示。由于空间限制，该图以截断形式显示了前面输入张量的值；例如，0.87被截断为0.8。在这个截断版本中，“journey”（旅程）和“starts”（开始）这两个词的嵌入可能因随机因素而看起来相似。

图3.8 总体目标是说明如何以第二个输入元素 x⁽²⁾作为查询值来计算上下文向量 z⁽²⁾ 。该图展示了第一个中间步骤，即通过点积计算查询值 x⁽²⁾ 与所有其他输入元素之间的注意力分数 ω \omega ω 。（注意，为减少视觉干扰，数字均保留到小数点后一位 。）
图 3.8 展示了我们如何计算查询词元与每个输入词元之间的中间注意力分数。我们通过计算查询值 与其他每个输入词元的点积来确定这些分数：

query = inputs[1] 
attn_scores_2 = torch.empty(inputs.shape[0]) 
for i, x_i in enumerate(inputs): 
	attn_scores_2[i] = torch.dot(x_i, query) 
	print(attn_scores_2)

计算得到的注意力分数为：

tensor([0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865])

理解点积

点积本质上是一种简洁的方法，先将两个向量按元素相乘，然后对乘积进行求和，如下所示：

res = 0. 
for idx, element in enumerate(inputs[0]): 
	res += inputs[0][idx] * query[idx] 
print(res) 
print(torch.dot(inputs[0], query))

输出结果证实，按元素相乘再求和，所得结果与点积运算结果相同：

tensor(0.9544) 
tensor(0.9544)

除了将点积运算视为一种把两个向量结合以产生一个标量值的数学工具外，点积还是一种相似性度量，因为它量化了两个向量的对齐程度：点积越高，表明向量之间的对齐程度或相似性越高。在自注意力机制的背景下，点积决定了序列中每个元素对其他任何元素的关注程度：点积越高，两个元素之间的相似性和注意力分数就越高。
接下来，如图3.9所示，我们对之前计算出的每个注意力分数进行归一化处理。归一化的主要目的是获得总和为1的注意力权重。这种归一化是一种惯例，有助于在大语言模型（LLM）中进行结果解读，并维持训练的稳定性。以下是实现这一归一化步骤的一种简单方法：

attn_weights_2_tmp = attn_scores_2 / attn_scores_2.sum() 
print("Attention weights:", attn_weights_2_tmp) 
print("Sum:", attn_weights_2_tmp.sum())

图 3.9 针对输入查询x⁽²⁾计算出注意力分数 ω \omega ω ₂₁至 ω \omega ω _2T 后，下一步是通过对注意力分数进行归一化，得到注意力权重 α \alpha α₂₁ 至 α \alpha α_2T 。

Attention weights: tensor([0.1455, 0.2278, 0.2249, 0.1285, 0.1077, 0.1656]) 
Sum: tensor(1.0000)

在实际应用中，使用softmax函数进行归一化更为常见且可取。这种方法在处理极端值方面表现更优，并且在训练过程中具有更良好的梯度特性。以下是用于对注意力分数进行归一化的softmax函数的基本实现：

def softmax_naive(x): 
	return torch.exp(x) / torch.exp(x).sum(dim=0) 
attn_weights_2_naive = softmax_naive(attn_scores_2) 
print("Attention weights:", attn_weights_2_naive) 
print("Sum:", attn_weights_2_naive.sum())

正如输出结果所示，softmax函数同样达到了目标，将注意力权重进行归一化，使其总和为1 ：

Attention weights: tensor([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581]) 
Sum: tensor(1.)

此外，softmax函数确保注意力权重始终为正。这使得输出可以被解释为概率或相对重要性，权重越高表示重要性越大。
需要注意的是，这种简单的softmax实现（softmax_naive）在处理较大或较小的输入值时，可能会遇到数值不稳定的问题，比如上溢和下溢。因此，在实际应用中，建议使用PyTorch实现的softmax函数，它在性能方面经过了大量优化：

attn_weights_2 = torch.softmax(attn_scores_2, dim=0) 
print("Attention weights:", attn_weights_2) 
print("Sum:", attn_weights_2.sum())

在这种情况下，它产生的结果与我们之前的 softmax_naive 函数相同：

Attention weights: tensor([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581]) 
Sum: tensor(1.)

既然我们已经计算出了归一化后的注意力权重，那么就可以进行最后一步了，如图 3.10 所示：通过将嵌入后的输入词元 x⁽ⁱ⁾ 与相应的注意力权重相乘，然后对所得向量求和，来计算上下文向量 z⁽²⁾。因此，上下文向量 z⁽²⁾ 是所有输入向量的加权和，通过将每个输入向量乘以其对应的注意力权重得到：

query = inputs[1] 
context_vec_2 = torch.zeros(query.shape) 
for i,x_i in enumerate(inputs): 
	context_vec_2 += attn_weights_2[i]*x_i 
print(context_vec_2)

此次计算的结果如下：

tensor([0.4419, 0.6515, 0.5683])

![[Pasted image 20250123135151.png]]
图3.10 在计算并归一化注意力分数以获得针对查询 x⁽²⁾ 的注意力权重之后，最后一步是计算上下文向量 z⁽²⁾。这个上下文向量是所有输入向量 x⁽¹⁾ 到 x^(T) 以注意力权重进行加权后的组合。
接下来，我们要将计算上下文向量的这一过程进行推广，从而实现同时计算所有上下文向量。

3.3.2 计算所有输入词元的注意力权重

到目前为止，我们已经计算出了输入 2 的注意力权重和上下文向量，如图 3.11 中高亮行所示。现在，让我们将这个计算过程扩展到所有输入，以计算它们的注意力权重和上下文向量。

图3.11中高亮的行展示了以第二个输入元素作为查询时的注意力权重。现在我们将把这个计算过程进行泛化，以得到所有其他的注意力权重。（请注意，为减少视觉干扰，此图中的数字均保留到小数点后两位。每行的值相加应为1.0，即100%。）
我们遵循与之前相同的三个步骤（见图3.12），不过在代码上做了一些修改，目的是计算所有的上下文向量，而非仅仅计算第二个上下文向量 z⁽²⁾。

attn_scores = torch.empty(6, 6) 
for i, x_i in enumerate(inputs): 
	for j, x_j in enumerate(inputs): 
		attn_scores[i, j] = torch.dot(x_i, x_j) 
print(attn_scores)

图 3.12：在步骤 1 中，我们额外添加了一个 for 循环，用于计算所有输入对之间的点积。
由此得到的注意力分数如下：

tensor([[0.9995, 0.9544, 0.9422, 0.4753, 0.4576, 0.6310], 
	[0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865], 
	[0.9422, 1.4754, 1.4570, 0.8296, 0.7154, 1.0605], 
	[0.4753, 0.8434, 0.8296, 0.4937, 0.3474, 0.6565], 
	[0.4576, 0.7070, 0.7154, 0.3474, 0.6654, 0.2935], 
	[0.6310, 1.0865, 1.0605, 0.6565, 0.2935, 0.9450]])

正如我们在图3.11中所看到的，张量中的每个元素都代表了每对输入之间的注意力分数。要注意，图中的数值是经过归一化处理的，这就是它们与前面张量中未归一化的注意力分数有所不同的原因。我们之后会处理归一化的问题。
在计算前面的注意力分数张量时，我们在Python中使用了 for 循环。不过，for 循环通常速度较慢，我们可以通过矩阵乘法来获得相同的结果：

attn_scores = inputs @ inputs.T 
print(attn_scores)

我们可以直观地确认，结果和之前是一样的。

tensor([[0.9995, 0.9544, 0.9422, 0.4753, 0.4576, 0.6310], 
	[0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865], 
	[0.9422, 1.4754, 1.4570, 0.8296, 0.7154, 1.0605], 
	[0.4753, 0.8434, 0.8296, 0.4937, 0.3474, 0.6565],
	[0.4576, 0.7070, 0.7154, 0.3474, 0.6654, 0.2935], 
	[0.6310, 1.0865, 1.0605, 0.6565, 0.2935, 0.9450]])

在图3.12的步骤2中，我们对每一行进行归一化处理，使得每一行的值总和为1。

attn_weights = torch.softmax(attn_scores, dim=-1) 
print(attn_weights)

在使用 PyTorch 时，像 torch.softmax 这类函数中的 dim 参数，用于指定函数将沿着输入张量的哪个维度进行计算。通过设置 dim=-1，我们是在指示 softmax 函数沿着 attn_scores 张量的最后一个维度进行归一化操作。若 attn_scores 是一个二维张量（例如，形状为 [行数, 列数]），它会跨列进行归一化，从而使每一行的值（按列维度求和）总和为 1。
我们可以验证，每一行的值确实总和为 1：

row_2_sum = sum([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581]) 
print("Row 2 sum:", row_2_sum) 
print("All row sums:", attn_weights.sum(dim=-1))

结果是……

Row 2 sum: 1.0 All row sums: tensor([1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000])

在图3.12的第三个也是最后一个步骤中，我们利用这些注意力权重，通过矩阵乘法来计算所有的上下文向量。

all_context_vecs = attn_weights @ inputs 
print(all_context_vecs)

在得到的输出张量中，每一行都包含一个三维的上下文向量。

tensor([[0.4421, 0.5931, 0.5790], 
	[0.4419, 0.6515, 0.5683], 
	[0.4431, 0.6496, 0.5671], 
	[0.4304, 0.6298, 0.5510], 
	[0.4671, 0.5910, 0.5266], 
	[0.4177, 0.6503, 0.5645]])

我们可以通过将输出张量的第二行与我们在3.3.1节中计算出的上下文向量 (z^{(2)}) 进行比较，来再次验证代码是否正确。

print("Previous 2nd context vector:", context_vec_2)

根据结果，我们可以看到之前计算得到的 context_vec_2 与前面张量中的第二行完全匹配。

Previous 2nd context vector: tensor([0.4419, 0.6515, 0.5683])

至此，简单自注意力机制的代码讲解就结束了。接下来，我们将添加可训练的权重，让大语言模型（LLM）能够从数据中学习，并提升其在特定任务上的表现。