《从零构建大语言模型》学习笔记4，注意力机制1

《从零构建大语言模型》学习笔记4，自注意力机制1

---

---

前言

本书原项目地址：https://github.com/rasbt/LLMs-from-scratch
我们进入第三章，探讨自注意力机制——这是大语言模型的核心基础算法。自注意力中的"自"表示该机制能够分析输入序列内部不同位置之间的关联，动态计算注意力权重。它通过学习输入元素（如句子中的单词或图像中的像素）之间的相互关系和依赖模式来实现这一功能。

---

一、实现一个简单的无训练权重的自注意力机制

比如我们有6个词元的embeddings的向量，接下来看下怎么计算第二个词元与其它词元之间的注意力分数。

计算方法也很简单，就是把第二个词元向量分别与其它词元向量做点积，原理图如下：

点积后得到的就是每一个向量与第二向量的自注意分数，然后进行归一化就得到了注意力权重。
计算代码如下：

import torchinputs = 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 queryattn_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对边缘值也挺友好的

然后把得到的注意力权重又分别与对应的向量相乘后再累加，就得到了上下文向量，这个就是我们最后要求的输出。

代码如下：

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想成后累加，就得到了上下文向量。当然应用可以训练的参数矩阵，所以后面可以根据上下文向量结果来训练参数矩阵。
实现代码如下：

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=2torch.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矩阵按照上面讲到的顺序进行矩阵相乘，最后得到上下文向量矩阵。
代码如下（示例）：

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_vectorch.manual_seed(789)
sa_v2 = SelfAttention_v2(d_in, d_out)
print(sa_v2(inputs))

这个类还需进一步优化，需要引入掩码的概念。掩码的作用是遮盖注意力权重矩阵的特定部分，通常是对矩阵的上三角部分进行处理。因为在实际推理过程中，模型需要预测后续未知的词元，所以在训练阶段就要通过掩码将未来的权重信息遮盖掉，让模型学会对未知信息的合理拟合。如果不这样做，可能会导致模型过快收敛。原理如下图：

代码很简单，就是给权重矩阵乘一个三角矩阵：

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层也引用进来，就是随机丢弃权重矩阵中的一些权重。

实现代码如下：

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

把以上两个技巧应用到类里面后，重写的代码如下：

class CausalAttention(nn.Module):def __init__(self, d_in, d_out, context_length,dropout, qkv_bias=False):#初始化定义网络结构和参数super().__init__()self.d_out = d_outself.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) # Newself.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-placeself.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_sizeattn_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_vectorch.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循环做多次单注意力计算，然后再拼接就可以了，代码如下：

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代表注意力头的数量），导致计算资源的浪费；其次，多个独立的参数矩阵会导致内存访问不连续，降低缓存命中率。更常见且高效的做法是在模型初始化阶段就进行维度扩展。最后代码如下：

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_outself.num_heads = num_headsself.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.shapekeys = 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 maskattn_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 scoresattn_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 projectionreturn context_vectorch.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的基础知识。若有任何表述不清或理解有误之处，欢迎随时指正交流。