从‘单核’到‘多核’：用PyTorch代码实战，拆解Transformer中Self-Attention与Multi-Head Attention的性能差异

从‘单核’到‘多核’：用PyTorch代码实战拆解Transformer中Self-Attention与Multi-Head Attention的性能差异

当你在Jupyter Notebook中敲下第一行PyTorch代码时，可能从未想过一个简单的矩阵乘法背后隐藏着怎样的计算艺术。本文不是又一篇关于注意力机制的科普，而是一次深度技术潜水——我们将用可运行的代码，亲手揭开Self-Attention与Multi-Head Attention在计算效率、表征能力和实际应用中的本质区别。

1. 环境准备与基础实现

在开始对比实验前，我们需要搭建一个可复现的测试环境。建议使用Colab Pro或配备GPU的本地环境，因为注意力机制的计算会随着序列长度呈平方级增长。

import torch import torch.nn as nn import torch.nn.functional as F from time import time import matplotlib.pyplot as plt device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f"Using {device} device")

1.1 单头注意力实现

让我们从最基础的Self-Attention实现开始。以下代码展示了完整的单头注意力计算过程，特别注意形状变换的注释：

class SelfAttention(nn.Module): def __init__(self, embed_size): super().__init__() self.embed_size = embed_size self.query = nn.Linear(embed_size, embed_size) self.key = nn.Linear(embed_size, embed_size) self.value = nn.Linear(embed_size, embed_size) def forward(self, x): # x shape: (batch, seq_len, embed_size) Q = self.query(x) # (batch, seq_len, embed_size) K = self.key(x) # (batch, seq_len, embed_size) V = self.value(x) # (batch, seq_len, embed_size) scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.embed_size)) attention = F.softmax(scores, dim=-1) out = torch.matmul(attention, V) return out

关键计算步骤的时间复杂度分析：

注意：这里的n代表序列长度，d代表嵌入维度。实际应用中，当n较大时（如2048个token），QK^T的计算会成为性能瓶颈。

2. 多头注意力的并行化实现

真正的技术突破在于Multi-Head Attention的并行计算设计。下面这个实现展示了如何利用PyTorch的矩阵操作特性实现高效并行：

class MultiHeadAttention(nn.Module): def __init__(self, embed_size, num_heads=8): super().__init__() assert embed_size % num_heads == 0, "Embed size must be divisible by num_heads" self.embed_size = embed_size self.num_heads = num_heads self.head_dim = embed_size // num_heads self.query = nn.Linear(embed_size, embed_size) self.key = nn.Linear(embed_size, embed_size) self.value = nn.Linear(embed_size, embed_size) self.fc_out = nn.Linear(embed_size, embed_size) def split_heads(self, x): # x shape: (batch, seq_len, embed_size) batch_size = x.size(0) return x.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) def forward(self, x): Q = self.split_heads(self.query(x)) # (batch, num_heads, seq_len, head_dim) K = self.split_heads(self.key(x)) V = self.split_heads(self.value(x)) scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.head_dim)) attention = F.softmax(scores, dim=-1) out = torch.matmul(attention, V) # (batch, num_heads, seq_len, head_dim) out = out.transpose(1, 2).contiguous().view(x.size(0), -1, self.embed_size) return self.fc_out(out)

多头注意力的计算优势体现在三个维度：

1. 内存访问局部性：每个头的计算都在较小的head_dim空间进行，提高了缓存命中率
2. 并行计算潜力：不同头的计算相互独立，适合GPU的SIMD架构
3. 模型容量扩展：通过增加头数而非单一维度来提升模型表达能力

3. 性能基准测试

让我们设计一个严谨的对比实验。以下测试代码会测量两种注意力机制在不同序列长度下的表现：

def benchmark(attention_module, seq_lengths, embed_size=512, batch_size=32, warmup=5, repeat=10): times = [] for seq_len in seq_lengths: module = attention_module(embed_size).to(device) x = torch.rand(batch_size, seq_len, embed_size).to(device) # Warmup for _ in range(warmup): _ = module(x) torch.cuda.synchronize() # Timing start = time() for _ in range(repeat): _ = module(x) torch.cuda.synchronize() elapsed = (time() - start) / repeat times.append(elapsed * 1000) # convert to ms return times seq_lengths = [64, 128, 256, 512, 1024] single_head_times = benchmark(SelfAttention, seq_lengths) multi_head_times = benchmark(lambda x: MultiHeadAttention(x, num_heads=8), seq_lengths) plt.plot(seq_lengths, single_head_times, label='Single Head') plt.plot(seq_lengths, multi_head_times, label='Multi-Head (8 heads)') plt.xlabel('Sequence Length') plt.ylabel('Time (ms)') plt.title('Attention Computation Time Comparison') plt.legend() plt.show()

典型测试结果分析（NVIDIA V100 GPU）：

关键发现：随着序列长度增加，多头注意力的并行优势逐渐显现。在1024长度时，8头注意力比单头快1.7倍，这超出了简单的8倍理论值，说明除了并行化，内存访问模式的优化也带来了额外收益。

4. 表征能力实验

性能只是故事的一半，我们更关心两种机制学到的表征差异。设计以下对比实验：

def analyze_attention_patterns(text, tokenizer, model): inputs = tokenizer(text, return_tensors='pt').to(device) with torch.no_grad(): outputs = model(**inputs, output_attentions=True) # 可视化第一个头的注意力模式 first_head_attention = outputs.attentions[0][0, 0].cpu().numpy() fig, ax = plt.subplots(figsize=(10, 6)) im = ax.imshow(first_head_attention, cmap='viridis') ax.set_xticks(range(len(inputs.input_ids[0]))) ax.set_yticks(range(len(inputs.input_ids[0]))) ax.set_xticklabels(tokenizer.convert_ids_to_tokens(inputs.input_ids[0])) ax.set_yticklabels(tokenizer.convert_ids_to_tokens(inputs.input_ids[0])) plt.colorbar(im) plt.show() # 示例使用HuggingFace的BERT模型 from transformers import BertTokenizer, BertModel tokenizer = BertTokenizer.from_pretrained('bert-base-uncased') model = BertModel.from_pretrained('bert-base-uncased', output_attentions=True).to(device) analyze_attention_patterns("The cat sat on the mat because it was tired", tokenizer, model)

通过对比单头和多头模型的注意力模式，我们可以观察到：

• 单头注意力：倾向于学习全局的、综合的注意力模式，所有位置的关系被压缩到一个注意力矩阵中
• 多头注意力：不同头会自发地学习不同的关注模式。在BERT中常见：
  • 局部语法关系（相邻词关联）
  • 长程依赖（如代词指代）
  • 标点符号关注
  • 内容与位置的特殊交互

在具体任务中的表现差异：

这些实验证实了多头设计在捕捉复杂语言关系时的优势，特别是在需要同时处理多种依赖关系的任务上。