从0开始搭建大语言模型：实现Attention机制

接上文：【教程】从0开始搭建大语言模型：Word和位置Embedding

建模长序列存在的问题

思考一个问题，之前语言模型没有注意力机制，它有什么问题？

当将文本从一种语言翻译成另一种语言时，例如从德语翻译成英语，仅仅是逐字翻译是不可能的。相反，翻译过程需要上下文理解和语法对齐。

为了解决这个问题，通常使用具有两个子模块的深度神经网络，即所谓的编码器和解码器。编码器的工作首先是读入并处理整个文本，然后解码器生成翻译后的文本。

在transformer出现之前，循环神经网络(RNN)是语言翻译中最流行的编码器-解码器架构。RNN是一种神经网络，前一步的输出作为当前步骤的输入，使其非常适合于文本等顺序数据。

下图是RNN进行语言翻译的例子：

编码器将来自源语言的token序列作为输入，其中编码器的隐藏状态(中间神经网络层)对整个输入序列的压缩表示进行编码。然后，解码器使用其当前的隐藏状态开始一个词一个词的翻译。

上面这种结果存在不足：RNN不能在解码阶段直接从编码器获取早期的隐藏状态。因此，它仅依赖于当前隐藏状态，其中保存了所有相关信息。这可能会导致上下文的丢失，特别是在依赖关系可能跨越长距离的复杂句子中。

使用attention机制获得数据间的依赖

前面提到，RNN可以很好地翻译短句，但不能很好地翻译较长的文本，因为它们不能直接访问输入中之前的单词。

为了解决这个问题，研究人员为RNN开发了所谓的Bahdanau注意力机制，如下：

上图的网络的文本生成解码器部分可以有选择地访问所有输入token，因此在生成给定输出token时，某些输入token比其他标记更重要，类似于注意力权重。

受Bahdanau注意力的启发，自注意力被提出，它允许在计算序列表示时，输入序列中的每个位置都关注同一序列中的所有位置。

Self-attention

介绍

在自注意力中，“self”是指机制通过关联单个输入序列中的不同位置来计算注意力权重的能力。它评估和学习输入本身不同部分之间的关系和依赖关系，例如句子中的单词或图像中的像素。

自注意力的目标是为每个输入元素计算一个上下文向量，该向量结合了来自所有其他输入元素的信息，如下图所示：

在self-attention中，我们的目标是为输入序列中的每个元素x (i)计算上下文向量z(i)。上下文向量可以被解释为丰富的embedding向量。

在self-attention中，上下文向量起着至关重要的作用。它们的目的是通过合并来自序列中所有其他元素的信息，创建输入序列中每个元素(如句子)的丰富表示。

x(2)和其他元素的attention score的计算如下图所示，通过计算点积得到。

点积是一种相似性度量，因为它量化了两个向量之间的对齐程度：

• 点积越高，表示向量之间的对齐程度或相似性越高。
• 在自注意力机制的背景下，点积决定了序列中元素相互关注的程度
  • 点积越高，两个元素之间的相似性和注意力分数越高

得到attention score后，还需要归一化得到最终的attention weights，如下图：

归一化背后的主要目标是获得总和为1的注意力权重。这种归一化是一种约定，对于解释和保持LLM的训练稳定性很有用。

在实践中，使用softmax函数进行归一化更常见，也更可取。这种方法可以更好地管理极值，并在训练过程中提供更有利的梯度属性。此外，softmax函数确保注意力权重始终为正。这使得输出可以解释为概率或相对重要性，其中权重越大，说明重要性越高。

最后，将各个元素的信息通过attention weights结合其他，得到x(2)的上下文信息z(2)：

如果要对所有token计算它与其他token的信息，可以使用矩阵乘法，而不是for循环，因为for循环计算效率低。

带有可训练权重的self-attention

在LLM中，self-attention中权重矩阵会在模型训练过程中更新。这些可训练的权重矩阵至关重要，以便模型(特别是模型中的注意力模块)可以学习产生“好的”上下文向量。

1.生成Q，K，V变量

引入三个可训练的权重矩阵Wq、Wk和Wv，这三个矩阵用于将embedding的输入标记x(i)投影为query、key和value，如下：

需要注意，权重参数（权重矩阵的值）是定义网络连接的基本的学习系数，而注意力权重是动态的、特定于上下文的值。

