从零构建Transformer：模块拆解与PyTorch实战

1. Transformer核心思想与架构概述Transformer模型彻底改变了自然语言处理领域其核心在于完全依赖注意力机制来处理序列数据。我第一次接触Transformer时被它抛弃传统RNN/CNN的勇气所震撼——这种架构竟然仅靠注意力机制就能实现更好的长距离依赖建模。自注意力机制的本质就像让序列中的每个词元召开圆桌会议当处理苹果这个词时模型会同时关注句子中所有其他词比如红色的、吃的通过动态权重决定哪些信息更重要。这种并行处理方式相比RNN的逐步处理有着天然的优势计算效率可以并行处理所有位置时间复杂度O(n²)带来平方级增长问题建模能力直接建立任意两个词元的关系不受距离限制可解释性注意力权重可视化显示模型关注点典型的Transformer架构包含输入 - 词嵌入 - 位置编码 - [编码器×N] - [解码器×N] - 输出其中每个编码器包含多头自注意力层前馈神经网络层残差连接和层归一化而解码器额外增加了掩码多头注意力防止信息泄露编码器-解码器注意力层2. 词嵌入与位置编码实现2.1 词嵌入层(Embedding)词嵌入的本质是将离散的词语映射到连续的向量空间。我在实际项目中发现好的词嵌入能显著提升模型性能。PyTorch实现非常简单import torch.nn as nn class TokenEmbedding(nn.Module): def __init__(self, vocab_size, embed_dim): super().__init__() self.embedding nn.Embedding(vocab_size, embed_dim) self.embed_dim embed_dim def forward(self, tokens): # tokens形状: [batch_size, seq_len] # 输出形状: [batch_size, seq_len, embed_dim] return self.embedding(tokens) * math.sqrt(self.embed_dim) # 缩放为什么需要缩放初始阶段词嵌入的数值通常较小乘以√d_model可以让数值范围与位置编码匹配有利于训练稳定性。2.2 位置编码(Positional Encoding)由于Transformer没有递归结构必须显式注入位置信息。原始论文使用了一种巧妙的固定编码方式class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len5000): super().__init__() pe torch.zeros(max_len, d_model) position torch.arange(0, max_len).unsqueeze(1) div_term torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model)) pe[:, 0::2] torch.sin(position * div_term) # 偶数位置 pe[:, 1::2] torch.cos(position * div_term) # 奇数位置 self.register_buffer(pe, pe.unsqueeze(0)) def forward(self, x): # x形状: [batch_size, seq_len, d_model] return x self.pe[:, :x.size(1)]为什么选择正弦函数我在实验中尝试过可学习的位置编码发现固定编码在小数据上表现更好不易过拟合正弦编码可以外推到比训练更长的序列相对位置关系可以通过线性变换表示3. 注意力机制详解与实现3.1 缩放点积注意力def scaled_dot_product_attention(Q, K, V, maskNone): # Q, K, V形状: [batch_size, ..., seq_len, d_k] d_k Q.size(-1) scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k) if mask is not None: scores scores.masked_fill(mask 0, -1e9) attn_weights F.softmax(scores, dim-1) output torch.matmul(attn_weights, V) return output, attn_weights缩放因子(√d_k)的重要性当维度d_k较大时点积结果会变得很大导致softmax进入梯度很小的饱和区。我在调试模型时曾移除这个缩放因子结果模型完全无法收敛。3.2 多头注意力机制class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): 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 split_heads(self, x): batch_size x.size(0) return x.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) def forward(self, Q, K, V, maskNone): # 线性变换 Q self.split_heads(self.W_q(Q)) # [batch, heads, seq_len, d_k] K self.split_heads(self.W_k(K)) V self.split_heads(self.W_v(V)) # 计算注意力 attn_output, _ scaled_dot_product_attention(Q, K, V, mask) attn_output attn_output.transpose(1, 2).contiguous().view( Q.size(0), -1, self.num_heads * self.d_k) return self.W_o(attn_output)多头设计的优势在我的文本分类任务中可视化不同头的注意力模式时发现某些头专注于局部语法关系如形容词-名词另一些头捕捉长距离语义关系如指代消解还有头关注标点等特殊符号4. 编码器模块实现4.1 前馈网络class PositionwiseFFN(nn.Module): def __init__(self, d_model, d_ff, dropout0.1): super().__init__() self.linear1 nn.Linear(d_model, d_ff) self.linear2 nn.Linear(d_ff, d_model) self.dropout nn.Dropout(dropout) def forward(self, x): return self.linear2(self.dropout(F.relu(self.linear1(x))))为什么需要两个线性层这是一个小型瓶颈结构典型配置是d_ff4×d_model。我在消融实验中发现更窄的中间层会损害模型性能更宽的层提升有限但增加计算量ReLU激活比GELU略快但效果稍差4.2 编码器层完整实现class EncoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout): super().__init__() self.self_attn MultiHeadAttention(d_model, num_heads) self.ffn PositionwiseFFN(d_model, d_ff, dropout) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.dropout1 nn.Dropout(dropout) self.dropout2 nn.Dropout(dropout) def forward(self, x, mask): # 自注意力子层 attn_output self.self_attn(x, x, x, mask) x x self.dropout1(attn_output) x self.norm1(x) # 前馈子层 ffn_output self.ffn(x) x x self.dropout2(ffn_output) x self.norm2(x) return x残差连接的重要性当堆叠很多层时如12层没有残差连接的模型几乎无法训练。我在实现时曾忘记加残差连接导致梯度消失问题。5. 解码器模块实现5.1 掩码注意力解码器需要防止当前位置关注后续位置这通过注意力掩码实现def create_look_ahead_mask(size): mask torch.triu(torch.ones(size, size), diagonal1) return mask 0 # 下三角为True5.2 解码器层完整实现class DecoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout): super().__init__() self.self_attn MultiHeadAttention(d_model, num_heads) self.cross_attn MultiHeadAttention(d_model, num_heads) self.ffn PositionwiseFFN(d_model, d_ff, dropout) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.norm3 nn.LayerNorm(d_model) self.dropout1 nn.Dropout(dropout) self.dropout2 nn.Dropout(dropout) self.dropout3 nn.Dropout(dropout) def forward(self, x, enc_output, look_ahead_mask, padding_mask): # 自注意力带掩码 attn1 self.self_attn(x, x, x, look_ahead_mask) x x self.dropout1(attn1) x self.norm1(x) # 编码器-解码器注意力 attn2 self.cross_attn(x, enc_output, enc_output, padding_mask) x x self.dropout2(attn2) x self.norm2(x) # 前馈网络 ffn_output self.ffn(x) x x self.dropout3(ffn_output) x self.norm3(x) return x编码器-解码器注意力的特点这是连接两个模块的关键允许解码器关注输入序列的相关部分。在机器翻译任务中这种注意力往往能准确对齐源语言和目标语言的词语。6. 完整Transformer组装class Transformer(nn.Module): def __init__(self, src_vocab_size, tgt_vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_len, dropout0.1): super().__init__() self.encoder_embedding nn.Sequential( TokenEmbedding(src_vocab_size, d_model), PositionalEncoding(d_model, max_seq_len) ) self.decoder_embedding nn.Sequential( TokenEmbedding(tgt_vocab_size, d_model), PositionalEncoding(d_model, max_seq_len) ) self.encoder_layers nn.ModuleList( [EncoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)] ) self.decoder_layers nn.ModuleList( [DecoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)] ) self.fc_out nn.Linear(d_model, tgt_vocab_size) def encode(self, src, src_mask): src self.encoder_embedding(src) for layer in self.encoder_layers: src layer(src, src_mask) return src def decode(self, tgt, enc_output, look_ahead_mask, tgt_mask): tgt self.decoder_embedding(tgt) for layer in self.decoder_layers: tgt layer(tgt, enc_output, look_ahead_mask, tgt_mask) return self.fc_out(tgt) def forward(self, src, tgt, src_maskNone, tgt_maskNone): enc_output self.encode(src, src_mask) return self.decode(tgt, enc_output, create_look_ahead_mask(tgt.size(1)), tgt_mask)参数初始化技巧Transformer对初始化敏感我通常使用Xavier均匀初始化for p in model.parameters(): if p.dim() 1: nn.init.xavier_uniform_(p)7. 训练技巧与优化7.1 学习率调度Transformer使用特殊的学习率预热策略class TransformerScheduler: def __init__(self, optimizer, d_model, warmup_steps4000): 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] lr7.2 标签平滑class LabelSmoothingLoss(nn.Module): def __init__(self, classes, padding_idx, smoothing0.1): super().__init__() self.confidence 1.0 - smoothing self.smoothing smoothing self.classes classes self.padding_idx padding_idx def forward(self, pred, target): pred pred.log_softmax(dim-1) true_dist torch.zeros_like(pred) true_dist.fill_(self.smoothing / (self.classes - 2)) true_dist.scatter_(1, target.unsqueeze(1), self.confidence) true_dist[:, self.padding_idx] 0 mask (target self.padding_idx) if mask.any(): true_dist[mask] 0 return torch.mean(torch.sum(-true_dist * pred, dim-1))为什么需要标签平滑在机器翻译任务中过于自信的预测会导致模型泛化能力下降对错误标记敏感校准性差预测概率不代表真实概率8. 常见问题与调试经验问题1训练初期损失不下降检查嵌入层缩放验证注意力掩码是否正确确认学习率预热生效问题2验证集性能波动大增加标签平滑尝试更大的dropout(0.3-0.5)检查梯度裁剪是否生效问题3长序列表现差检查相对位置编码实现尝试Reformer等改进架构调整注意力头数量更多头适合长序列我在实际项目中发现Transformer对超参数相当敏感。一个实用的技巧是先用小模型如d_model128快速验证想法再扩展到完整规模。另外使用混合精度训练可以显著减少显存占用允许更大的batch size。