Transformer架构精讲 —— 从零构建编码器与解码器-尧图网络科技

1. Transformer架构概述

Transformer模型彻底改变了自然语言处理领域，它完全基于自注意力机制，摒弃了传统的循环神经网络（RNN）和卷积神经网络（CNN）。我第一次在实际项目中使用Transformer时，就被它的并行处理能力和长距离依赖捕捉能力震撼了。下面我们从一个实际案例开始，逐步拆解它的核心组件。

想象你正在构建一个翻译系统：输入"I love coding"，模型需要输出"我爱编程"。传统RNN需要逐个词处理，而Transformer可以同时处理所有单词，并通过自注意力机制自动学习"love"和"爱"之间的关联。这种并行性使得训练速度大幅提升，实测下来比LSTM快3倍以上。

Transformer的核心架构包含：

编码器堆栈（左侧）：处理输入序列
解码器堆栈（右侧）: 生成输出序列
注意力机制：连接两者的关键桥梁

每个编码器层包含两个子层：

多头自注意力机制
前馈神经网络

而解码器层则多出一个编码器-解码器注意力层，用于聚焦相关输入。我在实现时发现，这种分层结构让模型可以逐层提取不同抽象级别的特征。

2. 输入处理：嵌入与位置编码

2.1 词嵌入实现

词嵌入是将离散的单词转换为连续向量的过程。在PyTorch中，我们可以这样实现：

import torch.nn as nn class TokenEmbedding(nn.Module): def __init__(self, vocab_size, d_model): super().__init__() self.embedding = nn.Embedding(vocab_size, d_model) self.d_model = d_model def forward(self, x): # 乘以sqrt(d_model)来缩放嵌入值 return self.embedding(x) * math.sqrt(self.d_model)

这里有个坑我踩过：一定要对嵌入值进行缩放，否则初始阶段注意力分数会过大，导致softmax进入饱和区。通常设置d_model=512，这样不同位置的嵌入值方差相近。

2.2 位置编码详解

因为Transformer没有递归结构，必须显式注入位置信息。原始论文使用正弦函数：

class PositionalEncoding(nn.Module): def __init__(self, d_model, dropout, max_len=5000): super().__init__() self.dropout = nn.Dropout(p=dropout) 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) self.register_buffer('pe', pe) def forward(self, x): x = x + self.pe[:x.size(1)] return self.dropout(x)

有趣的是，我在实验中发现对于短文本（<50词），可学习的位置嵌入效果更好；但对于长文本，正弦版本更具泛化性。这是因为正弦函数的周期性可以外推到训练时未见的位置。

3. 编码器实现

3.1 多头注意力机制

这是Transformer最核心的组件。我们先用代码实现单头注意力：

def scaled_dot_product_attention(Q, K, V, mask=None): matmul_qk = torch.matmul(Q, K.transpose(-2, -1)) dk = K.size(-1) scaled_attention_logits = matmul_qk / math.sqrt(dk) if mask is not None: scaled_attention_logits += (mask * -1e9) attention_weights = F.softmax(scaled_attention_logits, dim=-1) output = torch.matmul(attention_weights, V) return output

多头注意力的关键是将输入拆分为多个子空间：

class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): super().__init__() self.num_heads = num_heads self.d_model = d_model self.depth = d_model // num_heads self.Wq = nn.Linear(d_model, d_model) self.Wk = nn.Linear(d_model, d_model) self.Wv = nn.Linear(d_model, d_model) self.dense = nn.Linear(d_model, d_model) def split_heads(self, x, batch_size): return x.view(batch_size, -1, self.num_heads, self.depth).transpose(1, 2) def forward(self, Q, K, V, mask): batch_size = Q.size(0) Q = self.split_heads(self.Wq(Q), batch_size) K = self.split_heads(self.Wk(K), batch_size) V = self.split_heads(self.Wv(V), batch_size) scaled_attention = scaled_dot_product_attention(Q, K, V, mask) concat_attention = scaled_attention.transpose(1, 2).contiguous() \ .view(batch_size, -1, self.d_model) return self.dense(concat_attention)

