共计 2481 个字符，预计需要花费 7 分钟才能阅读完成。

1. 为什么需要自注意力机制

传统 RNN 在处理长序列时存在两个致命缺陷：

• 梯度消失 / 爆炸 ：随着序列长度增加，反向传播时梯度会指数级衰减或增长
• 顺序计算无法并行 ：必须按时间步依次计算，训练效率低下

比如在句子 ”The animal didn’t cross the street because it was too tired” 中，RNN 很难建立 ”it” 与 ”animal” 的长距离依赖关系。

2. 单头 vs 多头注意力对比

2.1 单头注意力计算

公式表示：

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

其中：
– Q(Query): 查询向量
– K(Key): 键向量
– V(Value): 值向量
– d_k: 向量的维度

2.2 多头注意力优势

多头相当于多个 ” 视角 ” 观察数据：

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

实验表明，8 头注意力在 GLUE 基准上比单头高 3.2 个点。

3. PyTorch 完整实现

import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, num_heads=8):
        super().__init__()
        assert d_model % num_heads == 0

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        # 线性变换层
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

    def scaled_dot_product(self, Q, K, V, mask=None):
        # Q,K,V 形状: [batch_size, num_heads, seq_len, d_k]
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)

        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        attn = torch.softmax(scores, dim=-1)
        output = torch.matmul(attn, V)
        return output, attn

    def forward(self, Q, K, V, mask=None):
        batch_size = Q.size(0)

        # 线性变换并分头 [batch_size, seq_len, d_model] -> [batch_size, seq_len, num_heads, d_k]
        Q = self.W_q(Q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(K).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(V).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)

        # 计算注意力
        scores, attn = self.scaled_dot_product(Q, K, V, mask)

        # 拼接多头结果 [batch_size, seq_len, d_model]
        concat = scores.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)

        # 最终线性变换
        output = self.W_o(concat)
        return output, attn

4. 可视化注意力权重

使用 seaborn 绘制热力图：

import seaborn as sns
import matplotlib.pyplot as plt

# 假设 attn 是计算得到的注意力矩阵
plt.figure(figsize=(10,8))
sns.heatmap(attn[0,0].detach().numpy(), cmap="YlGnBu")
plt.xlabel("Key Position")
plt.ylabel("Query Position")
plt.show()

5. 调优与部署指南

5.1 头数选择经验

5.2 内存优化技巧

• 梯度检查点 ：用时间换空间

  from torch.utils.checkpoint import checkpoint
  output = checkpoint(self.scaled_dot_product, Q, K, V, mask)

• 混合精度训练 ：减少显存占用

  scaler = torch.cuda.amp.GradScaler()
  with torch.cuda.amp.autocast():
      output, attn = model(inputs)

6. 常见问题解决方案

1. 梯度消失 ：
2. 使用 Layer Normalization

3. 初始化权重时乘以 1 /√d_k

4. 低资源部署 ：

5. 使用知识蒸馏压缩模型
6. 量化到 INT8 精度

   model = torch.quantization.quantize_dynamic(model, {nn.Linear}, dtype=torch.qint8
   )

7. 延伸阅读方向

• 高效注意力 ：Reformer 的 LSH 注意力
• 长序列处理 ：Transformer-XL 的片段递归机制
• 稀疏注意力 ：BigBird 的块稀疏模式

通过本文的代码实践和原理分析，相信你已经对 BERT 的多头注意力机制有了直观认识。建议动手修改头数参数，观察模型在具体任务上的表现变化，这是掌握该机制的最佳途径。