从零开始构建大型语言模型—

本章内容：

使用注意力机制的原因
基本的自注意力框架，逐步深入到增强的自注意力机制
允许LLMs逐个生成词元的因果注意力模块
通过dropout随机屏蔽部分注意力权重以减少过拟合
将多个因果注意力模块堆叠为多头注意力模块

到目前为止，你已经了解了如何通过将文本拆分为单词和子词词元来准备LLM的输入文本，并将其编码为向量表示（嵌入）。现在，我们将介绍LLM架构中的一个重要部分——注意力机制，如图3.1所示。我们将主要独立地研究注意力机制，并在机制层面深入探讨。然后，我们将编写围绕自注意力机制的LLM其他部分的代码，以观察其实际效果，并构建一个用于生成文本的模型。

在这里插入图片描述

我们将实现四种不同的注意力机制变体，如图3.2所示。这些不同的注意力变体是逐步构建的，目标是最终实现一个紧凑且高效的多头注意力机制实现，然后可以将其嵌入到我们将在下一章编写的LLM架构中。

在这里插入图片描述

处理长序列建模的问题

在深入探讨LLM核心的自注意力机制之前，让我们先考虑一下在没有注意力机制的传统架构中遇到的问题。假设我们想开发一个将文本从一种语言翻译为另一种语言的翻译模型。如图3.3所示，由于源语言和目标语言中的语法结构不同，我们无法简单地逐字逐句进行翻译。

在这里插入图片描述

为了解决这个问题，通常使用包含两个子模块的深度神经网络，即编码器和解码器。编码器的任务是首先读取并处理整个文本，而解码器随后生成翻译后的文本。

在Transformer模型出现之前，循环神经网络（RNN）是用于语言翻译的最流行的编码器-解码器架构。RNN是一种神经网络，它将前一步的输出作为当前步骤的输入，因此非常适合处理像文本这样的序列数据。如果你不熟悉RNN，也不必担心，理解它的详细工作原理并不是理解本讨论的必要条件；这里我们主要关注编码器-解码器结构的一般概念。

在编码器-解码器RNN中，输入文本按顺序输入编码器，编码器逐步处理输入。在每一步，编码器会更新其隐藏状态（隐藏层的内部值），尝试在最终的隐藏状态中捕捉输入句子的完整意义，如图3.4所示。然后，解码器使用这个最终的隐藏状态开始逐字生成翻译句子。解码器也会在每一步更新其隐藏状态，这个状态应携带下一步预测单词所需的上下文信息。

虽然我们不需要深入了解编码器-解码器RNN的内部工作原理，但关键的思想是，编码器部分将整个输入文本处理为一个隐藏状态（记忆单元），然后解码器使用这个隐藏状态生成输出。你可以将这个隐藏状态看作是一个嵌入向量，这个概念我们在第2章中讨论过。

编码器-解码器RNN的一个主要限制是，在解码阶段，RNN无法直接访问编码器中的早期隐藏状态。因此，它只能依赖当前的隐藏状态，这个状态包含所有相关的信息。这可能会导致上下文丢失，尤其是在复杂句子中，依赖关系可能跨越较长距离。

幸运的是，构建LLM并不需要深入理解RNN。只需记住，编码器-解码器RNN的这一缺点促使了注意力机制的设计。

通过注意力机制捕捉数据依赖关系

虽然RNN在翻译短句时表现良好，但在处理较长文本时效果不佳，因为它无法直接访问输入中的前面部分。这个方法的一个主要缺点是，RNN必须在一个隐藏状态中记住整个编码的输入，然后再将其传递给解码器（如图3.4所示）。

因此，研究人员在2014年为RNN开发了Bahdanau注意力机制（以论文的第一作者命名；详情见附录B），这种机制修改了编码器-解码器RNN，使得解码器在每个解码步骤都可以有选择地访问输入序列的不同部分，如图3.5所示。

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

自注意力是一种机制，它允许输入序列中的每个位置在计算序列表示时考虑所有其他位置的相关性，或“关注”同一序列中的所有其他位置。自注意力是基于Transformer架构的现代LLM（如GPT系列）的关键组件。

本章将重点介绍如何编写和理解GPT类模型中使用的自注意力机制，如图3.6所示。在下一章中，我们将编写LLM的其他部分代码。

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

现在我们将介绍自注意力机制的内部工作原理，并从头开始学习如何编写相关代码。自注意力是基于Transformer架构的每个LLM的核心基础。这个主题可能需要高度集中注意力（双关意图），但一旦掌握了它的基础，你将攻克本书和LLM实现中最具挑战性的部分之一。

自注意力中的“自我”

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

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

由于自注意力机制可能看起来复杂，尤其是当你第一次接触它时，我们将从一个简化版本开始进行讲解。接着，我们将实现LLM中使用的带有可训练权重的自注意力机制。

没有可训练权重的简单自注意力机制

