当前位置：首页 > news >正文

图解Transformer：Self-Attention与多头注意力机制详解

news 2026/5/12 3:21:36

图解Transformer：Self-Attention与多头注意力机制详解

在自然语言处理领域，Transformer架构已经彻底改变了序列建模的范式。与传统的循环神经网络不同，Transformer完全基于注意力机制构建，特别是其核心组件——Self-Attention与多头注意力机制。本文将通过直观的可视化方式，逐步拆解这些关键技术的数学原理和实现细节，帮助开发者建立清晰的技术认知。

1. 从序列建模到注意力机制

传统RNN架构在处理长序列时面临两个根本性挑战：信息衰减和顺序计算瓶颈。当序列长度超过50个token时，早期位置的信息在传递过程中会逐渐稀释。而Transformer提出的Self-Attention机制，通过建立全连接注意力网络，实现了三个关键突破：

全局感知：每个位置可以直接访问序列中所有其他位置的信息
并行计算：摆脱了RNN的时序依赖，大幅提升训练效率
动态权重：根据内容相关性自动调整不同位置的关注强度

典型场景对比：在翻译"The animal didn't cross the street because it was too tired"时，传统模型可能混淆"it"的指代对象，而Self-Attention能准确建立"it"与"animal"的关联。

2. Self-Attention的数学解剖

Self-Attention的核心计算过程可以分为六个可解释的步骤，我们以输入序列"Thinking Machines"为例进行说明：

2.1 向量空间映射

首先将每个token的嵌入向量（假设维度d=4）通过可学习参数矩阵转换为三个基本向量：

向量类型	计算方式	维度	作用
Query	Q = X · W_Q	d × d_k	表示当前关注点
Key	K = X · W_K	d × d_k	表示待匹配特征
Value	V = X · W_V	d × d_v	携带实际信息内容

# PyTorch实现示例 class SelfAttention(nn.Module): def __init__(self, embed_size): super().__init__() self.W_q = nn.Linear(embed_size, embed_size) self.W_k = nn.Linear(embed_size, embed_size) self.W_v = nn.Linear(embed_size, embed_size)

2.2 注意力分数计算

计算Query与所有Key的点积，得到原始注意力分数。以第一个词"Thinking"为例：

Score_i = Q_think · K_i^T

这个操作本质上是在衡量当前查询与各个键的匹配程度。点积值越大，表明两个向量在语义空间中的方向越接近。

2.3 分数缩放与归一化

将原始分数除以√d_k（Key向量维度）进行缩放，然后应用Softmax：

scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k) attn_weights = F.softmax(scores, dim=-1)

缩放操作确保梯度稳定性，而Softmax将分数转换为概率分布。下图展示了归一化后的注意力权重：

Thinking → [0.6, 0.4] Machines → [0.2, 0.8]

2.4 加权求和

用注意力权重对Value向量进行加权求和，得到当前位置的输出表示：

Output = Σ (attn_weight_i × V_i)

这个步骤实现了信息的动态聚合——相关度高的位置贡献更大信息量。

3. 多头注意力机制解析

单一注意力头存在表征能力有限的缺陷。Transformer通过多头机制实现了三个维度的增强：

3.1 并行注意力头设计

典型配置包含8个独立的注意力头，每个头具有不同的参数矩阵：

class MultiHeadAttention(nn.Module): def __init__(self, num_heads, embed_size): super().__init__() self.heads = nn.ModuleList([ SelfAttention(embed_size) for _ in range(num_heads) ])

每个头学习不同的注意力模式，例如：

头1可能关注语法结构
头2可能捕捉指代关系
头3可能跟踪主题一致性

3.2 子空间投影

每个头的Q/K/V会先投影到低维子空间（通常d_k = d_model / h），这种设计带来两个优势：

计算复杂度保持与单头相同
强制不同头学习差异化特征

3.3 输出融合

各头的输出通过拼接和线性变换整合：

output = torch.cat([head(x) for head in self.heads], dim=-1) output = self.fc(output) # 投影回d_model维度

这种结构类似于卷积神经网络中的多通道设计，极大提升了模型的表征能力。

4. 工程实现关键技巧

在实际部署Transformer时，以下几个优化策略至关重要：

4.1 矩阵并行计算

利用矩阵运算一次性完成所有位置的注意力计算：

# 输入矩阵X形状: [batch_size, seq_len, embed_size] Q = torch.matmul(X, W_Q) # [bs, seq_len, d_k] K = torch.matmul(X, W_K) V = torch.matmul(X, W_V) scores = torch.matmul(Q, K.transpose(-2, -1)) # [bs, seq_len, seq_len]

这种实现相比循环计算可获得数百倍的加速比。

4.2 掩码机制

在解码器中需要防止信息泄露，通过注意力掩码实现：

mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool() scores.masked_fill_(mask, float('-inf'))

这确保每个位置只能关注之前的位置，符合自回归生成特性。

4.3 残差连接与层归一化

每个子层都采用残差连接和层归一化：

x = x + self.dropout(self.attention(x)) x = self.norm(x)

这种设计有效缓解了深层网络的梯度消失问题，在12层以上的Transformer中表现尤为关键。

5. 可视化理解注意力模式

通过热力图可以直观展示不同注意力头学到的模式。下图展示了一个翻译任务中的典型注意力分布：

Source: The cat sat on the mat ┌─────────┬─────┬───┬───┬────┐ │ the │ cat │sat│on│mat │ ┌───────┼─────────┼─────┼───┼───┼────┤ │ the │ ████████│ │ │ │ │ │ cat │ │████ │ │ │ │ │ sat │ │ │███│ │ │ │ on │ │ │ │██ │ │ │ mat │ │ │ │ │████│ └───────┴─────────┴─────┴───┴───┴────┘

可以看到：