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, VQKV 矩阵的数学本质查询矩阵 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. 多头注意力机制原理多头注意力通过并行运行多个独立的注意力头让模型能够同时关注输入序列的不同子空间单头注意力多头注意力单一注意力模式多个注意力模式并行可能遗漏复杂关系捕获多样化关系计算效率高表达能力更强多头注意力的计算过程可以分为四个关键阶段线性投影将输入映射到多个子空间缩放点积注意力在每个子空间独立计算注意力拼接输出合并所有头的输出最终投影将拼接结果映射回原始维度def scaled_dot_product_attention(Q, K, V, maskNone): 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_weights3. PyTorch 实现 8 头注意力下面我们实现一个完整的 8 头注意力模块重点关注张量形状的变化class MultiHeadAttention(nn.Module): def __init__(self, embed_dim512, num_heads8): 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, maskNone): 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关键张量形状变化输入 x: (batch_size, seq_len, embed_dim)qkv投影后: (batch_size, seq_len, 3*embed_dim)重塑后: (batch_size, seq_len, num_heads, 3*head_dim)分离QKV后: 每个形状为 (batch_size, num_heads, seq_len, head_dim)注意力输出: (batch_size, num_heads, seq_len, head_dim)拼接后: (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:ichunk_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, dim2)2. 混合精度训练with torch.cuda.amp.autocast(): output, attn mha(input_tensor)3. 注意力掩码应用掩码类型用途实现方式填充掩码忽略padding位置mask (x ! pad_idx).unsqueeze(1)因果掩码防止未来信息泄露torch.tril(torch.ones(seq_len, seq_len))局部注意力限制注意力范围滑动窗口掩码性能优化对比优化方法速度提升内存节省适用场景内存高效注意力1.5x3x长序列处理混合精度2x1.5x大规模训练稀疏注意力3x4x特定领域模型在实际项目中多头注意力的头数选择需要权衡模型容量和计算效率。对于大多数NLP任务8头注意力在表达能力和计算成本之间提供了良好的平衡。