PyTorch 2.0 多头自注意力机制实战：从零实现 8 头注意力模块

在自然语言处理和计算机视觉领域，Transformer 架构已经成为当前最强大的模型基础。而 Transformer 的核心组件——多头自注意力机制（Multi-Head Self-Attention），则是其成功的关键所在。本文将带你深入理解这一机制，并手把手教你用 PyTorch 2.0 实现一个完整的 8 头自注意力模块。

1. 自注意力机制基础

自注意力机制的核心思想是让序列中的每个元素都能直接关注到序列中的所有其他元素，通过动态计算注意力权重来决定关注的程度。这种机制打破了传统 RNN 按顺序处理数据的限制，能够更好地捕捉长距离依赖关系。

1.1 关键概念：Q, K, V

自注意力机制涉及三个核心矩阵：

• Query (Q): 表示当前关注的元素
• Key (K): 表示被查询的元素
• Value (V): 实际被提取的信息

计算过程可以表示为：

Attention(Q, K, V) = softmax(QK^T/√d_k)V

其中 d_k 是 Key 的维度，√d_k 的缩放是为了防止点积结果过大导致 softmax 梯度消失。

1.2 自注意力的 PyTorch 实现

import torch import torch.nn as nn import torch.nn.functional as F class SelfAttention(nn.Module): def __init__(self, embed_size): super(SelfAttention, self).__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_size, seq_len, embed_size) Q = self.query(x) # (batch_size, seq_len, embed_size) K = self.key(x) # (batch_size, seq_len, embed_size) V = self.value(x) # (batch_size, seq_len, embed_size) # 计算注意力分数 scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.embed_size, dtype=torch.float32)) attention = F.softmax(scores, dim=-1) # 应用注意力权重 out = torch.matmul(attention, V) return out

2. 多头自注意力机制

单一注意力头只能学习一种关注模式，而多头注意力通过并行多个注意力头，让模型能够同时关注不同方面的信息。

2.1 多头机制原理

多头注意力的关键步骤：

1. 将 Q, K, V 线性投影到多个子空间
2. 在每个子空间独立计算注意力
3. 拼接所有头的输出
4. 通过最终线性变换得到结果

数学表达式：

MultiHead(Q, K, V) = Concat(head_1, ..., head_h)W^O where head_i = Attention(QW_i^Q, KW_i^K, VW_i^V)

2.2 8 头注意力模块设计

我们将实现一个完整的 8 头注意力模块，包含以下组件：

• 独立的 Q, K, V 线性变换层
• 头拆分与合并机制
• 残差连接和层归一化
• 前馈网络子层

class MultiHeadAttention(nn.Module): def __init__(self, embed_size=512, num_heads=8): 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, "Embed size must be divisible by 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) # 层归一化 self.layer_norm = nn.LayerNorm(embed_size) def split_heads(self, x): # 将嵌入维度拆分为多个头 batch_size = x.size(0) x = x.view(batch_size, -1, self.num_heads, self.head_dim) return x.transpose(1, 2) # (batch_size, num_heads, seq_len, head_dim) def forward(self, x, mask=None): batch_size, seq_len, _ = x.size() # 保留原始输入用于残差连接 residual = x # 线性变换并拆分头 Q = self.split_heads(self.query(x)) 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, dtype=torch.float32)) # 应用掩码（如果有） if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) attention = F.softmax(scores, dim=-1) # 应用注意力权重并合并头 out = torch.matmul(attention, V) out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, self.embed_size) # 最终线性变换和残差连接 out = self.fc_out(out) out = self.layer_norm(out + residual) return out

3. 完整 Transformer 层实现

一个完整的 Transformer 层通常包含：

1. 多头自注意力子层
2. 前馈神经网络子层
3. 残差连接和层归一化

3.1 前馈网络子层

class FeedForward(nn.Module): def __init__(self, embed_size=512, ff_dim=2048): super(FeedForward, self).__init__() self.embed_size = embed_size self.ff_dim = ff_dim self.fc1 = nn.Linear(embed_size, ff_dim) self.fc2 = nn.Linear(ff_dim, embed_size) self.relu = nn.ReLU() self.layer_norm = nn.LayerNorm(embed_size) def forward(self, x): residual = x out = self.fc1(x) out = self.relu(out) out = self.fc2(out) out = self.layer_norm(out + residual) return out

3.2 完整 Transformer 层

class TransformerLayer(nn.Module): def __init__(self, embed_size=512, num_heads=8, ff_dim=2048): super(TransformerLayer, self).__init__() self.attention = MultiHeadAttention(embed_size, num_heads) self.ffn = FeedForward(embed_size, ff_dim) def forward(self, x, mask=None): x = self.attention(x, mask) x = self.ffn(x) return x

4. 模块测试与验证

为了验证我们的实现是否正确，我们需要构建测试用例并检查输出。

4.1 测试数据准备

def test_implementation(): # 参数设置 batch_size = 2 seq_len = 10 embed_size = 512 num_heads = 8 # 创建随机输入 x = torch.randn(batch_size, seq_len, embed_size) # 初始化模型 model = MultiHeadAttention(embed_size, num_heads) # 前向传播 output = model(x) # 检查输出形状 assert output.shape == (batch_size, seq_len, embed_size), \ f"Expected shape {(batch_size, seq_len, embed_size)}, got {output.shape}" print("测试通过！输出形状:", output.shape) test_implementation()

