### 多头注意力机制的工作原理 多头注意力机制（Multi-Head Attention）是Transformer模型的核心组件之一，它通过并行计算多个独立的注意力机制来增强模型的表现力。具体来说，输入向量被线性变换为查询（Query）、键（Key）和值（Value），并通过缩放点积注意力（Scaled Dot-Product Attention）计算加权平均值[^2]。 #### 查询、键和值的定义 在多头注意力机制中，每个位置的输入表示会被映射到三个不同的向量：查询（Q）、键（K）和值（V）。这些向量分别用于衡量不同位置之间的关系以及提取有用的信息。为了实现这一点，通常会使用矩阵乘法将输入嵌入转换成所需的维度： \[ Q = XW_Q, \quad K = XW_K, \quad V = XW_V \] 其中 \(X\) 是输入序列的嵌入矩阵，\(W_Q\), \(W_K\), 和 \(W_V\) 是可学习的参数矩阵。 #### 缩放点积注意力 一旦得到了查询、键和值，就可以通过以下公式计算注意力分数： \[ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V \] 这里，\(d_k\) 表示键向量的维度，分母中的平方根是为了防止点积过大而导致梯度消失或爆炸问题。 #### 并行化与头部组合 多头注意力的关键在于引入了多个平行运行的注意力层，即所谓的“头”。每一个头都拥有自己的一组权重矩阵，并且能够捕捉输入的不同方面特征。最终，各个头的结果会被拼接起来再经过一次线性投影得到输出： \[ \text{MultiHead}(Q, K, V) = [\text{head}_1; \dots ; \text{head}_h] W_O \] 这里的 \(h\) 表示头的数量，而 \(W_O\) 则是用来融合各头信息的最后一组共享权重。 以下是基于PyTorch框架的一个简单代码示例展示如何构建一个多头注意力模块: ```python import torch 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.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 scaled_dot_product_attention(self, q, k, v, mask=None): scores = torch.matmul(q, k.transpose(-2, -1)) / (k.size(-1)**0.5) if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) attention_weights = torch.softmax(scores, dim=-1) output = torch.matmul(attention_weights, v) return output def split_heads(self, x): batch_size, seq_len, _ = x.shape x = x.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) return x.flatten(end_dim=1) # shape: (batch * heads, seq_len, depth) def forward(self, query, key, value, mask=None): batch_size = query.size(0) q = self.split_heads(self.W_q(query)) k = self.split_heads(self.W_k(key)) v = self.split_heads(self.W_v(value)) context_vectors = self.scaled_dot_product_attention(q, k, v, mask) context_vectors = context_vectors.unflatten( 0, (batch_size, self.num_heads)).transpose(1, 2).contiguous() concat_context_vector = context_vectors.view(batch_size, -1, self.d_model) output = self.fc_out(concat_context_vector) return output ``` 此段代码实现了基本的多头注意力建构单元，包括分割头操作、执行点积注意力运算以及最后重新组装结果等功能[^3]。