### Transformer多头注意力机制简介 Transformer模型的核心组件之一是自注意力机制（Self-Attention），它允许模型关注输入序列的不同位置，捕获全局依赖关系[^1]。多头注意力机制（Multi-Head Attention, MHA）通过并行计算多个注意力头，分别学习不同的特征子空间表示，增强了模型表达能力。具体而言，输入向量被线性变换为查询（Q）、键（K）和值（V），并通过缩放点积注意力计算权重矩阵[^3]。最终，各头的结果拼接后再经过线性变换输出[^2]。 ```python import torch.nn as nn class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): super(MultiHeadAttention, self).__init__() assert d_model % num_heads == 0 self.d_k = d_model // num_heads self.num_heads = num_heads self.linear_qkv = nn.Linear(d_model, 3 * d_model) self.output_linear = nn.Linear(d_model, d_model) def forward(self, query, key, value, mask=None): batch_size = query.size(0) # Linear projections and split heads qkv = self.linear_qkv(query).view(batch_size, -1, self.num_heads, 3 * self.d_k) q, k, v = torch.chunk(qkv, chunks=3, dim=-1) # Scaled dot-product attention scores = torch.matmul(q, k.transpose(-2, -1)) / np.sqrt(self.d_k) if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) attn_weights = torch.softmax(scores, dim=-1) context = torch.matmul(attn_weights, v) # Concatenate heads and final linear projection output = self.output_linear(context.contiguous().view(batch_size, -1, self.num_heads * self.d_k)) return output ```