Transformer QKV 矩阵与多头注意力：从 PyTorch 代码到 8 头并行计算详解

Transformer QKV 矩阵与多头注意力：从 PyTorch 代码到 8 头并行计算详解

在深度学习领域，Transformer 架构已经成为自然语言处理、计算机视觉等多个领域的基石。其核心组件——自注意力机制（Self-Attention）和多头注意力（Multi-Head Attention）——赋予了模型捕捉长距离依赖关系的能力。本文将深入探讨 QKV 矩阵的生成、多头拆分、并行计算与最终融合的完整流程，并通过 PyTorch 实现一个完整的 MultiHeadAttention 类。

1. QKV 矩阵：自注意力的核心

自注意力机制的核心在于三个关键矩阵：查询矩阵（Query）、键矩阵（Key）和值矩阵（Value），简称 QKV。这三个矩阵通过线性变换从输入序列中提取不同特征：

import torch import torch.nn as nn class QKVProjection(nn.Module): def __init__(self, embed_dim): super().__init__() self.query = nn.Linear(embed_dim, embed_dim) self.key = nn.Linear(embed_dim, embed_dim) self.value = nn.Linear(embed_dim, embed_dim) def forward(self, x): Q = self.query(x) # (batch_size, seq_len, embed_dim) K = self.key(x) # (batch_size, seq_len, embed_dim) V = self.value(x) # (batch_size, seq_len, embed_dim) return Q, K, V

QKV 矩阵的数学本质：

• 查询矩阵 Q：表示当前词对上下文的"询问"，决定关注哪些位置
• 键矩阵 K：表示每个词的"身份标识"，用于与查询匹配
• 值矩阵 V：包含实际要传递的信息，根据注意力权重加权求和

计算注意力得分的标准公式为：

$$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V $$

其中 $d_k$ 是键向量的维度，缩放因子 $\sqrt{d_k}$ 用于防止点积结果过大导致梯度消失。

2. 多头注意力机制原理

多头注意力通过并行运行多个独立的注意力头，让模型能够同时关注输入序列的不同子空间：

多头注意力的计算过程可以分为四个关键阶段：

1. 线性投影：将输入映射到多个子空间
2. 缩放点积注意力：在每个子空间独立计算注意力
3. 拼接输出：合并所有头的输出
4. 最终投影：将拼接结果映射回原始维度

def scaled_dot_product_attention(Q, K, V, mask=None): d_k = Q.size(-1) scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k) if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) attn_weights = F.softmax(scores, dim=-1) output = torch.matmul(attn_weights, V) return output, attn_weights

3. PyTorch 实现 8 头注意力

下面我们实现一个完整的 8 头注意力模块，重点关注张量形状的变化：

class MultiHeadAttention(nn.Module): def __init__(self, embed_dim=512, num_heads=8): super().__init__() assert embed_dim % num_heads == 0, "Embedding dimension must be divisible by number of heads" self.embed_dim = embed_dim self.num_heads = num_heads self.head_dim = embed_dim // num_heads # 合并的QKV投影，提高计算效率 self.qkv_proj = nn.Linear(embed_dim, 3 * embed_dim) self.out_proj = nn.Linear(embed_dim, embed_dim) def forward(self, x, mask=None): batch_size, seq_len, _ = x.shape # 步骤1：合并计算QKV (batch_size, seq_len, 3*embed_dim) qkv = self.qkv_proj(x) # 步骤2：拆分QKV并重塑为多头 (batch_size, seq_len, num_heads, 3*head_dim) qkv = qkv.reshape(batch_size, seq_len, self.num_heads, 3 * self.head_dim) # 步骤3：分离Q,K,V (each shape: batch_size, num_heads, seq_len, head_dim) q, k, v = torch.chunk(qkv, 3, dim=-1) q = q.transpose(1, 2) # (batch_size, num_heads, seq_len, head_dim) k = k.transpose(1, 2) v = v.transpose(1, 2) # 步骤4：计算缩放点积注意力 attn_output, attn_weights = scaled_dot_product_attention(q, k, v, mask) # 步骤5：拼接多头输出 (batch_size, seq_len, embed_dim) attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, self.embed_dim) # 步骤6：最终投影 output = self.out_proj(attn_output) return output, attn_weights

关键张量形状变化：

1. 输入 x: (batch_size, seq_len, embed_dim)
2. qkv投影后: (batch_size, seq_len, 3*embed_dim)
3. 重塑后: (batch_size, seq_len, num_heads, 3*head_dim)
4. 分离QKV后: 每个形状为 (batch_size, num_heads, seq_len, head_dim)
5. 注意力输出: (batch_size, num_heads, seq_len, head_dim)
6. 拼接后: (batch_size, seq_len, embed_dim)

4. 多头注意力的高级应用与优化

在实际应用中，我们可以对基础的多头注意力进行多种优化：

1. 内存高效注意力：

class MemoryEfficientAttention(nn.Module): def forward(self, q, k, v): # 分块计算注意力，减少内存占用 chunk_size = 256 output = [] for i in range(0, q.size(2), chunk_size): q_chunk = q[:, :, i:i+chunk_size] scores = torch.einsum('bhid,bhjd->bhij', q_chunk, k) attn = F.softmax(scores / math.sqrt(k.size(-1)), dim=-1) output.append(torch.einsum('bhij,bhjd->bhid', attn, v)) return torch.cat(output, dim=2)

2. 混合精度训练：

with torch.cuda.amp.autocast(): output, attn = mha(input_tensor)

3. 注意力掩码应用：

性能优化对比：

在实际项目中，多头注意力的头数选择需要权衡模型容量和计算效率。对于大多数NLP任务，8头注意力在表达能力和计算成本之间提供了良好的平衡。