llms-from-scratch-cn/Translated_Book/ch03/3.4.ipynb
2024-05-16 22:22:49 +08:00

606 lines
26 KiB
Plaintext
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# 3.4 实现带有可训练权重的自注意力\n",
"\n",
"在本节中,我们将实现自注意力机制,这种机制被广泛应用于原始的 Transformer 架构、GPT 模型以及大多数其他流行的大语言模型中。这种自注意力机制也被称作缩放点积注意力scaled dot\n",
"product attention。图 3.13 描述了这种自注意力机制如何融入到构建大语言模型的更广泛的背景中。\n",
"\n",
"**图 3.13 展示了我们在本节编码的自注意力机制是如何与本书及本章的整体背景相融合的。在前一节中,我们实现了一个简化的注意力机制,以便理解注意力机制的基本原理。本节中,我们将在此基础上增加可训练的权重。在后续章节中,我们还将通过加入因果掩码和多头注意力来进一步扩展这种自注意力机制。**\n",
"\n",
"![3.13](../img/fig-3-13.jpg)\n",
"\n",
"如图 3.13 所示,具有可训练权重的自注意力机制是在之前概念的基础上建立的:我们希望根据特定输入元素,将输入向量的加权和计算为上下文向量。如你所见,它与我们在 3.3 节之前编写的基本自注意力机制相比,只有一些细微的差别。\n",
"\n",
"最显著的区别是引入了在模型训练期间会更新的权重矩阵。这些可训练的权重矩阵至关重要,因为它们使得模型(特别是模型内部的注意力模块)能够学习产生“良好”的上下文向量。(请注意,我们将在第 5 章训练大语言模型。)\n",
"\n",
"我们将在两个小节中讨论这种自注意力机制。首先,我们将像以前一样逐步编写代码。其次,我们将代码组织成一个紧凑的 Python 类,该类可以导入到我们在第 4 章的大语言模型架构中。\n",
"\n",
"## 3.4.1 逐步计算注意力权重\n",
"\n",
"我们将通过引入三个可训练的权重矩阵 Wq、Wk 和 Wv 逐步实现自注意力机制。这三个矩阵用于将嵌入的输入 Token x(i) 投影为查询向量、键向量和值向量,如图 3.14 所示。\n",
"\n",
"**图 3.14 展示了实现具备可训练权重矩阵的自注意力机制的第一步。在这一步中,我们针对每个输入元素 x计算其对应的查询q、键k和值v向量。如同之前章节所做的那样我们将第二个输入 x(2) 作为查询输入来处理。查询向量 q(2) 是通过将输入 x(2) 与查询权重矩阵 Wq 进行矩阵乘法得到的。类似地,我们也通过对应的矩阵乘法操作,使用权重矩阵 Wk 和 Wv 来分别计算键向量和值向量。**\n",
" \n",
"![3.14](../img/fig-3-14.jpg)\n",
"\n",
"在第 3.3.1 节中,我们定义第二个输入元素 x(2) 作为查询向量,以计算简化的注意力权重并得出上下文向量 z(2)。在第 3.3.2 节中我们将这一计算推广到了“Your journey starts with one step.” 这个六词输入句子的所有上下文向量 z(1) 到 z(T)。\n",
"\n",
"同样,出于示例说明的目的,我们将首先只计算一个上下文向量 z(2)。在下一节中,我们将修改这段代码以计算所有上下文向量。\n",
"\n",
"让我们开始定义一些变量:"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {
"trusted": true
},
"outputs": [],
"source": [
"x_2 = inputs[1] #A\n",
"d_in = inputs.shape[1] #B\n",
"d_out = 2 #C"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"请注意,在类似 GPT 的模型中输入和输出维度通常相同但为了更好地展示计算过程这里我们选择了不同的输入d_in=3和输出d_out=2维度。\n",
"\n",
"接下来,我们初始化图 3.14 中显示的三个权重矩阵 Wq、Wk 和 Wv\n"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {
"trusted": true
},
"outputs": [],
"source": [
"torch.manual_seed(123)\n",
"W_query = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)\n",
"W_key = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)\n",
"W_value = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"请注意,我们将 requires_grad 设置为 False 是为了示范时输出结果更清晰,但如果我们要将这些权重矩阵用于模型训练,我们会将 requires_grad 设置为 True以便在模型训练期间更新这些矩阵。\n",
"\n",
"接下来,如图 3.14 所示,我们计算查询、键和值向量:"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {
"trusted": true
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor([0.4306, 1.4551])\n"
]
}
],
"source": [
"query_2 = x_2 @ W_query \n",
"key_2 = x_2 @ W_key \n",
"value_2 = x_2 @ W_value\n",
"print(query_2)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"从查询的输出可以看到,这产生了一个二维向量,因为我们将相应权重矩阵的列数通过 d_out 设置为了 2\n",
"```python\n",
"tensor([0.4306, 1.4551])\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### 权重参数与注意力权重\n",
"\n",
"请注意,在权重矩阵 W 中,“权重”一词是“权重参数”的简称,这是在训练过程中被优化的神经网络的值。它不应该与注意力权重混淆。正如我们在前一节中已经看到的,注意力权重决定了上下文向量在多大程度上依赖输入的不同部分,即网络在多大程度上关注输入的不同部分。\n",
"\n",
"也就是说,权重参数是定义网络连接的基本学习系数,而注意力权重则是动态的、特定于上下文的值。\n",
"\n",
"尽管我们的临时目标只是计算一个上下文向量 z(2),我们仍然需要所有输入元素的键和值向量,因为它们参与了根据查询向量 q(2) 计算注意力权重,如图 3.14 所示。\n",
"\n",
"我们可以通过矩阵乘法获得所有键和值向量:\n"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {
"trusted": true
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"keys.shape: torch.Size([6, 2])\n",
"values.shape: torch.Size([6, 2])\n"
]
}
],
"source": [
"keys = inputs @ W_key \n",
"values = inputs @ W_value\n",
"print(\"keys.shape:\", keys.shape)\n",
"print(\"values.shape:\", values.shape)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"从输出中我们可以看出,我们成功地将 6 个输入 Token 从 3D 空间投影到了 2D 嵌入空间:\n",
"```python\n",
"keys.shape: torch.Size([6, 2])\n",
"values.shape: torch.Size([6, 2])\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"第二步是计算注意力得分,如图 3.15 所示。\n",
"\n",
"**图 3.15 注意力得分计算是一个点积计算,类似于我们在第 3.3 节中使用的简化自注意力机制。新的方面在于,我们不是直接计算输入元素之间的点积,而是使用通过各自的权重矩阵转换输入获得的查询向量和键向量。**\n",
"\n",
"![3.15](../img/fig-3-15.jpg)\n",
"\n",
"首先让我们来计算注意力得分ω22\n"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {
"trusted": true
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor(1.8524)\n"
]
}
],
"source": [
"keys_2 = keys[1] #A\n",
"attn_score_22 = query_2.dot(keys_2)\n",
"print(attn_score_22)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"得到未归一化的注意力得分结果:\n",
"```python\n",
"tensor(1.8524)\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"再次,我们可以通过矩阵乘法将此计算拓展到所有注意力得分上:"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {
"trusted": true
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor([1.2705, 1.8524, 1.8111, 1.0795, 0.5577, 1.5440])\n"
]
}
],
"source": [
"attn_scores_2 = query_2 @ keys.T # All attention scores for given query\n",
"print(attn_scores_2)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"从输出中我们很快可以发现,输出的第二个元素与我们之前计算的 attn_score_22 相匹配:\n",
"```python\n",
"tensor([1.2705, 1.8524, 1.8111, 1.0795, 0.5577, 1.5440])\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"第三步是从注意力得分转到注意力权重,如图 3.16 所示。\n",
"\n",
"**图 3.16 在计算注意力得分 ω 之后,下一步是使用 softmax 函数归一化这些得分以获得注意力权重 α。**\n",
"\n",
"![3.16](../img/fig-3-16.jpg)\n",
"\n",
"接下来,如图 3.16 所示,我们通过缩放注意力得分并使用我们之前使用的 softmax 函数来计算注意力权重。与之前不同的是,我们现在通过除以键的嵌入维度的平方根来缩放注意力得分,(注意,取平方根在数学上与指数化为 0.5 相同)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"trusted": true
},
"outputs": [],
"source": [
"d_k = keys.shape[-1]\n",
"attn_weights_2 = torch.softmax(attn_scores_2 / d_k**0.5, dim=-1)\n",
"print(attn_weights_2)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"得到的注意力权重如下:\n",
"```python\n",
"tensor([0.1500, 0.2264, 0.2199, 0.1311, 0.0906, 0.1820])\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### 缩放点积注意力背后的逻辑\n",
"\n",
"通过嵌入维度大小进行归一化的原因是为了通过避免小梯度来提高训练性能。例如,在扩大嵌入维度时,对于类似 GPT 的大语言模型,其维度通常超过千,较大的点积可能会因为应用了 softmax 函数而在反向传播过程中产生非常小的梯度。随着点积的增加softmax 函数的表现更像是一个阶跃函数,导致梯度接近零。这些小梯度可以极大地减缓学习速度或导致训练停滞。\n",
"\n",
"这种自注意力机制也被称为缩放点积注意力的原因,是它通过嵌入维度的平方根进行缩放。\n",
"\n",
"最后一步是计算上下文向量,如图 3.17 所示。 \n",
"\n",
"**图 3.17 在自注意力计算的最后一步中,我们通过注意力权重组合所有值向量来计算上下文向量。**\n",
"\n",
"![3.17](../img/fig-3-17.jpg)\n",
"\n",
"类似于第 3.3 节,我们通过输入向量的加权求和计算上下文向量,现在我们通过值向量的加权求和来计算上下文向量。在这里,注意力权重充当了衡量每个值向量相应重要性的权重因子。类似于第 3.3 节,我们可以使用矩阵乘法一步获得输出:"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {
"trusted": true
},
"outputs": [
{
"ename": "NameError",
"evalue": "name 'attn_weights_2' is not defined",
"output_type": "error",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mNameError\u001b[0m Traceback (most recent call last)",
"Input \u001b[0;32mIn [10]\u001b[0m, in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0m context_vec_2 \u001b[38;5;241m=\u001b[39m \u001b[43mattn_weights_2\u001b[49m \u001b[38;5;241m@\u001b[39m values\n\u001b[1;32m 2\u001b[0m \u001b[38;5;28mprint\u001b[39m(context_vec_2)\n",
"\u001b[0;31mNameError\u001b[0m: name 'attn_weights_2' is not defined"
]
}
],
"source": [
"context_vec_2 = attn_weights_2 @ values\n",
"print(context_vec_2)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"生成的张量内容如下:\n",
"```python\n",
"tensor([0.3061, 0.8210])\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"到目前为止我们只计算了一个上下文向量z(2)。在下一节中我们将拓展代码以计算输入序列中的所有上下文向量z(1) 到 z(T)中。"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### 为什么是查询、键和值?\n",
"\n",
"在注意力机制的上下文中,“键”、“查询”和“值”这些术语是从信息检索和数据库领域借鉴来的,在这些领域中,类似的概念被用于存储、搜索和检索信息。\n",
"\n",
"“查询”query类似于数据库中的搜索查询。它代表模型当前关注或试图理解的项目例如句子中的一个词或 Token。查询用于探查输入序列的其他部分以确定应该给予它们多少注意力。\n",
"\n",
"“键”key类似于数据库中用于索引和搜索的键。在注意力机制中输入序列中的每个项目例如句子中的每个词都有一个关联的键。这些键用于与查询匹配。\n",
"\n",
"“值”value在这个上下文中类似于数据库中键值对的值。它代表输入项目的实际内容或表示。一旦模型确定哪些键哪些输入部分与查询当前关注项目最相关它就检索相应的值。\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 3.4.2 实现一个紧凑的自注意力 Python 类\n",
"\n",
"在之前的章节里,我们详细介绍了计算自注意力输出的各个步骤,这样做主要是为了便于逐步讲解和展示。但在实际应用中,特别是考虑到下一章将要介绍的大语言模型的实现,把这些代码整合到一个 Python 类中会更加高效。如下所示:\n",
"\n",
"### 清单 3.1 紧凑型 self-attention 类"
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {
"trusted": true
},
"outputs": [],
"source": [
"import torch.nn as nn\n",
"class SelfAttention_v1(nn.Module):\n",
" def __init__(self, d_in, d_out):\n",
" super().__init__()\n",
" self.d_out = d_out\n",
" self.W_query = nn.Parameter(torch.rand(d_in, d_out))\n",
" self.W_key = nn.Parameter(torch.rand(d_in, d_out))\n",
" self.W_value = nn.Parameter(torch.rand(d_in, d_out))\n",
" \n",
" def forward(self, x):\n",
" keys = x @ self.W_key\n",
" queries = x @ self.W_query\n",
" values = x @ self.W_value\n",
" attn_scores = queries @ keys.T # omega\n",
" attn_weights = torch.softmax(\n",
" attn_scores / keys.shape[-1]**0.5, dim=-1)\n",
" context_vec = attn_weights @ values\n",
" return context_vec"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"在这段 PyTorch 代码中SelfAttention_v1 是继承自 nn.Module 的一个类,而 nn.Module 是构成 PyTorch 模型的基本单元,它提供了创建和管理模型层所需的功能。\n",
"\n",
"__init__方法负责初始化可训练的权重矩阵W_query、W_key 和 W_value这些矩阵分别用于查询、键和值每个矩阵都将输入维度 d_in 转换为输出维度 d_out。\n",
"\n",
"在前向传播过程中,通过 forward 方法我们通过查询和键的乘积计算注意力得分attn_scores并使用 softmax 函数对这些得分进行归一化处理。最后,我们通过这些归一化的注意力得分对值进行加权,以此创建一个上下文向量。\n",
"\n",
"我们可以按照以下方式使用这个类:\n"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {
"trusted": true
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor([[0.2996, 0.8053],\n",
" [0.3061, 0.8210],\n",
" [0.3058, 0.8203],\n",
" [0.2948, 0.7939],\n",
" [0.2927, 0.7891],\n",
" [0.2990, 0.8040]], grad_fn=<MmBackward>)\n"
]
}
],
"source": [
"torch.manual_seed(123)\n",
"sa_v1 = SelfAttention_v1(d_in, d_out)\n",
"print(sa_v1(inputs))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"由于输入包含六个嵌入向量,因此我们得到一个包含六个上下文向量的矩阵:\n",
"```python\n",
"tensor([[0.2996, 0.8053],\n",
" [0.3061, 0.8210],\n",
" [0.3058, 0.8203],\n",
" [0.2948, 0.7939],\n",
" [0.2927, 0.7891],\n",
" [0.2990, 0.8040]], grad_fn=<MmBackward0>)\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"快速核对一下,第二行 ([0.3061, 0.8210]) 实际上与上一节中的 context_vec_2 的内容相对应。\n",
"\n",
"图 3.18 概述了我们刚刚实现的自注意力机制。\n",
"\n",
"**图3.18 在自注意力机制中,我们通过三个权重矩阵 Wq、Wk 和 Wv 来转换输入矩阵 X 中的输入向量。然后,我们基于生成的查询 (Q) 和键 (K) 计算注意力权重矩阵。利用这些注意力权重和值 (V),我们计算出上下文向量 (Z)。为了视觉上的简洁,本图仅展示了一个包含 n 个 Token 的单一输入文本,而不是多个输入的批次。这种简化为 2D 矩阵的表示,使得过程的可视化和理解变得更为直观。**\n",
"\n",
"![3.18](../img/fig-3-18.jpg)\n",
"\n",
"如图 3.18 所示,自注意力涉及三个可训练的权重矩阵 Wq、Wk 和 Wv。这些矩阵将输入数据转化为查询、键和值它们是注意力机制的核心组成部分。随着模型在训练过程中接触更多数据这些可训练的权重会进行相应的调整我们将在后续章节中详细讨论。\n",
"\n",
"通过使用 PyTorch 的 nn.Linear 层,我们可以进一步改进 SelfAttention_v1 的实现。这些层在不使用偏置单元时可以有效地执行矩阵乘法。此外与手动实现nn.Parameter(torch.rand(...)) 相比,使用 nn.Linear 的一个显著优势是其具有优化的权重初始化方案,有助于实现更稳定和有效的模型训练。\n",
"\n",
"### 清单 3.2 使用 PyTorch 的 Linear 层的 self-attention 类"
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {
"trusted": true
},
"outputs": [],
"source": [
"class SelfAttention_v2(nn.Module):\n",
" def __init__(self, d_in, d_out, qkv_bias=False):\n",
" super().__init__()\n",
" self.d_out = d_out\n",
" self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
" self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
" self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
" \n",
" def forward(self, x):\n",
" keys = self.W_key(x)\n",
" queries = self.W_query(x)\n",
" values = self.W_value(x)\n",
" attn_scores = queries @ keys.T\n",
" attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)\n",
" context_vec = attn_weights @ values\n",
" return context_vec"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"你可以像使用 SelfAttention_v1 那样使用 SelfAttention_v2"
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {
"trusted": true
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor([[-0.0739, 0.0713],\n",
" [-0.0748, 0.0703],\n",
" [-0.0749, 0.0702],\n",
" [-0.0760, 0.0685],\n",
" [-0.0763, 0.0679],\n",
" [-0.0754, 0.0693]], grad_fn=<MmBackward>)\n"
]
}
],
"source": [
"torch.manual_seed(789)\n",
"sa_v2 = SelfAttention_v2(d_in, d_out)\n",
"print(sa_v2(inputs))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"输出结果是:\n",
"```python\n",
"tensor([[-0.0739, 0.0713],\n",
" [-0.0748, 0.0703],\n",
" [-0.0749, 0.0702],\n",
" [-0.0760, 0.0685],\n",
" [-0.0763, 0.0679],\n",
" [-0.0754, 0.0693]], grad_fn=<MmBackward0>)\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"注意SelfAttention_v1 和 SelfAttention_v2 的输出不同因为它们采用了不同的权重矩阵初始方案。nn.Linear 采用了比 nn.Parameter(torch.rand(d_in, d_out)) 更为复杂的权重初始化方案。"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### 练习 3.1 比较 SelfAttention_v1 和 SelfAttention_v2\n",
"\n",
"请注意SelfAttention_v2 中的 nn.Linear 使用了与 SelfAttention_v1 中的 nn.Parameter(torch.rand(d_in, d_out)) 不同的权重初始化方案,这导致两种机制产生不同的结果。为了检查 SelfAttention_v1 和 SelfAttention_v2 的实现在其他方面是否相似我们可以将SelfAttention_v2对象的权重矩阵转移到SelfAttention_v1对象上以便两个对象生成相同的结果。\n",
"\n",
"你的任务是正确地将SelfAttention_v2实例的权重分配给SelfAttention_v1实例。为此你需要理解两种版本中权重之间的关系。提示nn.Linear以转置形式存储权重矩阵。完成权重分配后你应该会发现两个实例生成相同的输出。\n",
"\n",
"在下一节中,我们将对自注意力机制进行增强,特别是我们将结合因果关系和多头注意力元素。因果关系方面的改进涉及修改注意力机制,以防止模型访问序列中的未来信息,这对于语言建模等任务至关重要,在这些任务中,每个词的预测只能依赖于之前的词。\n",
"\n",
"多头注意力的组件涉及将注意力机制分成多个“头”。每个头学习数据的不同方面,使模型能够同时关注来自不同表示子空间的不同位置的信息。这提高了模型在复杂任务中的表现。\n",
"\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"trusted": true
},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"trusted": true
},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "minitorch",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.8.12"
}
},
"nbformat": 4,
"nbformat_minor": 4
}