GPT - 多头注意力机制（Multi-Head Attention）模块

本节代码实现了一个多头注意力机制（Multi-Head Attention）模块，它是Transformer架构中的核心组件之一。
 

⭐关于多头自注意力机制的数学原理请见文章：
Transformer - 多头自注意力机制复现-CSDN博客
本节要求理解原理后手敲实现多头注意力机制
1. 初始化部门

1. class MultiHeadAttention(nn.Module):
2.     def __init__(self, d_model, num_heads, dropout):
3.         super().__init__()
4.         self.num_heads = num_heads
5.         self.d_k = d_model // num_heads
6.         self.q_project = nn.Linear(d_model, d_model)
7.         self.k_project = nn.Linear(d_model, d_model)
8.         self.v_project = nn.Linear(d_model, d_model)
9.         self.o_project = nn.Linear(d_model, d_model)
11.         self.dropout = nn.Dropout(dropout)

复制代码

• d_model：模型的维度，表现输入的特性维度。
  
• num_heads：注意力头的数量。多头注意力机制将输入分成多个差别的“头”，每个头学习差别的特性，末了再将这些特性归并起来。
  
• d_k：每个头的维度，计算公式为d_model // num_heads。比方，如果d_model=512，num_heads=8，则每个头的维度为512 // 8 = 64。
  
• q_project、k_project、v_project：这三个线性层分别用于将输入x投影到查询（Query）、键（Key）和值（Value）空间。投影后的维度仍然是d_model。
  
• o_project：输出投影层，将多头注意力的结果再次投影到d_model维度。
  
• dropout：用于防止过拟合的Dropout层。
  

2. 前向传播部门

1. def forward(self, x, attn_mask=None):
2.     batch_size, seq_len, d_model = x.shape
3.     Q = self.q_project(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
4.     K = self.k_project(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
5.     V = self.v_project(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)

复制代码

• 输入x：形状为(batch_size, seq_len, d_model)，其中seq_len是序列长度。
  
• 投影操作：
  
  
  
  • 利用q_project、k_project、v_project将输入x分别投影到查询（Q）、键（K）和值（V）空间。
    
  • 投影后的张量形状为(batch_size, seq_len, d_model)。
    
  
  
• 多头拆分：
  
  
  
  • 利用.view(batch_size, seq_len, self.num_heads, self.d_k)将投影后的张量拆分成多个头，形状变为(batch_size, seq_len, num_heads, d_k)。
    
  • 利用.transpose(1, 2)将头的维度提到前面，形状变为(batch_size, num_heads, seq_len, d_k)。
    
  
  

1.     atten_scores = Q @ K.transpose(2, 3) / math.sqrt(self.d_k)

复制代码

• 计算注意力分数：
  
  
  
  • 利用矩阵乘法@计算Q和K的点积，K.transpose(2, 3)将K的形状变为(batch_size, num_heads, d_k, seq_len)。
    
  • 点积结果的形状为(batch_size, num_heads, seq_len, seq_len)，表现每个位置之间的注意力分数。
    
  • 除以math.sqrt(self.d_k)是为了防止点积结果过大，导致梯度消失或爆炸。
    
  
  

1.     if attn_mask is not None:
2.         attn_mask = attn_mask.unsqueeze(1)
3.         atten_scores = atten_scores.masked_fill(attn_mask == 0, -1e9)

复制代码

• 注意力掩码（关于掩码的详细实现将在下一篇文章进行解说）：
  
  
  
  • 如果提供了注意力掩码attn_mask，则利用unsqueeze(1)将掩码的形状扩展为(batch_size, 1, seq_len, seq_len)。
    
  • 利用masked_fill将掩码为0的位置的注意力分数设置为一个非常小的值（如-1e9），这样在softmax计算时，这些位置的注意力权重会靠近0。
    
  
  

1.     atten_scores = torch.softmax(atten_scores, dim=-1)
2.     out = atten_scores @ V

复制代码

• 归一化注意力分数：
  
  
  
  • 利用torch.softmax对注意力分数进行归一化，形状仍为(batch_size, num_heads, seq_len, seq_len)。
    
  
  
• 计算加权和：
  
  
  
  • 利用矩阵乘法@将归一化后的注意力分数与V相乘，得到每个头的加权和，形状为(batch_size, num_heads, seq_len, d_k)。
    
  
  

1.     out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)
2.     out = self.o_project(out)
3.     return self.dropout(out)

复制代码

• 归并多头结果：
  
  
  
  • 利用.transpose(1, 2)将头的维度放回原来的位置，形状变为(batch_size, seq_len, num_heads, d_k)。
    
  • 利用.contiguous().view(batch_size, seq_len, d_model)将多头结果归并成一个张量，形状为(batch_size, seq_len, d_model)。
    
  
  
• 输出投影：
  
  
  
  • 利用o_project将归并后的结果再次投影到d_model维度。
    
  
  
• Dropout：
  
  
  
  • 利用dropout层对输出进行Dropout操作，防止过拟合。
    
  
  

需复现完整代码

1. class MultiHeadAttention(nn.Module):
2.     def __init__(self, d_model, num_heads, dropout):
3.         super().__init__()
4.         self.num_heads = num_heads
5.         self.d_k = d_model // num_heads
6.         self.q_project = nn.Linear(d_model, d_model)
7.         self.k_project = nn.Linear(d_model, d_model)
8.         self.v_project = nn.Linear(d_model, d_model)
9.         self.o_project = nn.Linear(d_model, d_model)
11.         self.dropout = nn.Dropout(dropout)    def forward(self, x, attn_mask=None):                batch_size, seq_len, d_model = x.shape        Q = self.q_project(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)        K = self.q_project(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)        V = self.q_project(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)        atten_scores = Q @ K.transpose(2, 3) / math.sqrt(self.d_k)        if attn_mask is not None:            attn_mask = attn_mask.unsqueeze(1)            atten_scores = atten_scores.masked_fill(attn_mask == 0, -1e9)        atten_scores = torch.softmax(atten_scores, dim=-1)        out = atten_scores @ V        out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)        out = self.o_project(out)        return self.dropout(out)

复制代码

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。