《从零构建大语言模型》学习笔记4，自注意力机制1
文章目录
- 《从零构建大语言模型》学习笔记4，自注意力机制1
- 前言
- 一、实现一个简单的无训练权重的自注意力机制
- 二、实现具有可训练权重的自注意力机制
-
- [1. 分步计算注意力权重](#1. 分步计算注意力权重)
- 2.实现自注意力Python类
- 三、将单头注意力扩展到多头注意力
- 总结
前言
本书原项目地址：https://github.com/rasbt/LLMs-from-scratch
我们进入第三章，探讨自注意力机制------这是大语言模型的核心基础算法。自注意力中的"自"表示该机制能够分析输入序列内部不同位置之间的关联，动态计算注意力权重。它通过学习输入元素（如句子中的单词或图像中的像素）之间的相互关系和依赖模式来实现这一功能。
一、实现一个简单的无训练权重的自注意力机制
比如我们有6个词元的embeddings的向量，接下来看下怎么计算第二个词元与其它词元之间的注意力分数。
计算方法也很简单，就是把第二个词元向量分别与其它词元向量做点积，原理图如下：
点积后得到的就是每一个向量与第二向量的自注意分数，然后进行归一化就得到了注意力权重。
计算代码如下：
python
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)
)
#对于一句话中的每个单词定义了一个三维的向量
query = inputs[1] # 2nd input token is the query
attn_scores_2 = torch.empty(inputs.shape[0])
#建立一个未初始化的张量来记录注意力得分
for i, x_i in enumerate(inputs):
attn_scores_2[i] = torch.dot(x_i, query)
# 相似性度量计算attention分数
# 从公式上看也就是点乘
print(attn_scores_2)
attn_weights_2 = torch.softmax(attn_scores_2, dim=0)
print("Attention weights:", attn_weights_2)
print("Sum:", attn_weights_2.sum())
#用torch优化过的softmax对边缘值也挺友好的然后把得到的注意力权重又分别与对应的向量相乘后再累加，就得到了上下文向量，这个就是我们最后要求的输出。
代码如下：
python
query = inputs[1] # 2nd input token is the query
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. 分步计算注意力权重
上述注意力计算过程可拆解为三个步骤，每个步骤都需要使用词向量。为此，我们为每个步骤的词向量分别乘以可训练的参数矩阵（Wq、Wk、Wv），从而得到对应的query、key和value向量。这样计算过程就如下图所示：
同样使用上面的例子，先用第二个词元向量点乘对应的Wq矩阵得到q2，然后用q2去点乘其它所有词元的key,就得到对应词元的注意力分数，同样归一化后得到对应的注意力权重。
最后用对应的注意力权重与value想成后累加，就得到了上下文向量。当然应用可以训练的参数矩阵，所以后面可以根据上下文向量结果来训练参数矩阵。
实现代码如下：
python
x_2 = inputs[1] # second input element
d_in = inputs.shape[1] # the input embedding size, d=3
d_out = 2 # the output embedding size, d=2
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)
#初始化三个矩阵来存放
#不要求梯度降低了复杂度
query_2 = x_2 @ W_query # _2 because it's with respect to the 2nd input element
key_2 = x_2 @ W_key
value_2 = x_2 @ W_value
#点积计算
print(query_2)
keys = inputs @ W_key
values = inputs @ W_value
print("keys.shape:", keys.shape)
print("values.shape:", values.shape)
#中途检验下
keys_2 = keys[1] # Python starts index at 0
attn_score_22 = query_2.dot(keys_2)
print(attn_score_22)
attn_scores_2 = query_2 @ keys.T # All attention scores for given query
print(attn_scores_2)
#计算注意力跟query值
d_k = keys.shape[1]
attn_weights_2 = torch.softmax(attn_scores_2 / d_k**0.5, dim=-1)
#压缩函数, 有利于储存与比较
print(attn_weights_2)
context_vec_2 = attn_weights_2 @ values
print(context_vec_2)2.实现自注意力Python类
把上面的过程集中到一个类里面，并且按照pytorch里面构建神经网络的方式来重写这个类，__init__函数是初始化一些参数，其中就包括用nn模块里面的线性层来初始化W_query，W_key ，W_value 这三个参数矩阵。forward函数是这个网络的前向运算过程，就是用初始化后的K，Q，V矩阵按照上面讲到的顺序进行矩阵相乘，最后得到上下文向量矩阵。
代码如下（示例）：
python
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.T
#Query跟Key的计算 得出初始的分数传递到后面进行归一化操作
attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
context_vec = attn_weights @ values
#直接基于注意力对于文本计算
return context_vec
torch.manual_seed(789)
sa_v2 = SelfAttention_v2(d_in, d_out)
print(sa_v2(inputs))这个类还需进一步优化，需要引入掩码的概念。掩码的作用是遮盖注意力权重矩阵的特定部分，通常是对矩阵的上三角部分进行处理。因为在实际推理过程中，模型需要预测后续未知的词元，所以在训练阶段就要通过掩码将未来的权重信息遮盖掉，让模型学会对未知信息的合理拟合。如果不这样做，可能会导致模型过快收敛。原理如下图：
代码很简单，就是给权重矩阵乘一个三角矩阵：
python
context_length = attn_scores.shape[0]
mask_simple = torch.tril(torch.ones(context_length, context_length))
#Mask矩阵,直接保留Diagonal下部分的,上部分掩盖掉
print(mask_simple)
masked_simple = attn_weights*mask_simple
print(masked_simple)
#简单的效果图
row_sums = masked_simple.sum(dim=-1, keepdim=True)
masked_simple_norm = masked_simple / row_sums
print(masked_simple_norm)
#掩码之后的softmax记得掩码后，要重新归一化。
为了加大训练难点，还会把卷积网络中比较成熟的drop层也引用进来，就是随机丢弃权重矩阵中的一些权重。
实现代码如下：
python
torch.manual_seed(123)
dropout = torch.nn.Dropout(0.5)
# dropout rate of 50%丢包率doge
example = torch.ones(6, 6)
# create a matrix of ones满的6*6矩阵被1包圆了
print(dropout(example))
#输出需要被放大相应的倍数,为了维持恒定
torch.manual_seed(123)
print(dropout(attn_weights))把以上两个技巧应用到类里面后，重写的代码如下：
python
class CausalAttention(nn.Module):
def __init__(self, d_in, d_out, context_length,
dropout, 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)
self.dropout = nn.Dropout(dropout) # New
self.register_buffer('mask', torch.triu(torch.ones(context_length, context_length), diagonal=1)) # New
#定义QKV并对进行dropout防止过拟合
#注册mask向量,对未来进行负无穷的拟合
def forward(self, x):
b, num_tokens, d_in = x.shape # New batch dimension b
#提取batch的大小、token的数量、跟宽度
keys = self.W_key(x)
queries = self.W_query(x)
values = self.W_value(x)
#进行运算计算
attn_scores = queries @ keys.transpose(1, 2) # Changed transpose
#通过点积来计算attention的数值
attn_scores.masked_fill_( # New, _ ops are in-place
self.mask.bool()[:num_tokens, :num_tokens], -torch.inf) # `:num_tokens` to account for cases where the number of tokens in the batch is smaller than the supported context_size
attn_weights = torch.softmax(
attn_scores / keys.shape[-1]**0.5, dim=-1## 缩放因子 √d，用于稳定梯度
)
#在时间顺序上进行mask确保信息不会被泄露
attn_weights = self.dropout(attn_weights) # New
#防止过拟合的dropout处理方式
context_vec = attn_weights @ values
# 根据注意力权重计算上下文向量
return context_vec
torch.manual_seed(123)
context_length = batch.shape[1]
ca = CausalAttention(d_in, d_out, context_length, 0.0)
context_vecs = ca(batch)
print(context_vecs)
print("context_vecs.shape:", context_vecs.shape)三、将单头注意力扩展到多头注意力
在卷积神经网络(CNN)中，通常采用不同尺寸的卷积核（如3×3、5×5等）来捕获图像的多尺度特征。这些不同尺寸卷积核提取的特征图经过通道维度的拼接（concat）后，能形成更全面的特征表示。类似地，注意力网络通过初始化多组q、k、v参数来获取不同的上下文向量并进行合并，这种机制被称为多头注意力。
比较简单的实现方式是，使用for循环做多次单注意力计算，然后再拼接就可以了，代码如下：
python
class MultiHeadAttentionWrapper(nn.Module):
def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
super().__init__()
#多个实例,每个都是一个头
self.heads = nn.ModuleList(
[CausalAttention(d_in, d_out, context_length, dropout, qkv_bias)
for _ in range(num_heads)]
)
def forward(self, x):
return torch.cat([head(x) for head in self.heads], dim=-1)
#模型的训练
torch.manual_seed(123)
context_length = batch.shape[1] # This is the number of tokens
d_in, d_out = 3, 2
mha = MultiHeadAttentionWrapper(
d_in, d_out, context_length, 0.0, num_heads=2
)
context_vecs = mha(batch)
print(context_vecs)
print("context_vecs.shape:", context_vecs.shape)然而，上述方法的效率较低，主要体现在以下两个方面：首先，这种方法需要重复进行n次参数初始化和前向传播计算（n代表注意力头的数量），导致计算资源的浪费；其次，多个独立的参数矩阵会导致内存访问不连续，降低缓存命中率。更常见且高效的做法是在模型初始化阶段就进行维度扩展。最后代码如下：
python
class MultiHeadAttention(nn.Module):
def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
super().__init__()
assert (d_out % num_heads == 0), \
"d_out must be divisible by num_heads"
#确保是可以被整除的
self.d_out = d_out
self.num_heads = num_heads
self.head_dim = d_out // num_heads # Reduce the projection dim to match desired output dim
#初始化头的维度、数量
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)
self.out_proj = nn.Linear(d_out, d_out) # Linear layer to combine head outputs
#头的输出结合线性层
self.dropout = nn.Dropout(dropout)
#进行dropout防止过拟合
self.register_buffer(
"mask",
torch.triu(torch.ones(context_length, context_length),
diagonal=1)
)
# 上三角掩码，确保因果性
def forward(self, x):
b, num_tokens, d_in = x.shape
keys = self.W_key(x) # Shape: (b, num_tokens, d_out)
queries = self.W_query(x)
values = self.W_value(x)
#把输出的维度拆成头*头大小
# We implicitly split the matrix by adding a `num_heads` dimension
# Unroll last dim: (b, num_tokens, d_out) -> (b, num_tokens, num_heads, head_dim)
keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)
values = values.view(b, num_tokens, self.num_heads, self.head_dim)
queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)
#转制维度,听说是为了更好的计算注意力
# Transpose: (b, num_tokens, num_heads, head_dim) -> (b, num_heads, num_tokens, head_dim)
keys = keys.transpose(1, 2)
queries = queries.transpose(1, 2)
values = values.transpose(1, 2)
# 计算缩放点积注意力
# Compute scaled dot-product attention (aka self-attention) with a causal mask
attn_scores = queries @ keys.transpose(2, 3) # Dot product for each head
# 将掩码缩减到当前 token 数量，并转换为布尔型
# 进而实现动态遮蔽,所以不用另开好几个数组
mask_bool = self.mask.bool()[:num_tokens, :num_tokens]
# 遮蔽矩阵
# Use the mask to fill attention scores
attn_scores.masked_fill_(mask_bool, -torch.inf)
#归一化
attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
attn_weights = self.dropout(attn_weights)
# Shape: (b, num_tokens, num_heads, head_dim)
context_vec = (attn_weights @ values).transpose(1, 2)
#头的合并
# Combine heads, where self.d_out = self.num_heads * self.head_dim
#对上下文向量的形状进行调整，确保输出的形状
context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)
context_vec = self.out_proj(context_vec) # optional projection
return context_vec
torch.manual_seed(123)
batch_size, context_length, d_in = batch.shape
d_out = 2
mha = MultiHeadAttention(d_in, d_out, context_length, 0.0, num_heads=2)
context_vecs = mha(batch)
print(context_vecs)
print("context_vecs.shape:", context_vecs.shape)总结
以上是关于注意力机制的讲解，重点在于理解Q、K、V三个参数。我一直在思考为什么是三个参数，是否能构造更多参数。从理论上看，增加更多参数是可行的，但从数学角度来说，过多的线性相乘参数可能对后续求导没有实质性帮助。此外，采用query、key、value的命名方式也使其含义更加直观易懂。在撰写过程中，已假设大家具备卷积神经网络和PyTorch的基础知识。若有任何表述不清或理解有误之处，欢迎随时指正交流。