我们首先来实现一个没有任何可训练权重的简化自注意力机制，如图3.7所示。这样做的目的是在引入可训练权重之前，先阐明自注意力机制中的一些关键概念。

图3.7显示了一个输入序列，记作x，由T个元素组成，表示为x(1)到x(T)。这个序列通常表示已转换为词元嵌入的文本，例如一个句子。

例如，考虑输入文本“Your journey starts with one step.” 在这种情况下，序列中的每个元素（如x(1)）对应一个d维嵌入向量，表示一个特定的词元，比如“Your”。图3.7显示了这些输入向量作为三维嵌入。

在自注意力机制中，我们的目标是为输入序列中的每个元素x(i)计算上下文向量z(i)。上下文向量可以解释为一个增强的嵌入向量。

为了说明这个概念，我们重点关注第二个输入元素x(2)（对应词元“journey”）的嵌入向量及其对应的上下文向量z(2)，如图3.7底部所示。这个增强的上下文向量z(2)是一个包含x(2)及所有其他输入元素x(1)到x(T)信息的嵌入。

上下文向量在自注意力机制中起着至关重要的作用。它们的目的是通过融合输入序列（如一个句子）中所有其他元素的信息，来创建每个元素的增强表示（图3.7）。这对LLMs至关重要，因为它们需要理解句子中词与词之间的关系和相关性。稍后，我们将添加可训练的权重，帮助LLM学习构建这些上下文向量，使它们与生成下一个词元相关。但首先，让我们实现一个简化的自注意力机制，一步步计算这些权重和相应的上下文向量。

考虑以下句子，它已经嵌入为三维向量（见第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)
)

实现自注意力机制的第一步是计算中间值w，即所谓的注意力分数，如图3.8所示。由于空间限制，图中显示了截断版的输入张量值；例如，0.87被截断为0.8。在这个截断版本中，“journey”和“starts”的嵌入可能由于随机原因看起来相似。

图3.8展示了如何计算查询词元与每个输入词元之间的中间注意力分数。我们通过计算查询词元 x(2)x(2)x(2) 与每个其他输入词元的点积来确定这些分数：

query = inputs[1]                            #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)

#1 第二个输入词元作为查询词元。

计算出的注意力分数为：

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

输出结果为：

Attention weights: tensor([0.1444, 0.2261, 0.2232, 0.1276, 0.1069, 0.1718])
Sum: tensor(1.)

输出显示，注意力权重的总和现在为1：

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(i)x(i)x(i) 与相应的注意力权重相乘并对得到的向量求和来计算上下文向量 z(2)z(2)z(2)。因此，上下文向量 z(2)z(2)z(2) 是所有输入向量的加权和，即通过将每个输入向量乘以其对应的注意力权重得到：

query = inputs[1]         #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)

#1 第二个输入词元作为查询词元。

该计算的结果为：

tensor([0.4419, 0.6515, 0.5683])

接下来，我们将对计算上下文向量的过程进行泛化，以便同时计算所有的上下文向量。

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

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

我们将遵循之前的三个步骤（参见图3.12），只是对代码进行了一些修改，以计算所有的上下文向量，而不仅仅是第二个 z(2)z(2)z(2)：

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)

这个代码计算了每个输入之间的注意力分数，形成了一个 6×6 的矩阵，其中每个元素表示输入序列中两个词元之间的相似性。

得到的注意力分数如下：

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, 06654, 0.2935],[0.6310, 1.0865, 1.0605, 0.6565, 0.2935, 0.9450]])

在图3.12的第二步中，我们对每一行进行归一化，使每行的值总和为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)z(2)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能够从数据中学习，并在特定任务上提高性能。

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

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

如图3.13所示，带有可训练权重的自注意力机制基于之前的概念：我们希望计算上下文向量，作为针对特定输入元素的输入向量的加权和。你将会发现，这与我们之前编写的基本自注意力机制相比，只有细微的区别。

最显著的区别是引入了在模型训练过程中更新的权重矩阵。这些可训练的权重矩阵至关重要，因为模型（特别是模型中的注意力模块）需要通过它们学习生成“良好”的上下文向量。（我们将在第5章训练LLM。）

我们将在两个小节中处理这个自注意力机制。首先，我们将像之前一样逐步编写代码。其次，我们将把代码组织成一个紧凑的Python类，便于导入LLM架构中。

逐步计算注意力权重

我们将逐步实现自注意力机制，并引入三个可训练的权重矩阵 WqW_qWq、WkW_kWk 和 WvW_vWv。这三个矩阵分别用于将嵌入的输入词元 x(i)x(i)x(i) 映射到查询向量（query）、键向量（key）和值向量（value），如图3.14所示。

之前，我们在计算简化的注意力权重以得出上下文向量 z(2)z(2)z(2) 时，将第二个输入元素 x(2)x(2)x(2) 作为查询向量。随后，我们将其推广以计算六个单词组成的输入句子 “Your journey starts with one step” 的所有上下文向量 z(1)…z(T)z(1) \dots z(T)z(1)…z(T)。

