引言2017年Google在《Attention Is All You Need》中提出Transformer彻底革新了序列建模的范式。它抛弃了RNN和CNN仅依赖注意力机制成为了BERT、GPT等大模型的基石。本文将从直观理解Self-Attention出发逐步拆解Transformer的核心组件并提供一个基于PyTorch的完整可运行示例带你亲手实现一个Transformer编码器-解码器完成序列到序列的复制任务。读完这篇文章你不仅能理解原理还能写出属于自己的Transformer。一、核心概念从注意力到多头自注意力1.1 缩放点积注意力Scaled Dot-Product Attention注意力机制可以看作一个查询Query去键Key库中查找相关值Value的过程。缩放点积注意力的公式如下[\text{Attention}(Q, K, V) \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V](Q, K, V) 分别是查询、键、值矩阵形状为序列长度维度。除以 (\sqrt{d_k}) 是为了缩放防止点积结果过大导致softmax梯度极小。Softmax后得到注意力权重再与V相乘得到最终的注意力输出。从几何角度看(QK^T) 计算了序列中每个位置相对于其他位置的“相关性分数”。缩放后通过Softmax归一化成概率分布然后用V的加权和得到每个位置的表示。1.2 多头注意力Multi-Head Attention单头注意力只能捕捉一种关系模式。多头注意力并行地进行多次注意力计算每个“头”关注不同子空间的信息最后拼接结果并线性变换。[\text{MultiHead}(Q, K, V) \text{Concat}(\text{head}_1, \dots, \text{head}_h)W^O][\text{head}_i \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)]其中 (W_i^Q, W_i^K, W_i^V) 是每个头的投影矩阵(W^O) 是输出投影。通过这种方式模型可以同时关注不同位置、不同表示层面的依赖。1.3 自注意力Self-Attention在Transformer中编码器和解码器的每一层都使用多头自注意力。所谓“自注意力”就是 (QKV)它们都来自同一个输入序列。这意味着每个词都能与序列中的所有词建立联系直接捕获长距离依赖而无需像RNN那样逐步传递信息。1.4 位置编码Positional Encoding由于注意力机制本身不包含位置信息Transformer需要显式地注入位置编码。原始论文使用正弦和余弦函数生成固定位置编码[PE_{(pos, 2i)} \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right)][PE_{(pos, 2i1)} \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right)]其中 (pos) 是位置(i) 是维度索引。这种编码使得模型可以轻易学到相对位置关系因为对于任意偏移k(PE_{posk}) 可以表示为 (PE_{pos}) 的线性函数。1.5 前馈网络与残差连接每个注意力子层后都有一个全连接的前馈网络FFN它由两个线性变换和一个激活函数组成通常为ReLU[\text{FFN}(x) \max(0, xW_1 b_1)W_2 b_2]此外每个子层都使用了残差连接和层归一化(LayerNorm(x Sublayer(x)))。这能有效缓解深层网络中的梯度消失问题。二、实战示例用PyTorch从零实现Transformer下面我们实现一个简化版的Transformer包含编码器和解码器并训练它完成序列复制任务输入一个随机序列期望输出相同的序列向右偏移一位作为经典的seq2seq任务。代码完整可运行注释详细。import torch import torch.nn as nn import torch.optim as optim import math # ------------------ 1. 位置编码 ------------------ 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).float() div_term torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)) pe[:, 0::2] torch.sin(position * div_term) pe[:, 1::2] torch.cos(position * div_term) pe pe.unsqueeze(0) # (1, max_len, d_model) self.register_buffer(pe, pe) def forward(self, x): # x: (batch_size, seq_len, d_model) return x self.pe[:, :x.size(1)] # ------------------ 2. 多头注意力 ------------------ class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads, dropout0.1): 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) self.dropout nn.Dropout(dropout) def scaled_dot_product_attention(self, Q, K, V, maskNone): scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: scores scores.masked_fill(mask 0, -1e9) attn torch.softmax(scores, dim-1) attn self.dropout(attn) return torch.matmul(attn, V), attn def forward(self, Q, K, V, maskNone): batch_size Q.size(0) # 线性投影并分割为多头 Q self.W_q(Q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) K self.W_k(K).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) V self.W_v(V).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) # 注意力 out, attn self.scaled_dot_product_attention(Q, K, V, mask) # 合并多头 out out.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.d_k) return self.W_o(out) # ------------------ 3. 编码器层 ------------------ class EncoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout0.1): super().__init__() self.self_attn MultiHeadAttention(d_model, num_heads, dropout) self.feed_forward nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model) ) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) def forward(self, x, maskNone): attn_out self.self_attn(x, x, x, mask) x self.norm1(x self.dropout(attn_out)) ff_out self.feed_forward(x) x self.norm2(x self.dropout(ff_out)) return x # ------------------ 4. 解码器层带掩码自注意力和交叉注意力 ------------------ class DecoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout0.1): super().__init__() self.self_attn MultiHeadAttention(d_model, num_heads, dropout) self.cross_attn MultiHeadAttention(d_model, num_heads, dropout) self.feed_forward nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model) ) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.norm3 nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) def forward(self, x, enc_output, src_maskNone, tgt_maskNone): # 带掩码的自注意力确保预测时不偷看未来 self_attn_out self.self_attn(x, x, x, tgt_mask) x self.norm1(x self.dropout(self_attn_out)) # 交叉注意力Q来自解码器K、V来自编码器输出 cross_attn_out self.cross_attn(x, enc_output, enc_output, src_mask) x self.norm2(x self.dropout(cross_attn_out)) # 前馈网络 ff_out self.feed_forward(x) x self.norm3(x self.dropout(ff_out)) return x # ------------------ 5. 完整Transformer ------------------ class Transformer(nn.Module): def __init__(self, src_vocab_size, tgt_vocab_size, d_model128, num_heads8, num_encoder_layers3, num_decoder_layers3, d_ff512, dropout0.1): super().__init__() self.encoder_embedding nn.Embedding(src_vocab_size, d_model) self.decoder_embedding nn.Embedding(tgt_vocab_size, d_model) self.positional_encoding PositionalEncoding(d_model) self.encoder_layers nn.ModuleList( [EncoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_encoder_layers)] ) self.decoder_layers nn.ModuleList( [DecoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_decoder_layers)] ) self.fc_out nn.Linear(d_model, tgt_vocab_size) def generate_mask(self, src, tgt): # 源序列掩码 (1表示保留0表示屏蔽此处简单设为全1) src_mask torch.ones(src.size(0), 1, 1, src.size(1)).to(src.device) # 目标序列掩码上三角为0防止看到未来信息 tgt_len tgt.size(1) subsequent_mask torch.triu(torch.ones(tgt_len, tgt_len), diagonal1).bool() tgt_mask (1 - subsequent_mask.int()).unsqueeze(0).unsqueeze(0).to(tgt.device) return src_mask, tgt_mask def forward(self, src, tgt): src_mask, tgt_mask self.generate_mask(src, tgt) # 编码器 src_emb self.positional_encoding(self.encoder_embedding(src)) enc_output src_emb for layer in self.encoder_layers: enc_output layer(enc_output, src_mask) # 解码器 tgt_emb self.positional_encoding(self.decoder_embedding(tgt)) dec_output tgt_emb for layer in self.decoder_layers: dec_output layer(dec_output, enc_output, src_mask, tgt_mask) return self.fc_out(dec_output) # ------------------ 6. 训练序列复制任务 ------------------ def prepare_batch(batch_size, seq_len, vocab_size): # 生成随机整数序列作为输入输出为相同序列向右偏移一位 src torch.randint(1, vocab_size, (batch_size, seq_len)) tgt_input torch.cat([torch.full((batch_size, 1), 2), src[:, :-1]], dim1) # 起始符为2 tgt_output src # 期望输出原始序列 return src, tgt_input, tgt_output def train(): device torch.device(cuda if torch.cuda.is_available() else cpu) vocab_size 50 d_model 128 num_heads 4 lr 0.001 epochs 2000 seq_len 10 batch_size 32 model Transformer(vocab_size, vocab_size, d_model, num_heads).to(device) optimizer optim.Adam(model.parameters(), lrlr) criterion nn.CrossEntropyLoss() for epoch in range(1, epochs 1): model.train() src, tgt_input, tgt_output prepare_batch(batch_size, seq_len, vocab_size) src, tgt_input, tgt_output src.to(device), tgt_input.to(device), tgt_output.to(device) optimizer.zero_grad() logits model(src, tgt_input) # (batch, seq_len, vocab_size) loss criterion(logits.view(-1, vocab_size), tgt_output.view(-1)) loss.backward() optimizer.step() if epoch % 200 0: print(fEpoch {epoch}, Loss: {loss.item():.4f})运行说明上述代码定义了一个Transformer模型并在随机生成的序列复制任务上进行了训练。损失会随着训练下降证明模型能学习到复制行为。代码亮点-PositionalEncoding使用正弦位置编码注册为buffer无需梯度更新。-MultiHeadAttention实现了缩放点积注意力并支持掩码。- 编码器和解码器层均包含残差连接和层归一化。- 解码器中的mask确保自回归预测时不会窥探未来。- 训练数据动态生成任务简单可以快速验证模型正确性。三、常见问题与注意事项为什么需要除以 (\sqrt{d_k})如果不缩放当 (d_k) 很大时点积结果的方差会变大导致Softmax后的分布极端近似one-hot梯度几乎为零模型无法学习。缩放后点积的方差被控制为1使得梯度稳定。自注意力中的Mask如何起作用在解码器中我们将未来位置的注意力权重设为 (-\infty)代码中 -1e9经过Softmax后这些位置的权重接近0从而禁止信息从未来流向当前时刻。多头注意力与单头区别多头允许模型在不同的表示子空间里关注不同特征类似于CNN中的多个滤波器。实验表明多头能显著提升性能。位置编码是否可学习原始论文使用固定的正弦位置编码但后续工作如BERT使用可学习的位置嵌入。二者效果相当固定编码的好处是可以外推到更长的序列可学习编码则需要重新训练。训练Transformer的难点- 需要较大的训练数据和合适的超参数学习率、预热策略。- 对输入序列长度敏感长序列的计算复杂度为 (O(n^2))。- 容易过拟合小数据集需添加Dropout和标签平滑。总结本文从Self-Attention的数学原理出发逐步构建了多头注意力、编码器层和解码器层最终给出了一个完整的Transformer PyTorch实现。通过序列复制任务我们验证了模型的正确性。理解Transformer不仅是掌握当前NLP主流范式的基础更是深入大模型设计的第一步。建议读者运行代码尝试修改超参数观察注意力权重相信你会对“Attention Is All You Need”有更深的体会。