共计 2650 个字符,预计需要花费 7 分钟才能阅读完成。
为什么需要 Transformer?
在自然语言处理领域,传统的 RNN(循环神经网络)在处理序列数据时存在两个主要问题:

- 顺序计算瓶颈:RNN 必须逐个处理序列中的元素,无法充分利用现代 GPU 的并行计算能力
- 长距离依赖衰减:随着序列长度的增加,早期信息在传递过程中会逐渐丢失或稀释
Transformer 通过完全基于注意力机制的架构,完美解决了这两个痛点。它允许模型同时处理整个序列,并直接建立任意两个位置之间的依赖关系。
自注意力机制详解
自注意力 (Self-Attention) 是 Transformer 的核心组件,其计算过程可以分为以下几步:
- 输入表示:对于输入序列中的每个词,我们都有一个对应的嵌入向量 $x_i$
- 线性变换 :通过三个可学习的权重矩阵 $W_Q$、$W_K$、$W_V$,将每个 $x_i$ 转换为查询(Query)、键(Key) 和值 (Value) 向量
数学表达式为:
$$
Q = XW_Q, \quad K = XW_K, \quad V = XW_V
$$
-
注意力分数计算:衡量查询与键的相似度,通常使用点积后缩放
$$
\text{Attention}(Q, K, V) = \text{softmax}(\frac{QK^T}{\sqrt{d_k}})V
$$ -
缩放与归一化:除以 $\sqrt{d_k}$(键向量的维度)防止梯度消失,然后应用 softmax 归一化
- 加权求和:用归一化的权重对值向量进行加权求和
多头自注意力实现
多头自注意力 (Multi-Head Attention) 将自注意力机制并行执行多次,允许模型在不同表示子空间中学习信息。PyTorch 实现关键代码如下:
import torch
import torch.nn as nn
class MultiHeadAttention(nn.Module):
def __init__(self, embed_dim, num_heads):
super().__init__()
self.embed_dim = embed_dim
self.num_heads = num_heads
self.head_dim = embed_dim // num_heads
# 线性变换层
self.qkv_proj = nn.Linear(embed_dim, 3*embed_dim)
self.out_proj = nn.Linear(embed_dim, embed_dim)
def forward(self, x, mask=None):
batch_size, seq_len, _ = x.shape
# 生成 Q,K,V [batch_size, seq_len, embed_dim]
qkv = self.qkv_proj(x)
q, k, v = torch.chunk(qkv, 3, dim=-1)
# 分头处理 [batch_size, num_heads, seq_len, head_dim]
q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
# 计算注意力分数 [batch_size, num_heads, seq_len, seq_len]
scores = torch.matmul(q, k.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.head_dim))
if mask is not None:
scores = scores.masked_fill(mask == 0, float('-inf'))
attn_weights = torch.softmax(scores, dim=-1)
# 加权求和 [batch_size, num_heads, seq_len, head_dim]
output = torch.matmul(attn_weights, v)
# 合并多头输出 [batch_size, seq_len, embed_dim]
output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
output = self.out_proj(output)
return output
头数选择与性能影响
多头注意力的头数选择需要权衡以下因素:
- 计算效率:头数增加会提升计算开销,但可以并行处理
- 模型容量:更多头数意味着更强的表达能力,但也更容易过拟合
- 任务特性:不同任务可能需要关注不同粒度的特征
实验表明,在大多数 NLP 任务中,8-16 个头通常能取得较好的平衡。对于小规模数据集,可以适当减少头数以避免过拟合。
训练调参实战建议
- 学习率设置:
- 使用学习率预热 (warmup) 策略,初期逐步增大学习率
-
典型配置:Adam 优化器,初始 lr=3e-5,warmup_steps=4000
-
梯度裁剪:
-
防止梯度爆炸,clip_norm 通常在 1.0-5.0 之间
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) -
批次大小:
- 在显存允许范围内尽可能使用大 batch
-
小 batch 时可使用梯度累积技巧
-
正则化:
- dropout 率通常设为 0.1-0.3
- 标签平滑 (label smoothing) 可提高模型泛化能力
工业部署优化策略
实际部署时需要特别关注内存和计算效率:
- 内存优化:
- 使用混合精度训练(FP16/FP32)
- 激活检查点技术 (checkpointing) 减少中间结果存储
-
序列长度较大时采用内存高效的注意力实现
-
计算优化:
- 利用 Flash Attention 等优化实现
- 对于超长序列,可采用稀疏注意力或分块处理
- 量化推理 (INT8) 可显著提升推理速度
总结与展望
Transformer 的自注意力机制通过全局依赖建模能力,彻底改变了序列建模的方式。多头设计让模型能够同时关注不同位置和不同表示子空间的信息。虽然 Transformer 在诸多任务上表现出色,但在处理超长序列时仍面临计算复杂度高的问题,这也是未来研究的重要方向之一。
对于初学者来说,建议从理解基础的自注意力计算开始,逐步扩展到多头实现,最后再考虑优化和部署问题。实践中多关注模型在不同配置下的表现,积累调参经验,才能真正掌握 Transformer 的精髓。
