Transformer 多头自注意力 PyTorch 实现从 QKV 矩阵到 8 头并行计算在自然语言处理领域Transformer 架构凭借其强大的并行计算能力和长距离依赖建模优势已成为现代深度学习模型的基石。本文将聚焦 Transformer 中最核心的多头自注意力机制通过 PyTorch 实现从基础矩阵运算到 8 头并行计算的完整过程。1. 自注意力机制基础自注意力机制的核心思想是通过计算序列内部元素间的相关性动态生成每个位置的上下文感知表示。其计算过程涉及三个关键矩阵查询矩阵Query、键矩阵Key和值矩阵Value。1.1 QKV 矩阵生成首先定义输入序列的嵌入表示。假设我们有一个包含 5 个词元的序列每个词元的嵌入维度为 512import torch import torch.nn as nn batch_size 1 seq_len 5 embed_dim 512 d_k d_v 64 # 键/值向量维度 # 随机初始化输入嵌入 inputs torch.randn(batch_size, seq_len, embed_dim)接下来创建线性变换层来生成 Q、K、V 矩阵class QKVProjection(nn.Module): def __init__(self, embed_dim, d_k, d_v): super().__init__() self.W_q nn.Linear(embed_dim, d_k) self.W_k nn.Linear(embed_dim, d_k) self.W_v nn.Linear(embed_dim, d_v) def forward(self, x): Q self.W_q(x) # (batch_size, seq_len, d_k) K self.W_k(x) # (batch_size, seq_len, d_k) V self.W_v(x) # (batch_size, seq_len, d_v) return Q, K, V qkv_proj QKVProjection(embed_dim, d_k, d_v) Q, K, V qkv_proj(inputs)1.2 注意力分数计算注意力分数的计算分为三个关键步骤点积计算查询与键的相似度缩放处理防止梯度消失Softmax归一化生成注意力权重def scaled_dot_product_attention(Q, K, V, maskNone): # 计算点积注意力分数 scores torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k)) # 应用掩码如有 if mask is not None: scores scores.masked_fill(mask 0, -1e9) # Softmax归一化 attn_weights torch.softmax(scores, dim-1) # 加权求和 output torch.matmul(attn_weights, V) return output, attn_weights attention_output, attn_weights scaled_dot_product_attention(Q, K, V)下表展示了单头注意力机制中各张量的维度变化操作输入维度输出维度QKV投影(1,5,512)Q/K:(1,5,64), V:(1,5,64)点积计算Q:(1,5,64), K:(1,5,64)(1,5,5)加权求和(1,5,5), V:(1,5,64)(1,5,64)2. 多头注意力机制实现多头注意力的核心思想是将输入序列投影到多个子空间在不同表示空间中学习多样化的特征关系。2.1 多头投影层我们实现一个支持多头投影的线性变换层class MultiHeadProjection(nn.Module): def __init__(self, embed_dim, d_k, d_v, num_heads): super().__init__() self.num_heads num_heads self.d_k d_k self.d_v d_v # 合并所有头的投影矩阵 self.W_q nn.Linear(embed_dim, d_k * num_heads) self.W_k nn.Linear(embed_dim, d_k * num_heads) self.W_v nn.Linear(embed_dim, d_v * num_heads) def forward(self, x): batch_size x.size(0) # 投影并分割多头 Q self.W_q(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) K self.W_k(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) V self.W_v(x).view(batch_size, -1, self.num_heads, self.d_v).transpose(1, 2) return Q, K, V2.2 8头并行计算现在我们实现完整的 8 头注意力机制class MultiHeadAttention(nn.Module): def __init__(self, embed_dim, num_heads8): super().__init__() assert embed_dim % num_heads 0, embed_dim必须能被num_heads整除 self.embed_dim embed_dim self.num_heads num_heads self.d_k embed_dim // num_heads self.d_v self.d_k # 投影层 self.qkv_proj MultiHeadProjection(embed_dim, self.d_k, self.d_v, num_heads) # 输出线性层 self.out_proj nn.Linear(num_heads * self.d_v, embed_dim) def forward(self, x, maskNone): batch_size x.size(0) # 生成Q,K,V (每个头的维度为d_k/d_v) Q, K, V self.qkv_proj(x) # 缩放点积注意力 scores torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.d_k)) if mask is not None: scores scores.masked_fill(mask 0, -1e9) attn_weights torch.softmax(scores, dim-1) attn_output torch.matmul(attn_weights, V) # 拼接多头输出 attn_output attn_output.transpose(1, 2).contiguous().view( batch_size, -1, self.num_heads * self.d_v) # 最终线性变换 output self.out_proj(attn_output) return output, attn_weights2.3 8头注意力可视化为了直观理解多头注意力的工作原理我们可以可视化不同注意力头的关注模式import matplotlib.pyplot as plt def plot_attention_weights(attn_weights, head_idx0): # attn_weights形状: (batch_size, num_heads, seq_len, seq_len) plt.figure(figsize(10, 8)) plt.imshow(attn_weights[0, head_idx].detach().numpy(), cmapviridis) plt.colorbar() plt.xlabel(Key Positions) plt.ylabel(Query Positions) plt.title(fAttention Head {head_idx1}) plt.show() # 测试8头注意力 mha MultiHeadAttention(embed_dim, num_heads8) output, attn_weights mha(inputs) # 可视化第一个注意力头 plot_attention_weights(attn_weights, head_idx0)提示在实际应用中不同注意力头往往会学习到不同的关注模式有的关注局部上下文有的关注特定语法关系这种多样性正是多头注意力的优势所在。3. 完整Transformer层实现将多头注意力机制与位置前馈网络结合构建完整的Transformer编码层3.1 位置前馈网络class PositionwiseFeedForward(nn.Module): def __init__(self, embed_dim, ff_dim2048): super().__init__() self.linear1 nn.Linear(embed_dim, ff_dim) self.linear2 nn.Linear(ff_dim, embed_dim) self.activation nn.GELU() def forward(self, x): return self.linear2(self.activation(self.linear1(x)))3.2 残差连接与层归一化class TransformerEncoderLayer(nn.Module): def __init__(self, embed_dim, num_heads, ff_dim2048): super().__init__() self.self_attn MultiHeadAttention(embed_dim, num_heads) self.ffn PositionwiseFeedForward(embed_dim, ff_dim) self.norm1 nn.LayerNorm(embed_dim) self.norm2 nn.LayerNorm(embed_dim) def forward(self, x, maskNone): # 自注意力子层 attn_output, _ self.self_attn(x, mask) x x attn_output # 残差连接 x self.norm1(x) # 层归一化 # 前馈网络子层 ffn_output self.ffn(x) x x ffn_output # 残差连接 x self.norm2(x) # 层归一化 return x3.3 测试完整编码层encoder_layer TransformerEncoderLayer(embed_dim, num_heads8) output encoder_layer(inputs) print(f输入形状: {inputs.shape}) print(f输出形状: {output.shape})4. 工程实践技巧在实际项目中实现高效的多头注意力需要考虑以下关键因素4.1 内存优化策略处理长序列时注意力矩阵会消耗大量内存。对于序列长度L注意力矩阵的空间复杂度为O(L²)。以下是一些优化方法分块计算将大矩阵拆分为小块处理混合精度训练使用FP16减少内存占用稀疏注意力只计算关键位置对的注意力分数# 示例内存高效的注意力计算 def memory_efficient_attention(Q, K, V, chunk_size64): batch, heads, seq_len, dim Q.shape output torch.zeros_like(V) for i in range(0, seq_len, chunk_size): Q_chunk Q[:, :, i:ichunk_size] scores torch.matmul(Q_chunk, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(dim)) attn torch.softmax(scores, dim-1) output[:, :, i:ichunk_size] torch.matmul(attn, V) return output4.2 并行计算加速利用PyTorch的并行计算特性可以显著提升多头注意力的计算效率import torch.nn.functional as F def parallel_multihead_attention(Q, K, V, maskNone): # Q,K,V形状: (batch, heads, seq_len, dim) scores torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(Q.size(-1))) if mask is not None: scores scores.masked_fill(mask 0, -1e9) attn F.softmax(scores, dim-1) return torch.matmul(attn, V)4.3 梯度检查点技术对于极深模型可以使用梯度检查点来节省内存from torch.utils.checkpoint import checkpoint class CheckpointedMHA(nn.Module): def __init__(self, embed_dim, num_heads): super().__init__() self.mha MultiHeadAttention(embed_dim, num_heads) def forward(self, x): return checkpoint(self.mha, x) # 使用示例 checkpointed_mha CheckpointedMHA(embed_dim, 8) output checkpointed_mha(inputs)下表对比了不同优化技术的效果优化技术内存占用计算速度实现复杂度原始实现高中等低分块计算中慢中混合精度低快低梯度检查点很低慢高