4.2 梯度检查

def check_gradients(): model = MultiHeadAttention(embed_size=64, num_heads=8) x = torch.randn(1, 5, 64, requires_grad=True) # 前向传播 output = model(x) # 反向传播 output.mean().backward() # 检查梯度 for name, param in model.named_parameters(): if param.grad is None: print(f"参数 {name} 没有梯度！") elif torch.all(param.grad == 0): print(f"参数 {name} 梯度为零！") else: print(f"参数 {name} 梯度正常") print("梯度检查完成") check_gradients()

5. 性能优化技巧

PyTorch 2.0 引入了多项性能优化，我们可以利用这些特性来提升多头注意力的计算效率。

5.1 使用 torch.nn.functional.scaled_dot_product_attention

PyTorch 2.0 提供了优化后的注意力计算函数：

def optimized_attention(Q, K, V, mask=None): return F.scaled_dot_product_attention(Q, K, V, attn_mask=mask)

5.2 混合精度训练

from torch.cuda.amp import autocast def train_with_mixed_precision(model, data_loader): scaler = torch.cuda.amp.GradScaler() for batch in data_loader: inputs = batch['input'].to('cuda') with autocast(): outputs = model(inputs) loss = compute_loss(outputs, batch['target']) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

5.3 内存高效注意力

对于长序列，可以使用内存高效的注意力实现：

from torch.nn.attention import SDPBackend, sdpa_kernel with sdpa_kernel(SDPBackend.MATH): # 使用内存高效但计算较慢的实现 output = model(inputs) with sdpa_kernel(SDPBackend.FLASH_ATTENTION): # 使用Flash Attention（如果可用） output = model(inputs)

6. 实际应用示例

让我们看一个简单的文本分类任务中如何使用我们实现的多头注意力模块。

6.1 文本分类模型架构

class TextClassifier(nn.Module): def __init__(self, vocab_size, embed_size=128, num_heads=8, num_classes=2): super(TextClassifier, self).__init__() self.embedding = nn.Embedding(vocab_size, embed_size) self.attention = MultiHeadAttention(embed_size, num_heads) self.fc = nn.Linear(embed_size, num_classes) def forward(self, x): # 嵌入层 x = self.embedding(x) # (batch_size, seq_len, embed_size) # 多头注意力 x = self.attention(x) # 全局平均池化 x = x.mean(dim=1) # 分类层 x = self.fc(x) return x

6.2 训练循环

def train_model(model, train_loader, val_loader, epochs=10): criterion = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=0.001) for epoch in range(epochs): model.train() train_loss = 0.0 for batch in train_loader: inputs = batch['input'].to('cuda') labels = batch['label'].to('cuda') optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() train_loss += loss.item() # 验证阶段 model.eval() val_loss = 0.0 correct = 0 total = 0 with torch.no_grad(): for batch in val_loader: inputs = batch['input'].to('cuda') labels = batch['label'].to('cuda') outputs = model(inputs) loss = criterion(outputs, labels) val_loss += loss.item() _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() print(f"Epoch {epoch+1}/{epochs} | " f"Train Loss: {train_loss/len(train_loader):.4f} | " f"Val Loss: {val_loss/len(val_loader):.4f} | " f"Val Acc: {100*correct/total:.2f}%")

7. 高级主题与扩展

7.1 相对位置编码

原始 Transformer 使用绝对位置编码，但相对位置编码通常表现更好：

class RelativePositionEmbedding(nn.Module): def __init__(self, max_len=512, embed_size=512): super().__init__() self.embedding = nn.Embedding(2*max_len-1, embed_size) def forward(self, seq_len): positions = torch.arange(seq_len, dtype=torch.long) relative_positions = positions[:, None] - positions[None, :] relative_positions += seq_len - 1 # 使索引非负 return self.embedding(relative_positions)

7.2 稀疏注意力

对于长序列，完全注意力计算代价高昂，可以使用稀疏注意力：

class SparseAttention(nn.Module): def __init__(self, embed_size, num_heads, window_size=32): super().__init__() self.embed_size = embed_size self.num_heads = num_heads self.window_size = window_size 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 forward(self, x): batch_size, seq_len, _ = x.size() # 线性变换 Q = self.query(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) K = self.key(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) V = self.value(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) # 创建局部注意力掩码 mask = torch.ones(seq_len, seq_len, dtype=torch.bool) for i in range(seq_len): start = max(0, i - self.window_size // 2) end = min(seq_len, i + self.window_size // 2 + 1) mask[i, start:end] = 0 # 计算注意力分数 scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32)) scores = scores.masked_fill(mask, -1e9) attention = F.softmax(scores, dim=-1) out = torch.matmul(attention, V) # 合并头 out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, self.embed_size) out = self.fc_out(out) return out

7.3 跨模态注意力

多头注意力也可以用于跨模态任务，如图文匹配：

class CrossModalAttention(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.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 forward(self, x1, x2): # x1: (batch_size, seq_len1, embed_size) # x2: (batch_size, seq_len2, embed_size) batch_size = x1.size(0) # 线性变换 Q = self.query(x1).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) K = self.key(x2).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) V = self.value(x2).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) # 计算注意力分数 scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32)) attention = F.softmax(scores, dim=-1) # 应用注意力 out = torch.matmul(attention, V) out = out.transpose(1, 2).contiguous().view(batch_size, -1, self.embed_size) out = self.fc_out(out) return out