从自注意力到多头自注意力：Transformer架构核心机制解析

2026年1月8日 互联网

一、自注意力机制：重新定义序列建模

1.1 传统序列模型的局限性

在Transformer诞生之前，序列建模主要依赖循环神经网络（RNN）及其变体（如LSTM、GRU）。这些模型通过隐状态传递信息，但存在两个核心问题：

  • 长程依赖失效：梯度消失/爆炸导致远距离信息难以传递；
  • 并行化困难：序列依赖要求逐token处理，计算效率低下。

自注意力机制通过直接建模token间的全局关系，突破了上述限制。其核心思想是：每个token的表示应综合所有token的信息，权重由token间的相关性动态决定。

1.2 自注意力的数学表达

给定输入序列 ( X \in \mathbb{R}^{n \times d} )（( n )为序列长度，( d )为特征维度），自注意力通过以下步骤计算：

  1. 线性变换：生成Query、Key、Value矩阵：
    [
    Q = XW_Q, \quad K = XW_K, \quad V = XW_V
    ]
    其中 ( W_Q, W_K, W_V \in \mathbb{R}^{d \times d_k} ) 为可学习参数。

  2. 相似度计算：计算Query与Key的点积，并缩放以避免梯度消失：
    [
    \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V
    ]
    缩放因子 ( \sqrt{d_k} ) 用于稳定训练。

  3. 加权求和：通过softmax归一化的权重对Value进行加权，得到每个token的上下文表示。

1.3 代码实现示例

  1. import torch
  2. import torch.nn as nn
  4. class SelfAttention(nn.Module):
  5.     def __init__(self, d_model):
  6.         super().__init__()
  7.         self.d_k = d_model // 8  # 典型缩放因子
  8.         self.W_Q = nn.Linear(d_model, self.d_k)
  9.         self.W_K = nn.Linear(d_model, self.d_k)
  10.         self.W_V = nn.Linear(d_model, self.d_k)
  11.         self.scale = 1.0 / torch.sqrt(torch.tensor(self.d_k, dtype=torch.float32))
  13.     def forward(self, x):
  14.         Q = self.W_Q(x)  # [n, d_k]
  15.         K = self.W_K(x)
  16.         V = self.W_V(x)
  18.         # 计算注意力分数
  19.         scores = torch.bmm(Q, K.transpose(1, 2)) * self.scale  # [n, n]
  20.         attn_weights = torch.softmax(scores, dim=-1)
  22.         # 加权求和
  23.         output = torch.bmm(attn_weights, V)  # [n, d_k]
  24.         return output

二、多头自注意力：并行捕捉多样关系

2.1 为什么需要多头？

单头自注意力可能无法同时捕捉多种语义关系（如语法、语义、指代等）。多头机制通过以下方式增强表达能力：

  • 并行化：将Query、Key、Value投影到多个子空间（头），每个头独立计算注意力；
  • 特征分离：不同头关注不同模式的信息，最终拼接融合。

2.2 多头自注意力的计算流程

  1. 头划分：将 ( d )-维特征划分为 ( h ) 个头，每个头维度为 ( d_h = d/h )。
  2. 并行计算：对每个头 ( i )，独立计算自注意力：
    [
    \text{head}_i = \text{Attention}(Q_i, K_i, V_i)
    ]
  3. 拼接与投影：将所有头的输出拼接，并通过线性层融合：
    [
    \text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h)W^O
    ]
    其中 ( W^O \in \mathbb{R}^{h \cdot d_h \times d} )。

2.3 代码实现示例

  1. class MultiHeadAttention(nn.Module):
  2.     def __init__(self, d_model, num_heads):
  3.         super().__init__()
  4.         self.num_heads = num_heads
  5.         self.d_h = d_model // num_heads
  6.         assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
  8.         self.W_Q = nn.Linear(d_model, d_model)
  9.         self.W_K = nn.Linear(d_model, d_model)
  10.         self.W_V = nn.Linear(d_model, d_model)
  11.         self.W_O = nn.Linear(d_model, d_model)
  13.     def forward(self, x):
  14.         batch_size = x.size(0)
  15.         n = x.size(1)
  17.         # 线性变换并拆分多头
  18.         Q = self.W_Q(x).view(batch_size, n, self.num_heads, self.d_h).transpose(1, 2)
  19.         K = self.W_K(x).view(batch_size, n, self.num_heads, self.d_h).transpose(1, 2)
  20.         V = self.W_V(x).view(batch_size, n, self.num_heads, self.d_h).transpose(1, 2)
  22.         # 计算每个头的注意力
  23.         scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.d_h, dtype=torch.float32))
  24.         attn_weights = torch.softmax(scores, dim=-1)
  25.         output = torch.matmul(attn_weights, V)
  27.         # 拼接多头并投影
  28.         output = output.transpose(1, 2).contiguous().view(batch_size, n, -1)
  29.         return self.W_O(output)

三、工程实践与优化技巧

3.1 性能优化策略

  • 头维度选择：通常设置 ( d_h = 64 )，头数 ( h ) 根据任务调整（如BERT-base使用12头）。
  • 稀疏注意力：对于长序列（如文档级任务），可采用局部窗口或块稀疏模式减少计算量。
  • 内存效率：使用梯度检查点（Gradient Checkpointing）降低显存占用。

3.2 初始化与正则化

  • 参数初始化：Query/Key矩阵的权重建议使用Xavier初始化，避免初始阶段梯度消失。
  • 注意力正则化：可添加注意力权重熵正则项，防止模型过度依赖少数token。

3.3 实际应用场景

  • 自然语言处理：多头自注意力是BERT、GPT等模型的核心，适用于文本分类、生成等任务。
  • 多模态建模：在视觉Transformer（ViT）中，自注意力可捕捉图像区域间的空间关系。
  • 推荐系统：通过自注意力建模用户历史行为的时序依赖，提升点击率预测精度。

四、总结与展望

自注意力与多头自注意力机制通过动态权重分配和并行特征提取，重新定义了序列建模的范式。其成功不仅体现在NLP领域，更推动了计算机视觉、语音识别等任务的革新。未来研究方向包括：

  • 高效自注意力变体：如线性注意力（Linear Attention）、核方法等，降低时间复杂度；
  • 硬件友好实现：针对GPU/TPU架构优化计算图，提升吞吐量；
  • 跨模态融合：探索自注意力在多模态大模型中的统一表示能力。

对于开发者而言，深入理解自注意力机制的设计哲学，有助于在定制模型时灵活调整头数、维度等超参数，平衡性能与效率。例如，在资源受限场景下，可减少头数并增大 ( d_h )；而在需要捕捉复杂关系的任务中，增加头数通常能带来收益。

最新文章