多头自注意力机制

### 多头自注意力机制概述 多头自注意力机制是一种增强神经网络模型表示能力的技术，在自然语言处理和其他序列建模任务中表现出色。该方法允许模型并行地学习不同位置之间的关系，从而捕捉更丰富的上下文信息[^1]。 ### 原理详解 在实现多头自注意力时，输入数据被转换成三个不同的向量集：查询（Query）、键（Key）以及值（Value）。这些变换通过线性映射完成，具体来说就是将原始特征矩阵 \(X\) 分别乘上对应的权重参数得到新的矩阵： \[Q_i=XW^Q_i, \quad K_i=XW^K_i,\quad V_i=XW^V_i\] 其中下标\(i\)代表第\(i\)个注意力头；\(W^{Q}_i,W^{K}_i,W^{V}_i\)分别是用于计算查询、键和值得到的可训练参数矩阵。这样做之后，每个头部都能独立工作，形成各自的注意力分布，并最终汇总起来作为整体输出的一部分[^3]。 对于每一个注意力头而言，其核心运算遵循标准缩放点积注意公式: \[Attention(Q,K,V)=softmax(\frac{QK^\top}{\sqrt{d_k}})V\] 这里\(d_k\)指的是键维度大小，用来稳定梯度传播过程中的数值范围。经过上述操作后获得的结果再经由全连接层投影回原空间尺寸，最后叠加所有单独头产生的贡献构成完整的多头自注意力表达形式。 ```python import torch.nn as nn import math class MultiHeadSelfAttention(nn.Module): def __init__(self, d_model, num_heads): super(MultiHeadSelfAttention, self).__init__() assert d_model % num_heads == 0 self.d_k = d_model // num_heads self.num_heads = num_heads # 定义权重组件 self.W_Q = nn.Linear(d_model, d_model) self.W_K = nn.Linear(d_model, d_model) self.W_V = nn.Linear(d_model, d_model) self.fc_out = nn.Linear(d_model, d_model) def forward(self, x): batch_size = x.size(0) # 计算QKV矩阵 Q = self.W_Q(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) K = self.W_K(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) V = self.W_V(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) attention_weights = F.softmax(scores, dim=-1) context_vector = torch.matmul(attention_weights, V) output = context_vector.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.d_k) return self.fc_out(output) ``` 此代码片段展示了如何构建一个多头自注意力模块，它接收一个形状为 `[batch_size, seq_len, embed_dim]` 的张量 `x` 并返回相同形状的新嵌入表示。内部实现了对输入序列元素间相互依赖性的编码，使得下游任务能够更好地利用全局信息进行预测或生成[^2]。