### Transformer模型中多头注意力机制的工作原理 在Transformer模型中，多头注意力机制是一种扩展自注意力的方法，旨在通过不同表示子空间的不同位置来捕获更丰富的上下文信息[^1]。 #### 查询、键和值向量的生成 对于给定的输入序列，首先将其映射到查询（Query）、键（Key）以及值（Value）三个不同的线性变换矩阵。这些矩阵通常由可训练参数构成，用于将原始输入转换成适合计算注意力分数的形式： ```python import torch.nn as nn class MultiHeadAttention(nn.Module): def __init__(self, embed_size, num_heads): super(MultiHeadAttention, self).__init__() self.embed_size = embed_size self.num_heads = num_heads # 定义Q,K,V的线性层 self.q_linear = nn.Linear(embed_size, embed_size) self.k_linear = nn.Linear(embed_size, embed_size) self.v_linear = nn.Linear(embed_size, embed_size) self.fc_out = nn.Linear(embed_size, embed_size) ``` #### 向量分割与独立处理 为了实现多头特性，在实际操作过程中会将上述得到的`embed_size`维度拆分为多个较小的部分——即所谓的“头部”。每个部分都对应着一组新的查询、键和值向量，并且可以在各自的子空间内独立执行自注意力运算[^3]。 假设嵌入尺寸为512维，而有8个注意头，则每个头负责64维的数据流。这种设计使得网络可以并行学习多种模式下的特征关联，增强了表达能力。 #### 自注意力计算流程 一旦获得了针对各个头部调整后的查询、键和值之后，就可以应用标准的缩放点积注意力公式来进行加权求和操作。此过程涉及两步：一是计算注意力权重；二是利用这些权重对相应的值进行加权平均以形成最终输出。 ```python def forward(self, query, key, value, mask=None): N = query.shape[0] q = self.q_linear(query).view(N, -1, self.num_heads, self.embed_size // self.num_heads) k = self.k_linear(key).view(N, -1, self.num_heads, self.embed_size // self.num_heads) v = self.v_linear(value).view(N, -1, self.num_heads, self.embed_size // self.num_heads) scores = torch.einsum("nqhd,nkhd->nhqk", [q, k]) / (self.embed_size ** 0.5) if mask is not None: scores = scores.masked_fill(mask == 0, float("-1e20")) attention = torch.softmax(scores, dim=-1) out = torch.einsum("nhql,nlhd->nqhd", [attention, v]).reshape( N, -1, self.embed_size ) return self.fc_out(out) ``` 以上代码展示了如何在一个批次内的样本间高效地完成多头注意力机制中的核心计算步骤。 #### 输入序列表示与上下文信息获取 经过多轮这样的自我关注后，原本单一方向上的词项现在拥有了关于整个句子结构更加全面的理解。这不仅有助于提高语言建模的效果，也促进了诸如机器翻译等任务的表现提升，因为模型学会了更好地编码远距离依赖关系及其间的复杂交互作用[^2]。