在Transformer架构里，多头注意力机制（Multi - Head Attention）中涉及到`kvahead`（多头的`K`和`V`）的概念，下面详细阐述其在注意力公式里的应用。 ### 基本概念 在注意力机制中，输入序列会通过线性变换得到查询（Query，`Q`）、键（Key，`K`）和值（Value，`V`）三个矩阵。多头注意力机制把`Q`、`K`和`V`分别划分成多个头（head），每个头独立计算注意力分数，最后再把结果拼接起来。 ### 注意力公式 单头注意力的计算公式如下： \[ \text{Attention}(Q, K, V)=\text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V \] 其中，$Q$是查询矩阵，$K$是键矩阵，$V$是值矩阵，$d_k$是键向量的维度。 ### `kvahead`（`K`和`V`多头）的应用步骤 1. **线性变换与多头划分**： 输入的序列$X$会通过三个不同的线性变换得到$Q$、$K$和$V$。假设输入维度是$d_{model}$，头的数量是$h$，每个头的维度是$d_k = \frac{d_{model}}{h}$。 \[ \begin{align*} Q&=XW^Q\\ K&=XW^K\\ V&=XW^V \end{align*} \] 其中，$W^Q$、$W^K$和$W^V$是可学习的权重矩阵。然后把$Q$、$K$和$V$分别划分成$h$个头： \[ \begin{align*} Q_i&=Q[:, :, i\times d_k:(i + 1)\times d_k]\\ K_i&=K[:, :, i\times d_k:(i + 1)\times d_k]\\ V_i&=V[:, :, i\times d_k:(i + 1)\times d_k] \end{align*} \] 这里的$i$表示第$i$个头。 2. **每个头计算注意力**： 对每个头分别应用注意力公式： \[ \text{Attention}_i(Q_i, K_i, V_i)=\text{softmax}\left(\frac{Q_iK_i^T}{\sqrt{d_k}}\right)V_i \] 3. **拼接结果**： 将所有头的注意力结果拼接起来： \[ \text{MultiHead}(Q, K, V)=\text{Concat}(\text{Attention}_1, \text{Attention}_2, \cdots, \text{Attention}_h)W^O \] 其中，$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, "d_model must be divisible by num_heads" self.d_model = d_model self.num_heads = num_heads self.d_k = d_model // 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.W_o = nn.Linear(d_model, d_model) def forward(self, Q, K, V): batch_size = Q.size(0) # 线性变换 Q = self.W_q(Q) K = self.W_k(K) V = self.W_v(V) # 划分成多个头 Q = Q.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) K = K.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) V = V.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) # 计算注意力分数 attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.d_k, dtype=torch.float32)) attn_probs = torch.softmax(attn_scores, dim=-1) # 计算注意力输出 attn_output = torch.matmul(attn_probs, V) # 拼接结果 attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) # 最终线性变换 output = self.W_o(attn_output) return output ```