从0开始LLM-注意力机制-2

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

自注意力机制，被广泛应用于原始的 Transformer 架构、GPT 模型以及大多数其他流行的大语言模型中。这种自注意力机制也被称作缩放点积注意力（scaled dot product attention）。

与前文最显著的区别是引入了在模型训练期间会更新的权重矩阵。这些可训练的权重矩阵至关重要，因为它们使得模型（特别是模型内部的注意力模块）能够学习产生"良好"的上下文向量。

逐步计算注意力权重

将通过引入三个可训练的权重矩阵 Wq、Wk 和 Wv 逐步实现自注意力机制。这三个矩阵用于将嵌入的输入 Token x(i) 投影为查询向量、键向量和值向量。

图展示了实现具备可训练权重矩阵的自注意力机制的第一步。在这一步中，我们针对每个输入元素 x，计算其对应的查询（q）、键（k）和值（v）向量。如同之前章节所做的那样，我们将第二个输入 x(2) 作为查询输入来处理。查询向量 q(2) 是通过将输入 x(2) 与查询权重矩阵 Wq 进行矩阵乘法得到的。类似地，我们也通过对应的矩阵乘法操作，使用权重矩阵 Wk 和 Wv 来分别计算键向量和值向量。

定义一些变量：

ini 复制代码

x_2 = inputs[1] #A
d_in = inputs.shape[1] #B
d_out = 2 #C

初始化图中显示的三个权重矩阵 Wq、Wk 和 Wv：

ini 复制代码

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，以便在模型训练期间更新这些矩阵。

计算查询、键和值向量：

ini 复制代码

query_2 = x_2 @ W_query 
key_2 = x_2 @ W_key 
value_2 = x_2 @ W_value
print(query_2)
#结果
tensor([0.4306, 1.4551])

权重参数与注意力权重

在权重矩阵 W 中，"权重"一词是"权重参数"的简称，这是在训练过程中被优化的神经网络的值。它不应该与注意力权重混淆。正如我们在前一节中已经看到的，注意力权重决定了上下文向量在多大程度上依赖输入的不同部分，即网络在多大程度上关注输入的不同部分。

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

通过矩阵乘法获得所有键和值向量：

ini 复制代码

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

2.计算注意力得分

注意力得分计算是一个点积计算，类似于简化不是直接计算输入元素之间的点积，而是使用通过各自的权重矩阵转换输入获得的查询向量和键向量。

ini 复制代码

keys_2 = keys[1] #A
attn_score_22 = query_2.dot(keys_2)
print(attn_score_22)
#结果
tensor(1.8524)

通过矩阵乘法将此计算拓展到所有注意力得分上：

ini 复制代码

attn_scores_2 = query_2 @ keys.T # All attention scores for given query
print(attn_scores_2)
#结果
tensor([1.2705, 1.8524, 1.8111, 1.0795, 0.5577, 1.5440])

从注意力得分转到注意力权重 在计算注意力得分 ω 之后，下一步是使用 softmax 函数归一化这些得分以获得注意力权重 α。 与之前不同的是，我们现在通过除以键的嵌入维度的平方根来缩放注意力得分，(注意，取平方根在数学上与指数化为 0.5 相同)：

ini 复制代码

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 的大语言模型，其维度通常超过千，较大的点积可能会因为应用了 softmax 函数而在反向传播过程中产生非常小的梯度。随着点积的增加，softmax 函数的表现更像是一个阶跃函数，导致梯度接近零。这些小梯度可以极大地减缓学习速度或导致训练停滞。

这种自注意力机制也被称为缩放点积注意力的原因，是它通过嵌入维度的平方根进行缩放。 4. 最后一步是计算上下文向量

通过输入向量的加权求和计算上下文向量，现在我们通过值向量的加权求和来计算上下文向量。在这里，注意力权重充当了衡量每个值向量相应重要性的权重因子。

为什么是查询、键和值？

在注意力机制的上下文中，"键"、"查询"和"值"这些术语是从信息检索和数据库领域借鉴来的，在这些领域中，类似的概念被用于存储、搜索和检索信息。

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

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

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

实现一个紧凑的自注意力 Python 类

ini 复制代码

import torch.nn as nn
class SelfAttention_v1(nn.Module):
    def __init__(self, d_in, d_out):
        super().__init__()
        self.d_out = d_out
        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_key
        queries = x @ self.W_query
        values = x @ self.W_value
        attn_scores = queries @ keys.T # omega
        attn_weights = torch.softmax(
            attn_scores / keys.shape[-1]**0.5, dim=-1)
        context_vec = attn_weights @ values
        return 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 函数对这些得分进行归一化处理。最后，我们通过这些归一化的注意力得分对值进行加权，以此创建一个上下文向量。

我们可以按照以下方式使用这个类：

scss 复制代码

torch.manual_seed(123)
sa_v1 = SelfAttention_v1(d_in, d_out)
print(sa_v1(inputs))

在自注意力机制中，通过三个权重矩阵 Wq、Wk 和 Wv 来转换输入矩阵 X 中的输入向量。然后，我们基于生成的查询 (Q) 和键 (K) 计算注意力权重矩阵。利用这些注意力权重和值 (V)，我们计算出上下文向量 (Z)。为了视觉上的简洁，本图仅展示了一个包含 n 个 Token 的单一输入文本，而不是多个输入的批次。这种简化为 2D 矩阵的表示，使得过程的可视化和理解变得更为直观。自注意力涉及三个可训练的权重矩阵 Wq、Wk 和 Wv。这些矩阵将输入数据转化为查询、键和值，它们是注意力机制的核心组成部分。随着模型在训练过程中接触更多数据，这些可训练的权重会进行相应的调整。改进：通过使用 PyTorch 的 nn.Linear 层，我们可以进一步改进 SelfAttention_v1 的实现。这些层在不使用偏置单元时，可以有效地执行矩阵乘法。此外，与手动实现nn.Parameter(torch.rand(...)) 相比，使用 nn.Linear 的一个显著优势是其具有优化的权重初始化方案，有助于实现更稳定和有效的模型训练。

ini 复制代码

class SelfAttention_v2(nn.Module):
    def __init__(self, d_in, d_out, qkv_bias=False):
        super().__init__()
        self.d_out = d_out
        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.T
        attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
        context_vec = attn_weights @ values
        return context_vec

使用：

scss 复制代码

torch.manual_seed(789)
sa_v2 = SelfAttention_v2(d_in, d_out)
print(sa_v2(inputs))