1. Encoder-Decoder Transformer架构概述
2017年那篇《Attention is All You Need》论文提出的Transformer架构,彻底改变了自然语言处理领域的游戏规则。与传统RNN不同,这个架构完全基于自注意力机制,实现了前所未有的并行计算能力。我在实际项目中使用过各种变体的Transformer,发现原始版本的设计至今仍是理解后续所有改进的基础。
经典Encoder-Decoder结构包含三个核心部分:左侧的Encoder堆栈、右侧的Decoder堆栈,以及中间的注意力连接层。Encoder负责将输入序列编码为上下文相关的表示,Decoder则利用这些表示逐步生成输出序列。这种设计在机器翻译任务中表现尤为出色——比如把"Hello world"编码为向量表示,再解码为"你好世界"的过程。
关键提示:虽然后续出现了仅用Encoder(如BERT)或仅用Decoder(如GPT)的模型,但理解完整架构对掌握Transformer精髓至关重要。我在初学时曾忽视Decoder部分,后来在实现文本生成任务时吃了大亏。
2. 输入编码与位置嵌入
2.1 词嵌入层实现
假设我们处理英文句子"The cat sat on the mat",首先需要通过词嵌入层将单词转换为向量。以512维的嵌入空间为例:
python复制import torch
import torch.nn as nn
vocab_size = 10000 # 假设词表大小
d_model = 512 # 嵌入维度
embedding = nn.Embedding(vocab_size, d_model)
# 输入句子编码为索引
input_ids = torch.LongTensor([1, 5, 10, 8, 1, 20]) # 假设的token索引
embedded_input = embedding(input_ids) # 输出形状: [6, 512]
这里有个容易踩坑的地方:嵌入层通常需要乘以√d_model来缩放。论文中解释这是为了让位置编码的相对影响保持稳定。
2.2 位置编码详解
Transformer不使用RNN那样的序列处理,因此必须显式注入位置信息。原始论文使用正弦和余弦函数的固定编码:
python复制def positional_encoding(max_len, d_model):
position = torch.arange(max_len).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
pe = torch.zeros(max_len, d_model)
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
return pe
我在实现时发现几个关键点:
- 偶数维度用sin,奇数维度用cos
- 10000.0这个魔法数决定了波长范围
- 实际应用中通常将位置编码与词嵌入直接相加
3. Encoder层实现细节
3.1 多头注意力机制
这是Transformer最核心的创新点。以8个头、512维模型为例:
python复制class MultiHeadAttention(nn.Module):
def __init__(self, d_model=512, num_heads=8):
super().__init__()
assert d_model % num_heads == 0
self.d_k = d_model // num_heads
self.num_heads = 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 forward(self, x):
batch_size = x.size(0)
# 线性变换并分头 [batch_size, seq_len, d_model] -> [batch_size, num_heads, seq_len, d_k]
q = self.W_q(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
k = self.W_k(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
v = self.W_v(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
# 计算注意力分数
scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
attn = torch.softmax(scores, dim=-1)
# 应用注意力权重
context = torch.matmul(attn, v)
# 合并头并输出
context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
return self.W_o(context)
实际调试时我发现:
- 分头操作需要特别注意维度变换顺序
- 注意力分数除以√d_k对稳定训练至关重要
- 最后输出的线性层W_o不能省略
3.2 前馈网络与残差连接
每个Encoder层还包含一个前馈网络(FFN):
python复制class FeedForward(nn.Module):
def __init__(self, d_model=512, d_ff=2048):
super().__init__()
self.linear1 = nn.Linear(d_model, d_ff)
self.linear2 = nn.Linear(d_ff, d_model)
self.dropout = nn.Dropout(0.1)
def forward(self, x):
return self.linear2(self.dropout(torch.relu(self.linear1(x))))
这里有几个设计要点:
- 论文使用d_ff=2048作为中间维度
- ReLU激活函数后需要加Dropout
- 每子层都有残差连接和LayerNorm
完整的Encoder层实现:
python复制class EncoderLayer(nn.Module):
def __init__(self, d_model=512, num_heads=8, d_ff=2048):
super().__init__()
self.self_attn = MultiHeadAttention(d_model, num_heads)
self.ffn = FeedForward(d_model, d_ff)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(0.1)
def forward(self, x):
# 自注意力子层
attn_output = self.self_attn(x)
x = self.norm1(x + self.dropout(attn_output))
# 前馈子层
ffn_output = self.ffn(x)
x = self.norm2(x + self.dropout(ffn_output))
return x
4. Decoder层特殊设计
4.1 掩码多头注意力
Decoder有两个关键不同点:
- 第一层是掩码自注意力,防止看到未来信息
- 第二层是Encoder-Decoder注意力
python复制class DecoderLayer(nn.Module):
def __init__(self, d_model=512, num_heads=8, d_ff=2048):
super().__init__()
self.self_attn = MultiHeadAttention(d_model, num_heads)
self.cross_attn = MultiHeadAttention(d_model, num_heads)
self.ffn = FeedForward(d_model, d_ff)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.norm3 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(0.1)
def forward(self, x, encoder_output, tgt_mask):
# 掩码自注意力
attn_output = self.self_attn(x, x, x, tgt_mask)
x = self.norm1(x + self.dropout(attn_output))
# Encoder-Decoder注意力
attn_output = self.cross_attn(x, encoder_output, encoder_output)
x = self.norm2(x + self.dropout(attn_output))
# 前馈网络
ffn_output = self.ffn(x)
x = self.norm3(x + self.dropout(ffn_output))
return x
生成掩码的典型方法:
python复制def generate_square_subsequent_mask(sz):
mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
return mask
4.2 输出层设计
Decoder输出经过线性层和softmax:
python复制class Transformer(nn.Module):
def __init__(self, src_vocab_size, tgt_vocab_size, d_model=512, num_heads=8, num_layers=6, d_ff=2048):
super().__init__()
self.encoder = Encoder(src_vocab_size, d_model, num_heads, num_layers, d_ff)
self.decoder = Decoder(tgt_vocab_size, d_model, num_heads, num_layers, d_ff)
self.linear = nn.Linear(d_model, tgt_vocab_size)
def forward(self, src, tgt, src_mask=None, tgt_mask=None):
encoder_output = self.encoder(src, src_mask)
decoder_output = self.decoder(tgt, encoder_output, tgt_mask)
return torch.softmax(self.linear(decoder_output), dim=-1)
5. 矩阵维度全流程追踪
让我们以具体例子跟踪维度变化:
假设:
- 批量大小=32
- 源序列长度=10
- 目标序列长度=12
- d_model=512
- num_heads=8
- d_k = d_v = 512/8=64
5.1 Encoder流程
- 输入嵌入:[32, 10] -> [32, 10, 512]
- 位置编码相加:[32, 10, 512]
- 经过6个Encoder层:
- 每层多头注意力:
- Q/K/V投影:[32, 10, 512] -> [32, 10, 512]
- 分头:[32, 8, 10, 64]
- 注意力分数:[32, 8, 10, 10]
- 输出合并:[32, 10, 512]
- 前馈网络:[32, 10, 512] -> [32, 10, 2048] -> [32, 10, 512]
- 每层多头注意力:
5.2 Decoder流程
- 输出嵌入:[32, 12] -> [32, 12, 512]
- 位置编码相加:[32, 12, 512]
- 经过6个Decoder层:
- 掩码自注意力:
- 使用三角掩码:[12, 12]
- 输出:[32, 12, 512]
- Encoder-Decoder注意力:
- Q来自Decoder:[32, 12, 512]
- K/V来自Encoder:[32, 10, 512]
- 输出:[32, 12, 512]
- 前馈网络:[32, 12, 512] -> [32, 12, 2048] -> [32, 12, 512]
- 掩码自注意力:
- 输出层:[32, 12, 512] -> [32, 12, tgt_vocab_size]
6. 训练技巧与实战经验
6.1 学习率调度
论文使用Adam优化器配合特殊的学习率warmup:
python复制class TransformerScheduler:
def __init__(self, optimizer, d_model, warmup_steps=4000):
self.optimizer = optimizer
self.d_model = d_model
self.warmup_steps = warmup_steps
self.step_num = 0
def step(self):
self.step_num += 1
lr = self.d_model ** -0.5 * min(self.step_num ** -0.5, self.step_num * self.warmup_steps ** -1.5)
for param_group in self.optimizer.param_groups:
param_group['lr'] = lr
6.2 标签平滑
实践中发现标签平滑能显著提升模型泛化能力:
python复制criterion = nn.KLDivLoss(reduction='batchmean')
smoothing = 0.1
def smooth_label(targets, n_classes, smoothing):
return targets * (1 - smoothing) + smoothing / n_classes
6.3 梯度裁剪
Transformer训练需要梯度裁剪:
python复制torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
7. 常见问题排查
-
注意力分数爆炸:
- 检查是否除以√d_k
- 验证初始化方法(通常使用xavier初始化)
-
训练不稳定:
- 增加warmup步数
- 尝试更大的batch size
- 检查残差连接实现
-
验证集表现差:
- 调整dropout率(0.1-0.3)
- 尝试不同的标签平滑值
- 检查数据预处理是否一致
-
GPU内存不足:
- 减小batch size
- 使用梯度累积
- 尝试混合精度训练
我在实际项目中遇到过最棘手的问题是梯度消失,最终发现是LayerNorm的实现有误。建议新手特别注意各层的输入输出尺度,使用torchviz可视化计算图辅助调试。
