多头注意力（MultiHeadAttention）python实现

### 多头注意力机制 (Multi-Head Attention) 的 Python 实现 以下是基于 Transformer 架构中的多头注意力机制的一个标准 PyTorch 实现。该实现涵盖了前向传播过程，并通过梯度计算支持反向传播。 #### 前向传播逻辑 在 Multi-Head Attention 中，输入张量被分为多个头部（heads），每个头部独立计算 Query、Key 和 Value 的线性变换[^1]。随后，这些头部的结果会被拼接并再次经过一个线性变换得到最终输出。 ```python import torch import torch.nn as nn import math class MultiHeadAttention(nn.Module): def __init__(self, embed_dim, num_heads, dropout=0.1): super(MultiHeadAttention, self).__init__() assert embed_dim % num_heads == 0, "Embedding dimension must be divisible by number of heads." self.embed_dim = embed_dim self.num_heads = num_heads self.head_dim = embed_dim // num_heads self.q_linear = nn.Linear(embed_dim, embed_dim) self.k_linear = nn.Linear(embed_dim, embed_dim) self.v_linear = nn.Linear(embed_dim, embed_dim) self.dropout = nn.Dropout(dropout) self.out_projection = nn.Linear(embed_dim, embed_dim) def scaled_dot_product_attention(self, q, k, v, mask=None): d_k = q.size(-1) scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k) # Scaled dot-product if mask is not None: scores = scores.masked_fill(mask == 0, float('-inf')) attn_weights = torch.softmax(scores, dim=-1) attn_weights = self.dropout(attn_weights) output = torch.matmul(attn_weights, v) return output, attn_weights def split_heads(self, x): batch_size, seq_len, _ = x.size() x = x.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) return x def combine_heads(self, x): batch_size, _, seq_len, head_dim = x.size() x = x.transpose(1, 2).contiguous().view(batch_size, seq_len, self.embed_dim) return x def forward(self, query, key, value, mask=None): batch_size = query.size(0) # Linear projections Q = self.split_heads(self.q_linear(query)) K = self.split_heads(self.k_linear(key)) V = self.split_heads(self.v_linear(value)) # Scaled Dot Product Attention outputs, attn_weights = self.scaled_dot_product_attention(Q, K, V, mask=mask) # Concatenate and apply final linear transformation outputs = self.combine_heads(outputs) outputs = self.out_projection(outputs) return outputs, attn_weights ``` #### 关键点解析 1. **线性变换**: 输入 `query`、`key` 和 `value` 都会分别通过三个不同的线性层进行转换[^3]。 2. **分头操作**: 使用 `split_heads` 方法将嵌入维度拆分成多个头部，以便并行处理不同子空间的信息。 3. **缩放点积注意力**: 计算查询和键之间的相似度分数，并对其进行缩放以稳定梯度[^1]。 4. **掩码机制**: 如果存在填充标记，则可以通过掩码屏蔽掉无意义的位置。 5. **组合头部**: 将各头部的输出重新组合成单一表示形式，并通过最后一个线性层调整回原始嵌入维度。 --- ###