亚马逊做品牌备案自有网站必应搜索引擎网址
本文是github上的大模型教程LLMs-from-scratch的学习笔记,教程地址:教程链接

Chapter 3: Attention Mechanism
本文首先从固定参数的注意力机制说起,然后拓展到可以训练的注意力机制,然后加入掩码mask,最后拓展到多头注意力机制。
1. 注意力机制
一个句子中的每一个token,都会受到其他token的影响(这里先不考虑忽略未来的单词,掩码的问题后面再说),注意力机制可以让一个token收到其他token的影响,生成一个最终我们想要的embedding。即每个token有一个原始的embedding,通过注意力机制后,得到了一个新的embedding,这个embedding是结合了上下文语义得到的。
举个简单的例子,我们直接使用tokens的embedding之间两两点乘,得到互相之间的点乘结果,然后将点乘结果归一化,得到embeddings之间的注意力得分。
归一化一般使用softmax函数,通过取指数,除以求和得到归一化结果
torch.exp(x) / torch.exp(x).sum(dim=0)
得到token之间的相关权重后,我们就可以加权求和,得到每一个token的最终embedding。
2. 可以训练的注意力头
在上面的例子中,我们直接使用token对应的embedding来计算相关系数,以及最终的加权求和,这显然是不合理的,如果这样的话,那么我们只能训练token对应的词嵌入来学习模型,或者是一些全连接层,因此我们需要引入新的矩阵,来学习到更多的参数,这就是transformer的QKV矩阵。
QKV都是对原始的embedding做线性变换,得到新的向量,但是模型就可以通过训练QKV,学习海量知识。
QKV的维度不固定,可以与原始嵌入相同,也可以不同。总之,通过QKV三个矩阵,我们将原始token的embedding转换成了3个新的向量。
- Query vector: q ( i ) = W q x ( i ) q^{(i)} = W_q \,x^{(i)} q(i)=Wqx(i)
- Key vector: k ( i ) = W k x ( i ) k^{(i)} = W_k \,x^{(i)} k(i)=Wkx(i)
- Value vector: v ( i ) = W v x ( i ) v^{(i)} = W_v \,x^{(i)} v(i)=Wvx(i)
可以使用矩阵乘法实现:
keys = inputs @ W_key
values = inputs @ W_value
然后我们计算K
和Q
之间的点积,作为两两token之间的关联度。为什么要用两个不一样的矩阵,我的猜测是,如果使用的是一个矩阵计算相似度,那么关于对角线对称
的元素就会完全相同,但是使用两个不同的矩阵计算,就不会存在这样的情况,可以学习到的内容更多。

我们使用K和Q的点积得到了两两之间的注意力得分,同样使用softmax进行归一化,得到最终的注意力权重。
注意到没有直接对注意力得分softmax,而是除以维度的方根后再softmax,这是因为在计算注意力权重时,如果直接将Query和Key的点积结果用于softmax函数,当Key的维度较高时,点积的结果会变得非常大。这可能导致softmax函数在梯度下降过程中学习困难,因为大的数值会使softmax的梯度变得非常小(接近于0),这在数值稳定性上是一个问题,称为“梯度消失”。
最后一步,不再使用原始的embedding加权,我们使用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.Tattn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)context_vec = attn_weights @ valuesreturn context_vec
3. 隐藏未来的单词

对语言任务来说,在训练模型的时候不能使用未来的文本来预测之前的文本。因此我们需要屏蔽未见文本对先前文本的影响。在我们计算得到注意力权重后,我们人为地将上三角矩阵的权重置为0。
有一种naive的方法,就是将上注意力权重都置为0后,重新对剩下的元素归一化。但是我们要介绍的是一般使用的方法:
我们在计算出注意力得分后,对右上角元素都赋值为负无穷大,负无穷大在经过softmax后就变为了0。
mask = torch.triu(torch.ones(context_length, context_length), diagonal=1)
masked = attn_scores.masked_fill(mask.bool(), -torch.inf)
attn_weights = torch.softmax(masked / keys.shape[-1]**0.5, dim=-1)
最后为了防止过拟合,一般会使用dropout,对注意力权重矩阵进行随机丢弃,加强模型泛化性能。
总结以上的所有内容,我们现在就可以写出一个单头的注意力机制了,并且加入了对batch
输入的处理:
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)) # Newdef forward(self, x):b, num_tokens, d_in = x.shape # New batch dimension bkeys = self.W_key(x)queries = self.W_query(x)values = self.W_value(x)attn_scores = queries @ keys.transpose(1, 2) # Changed transposeattn_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)attn_weights = self.dropout(attn_weights) # Newcontext_vec = attn_weights @ valuesreturn context_vec
4. 多头注意力机制


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)
这是一种最简单直白的写法,直接声明num_heads
个注意力单元,然后在前向传播的时候,依次调用这num_heads
个注意力头,然后将输出拼接起来。(dim=-1
代表最后一维拼接)
问题是,这样的话需要循环num_heads
次得到结果,并且需要声明num_heads
个注意力头,相信熟悉线性代数的朋友已经想到了,可以通过曾广矩阵来拓展注意力头。比如单个的注意力头是(d_in, d_out)
,那么有n个头的注意力机制就是(d_in, n*d_out)
。
假设输入是(tokens, d_in)
,那么(tokens, d_in) @ (d_in, n*d_out) --> (tokens, n*d_out)
,输出的结果完美得到了n个头对应的输出,我们只需要按照每d_out
列拆开,得到n
个(tokens, d_out)
的矩阵,就能还原出n
个头对应的结果,进行后续的attention score计算。这样写起来虽然麻烦一些,但是效率更高。
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 dimself.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 outputsself.dropout = nn.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# Original mask truncated to the number of tokens and converted to booleanmask_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_dimcontext_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)context_vec = self.out_proj(context_vec) # optional projectionreturn context_vec

不同于前者,将不同的注意力头分开计算,第二种方法直接扩展query,key和value矩阵的列数,将多个矩阵运算简化为一个矩阵运算,计算完再更改维度还原成一个个注意力头,效率更高。这样,我们就完成了一个完整的注意力机制。