LLM(17)：计算所有输入 token 的注意力权重

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

如图 3.11 所示，到目前位置，已经计算了输入中的第二个词元的注意力权重和上下文向量，接下来将扩展这一计算过程，为所有输入计算注意力权重和上下文向量。

图 3.11 高亮的行展示了以第二个输入元素作为查询时的注意力权重。接下来把计算过程推广，以获得所有其他的注意力权重。（请注意，此图中的数字取到小数点后两位。每行中的数值总和应为 1.0 或 100%。）

与之前的步骤一样（见图 3.12），只是修改了一些代码，从而计算所有输入的上下文向量。

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 在第一步中，增加了 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)

会得到与之前一样的输出结果。

在图 3.12 的第 2 步中，对每一行进行归一化，使得每一行的数值之和为 1：

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

返回的是注意力权重张量，其数值与图 3.10 中显示的值一致：

tensor([[0.2098, 0.2006, 0.1981, 0.1242, 0.1220, 0.1452],
        [0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581],
        [0.1390, 0.2369, 0.2326, 0.1242, 0.1108, 0.1565],
        [0.1435, 0.2074, 0.2046, 0.1462, 0.1263, 0.1720],
        [0.1526, 0.1958, 0.1975, 0.1367, 0.1879, 0.1295],
        [0.1385, 0.2184, 0.2128, 0.1420, 0.0988, 0.1896]])

在使用 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）能够从数据中学习，并在特定任务上提升其性能。