实际调试时，我发现头数不是越多越好。对于d_model=512，8个头（每个头64维）效果最佳。头数过多会导致每个头的表达能力不足。

3.2 前馈网络与残差连接

编码器的第二个关键组件是位置式前馈网络：

class FeedForward(nn.Module): def __init__(self, d_model, dff): super().__init__() self.linear1 = nn.Linear(d_model, dff) self.linear2 = nn.Linear(dff, d_model) self.dropout = nn.Dropout(0.1) def forward(self, x): return self.linear2(self.dropout(F.relu(self.linear1(x))))

这里dff（前馈网络维度）通常设为4*d_model。我在一个文本分类任务中做过对比实验：

dff=2048（标准设置）准确率89.2%
dff=1024 准确率88.7%
dff=4096 准确率89.3%但训练速度慢2倍

残差连接和层归一化是稳定深层网络的关键：

class EncoderLayer(nn.Module): def __init__(self, d_model, num_heads, dff, dropout=0.1): super().__init__() self.mha = MultiHeadAttention(d_model, num_heads) self.ffn = FeedForward(d_model, dff) self.layernorm1 = nn.LayerNorm(d_model) self.layernorm2 = nn.LayerNorm(d_model) self.dropout1 = nn.Dropout(dropout) self.dropout2 = nn.Dropout(dropout) def forward(self, x, mask): attn_output = self.mha(x, x, x, mask) out1 = self.layernorm1(x + self.dropout1(attn_output)) ffn_output = self.ffn(out1) return self.layernorm2(out1 + self.dropout2(ffn_output))

4. 解码器实现

4.1 带掩码的多头注意力

解码器的第一个自注意力层需要防止信息泄露：

def create_look_ahead_mask(size): mask = torch.triu(torch.ones(size, size), diagonal=1) return mask # 上三角矩阵，1表示需要mask的位置

在训练时，即使模型预测第5个词，它也只能看到前4个词。这种因果掩码确保模型不能"作弊"。

4.2 编码器-解码器注意力

这是连接两个模块的关键：

class DecoderLayer(nn.Module): def __init__(self, d_model, num_heads, dff, dropout=0.1): super().__init__() self.mha1 = MultiHeadAttention(d_model, num_heads) self.mha2 = MultiHeadAttention(d_model, num_heads) self.ffn = FeedForward(d_model, dff) self.layernorm1 = nn.LayerNorm(d_model) self.layernorm2 = nn.LayerNorm(d_model) self.layernorm3 = 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.mha1(x, x, x, look_ahead_mask) out1 = self.layernorm1(x + self.dropout1(attn1)) # 编码器-解码器注意力 attn2 = self.mha2(out1, enc_output, enc_output, padding_mask) out2 = self.layernorm2(out1 + self.dropout2(attn2)) ffn_output = self.ffn(out2) return self.layernorm3(out2 + self.dropout3(ffn_output))

第二注意力层的Query来自解码器，而Key和Value来自编码器输出。这种设计让解码器可以动态关注输入的不同部分。

5. 完整模型组装

将所有组件组合起来：

class Transformer(nn.Module): def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size, target_vocab_size, dropout=0.1): super().__init__() self.encoder = Encoder(num_layers, d_model, num_heads, dff, input_vocab_size, dropout) self.decoder = Decoder(num_layers, d_model, num_heads, dff, target_vocab_size, dropout) self.final_layer = nn.Linear(d_model, target_vocab_size) def forward(self, inp, tar, enc_padding_mask, look_ahead_mask, dec_padding_mask): enc_output = self.encoder(inp, enc_padding_mask) dec_output = self.decoder(tar, enc_output, look_ahead_mask, dec_padding_mask) return self.final_layer(dec_output)

训练时需要注意三个关键技巧：