2.计算attention score

注意力分数计算是一种点积计算，需要注意的是：我们不直接计算输入元素之间的点积，而是使用通过各自的权重矩阵转换输入得到的query和key。

3.attention weight的获得

如下图，在计算出注意力分数ω之后，下一步是使用softmax函数对这些分数进行归一化，以获得注意力权重：

需要注意的是，我们通过除以key的embedding维度的平方根来缩放注意力分数。缩放注意力的原因：

• 通过避免小梯度来提高训练性能。例如，当扩大嵌入维度时(对于类似gpt的LLM通常大于1000)，由于对其应用了softmax函数，在反向传播过程中，较大的点积可能会导致非常小的梯度。
• 随着点积的增加，softmax函数的行为更像一个阶梯函数，导致梯度接近于零。这些小的梯度可能会大大减慢学习速度或导致训练停滞、

对embedding维度平方根的缩放是这种自注意力机制也被称为scaled-dot product attention的原因。

4.计算context vector

在自注意力计算的最后一步，我们通过注意力权重组合所有value向量来计算上下文向量。注意力权重作为一个权重因子，对每个值向量的重要性进行加权。

5.对于query, key, value的理解

注意力机制上下文中的术语“query”、“key”和“value”是从信息检索和数据库领域借用的。

“query”类似于数据库中的搜索查询。它表示模型关注或试图理解的当前项(例如，句子中的一个单词或标记)。query用于探测输入序列的其他部分，以确定对它们的关注程度。

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

这里的“value”类似于数据库中键值对中的值。它表示输入项的实际内容或表示。一旦模型确定了哪些键(以及输入的哪些部分)与查询(当前焦点项)最相关，它就检索相应的值。

上面过程的总体流程为：

在self-attention中，我们用三个权重矩阵Wq、Wk和Wv对输入矩阵X中的输入向量进行变换。然后，我们根据结果查询(Q)和键(K)计算注意力权重矩阵。使用注意力权重和值(V)，然后计算上下文向量(Z)。

使用权重矩阵的代码为：

import torch.nn as nn
class SelfAttention_v1(nn.Module):def __init__(self, d_in, d_out):super().__init__()self.d_out = d_outself.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_keyqueries = x @ self.W_queryvalues = x @ self.W_valueattn_scores = queries @ keys.T # omegaattn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)context_vec = attn_weights @ valuesreturn context_vecimport 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)
)
torch.manual_seed(123)
d_in = inputs.shape[1] #B
d_out = 2 #C
sa_v1 = SelfAttention_v1(d_in, d_out)
print(sa_v1(inputs))

我们也可以使用nn.Linear实现上述过程，代码为：

class SelfAttention_v2(nn.Module):def __init__(self, d_in, d_out, 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)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
torch.manual_seed(789)
sa_v2 = SelfAttention_v2(d_in, d_out)
print(sa_v2(inputs))

需要注意的是，Linear将权重矩阵存储为转置形式。

6.masked attention的应用

masked attention限制模型在处理任何给定token时只考虑序列中以前和当前的输入。这与标准的自注意力机制不一样，后者允许一次访问整个输入序列。在计算注意力分数时，masked attention机制确保模型只考虑序列中出现在当前token之前或之前的token。

如下图，在masked attention中，我们屏蔽了对角线上的注意力权重，使得对于给定的输入，LLM在使用注意力权重计算上下文向量时无法访问未来的token。

在上图中，我们mask掉对角线上的注意力权重，并对未mask的注意力权重进行归一化，使每行的注意力权重之和为1。

获得masked attention weights的一种方式如下：

需要注意：

• 当我们应用掩码，然后重新规范化注意力权重时，最初可能看起来来自未来token(打算mask)的信息仍然可能影响当前token，因为它们的值是softmax计算的一部分。
• 当我们重新规范化掩码后的注意力权重时，我们本质上所做的是在较小的子集上重新计算softmax(因为掩码位置对softmax值没有贡献)

softmax的数学优雅之处在于，尽管最初在分母中包括所有位置，但在掩码和重归一化之后，掩码位置的影响被消除——它们不会以任何有意义的方式对softmax分数做出贡献。

另外一种更高效的实现方法如下：

7.Dropout的使用