同样地，为了便于说明，我们从计算一个上下文向量 z(2)z(2)z(2) 开始。之后，我们将修改代码以计算所有上下文向量。

首先，定义几个变量：

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

#1 第二个输入元素
#2 输入嵌入维度，d=3d = 3d=3
#3 输出嵌入维度，dout=2d_{out} = 2dout=2

请注意，在类似GPT的模型中，输入和输出的维度通常是相同的，但为了更好地理解计算过程，这里我们使用不同的输入维度 din=3d_{in} = 3din=3 和输出维度 dout=2d_{out} = 2dout=2。

接下来，我们初始化三个权重矩阵 WqW_qWq、WkW_kWk 和 WvW_vWv，如图3.14所示：

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)

查询向量的输出结果是一个二维向量，因为我们通过 doutd_{out}dout 将相应的权重矩阵的列数设置为2：

tensor([0.4306, 1.4551])

权重参数与注意力权重

在权重矩阵 WWW 中，“权重”是“权重参数”的缩写，指的是在训练过程中优化的神经网络值。这与注意力权重不同。正如我们之前看到的，注意力权重决定了上下文向量在多大程度上依赖于输入的不同部分（即网络在多大程度上关注输入的不同部分）。

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

尽管我们当前的目标只是计算一个上下文向量 z(2)z(2)z(2)，但我们仍然需要所有输入元素的键向量和值向量，因为它们参与计算与查询 q(2)q(2)q(2) 相关的注意力权重（见图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所示。

首先，让我们计算注意力分数 ω22\omega_{22}ω22：

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

#1 请记住，Python的索引从0开始。

未归一化的注意力分数结果为：

tensor(1.8524)

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

attn_scores_2 = query_2 @ keys.T       #1
print(attn_scores_2)

#1 针对给定查询的所有注意力分数。

可以看到，作为快速检查，输出中的第二个元素与我们之前计算的 attn_score_22\text{attn_score_22}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)

得到的注意力权重为：

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

缩放点积注意力的原理

通过嵌入维度大小进行归一化的原因是为了提高训练性能，避免小梯度的出现。例如，在GPT类的LLM中，嵌入维度通常大于1000，缩放后的大点积在反向传播过程中可能导致非常小的梯度，这是由于对它们应用的softmax函数所致。随着点积的增加，softmax函数的表现越来越像一个阶跃函数，导致梯度接近于零。这些小梯度可能会极大地减缓学习速度，或导致训练停滞。

通过嵌入维度的平方根进行缩放是自注意力机制被称为缩放点积注意力的原因。

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

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

context_vec_2 = attn_weights_2 @ values
print(context_vec_2)

得到的向量内容如下：

tensor([0.3061, 0.8210])

到目前为止，我们只计算了一个上下文向量 z(2)。接下来，我们将对代码进行泛化，以计算输入序列中的所有上下文向量 z(1) 到 z(T)。

为什么使用查询、键和值？

在注意力机制的上下文中，“键”、“查询”和“值”这三个术语借鉴自信息检索和数据库领域，其中类似的概念用于存储、搜索和检索信息。

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

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

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

实现紧凑的自注意力Python类

到目前为止，我们经历了许多步骤来计算自注意力输出。这主要是为了说明目的，使我们能够逐步进行。在实际应用中，考虑到下一章中的LLM实现，将这些代码组织成一个Python类是很有帮助的，如下所示。

代码清单3.1：紧凑的自注意力类

import torch.nn as nnclass 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_keyqueries = x @ self.W_queryvalues = x @ self.W_valueattn_scores = queries @ keys.T  # omegaattn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)context_vec = attn_weights @ valuesreturn 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))

由于 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=<MmBackward0>)

作为快速检查，请注意第二行（[0.3061, 0.8210]）与上一节中的 context_vec_2 的内容相匹配。图3.18总结了我们刚刚实现的自注意力机制。

自注意力机制涉及可训练的权重矩阵 WqW_qWq、WkW_kWk 和 WvW_vWv。这些矩阵分别将输入数据转换为查询、键和值，这些都是注意力机制的重要组成部分。随着模型在训练过程中接触到更多数据，它会调整这些可训练的权重，正如我们将在后面的章节中看到的。

我们可以进一步改进 SelfAttention_v1 的实现，利用 PyTorch 的 nn.Linear 层，这样在禁用偏置单元时可以有效地执行矩阵乘法。此外，使用 nn.Linear 而不是手动实现 nn.Parameter(torch.rand(...)) 的一个显著优势是，nn.Linear 具有优化的权重初始化方案，有助于更稳定和有效的模型训练。

代码清单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.Tattn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)context_vec = attn_weights @ valuesreturn 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 以转置形式存储权重矩阵。）在分配之后，你应该观察到两个实例产生相同的输出。