Attention 机制 PyTorch 手写实现从 Q/K/V 矩阵到 8 头注意力完整代码1. 理解 Attention 机制的核心概念在自然语言处理领域Attention 机制已经成为现代深度学习架构的核心组件。它的核心思想是让模型能够动态地关注输入序列中最相关的部分而不是像传统 RNN 那样强制按顺序处理信息。想象一下你在阅读一段技术文档时大脑会自然地聚焦于当前最相关的词汇和概念同时保留对上下文的理解。Attention 机制正是模拟了这一认知过程通过计算不同位置之间的相关性权重来决定在处理当前信息时应该注意哪些其他位置的信息。Attention 机制主要涉及三个关键矩阵Query (Q): 表示当前需要关注的内容Key (K): 表示可供关注的内容索引Value (V): 表示实际被关注的内容这三个矩阵通常通过线性变换从同一输入序列得到它们的关系可以用一个简单的比喻来理解Query 像是你在搜索引擎中输入的问题Key 是网页的标题而 Value 则是网页的实际内容。2. 单头 Self-Attention 的实现让我们从基础的 Self-Attention 开始逐步构建完整的实现。Self-Attention 的特点是 Q、K、V 都来自同一个输入序列。import torch import torch.nn as nn import math class SelfAttention(nn.Module): def __init__(self, embed_size): super(SelfAttention, self).__init__() self.embed_size embed_size # 初始化 Q、K、V 的线性变换层 self.values nn.Linear(embed_size, embed_size, biasFalse) self.keys nn.Linear(embed_size, embed_size, biasFalse) self.queries nn.Linear(embed_size, embed_size, biasFalse) def forward(self, x): # x 的形状: (batch_size, seq_length, embed_size) batch_size, seq_length, embed_size x.shape # 计算 Q, K, V Q self.queries(x) # (batch_size, seq_length, embed_size) K self.keys(x) # (batch_size, seq_length, embed_size) V self.values(x) # (batch_size, seq_length, embed_size) # 计算注意力分数 attention_scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.embed_size) # 应用 softmax 获取注意力权重 attention_weights torch.softmax(attention_scores, dim-1) # 应用注意力权重到 V out torch.matmul(attention_weights, V) return out注意在实际应用中我们通常会对注意力分数进行缩放除以√d_k以防止 softmax 函数的梯度消失问题。这是因为当输入维度较大时点积的结果可能会变得非常大导致 softmax 的梯度变得非常小。这个基础实现展示了 Attention 机制的核心计算步骤通过线性变换生成 Q、K、V计算 Q 和 K 的点积得到注意力分数应用 softmax 归一化得到注意力权重用注意力权重加权求和 V 得到输出3. 多头注意力机制单头注意力虽然有效但限制了模型同时关注不同位置信息的能力。多头注意力通过并行运行多个注意力头每个头学习不同的关注模式从而增强模型的表达能力。class MultiHeadAttention(nn.Module): def __init__(self, embed_size, num_heads): super(MultiHeadAttention, self).__init__() self.embed_size embed_size self.num_heads num_heads self.head_dim embed_size // num_heads assert ( self.head_dim * num_heads embed_size ), Embedding size needs to be divisible by number of heads # 线性变换层 self.values nn.Linear(embed_size, embed_size) self.keys nn.Linear(embed_size, embed_size) self.queries nn.Linear(embed_size, embed_size) self.fc_out nn.Linear(embed_size, embed_size) def forward(self, x): batch_size, seq_length, embed_size x.shape # 分割输入到多个头 Q self.queries(x).view(batch_size, seq_length, self.num_heads, self.head_dim) K self.keys(x).view(batch_size, seq_length, self.num_heads, self.head_dim) V self.values(x).view(batch_size, seq_length, self.num_heads, self.head_dim) # 转置以获得形状 (batch_size, num_heads, seq_length, head_dim) Q Q.transpose(1, 2) K K.transpose(1, 2) V V.transpose(1, 2) # 计算注意力分数 attention_scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim) # 应用 softmax attention_weights torch.softmax(attention_scores, dim-1) # 应用注意力权重到 V out torch.matmul(attention_weights, V) # 拼接多个头的输出 out out.transpose(1, 2).contiguous().view(batch_size, seq_length, self.embed_size) # 通过最后的线性层 out self.fc_out(out) return out多头注意力的关键改进在于将输入分割到多个头每个头学习不同的注意力模式并行计算多个注意力头拼接多个头的输出并通过线性层融合4. 完整的 8 头注意力实现现在我们将实现一个完整的 8 头注意力模块包含必要的层归一化和残差连接这是 Transformer 架构中的标准配置。class MultiHeadAttentionBlock(nn.Module): def __init__(self, embed_size512, num_heads8, dropout0.1): super(MultiHeadAttentionBlock, self).__init__() self.attention MultiHeadAttention(embed_size, num_heads) self.norm1 nn.LayerNorm(embed_size) self.norm2 nn.LayerNorm(embed_size) # 前馈网络 self.feed_forward nn.Sequential( nn.Linear(embed_size, 4 * embed_size), nn.ReLU(), nn.Linear(4 * embed_size, embed_size) ) self.dropout nn.Dropout(dropout) def forward(self, x): # 第一子层多头注意力 残差连接 层归一化 attention self.attention(x) x self.norm1(attention x) x self.dropout(x) # 第二子层前馈网络 残差连接 层归一化 forward self.feed_forward(x) x self.norm2(forward x) x self.dropout(x) return x这个完整实现包含了 Transformer 中的标准组件多头注意力机制层归一化LayerNorm残差连接前馈神经网络Dropout 正则化5. 可视化注意力权重理解模型关注什么是解释模型行为的关键。我们可以扩展我们的实现使其能够返回并可视化注意力权重。class VisualizableMultiHeadAttention(MultiHeadAttention): def forward(self, x): batch_size, seq_length, embed_size x.shape Q self.queries(x).view(batch_size, seq_length, self.num_heads, self.head_dim) K self.keys(x).view(batch_size, seq_length, self.num_heads, self.head_dim) V self.values(x).view(batch_size, seq_length, self.num_heads, self.head_dim) Q Q.transpose(1, 2) K K.transpose(1, 2) V V.transpose(1, 2) attention_scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim) attention_weights torch.softmax(attention_scores, dim-1) out torch.matmul(attention_weights, V) out out.transpose(1, 2).contiguous().view(batch_size, seq_length, self.embed_size) out self.fc_out(out) return out, attention_weights使用这个可视化版本我们可以获取注意力权重并绘制热力图import matplotlib.pyplot as plt import seaborn as sns def plot_attention_weights(attention_weights, sentence): # attention_weights 形状: (batch_size, num_heads, seq_length, seq_length) # 我们取第一个样本和第一个注意力头 attn attention_weights[0, 0].detach().numpy() plt.figure(figsize(10, 10)) sns.heatmap(attn, cmapviridis, xticklabelssentence, yticklabelssentence) plt.title(Attention Weights) plt.show()6. 梯度检查与调试为了确保我们的实现正确梯度检查是必不可少的。我们可以编写一个简单的梯度检查函数def gradient_check(module, input_tensor): input_tensor.requires_grad_(True) output module(input_tensor) # 创建一个与输出形状相同的随机梯度 grad_output torch.randn_like(output) # 反向传播 output.backward(grad_output) # 检查输入梯度是否存在 if input_tensor.grad is not None: print(梯度检查通过: 输入梯度存在) print(f输入梯度范数: {torch.norm(input_tensor.grad)}) else: print(梯度检查失败: 输入梯度不存在) # 检查参数梯度 for name, param in module.named_parameters(): if param.grad is None: print(f参数 {name} 没有梯度) else: print(f参数 {name} 梯度范数: {torch.norm(param.grad)})在实际项目中我们可以使用 PyTorch 的gradcheck工具进行更严格的数值梯度检查from torch.autograd import gradcheck def test_gradient(): embed_size 16 num_heads 4 model MultiHeadAttention(embed_size, num_heads).double() # 创建一个测试输入 input_tensor torch.randn(1, 5, embed_size, dtypetorch.double, requires_gradTrue) # 运行梯度检查 test gradcheck(model, input_tensor, eps1e-6, atol1e-4) print(梯度检查结果:, test)7. 性能优化技巧在实际应用中我们需要考虑实现的效率。以下是一些优化技巧矩阵乘法优化使用torch.bmm进行批量矩阵乘法内存效率使用原地操作减少内存使用半精度训练使用torch.cuda.amp进行混合精度训练class OptimizedMultiHeadAttention(nn.Module): def __init__(self, embed_size, num_heads): super().__init__() self.embed_size embed_size self.num_heads num_heads self.head_dim embed_size // num_heads # 合并所有线性变换以减少矩阵乘法次数 self.qkv_proj nn.Linear(embed_size, 3 * embed_size) self.fc_out nn.Linear(embed_size, embed_size) def forward(self, x): batch_size, seq_length, _ x.shape # 一次性计算 Q, K, V qkv self.qkv_proj(x) # 分割为 Q, K, V qkv qkv.reshape(batch_size, seq_length, 3, self.num_heads, self.head_dim) qkv qkv.permute(2, 0, 3, 1, 4) # (3, batch_size, num_heads, seq_length, head_dim) Q, K, V qkv[0], qkv[1], qkv[2] # 高效矩阵乘法 attention_scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim) attention_weights torch.softmax(attention_scores, dim-1) # 使用 torch.bmm 进行批量矩阵乘法 out torch.matmul(attention_weights, V) out out.transpose(1, 2).contiguous().view(batch_size, seq_length, self.embed_size) return self.fc_out(out)8. 实际应用示例让我们看一个完整的应用示例使用我们实现的注意力模块构建一个简单的文本分类器class AttentionClassifier(nn.Module): def __init__(self, vocab_size, embed_size, num_heads, num_classes, num_layers3): super().__init__() self.embedding nn.Embedding(vocab_size, embed_size) # 堆叠多个注意力层 self.attention_layers nn.ModuleList([ MultiHeadAttentionBlock(embed_size, num_heads) for _ in range(num_layers) ]) # 分类头 self.classifier nn.Sequential( nn.Linear(embed_size, embed_size // 2), nn.ReLU(), nn.Linear(embed_size // 2, num_classes) ) def forward(self, x): # 嵌入层 x self.embedding(x) # (batch_size, seq_length, embed_size) # 通过多个注意力层 for layer in self.attention_layers: x layer(x) # 取序列第一个位置的输出作为分类特征 # (也可以使用平均池化或最大池化) x x[:, 0, :] # 分类 return self.classifier(x)这个简单的分类器展示了如何将我们的注意力模块应用于实际任务。在实践中你可能还需要添加位置编码、更复杂的池化策略等组件。9. 测试与验证为了验证我们的实现是否正确我们可以编写单元测试并与 PyTorch 官方实现进行比较def test_against_official_implementation(): # 官方实现 official_attention nn.MultiheadAttention(embed_dim512, num_heads8) # 我们的实现 our_attention MultiHeadAttention(embed_size512, num_heads8) # 复制参数 our_attention.queries.weight.data official_attention.in_proj_weight[:512, :].data our_attention.keys.weight.data official_attention.in_proj_weight[512:1024, :].data our_attention.values.weight.data official_attention.in_proj_weight[1024:, :].data our_attention.fc_out.weight.data official_attention.out_proj.weight.data our_attention.fc_out.bias.data official_attention.out_proj.bias.data # 测试输入 x torch.randn(1, 10, 512) # (batch_size, seq_length, embed_size) # 官方输出 official_out, _ official_attention(x.transpose(0, 1), x.transpose(0, 1), x.transpose(0, 1)) official_out official_out.transpose(0, 1) # 我们的输出 our_out our_attention(x) # 比较 diff torch.abs(official_out - our_out).max().item() print(f与官方实现的最大差异: {diff}) if diff 1e-5: print(测试通过: 实现与官方版本一致) else: print(测试失败: 实现与官方版本存在显著差异)10. 扩展功能掩码注意力在许多序列任务中我们需要防止当前位置关注到未来的位置如在解码器中。这可以通过注意力掩码实现class MaskedMultiHeadAttention(MultiHeadAttention): def forward(self, x, maskNone): batch_size, seq_length, embed_size x.shape Q self.queries(x).view(batch_size, seq_length, self.num_heads, self.head_dim) K self.keys(x).view(batch_size, seq_length, self.num_heads, self.head_dim) V self.values(x).view(batch_size, seq_length, self.num_heads, self.head_dim) Q Q.transpose(1, 2) K K.transpose(1, 2) V V.transpose(1, 2) attention_scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim) # 应用掩码 if mask is not None: attention_scores attention_scores.masked_fill(mask 0, float(-1e20)) attention_weights torch.softmax(attention_scores, dim-1) out torch.matmul(attention_weights, V) out out.transpose(1, 2).contiguous().view(batch_size, seq_length, self.embed_size) out self.fc_out(out) return out使用示例# 创建一个下三角掩码防止关注未来位置 def create_look_ahead_mask(size): mask torch.triu(torch.ones(size, size), diagonal1) return mask 0 # 转换为布尔掩码 mask create_look_ahead_mask(10) masked_attention MaskedMultiHeadAttention(embed_size512, num_heads8) output masked_attention(x, maskmask)