深度学习中的Dropout是一种在训练过程中忽略随机选择的隐藏层单元的技术，有效地“丢弃”它们。这种方法通过确保模型不过度依赖任何特定的隐藏层单元集，有助于防止过拟合。需要强调的是，dropout只在训练过程中使用，并在训练结束后停用。

在transformer架构中，包括像GPT这样的模型，注意力机制中的dropout通常应用于两个特定领域：

• 计算注意力分数后或
• 注意力权重应用于value向量后

随机去掉一些attention weights的示例如下：

当将dropout应用于一个dropout率为50%的注意力权重矩阵时，矩阵中有一半元素随机设置为0。为了补偿活跃元素的减少，矩阵中剩余元素的值按比例放大1/0.5 =2。这种缩放对于保持注意力权重的整体平衡至关重要，确保注意力机制的平均影响在训练和推理阶段保持一致。

结合了mask attention和dropout的注意力代码为：

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) #A# 下三角,context_length为文本长度self.register_buffer('mask',torch.triu(torch.ones(context_length, context_length), diagonal=1)) #Bdef forward(self, x):b, num_tokens, d_in = x.shape #C# 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) #C# mask操作attn_scores.masked_fill_(self.mask.bool()[:num_tokens, :num_tokens], -torch.inf)attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)attn_weights = self.dropout(attn_weights)context_vec = attn_weights @ valuesreturn context_vec
torch.manual_seed(123)
batch = torch.stack((inputs, inputs), dim=0)
print(batch.shape)
context_length = batch.shape[1]
ca = CausalAttention(d_in, d_out, context_length, 0.0)
context_vecs = ca(batch)
print("context_vecs.shape:", context_vecs.shape)

8.多头注意力的应用

“多头”是指将注意力机制划分为多个“头”，每个“头”独立操作。在这种情况下，单个注意力模块可以被认为是单头注意力，其中只有一组注意力权重按顺序处理输入，而多头是有多组权重处理输入。

实现多头可以通过堆叠多个attention模块来实现，如下图：

如上图，在一个有两个头的多头注意力模块中，query, key和value各有两个权重矩阵，该部分用代码实现为：

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)

需要注意，最后需要将每个头得到的context vector拼接，如下图：

上面实现多头的方法不能并行处理每一个头，事实上，可以通过矩阵乘法同时计算所有注意力头的输出。

具体来说，可以通过reshape投影的query、key和value张量，将输入分成多个头，然后在计算注意力后结合这些头的结果。

代码为：

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_heads# 每个头中的输出维度self.head_dim = d_out // num_heads #Aself.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) #Bself.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) #Cqueries = self.W_query(x) #Cvalues = self.W_value(x) #C# 将映射的变量reshape成多个头keys = keys.view(b, num_tokens, self.num_heads, self.head_dim) #Dvalues = values.view(b, num_tokens, self.num_heads, self.head_dim) #Dqueries = queries.view(b, num_tokens, self.num_heads, self.head_dim)#D# 交换矩阵的维度keys = keys.transpose(1, 2) #Equeries = queries.transpose(1, 2) #Evalues = values.transpose(1, 2) #Eattn_scores = queries @ keys.transpose(2, 3) #Fmask_bool = self.mask.bool()[:num_tokens, :num_tokens] #Gattn_scores.masked_fill_(mask_bool, -torch.inf) #Hattn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)attn_weights = self.dropout(attn_weights)context_vec = (attn_weights @ values).transpose(1, 2) #I#Jcontext_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)context_vec = self.out_proj(context_vec) #Kreturn context_vec

将query,key,value变成(b, num_heads,num_tokens, head_dim)是很重要的，矩阵乘法在最后2个维度(num_tokens, head_dim)之间进行，然后对每个头重复。

除此之外，代码中还加入了output projection layer，这在LLM的代码中经常能看到。

该方法比第一种方法更加高效，因为只需要一次矩阵乘法就可以计算出key，而在第一种方法中，我们需要为每个注意力头重复这个矩阵乘法，这是计算成本最高的步骤之一。它们两个的区别如下：

需要注意的是，最小的GPT-2模型(1.17亿个参数)具有12个注意力头和768个上下文向量embedding大小。最大的GPT-2模型(15亿个参数)有25个注意力头和1600个上下文向量embedding大小。token输入的embedding大小和上下文embedding在GPT模型中是相同的(d_in = d_out)。