动画图解Transformer：从自注意力到编码器-解码器架构的完整实现

大家好我是专注于AI与深度学习领域的技术博主。Transformer模型自2017年横空出世以来已成为驱动当今AI大模型浪潮的绝对核心。无论是ChatGPT背后的GPT系列还是文生图领域的Stable Diffusion其底层架构都离不开Transformer。然而其论文《Attention Is All You Need》中的复杂公式和抽象描述常常让初学者望而却步。本文旨在用最直观的“动画”式思维结合清晰的代码实战带你从零开始彻底搞懂Transformer。我们将摒弃晦涩的数学推导专注于理解其核心思想、数据流动和工程实现。无论你是刚入门AI的学生还是希望深入理解大模型原理的开发者都能通过本文构建起对Transformer的完整认知框架并亲手实现一个可运行的简化版模型。1. Transformer 核心思想为什么要“注意力”在Transformer出现之前处理序列数据如文本、语音、时间序列的主流模型是RNN循环神经网络及其变体LSTM、GRU。RNN的核心问题是顺序计算必须一个接一个地处理序列中的元素这导致了训练速度慢无法并行和长距离依赖信息容易丢失。Transformer的革命性在于它完全抛弃了循环结构转而采用一种名为“自注意力机制Self-Attention”的架构实现了对序列数据的并行化全局建模。1.1 直观理解“注意力”想象一下你在阅读这句话“这只猫因为很久没吃饭所以它很饿。” 当你理解“它”指代谁时你的大脑会瞬间将“它”与前面的“猫”关联起来而忽略“吃饭”等词。这种聚焦于相关信息的能力就是“注意力”。在模型中自注意力机制让序列中的每一个词都能直接与序列中的所有其他词进行交互计算一个“关联度分数”从而动态地为每个词生成一个包含全局信息的新表示。1.2 Transformer 整体架构鸟瞰Transformer是一个“编码器-解码器Encoder-Decoder”架构最初为机器翻译任务设计。编码器Encoder负责理解并压缩输入序列如源语言句子的信息输出一个蕴含丰富语义的上下文矩阵。解码器Decoder基于编码器的输出和已生成的部分结果自回归地一个接一个生成目标序列如目标语言句子。其最核心的单元是Transformer Block编码器和解码器都由多个这样的Block堆叠而成。一个Transformer Block的核心包含两层多头自注意力层Multi-Head Self-Attention实现上文所述的全局信息交互。前馈神经网络层Feed-Forward Network对每个位置的表示进行独立、非线性的变换。每一层周围都包裹着残差连接Residual Connection和层归一化Layer Normalization这是训练深层网络稳定性的关键技巧。下面我们通过一个数据流的视角看看一个句子是如何穿越Transformer的。2. 从词到向量输入处理三部曲计算机无法直接理解文字所以第一步是将文本转化为数字——向量。2.1 词嵌入Word Embedding每个词或子词被映射为一个固定长度的稠密向量例如512维。这个向量空间具有语义特性语义相似的词其向量在空间中的距离也更近。import torch import torch.nn as nn # 假设我们的词汇表大小为10000嵌入维度为512 vocab_size 10000 d_model 512 embedding_layer nn.Embedding(vocab_size, d_model) # 输入一个句子每个词用索引表示 # 例如“Hello world” - [101, 102] input_indices torch.LongTensor([[101, 102]]) # shape: [batch_size, seq_len] word_embeddings embedding_layer(input_indices) # shape: [batch_size, seq_len, d_model] print(f词嵌入后形状{word_embeddings.shape})2.2 位置编码Positional Encoding自注意力机制本身没有顺序概念。“猫吃鱼”和“鱼吃猫”在它看来是一样的。为了注入序列的顺序信息我们需要加入位置编码。Transformer使用了一组固定公式正弦和余弦函数来生成位置编码向量其维度与词嵌入相同然后直接加到词嵌入上。class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len5000): super(PositionalEncoding, self).__init__() pe torch.zeros(max_len, d_model) position torch.arange(0, max_len).unsqueeze(1) # shape: [max_len, 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) # 偶数维用sin pe[:, 1::2] torch.cos(position * div_term) # 奇数维用cos pe pe.unsqueeze(0) # shape: [1, max_len, d_model] 便于广播 self.register_buffer(pe, pe) # 将其注册为缓冲区不参与训练 def forward(self, x): # x shape: [batch_size, seq_len, d_model] seq_len x.size(1) x x self.pe[:, :seq_len] # 将位置编码加到输入上 return x # 使用示例 pos_encoder PositionalEncoding(d_model) encoder_input pos_encoder(word_embeddings) # 此时输入已包含词义和位置信息2.3 进入Transformer Block前的准备经过词嵌入和位置编码相加后我们得到了编码器的初始输入X其形状为[batch_size, seq_len, d_model]。这个X将流入第一个编码器Block。3. 核心中的核心自注意力机制详解这是理解Transformer最关键的一步。我们通过“动画”思维来拆解。3.1 Q, K, V 的由来对于输入矩阵X我们通过三个不同的线性变换层为每个词生成三个新的向量查询向量Query代表当前词“想要寻找什么”。键向量Key代表当前词“能提供什么”。值向量Value代表当前词“实际的内容信息”。def generate_qkv(x, d_model, d_k): x: 输入形状 [batch_size, seq_len, d_model] 返回 Q, K, V batch_size, seq_len, _ x.shape # 简化起见假设我们有一个线性层 W_q nn.Linear(d_model, d_k) W_k nn.Linear(d_model, d_k) W_v nn.Linear(d_model, d_k) Q W_q(x) # [batch_size, seq_len, d_k] K W_k(x) # [batch_size, seq_len, d_k] V W_v(x) # [batch_size, seq_len, d_k] return Q, K, V3.2 计算注意力分数与权重注意力分数的计算可以类比为“匹配度查询”。对于序列中的第i个词Query我们计算它与序列中所有词包括自己的 Key 的点积来衡量它们之间的相关性。计算分数Score Q * K^T。结果是一个[seq_len, seq_len]的矩阵第i行第j列表示第i个词对第j个词的关注程度。缩放Scale为了防止点积结果过大导致Softmax梯度消失将分数除以sqrt(d_k)。Softmax归一化对每一行对应一个Query进行Softmax使得该行所有分数之和为1转化为概率分布形式的注意力权重。权重高的位置表示当前词应该更“关注”那个位置的词。import torch.nn.functional as F def scaled_dot_product_attention(Q, K, V, maskNone): Q, K, V: 形状均为 [batch_size, seq_len, d_k] mask: 可选用于在解码器屏蔽未来信息 d_k Q.size(-1) # 1. 计算分数 scores torch.matmul(Q, K.transpose(-2, -1)) # [batch_size, seq_len, seq_len] # 2. 缩放 scores scores / torch.sqrt(torch.tensor(d_k, dtypetorch.float32)) # 3. 可选应用掩码解码器用 if mask is not None: scores scores.masked_fill(mask 0, -1e9) # 将mask为0的位置置为负无穷 # 4. Softmax得到权重 attention_weights F.softmax(scores, dim-1) # [batch_size, seq_len, seq_len] # 5. 加权求和 output torch.matmul(attention_weights, V) # [batch_size, seq_len, d_k] return output, attention_weights # 示例 d_k 64 Q, K, V generate_qkv(encoder_input, d_model, d_k) context_vector, attn_weights scaled_dot_product_attention(Q, K, V) print(f上下文向量形状{context_vector.shape}) print(f注意力权重形状{attn_weights.shape})动画想象你可以把attention_weights矩阵想象成一个热力图。对于句子中的每个词行都有一列权重值颜色越亮值越大代表当前词在生成新表示时越“注意”那个位置的原始词。3.3 为什么要“多头”Multi-Head单一的注意力机制可能只关注一种类型的依赖关系例如语法依赖。多头注意力并行地运行多个例如8个独立的注意力层每个“头”可以将输入投影到不同的子空间从而学习到不同方面的依赖关系例如语法、语义、指代等。最后将所有头的输出拼接起来再经过一个线性变换得到最终的多头注意力输出。class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): super(MultiHeadAttention, self).__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, maskNone): batch_size, seq_len, d_model x.shape # 1. 线性投影并分割成多头 Q self.W_q(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) K self.W_k(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) V self.W_v(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) # Q,K,V 形状: [batch_size, num_heads, seq_len, d_k] # 2. 每个头独立计算缩放点积注意力 scores torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.d_k, dtypetorch.float32)) if mask is not None: scores scores.masked_fill(mask 0, -1e9) attention_weights F.softmax(scores, dim-1) context torch.matmul(attention_weights, V) # [batch_size, num_heads, seq_len, d_k] # 3. 合并多头 context context.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model) # 4. 最终线性投影 output self.W_o(context) return output, attention_weights4. 构建完整的Transformer编码器块一个编码器块 多头自注意力 残差层归一化 前馈网络 残差层归一化。4.1 前馈网络FFN这是一个简单的两层全连接网络通常中间有一个ReLU激活函数用于对每个位置的表示进行非线性变换。class PositionwiseFeedForward(nn.Module): def __init__(self, d_model, d_ff): super(PositionwiseFeedForward, self).__init__() self.linear1 nn.Linear(d_model, d_ff) # 通常 d_ff 4 * d_model self.linear2 nn.Linear(d_ff, d_model) self.activation nn.ReLU() def forward(self, x): return self.linear2(self.activation(self.linear1(x)))4.2 编码器块实现class EncoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout0.1): super(EncoderLayer, self).__init__() self.self_attn MultiHeadAttention(d_model, num_heads) self.feed_forward PositionwiseFeedForward(d_model, d_ff) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) def forward(self, x, maskNone): # 子层1: 多头自注意力 残差 层归一化 attn_output, _ self.self_attn(x, mask) x self.norm1(x self.dropout(attn_output)) # 残差连接后归一化 # 子层2: 前馈网络 残差 层归一化 ff_output self.feed_forward(x) x self.norm2(x self.dropout(ff_output)) return x残差连接将子层的输入直接加到其输出上x sublayer(x)。这有助于缓解深层网络中的梯度消失问题让模型更容易学习恒等映射。层归一化对单个样本的所有特征进行归一化稳定训练过程。5. 解码器与掩码多头注意力解码器比编码器稍复杂因为它需要自回归生成。5.1 解码器的两个注意力层掩码多头自注意力层Masked Multi-Head Self-Attention确保在预测第t个位置时只能看到1到t-1位置的信息即已生成的部分而不能“偷看”未来的信息。这是通过一个上三角掩码矩阵实现的。编码器-解码器注意力层Encoder-Decoder Attention其 Query 来自解码器上一层的输出而 Key 和 Value 来自编码器的最终输出。这让解码器在生成每一个词时都能有选择地聚焦于输入序列的不同部分。5.2 解码器块实现class DecoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout0.1): super(DecoderLayer, self).__init__() self.self_attn MultiHeadAttention(d_model, num_heads) self.cross_attn MultiHeadAttention(d_model, num_heads) # 编码器-解码器注意力 self.feed_forward PositionwiseFeedForward(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(dropout) def forward(self, x, encoder_output, src_maskNone, tgt_maskNone): # 子层1: 掩码自注意力 attn1_output, _ self.self_attn(x, tgt_mask) x self.norm1(x self.dropout(attn1_output)) # 子层2: 编码器-解码器注意力 # Query来自解码器Key和Value来自编码器输出 attn2_output, _ self.cross_attn(x, encoder_output, encoder_output, src_mask) x self.norm2(x self.dropout(attn2_output)) # 子层3: 前馈网络 ff_output self.feed_forward(x) x self.norm3(x self.dropout(ff_output)) return x6. 项目实战构建一个简化的Transformer模型现在我们将上述部件组装起来构建一个用于序列到序列任务如机器翻译的简化版Transformer。6.1 模型定义import torch import torch.nn as nn import math import copy class Transformer(nn.Module): def __init__(self, src_vocab_size, tgt_vocab_size, d_model512, num_heads8, num_encoder_layers6, num_decoder_layers6, d_ff2048, max_seq_len5000, dropout0.1): super(Transformer, self).__init__() self.d_model d_model # 1. 词嵌入和位置编码 self.src_embedding nn.Embedding(src_vocab_size, d_model) self.tgt_embedding nn.Embedding(tgt_vocab_size, d_model) self.positional_encoding PositionalEncoding(d_model, max_seq_len) # 2. 编码器堆叠 encoder_layer EncoderLayer(d_model, num_heads, d_ff, dropout) self.encoder nn.ModuleList([copy.deepcopy(encoder_layer) for _ in range(num_encoder_layers)]) # 3. 解码器堆叠 decoder_layer DecoderLayer(d_model, num_heads, d_ff, dropout) self.decoder nn.ModuleList([copy.deepcopy(decoder_layer) for _ in range(num_decoder_layers)]) # 4. 输出层线性变换 Softmax self.output_linear nn.Linear(d_model, tgt_vocab_size) self.dropout nn.Dropout(dropout) def generate_src_mask(self, src): # 简单示例假设padding的索引为0创建掩码 (batch_size, 1, 1, src_len) src_mask (src ! 0).unsqueeze(1).unsqueeze(2) return src_mask def generate_tgt_mask(self, tgt): # 创建后续位置掩码上三角矩阵和padding掩码 batch_size, tgt_len tgt.shape # 后续位置掩码 subsequent_mask torch.tril(torch.ones(tgt_len, tgt_len)).bool().unsqueeze(0).unsqueeze(0) # (1,1,tgt_len,tgt_len) # padding掩码 padding_mask (tgt ! 0).unsqueeze(1).unsqueeze(2) # (batch_size,1,1,tgt_len) # 合并掩码 tgt_mask subsequent_mask padding_mask return tgt_mask def forward(self, src, tgt): # 生成掩码 src_mask self.generate_src_mask(src) tgt_mask self.generate_tgt_mask(tgt) # 编码器前向传播 src_emb self.dropout(self.positional_encoding(self.src_embedding(src) * math.sqrt(self.d_model))) enc_output src_emb for enc_layer in self.encoder: enc_output enc_layer(enc_output, src_mask) # 解码器前向传播 tgt_emb self.dropout(self.positional_encoding(self.tgt_embedding(tgt) * math.sqrt(self.d_model))) dec_output tgt_emb for dec_layer in self.decoder: dec_output dec_layer(dec_output, enc_output, src_mask, tgt_mask) # 输出层 output self.output_linear(dec_output) return output # 实例化一个微型模型用于演示 src_vocab_size 5000 tgt_vocab_size 5000 model Transformer(src_vocab_size, tgt_vocab_size, d_model128, num_heads4, num_encoder_layers2, num_decoder_layers2, d_ff512) print(model)6.2 训练与推理流程简述训练阶段准备源语言句子src和目标语言句子tgt。将tgt整体输入解码器使用掩码确保自回归属性。模型输出预测的每个位置的下一个词的概率分布。计算预测分布与真实目标之间的交叉熵损失反向传播更新参数。推理阶段预测编码器处理源句子得到上下文表示。解码器从起始符如sos开始每次生成一个词。将已生成的序列作为新的tgt输入解码器预测下一个词的概率通常取概率最高的词贪婪搜索或使用束搜索Beam Search。重复步骤3直到生成结束符如eos或达到最大长度。7. 常见问题与排查思路问题现象可能原因解决思路训练时Loss为NaN或不收敛1. 学习率过高。2. 梯度爆炸。3. 数据中存在异常值或未进行归一化。1. 使用更小的学习率或使用学习率预热Warmup。2. 使用梯度裁剪Gradient Clipping。3. 检查数据预处理确保词嵌入和位置编码的数值范围合理。模型参数量巨大内存不足d_model,num_heads,num_layers设置过大。1. 减小模型尺寸或使用模型并行、梯度累积。2. 使用混合精度训练AMP。3. 检查批次大小Batch Size。推理速度慢1. 解码器自回归生成无法并行。2. 模型层数多计算量大。1. 使用缓存KV Cache避免重复计算。2. 考虑模型量化、剪枝或使用更高效的推理框架如ONNX Runtime, TensorRT。3. 对于特定任务可尝试仅使用编码器如BERT或仅使用解码器如GPT的架构。注意力权重过于均匀或稀疏1. 初始化问题。2. 模型深度不够或过拟合。1. 检查参数初始化方法如Xavier, Kaiming。2. 增加/减少Dropout率。3. 可视化注意力图辅助分析。位置编码效果不明显1. 序列长度远超训练时见过的最大长度。2. 可学习的位置编码可能在某些任务上更好。1. 使用相对位置编码如RoPE, ALiBi来外推更长的序列。2. 尝试可学习的位置嵌入Learnable Positional Embedding。8. 最佳实践与工程建议学习率策略是关键Transformer模型通常需要配合学习率预热Warmup和学习率衰减Decay。Warmup让模型在训练初期稳定避免震荡。AdamW优化器是标配。重视数据预处理对于文本子词切分如BPE、WordPiece、SentencePiece比整词更有效能更好地处理未登录词和稀有词。确保源语言和目标语言的词汇表对齐。批量训练与填充序列长度不一需要填充Padding到同一长度。务必使用**掩码Mask**来忽略填充位置对注意力计算和损失函数的影响。层归一化的位置原始Transformer使用“后归一化”Post-LN即LayerNorm(x Sublayer(x))。现在许多研究如GPT采用“前归一化”Pre-LN即x Sublayer(LayerNorm(x))通常训练更稳定。解码策略的选择贪婪搜索Greedy Search速度快但结果可能平庸束搜索Beam Search能平衡质量和多样性是机器翻译的常用选择对于创意生成可采样Sampling或核采样Top-p, Top-k能产生更多样化的结果。可视化工具使用torchviz或netron可视化计算图使用tensorboard监控训练过程使用seaborn或matplotlib绘制注意力权重热力图这对调试和理解模型行为至关重要。从预训练模型开始除非有特定研究目的否则在大多数下游任务上微调Fine-tune一个大规模预训练的Transformer模型如Hugging Face Transformers库中的BERT、GPT、T5等是最高效的路径。理解Transformer的架构是深入现代AI世界的基石。它优雅地将注意力机制发挥到极致通过并行计算和全局建模能力解决了序列建模的根本瓶颈。虽然本文实现的模型是简化版但它完整地勾勒出了Transformer的数据流和核心组件。建议你动手运行代码尝试在小数据集如数字加法、简单翻译对上训练观察其学习过程。之后可以深入探索BERT仅编码器、GPT仅解码器、Vision TransformerViT等变体以及更高级的技巧如稀疏注意力、线性注意力等。