Attention 机制 PyTorch 手写实现：从 Q/K/V 矩阵到 8 头注意力完整代码-尧图网络科技

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, bias=False) self.keys = nn.Linear(embed_size, embed_size, bias=False) self.queries = nn.Linear(embed_size, embed_size, bias=False) 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_size=512, num_heads=8, dropout=0.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, cmap="viridis", xticklabels=sentence, yticklabels=sentence) 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, dtype=torch.double, requires_grad=True) # 运行梯度检查 test = gradcheck(model, input_tensor, eps=1e-6, atol=1e-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_layers=3): 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_dim=512, num_heads=8) # 我们的实现 our_attention = MultiHeadAttention(embed_size=512, num_heads=8) # 复制参数 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, mask=None): 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), diagonal=1) return mask == 0 # 转换为布尔掩码 mask = create_look_ahead_mask(10) masked_attention = MaskedMultiHeadAttention(embed_size=512, num_heads=8) output = masked_attention(x, mask=mask)