这次我们来看一个来自CCF-LMCC-10的技术分享主题直击Transformer架构的核心。很多人学习Transformer时注意力机制Attention往往是焦点但理解了多头注意力Multi-Head Attention后面对整个Transformer的架构图依然会感到困惑各个模块之间到底如何连接数据流如何传递模型究竟是怎么“搭”起来的这篇文章将彻底拆解Transformer的构建过程。我们不只停留在注意力这个“零件”上而是从零开始一步步组装出完整的编码器Encoder和解码器Decoder。重点在于理解架构层面的设计逻辑、模块间的协作关系以及如何用代码实现一个可运行的简化版Transformer。这对于想深入理解BERT、GPT、T5等现代大模型底层原理甚至有志于进行模型架构创新的开发者来说是必不可少的一课。本文会以“搭积木”的方式展开先回顾核心组件然后分别构建编码器和解码器最后将它们组合成一个完整的Transformer模型。我们将重点关注模块的接口定义、数据维度的变化以及前向传播的完整流程并提供可执行的PyTorch代码片段。无论你是想巩固理论基础还是为后续的模型微调或魔改打下坚实基础这篇内容都能提供清晰的路径。1. 核心能力速览Transformer架构拆解目标在深入细节之前我们先明确本次“搭建”工程的目标和范围。能力项说明项目类型技术原理剖析与代码实现指南核心目标从模块层面理解Transformer如何组装掌握编码器-解码器数据流关键组件嵌入层、位置编码、多头注意力、前馈网络、层归一化、残差连接输出成果一个结构清晰、可执行训练的简化版Transformer模型代码硬件门槛极低。本文侧重于CPU上的原理验证模型规模小无需GPU。适合读者已了解注意力机制基础希望深入Transformer架构细节的开发者准备面试需要梳理模型结构的求职者。前置知识基本的Python/PyTorch知识对神经网络、矩阵运算、注意力机制有初步了解。2. 适用场景与使用边界理解Transformer的搭建过程其价值远不止于复现一个模型。它最适合以下场景深度学习面试准备面试官常问“请画出Transformer结构并解释数据流向”。通过亲手搭建你能对每个模块的位置和作用了如指掌。大模型技术栈深耕BERT、GPT、T5、ViT等模型都是Transformer的变体。掌握原版架构是理解这些变体如仅编码器、仅解码器、Encoder-Decoder的基础。自定义模型研发当你想修改Transformer的某个部分如设计新的注意力机制、位置编码方式或层间连接时清晰的模块化认知能让你快速定位修改点。教学与知识梳理用代码将理论具象化是检验和深化理解的最佳方式。需要明确的使用边界非生产级实现本文实现的将是简化版例如忽略Masked Attention的细节、使用固定位置编码旨在阐明原理。生产级的Transformer库如Hugging Facetransformers经过了大量优化和细节完善。侧重架构而非训练重点在于模型的前向传播构建不会深入讨论复杂的训练技巧如学习率预热、标签平滑等。计算规模有限为了演示清晰我们使用的模型维度d_model和头数num_heads都很小与实际的大规模预训练模型有数量级差异。3. 环境准备与前置条件搭建环境非常简单只需要基础的Python科学计算环境。基础环境清单操作系统Windows / macOS / Linux 均可。Python版本 3.7 或以上。核心库PyTorch 1.9.0。我们将使用其提供的nn.Module,nn.Linear,nn.Dropout等基础模块以及torch.matmul等张量操作。NumPy用于一些辅助计算如位置编码的生成。开发工具任意你熟悉的代码编辑器或IDE如VSCode、PyCharm。硬件普通CPU即可。我们的演示模型参数量极小不消耗显存。环境搭建步骤确保已安装Python。使用pip安装必要的包。建议使用虚拟环境。# 创建并激活虚拟环境可选但推荐 python -m venv transformer_env source transformer_env/bin/activate # Linux/macOS # transformer_env\Scripts\activate # Windows # 安装PyTorch (以CPU版本为例访问官网 https://pytorch.org/ 获取最适合你系统的命令) pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu # 安装numpy pip install numpy验证安装import torch import numpy as np print(torch.__version__) # 应输出类似 1.13.0 的版本号 print(np.__version__) # 应输出类似 1.21.0 的版本号4. 回顾核心“零件”注意力与FFN在组装整机前先确保我们手头的“零件”功能清晰。Transformer的核心零件有两个多头注意力机制MHA和前馈网络FFN。4.1 多头注意力机制Multi-Head Attention, MHA这是Transformer的标志性组件。其核心思想是将输入序列的表示空间投影到多个“子空间”即多个头中在每个头里独立进行注意力计算最后将结果合并。关键步骤线性投影对输入X分别进行三次线性变换得到Query(Q)、Key(K)、Value(V)矩阵。分割头将Q、K、V在特征维度上分割成num_heads份。缩放点积注意力对每个头计算Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V。合并头将多个头的输出在特征维度上拼接起来。输出投影通过一个线性层将拼接后的特征映射回原始维度。一个简化的MHA实现框架import torch import torch.nn as nn import torch.nn.functional as F class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): super().__init__() assert d_model % num_heads 0, d_model must be divisible by num_heads self.d_model d_model self.num_heads num_heads self.d_k d_model // num_heads # 定义Q, K, V的线性变换层和最后的输出层 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, query, key, value, maskNone): batch_size query.size(0) # 1. 线性投影并分头 Q self.W_q(query).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) K self.W_k(key).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) V self.W_v(value).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) # 2. 计算缩放点积注意力 scores torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5) if mask is not None: scores scores.masked_fill(mask 0, -1e9) attn_weights F.softmax(scores, dim-1) attn_output torch.matmul(attn_weights, V) # 3. 合并头并输出投影 attn_output attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) output self.W_o(attn_output) return output4.2 前馈网络Feed-Forward Network, FFN这是一个简单的两层全连接网络通常在每个注意力层之后应用。它为模型引入了非线性变换和维度变换的能力。典型结构FFN(x) max(0, x W1 b1) W2 b2中间层的维度d_ff通常比输入维度d_model更大例如d_ff 4 * d_model。实现代码class PositionwiseFeedForward(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) self.activation nn.ReLU() # 原论文使用ReLU def forward(self, x): return self.linear2(self.dropout(self.activation(self.linear1(x))))有了这两个核心零件再加上层归一化LayerNorm和残差连接Residual Connection我们就可以开始搭建编码器和解码器层了。5. 搭建编码器Encoder编码器由N个完全相同的层堆叠而成。每一层都包含两个子层多头自注意力子层Multi-Head Self-Attention输入序列自身做注意力捕捉序列内部元素间的关系。前馈网络子层Position-wise FFN对每个位置独立进行变换。每个子层周围都应用了“残差连接”和“层归一化”即LayerNorm(x Sublayer(x))。这是稳定训练深度网络的关键。5.1 编码器层实现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) self.feed_forward PositionwiseFeedForward(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, src, src_maskNone): # 子层1: 多头自注意力 Add Norm attn_output self.self_attn(src, src, src, src_mask) src self.norm1(src self.dropout1(attn_output)) # 子层2: 前馈网络 Add Norm ff_output self.feed_forward(src) src self.norm2(src self.dropout2(ff_output)) return src5.2 编码器实现编码器需要将输入符号序列转换为连续表示。这包括两个步骤输入嵌入Input Embedding将离散的token ID映射为d_model维的向量。位置编码Positional Encoding为序列中的每个位置注入位置信息因为自注意力机制本身不具备序列顺序感知能力。class Encoder(nn.Module): def __init__(self, vocab_size, d_model, num_layers, num_heads, d_ff, max_seq_len, dropout0.1): super().__init__() self.token_embedding nn.Embedding(vocab_size, d_model) self.position_encoding self._create_position_encoding(max_seq_len, d_model) self.layers nn.ModuleList([EncoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)]) self.norm nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) def _create_position_encoding(self, max_len, d_model): 创建正弦余弦位置编码 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) return pe def forward(self, src_tokens, src_maskNone): # src_tokens: (batch_size, src_len) batch_size, src_len src_tokens.size() # 1. 词嵌入 x self.token_embedding(src_tokens) # (batch_size, src_len, d_model) # 2. 加位置编码 x x self.position_encoding[:, :src_len, :].to(x.device) x self.dropout(x) # 3. 通过N个编码器层 for layer in self.layers: x layer(x, src_mask) return self.norm(x)至此一个完整的编码器就搭建好了。它接收一个整数ID序列输出一个融合了上下文信息的序列表示。6. 搭建解码器Decoder解码器同样由N个相同的层堆叠而成但每层比编码器层多一个子层共三个子层掩码多头自注意力子层Masked Multi-Head Self-Attention防止当前位置关注到未来的位置确保自回归生成特性。编码器-解码器注意力子层Encoder-Decoder Attention让解码器关注编码器的输出。这里的Query来自解码器上一子层的输出而Key和Value来自编码器的最终输出。前馈网络子层Position-wise FFN与编码器相同。每个子层同样遵循“Add Norm”模式。6.1 解码器层实现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) self.cross_attn MultiHeadAttention(d_model, num_heads) # 编码器-解码器注意力 self.feed_forward PositionwiseFeedForward(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, tgt, memory, tgt_maskNone, memory_maskNone): # memory: 编码器的输出 # 子层1: 掩码自注意力 Add Norm attn1_output self.self_attn(tgt, tgt, tgt, tgt_mask) tgt self.norm1(tgt self.dropout1(attn1_output)) # 子层2: 编码器-解码器注意力 Add Norm # Query来自解码器Key和Value来自编码器输出(memory) attn2_output self.cross_attn(tgt, memory, memory, memory_mask) tgt self.norm2(tgt self.dropout2(attn2_output)) # 子层3: 前馈网络 Add Norm ff_output self.feed_forward(tgt) tgt self.norm3(tgt self.dropout3(ff_output)) return tgt6.2 解码器实现解码器的输入处理与编码器类似也需要嵌入和位置编码。其输出需要通过一个线性层和Softmax映射回词汇表以预测下一个token。class Decoder(nn.Module): def __init__(self, vocab_size, d_model, num_layers, num_heads, d_ff, max_seq_len, dropout0.1): super().__init__() self.token_embedding nn.Embedding(vocab_size, d_model) self.position_encoding self._create_position_encoding(max_seq_len, d_model) # 复用编码器的方法 self.layers nn.ModuleList([DecoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)]) self.norm nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) # 输出投影层将d_model维映射回词汇表大小 self.output_projection nn.Linear(d_model, vocab_size) def _create_position_encoding(self, max_len, d_model): # ... 与Encoder中相同略 ... pass def forward(self, tgt_tokens, memory, tgt_maskNone, memory_maskNone): # tgt_tokens: (batch_size, tgt_len) batch_size, tgt_len tgt_tokens.size() # 1. 词嵌入 x self.token_embedding(tgt_tokens) # 2. 加位置编码 x x self.position_encoding[:, :tgt_len, :].to(x.device) x self.dropout(x) # 3. 通过N个解码器层 for layer in self.layers: x layer(x, memory, tgt_mask, memory_mask) x self.norm(x) # 4. 投影到词汇表维度并计算logits logits self.output_projection(x) # (batch_size, tgt_len, vocab_size) return logits7. 组装完整Transformer现在我们将编码器和解码器组合起来并添加必要的掩码生成逻辑就得到了完整的Transformer模型。关键掩码Mask源序列填充掩码src_mask用于在编码器自注意力中忽略padtoken。目标序列填充掩码tgt_mask用于在解码器自注意力中忽略padtoken。前瞻掩码look-ahead mask用于解码器掩码自注意力防止看到未来信息。通常与目标序列填充掩码结合使用。import math import torch import torch.nn as nn class Transformer(nn.Module): def __init__(self, src_vocab_size, tgt_vocab_size, d_model512, num_layers6, num_heads8, d_ff2048, max_seq_len100, dropout0.1): super().__init__() self.encoder Encoder(src_vocab_size, d_model, num_layers, num_heads, d_ff, max_seq_len, dropout) self.decoder Decoder(tgt_vocab_size, d_model, num_layers, num_heads, d_ff, max_seq_len, dropout) # 初始化参数 for p in self.parameters(): if p.dim() 1: nn.init.xavier_uniform_(p) def forward(self, src_tokens, tgt_tokens, src_maskNone, tgt_maskNone, memory_maskNone): # 编码 memory self.encoder(src_tokens, src_mask) # 解码 logits self.decoder(tgt_tokens, memory, tgt_mask, memory_mask) return logits def generate_square_subsequent_mask(self, sz): 生成前瞻掩码上三角矩阵为True/1 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 def create_padding_mask(self, seq): 创建填充掩码seq中为0的位置是pad需要被mask掉 # 假设pad_id为0 padding_mask (seq ! 0).unsqueeze(1).unsqueeze(2) # (batch_size, 1, 1, seq_len) return padding_mask8. 功能测试与效果验证现在我们用一个极简的例子来验证我们搭建的Transformer是否能正确运行。我们将模拟一个机器翻译任务将一句简短的英文“编码”成一句简短的德文。8.1 模拟数据与参数设置def test_transformer_forward(): # 超参数 (使用极小规模以便快速测试) src_vocab_size 100 # 假设源语言词汇表大小 tgt_vocab_size 100 # 假设目标语言词汇表大小 d_model 64 # 模型维度 num_heads 4 # 注意力头数 num_layers 2 # 编码器/解码器层数 d_ff 256 # 前馈网络中间层维度 max_seq_len 10 # 最大序列长度 batch_size 2 # 批大小 # 1. 初始化模型 model Transformer(src_vocab_size, tgt_vocab_size, d_model, num_layers, num_heads, d_ff, max_seq_len, dropout0.0) # dropout设为0便于测试 model.eval() # 切换到评估模式 # 2. 创建模拟输入数据 (假设pad_id0, 序列长度不足max_seq_len的部分用0填充) # 源序列: 假设是英文句子对应的ID序列 src_tokens torch.LongTensor([[3, 5, 7, 2, 0, 0, 0], # 第一句长度4 [1, 4, 6, 8, 9, 2, 0]]) # 第二句长度6 # 目标序列: 假设是德文句子对应的ID序列 (训练时输入是右移的) tgt_tokens torch.LongTensor([[1, 8, 4, 2, 0, 0], # 第一句 [1, 3, 5, 7, 2, 0]]) # 第二句 # 3. 创建掩码 # 源序列填充掩码 src_padding_mask model.create_padding_mask(src_tokens) # (batch, 1, 1, src_len) # 目标序列填充掩码 tgt_padding_mask model.create_padding_mask(tgt_tokens) # (batch, 1, 1, tgt_len) # 目标序列前瞻掩码 (防止看到未来信息) tgt_len tgt_tokens.size(1) look_ahead_mask model.generate_square_subsequent_mask(tgt_len) # (tgt_len, tgt_len) # 组合目标序列掩码填充掩码和前瞻掩码取或 tgt_mask tgt_padding_mask (look_ahead_mask 0) # 广播处理 # 4. 前向传播 with torch.no_grad(): logits model(src_tokens, tgt_tokens, src_masksrc_padding_mask, tgt_masktgt_mask) # 5. 检查输出 print(f模型输出logits的形状: {logits.shape}) # 应为 (batch_size, tgt_len, tgt_vocab_size) print(f示例输出 (第一个batch, 最后一个位置): {logits[0, -1, :5]}) # 查看前5个词的概率 # 计算损失 (例如交叉熵损失) # 注意实际训练时需要将目标序列左移作为decoder输入右移作为标签。 # 这里仅为演示前向过程。 criterion nn.CrossEntropyLoss(ignore_index0) loss criterion(logits.view(-1, tgt_vocab_size), tgt_tokens.view(-1)) print(f计算得到的损失值 (仅作参考): {loss.item():.4f}) return model, logits if __name__ __main__: model, logits test_transformer_forward()预期输出与验证运行上述测试代码你应该能看到类似以下的输出模型输出logits的形状: torch.Size([2, 6, 100]) 示例输出 (第一个batch, 最后一个位置): tensor([ 0.1234, -0.5678, 0.9012, -0.3456, 0.7890]) 计算得到的损失值 (仅作参考): 4.6052形状验证输出形状(2, 6, 100)符合预期表示2个样本每个目标序列长度为6每个位置对应100个词汇的logits。数值验证输出的是logits未归一化的分数有正有负是正常的。初始化的模型输出是随机的。损失值初始随机模型的损失值通常在-ln(1/vocab_size)附近对于词汇表100约为4.605这与输出相符说明前向传播计算是连贯的。测试成功标准代码能正常运行不报错。输出张量的形状与预期完全一致。损失函数能正常计算出一个标量值。整个过程在CPU上快速完成几秒内。9. 资源占用与性能观察由于我们搭建的是用于原理验证的极小模型资源占用可以忽略不计。但理解其扩展性至关重要。参数规模估算一个Transformer模型的参数量主要来源于词嵌入层(src_vocab_size tgt_vocab_size) * d_model注意力层的QKV和输出投影4 * num_layers * d_model * d_model(简化估算忽略偏置)前馈网络的两层线性变换2 * num_layers * d_model * d_ff层归一化和输出投影占比很小。以我们测试的微型模型为例 (d_model64, num_layers2, d_ff256, vocab100)参数量仅约几万。而原始论文的base模型 (d_model512, num_layers6, d_ff2048, vocab37000) 参数量约为65M。大型模型如GPT-3 (175B) 则是这个数字的数千倍。性能关键点自注意力复杂度序列长度L的注意力计算复杂度为O(L²)这是处理长文本的主要瓶颈。后续的改进如Longformer、BigBird的核心就是优化这一点。内存占用训练时需要存储中间激活值用于反向传播显存占用远大于模型参数本身。使用梯度检查点Gradient Checkpointing是一种用计算换内存的策略。批处理Batch增大批处理大小能提升GPU利用率但受限于显存。混合精度训练使用FP16/BF16精度可以显著减少显存占用并加速计算。对于本地学习的建议始终从微型模型开始确保前向/反向传播正确。逐步放大d_model,num_layers等超参数观察资源消耗。使用torch.cuda.memory_allocated()监控显存使用情况。10. 常见问题与排查方法在搭建和运行Transformer时你可能会遇到以下问题问题现象可能原因排查方式解决方案运行时错误矩阵维度不匹配1. 张量reshape或transpose操作维度错误。2. 多头注意力分头时d_model不能被num_heads整除。检查出错行的张量形状使用print(tensor.shape)调试。1. 仔细核对view/transpose的维度顺序。2. 确保assert d_model % num_heads 0。输出全是NaN或数值异常大1. 学习率过高。2. 没有进行梯度裁剪。3. 层归一化或残差连接实现有误。检查前几轮训练中权重和梯度的值。1. 使用更小的学习率。2. 添加梯度裁剪torch.nn.utils.clip_grad_norm_。3. 复查Add Norm的实现确保是LayerNorm(x Sublayer(x))。训练损失不下降1. 模型太小表达能力不足。2. 数据预处理或掩码错误。3. 优化器或学习率设置不当。1. 在极小的数据集上过拟合看模型能否记住。2. 可视化注意力权重看模型是否关注了合理的位置。1. 增加模型容量 (d_model,num_layers)。2. 仔细检查数据加载和掩码生成代码。3. 尝试不同的优化器如AdamW和学习率调度。GPU显存溢出OOM1. 模型太大或序列太长。2. 批处理大小batch size过大。使用nvidia-smi或torch.cuda.memory_summary()监控显存。1. 减小模型规模、序列长度或批大小。2. 使用梯度累积来模拟大batch。3. 启用激活检查点。解码时生成重复或无意义内容1. 训练不充分。2. 解码策略如贪婪搜索、束搜索参数不佳。3. 缺乏重复惩罚等后处理。检查验证集上的损失是否已收敛。分析生成样本。1. 延长训练时间。2. 调整束搜索的宽度beam width或长度惩罚length penalty。3. 在生成时引入n-gram重复惩罚。11. 最佳实践与使用建议基于这次从零搭建的经历为你后续深入使用Transformer提供一些建议理解优先于复现第一次实现时可以像本文一样简化如使用固定位置编码、简化掩码逻辑。关键是打通数据流理解每个模块的输入输出。之后再去研究优化版如可学习的位置编码、更高效的注意力实现。善用现有库进行对比在完成自己的实现后用相同的输入调用Hugging Facetransformers库的对应模型对比中间层输出或最终输出这是验证正确性的好方法。模块化保存与测试将MultiHeadAttention,EncoderLayer,DecoderLayer等类单独保存在文件中并为其编写单元测试。这有利于后续复用和调试。可视化辅助理解在调试时可视化注意力权重图。这能直观地看到模型在关注什么对于发现掩码错误或模型行为异常非常有帮助。从文本扩展到其他模态Transformer不仅是NLP的利器。理解了基本架构后可以尝试阅读Vision Transformer (ViT)、Audio Transformer等工作的代码看它们如何将图像、音频信号处理成序列形式输入Transformer。关注变体与进化掌握原版架构后应系统学习其重要变体如仅编码器Encoder-Only如BERT适用于理解类任务分类、标注。仅解码器Decoder-Only如GPT系列适用于生成类任务。高效注意力机制如Linformer、Performer、FlashAttention解决长序列的复杂度问题。结构改进如Pre-LN (层归一化放在残差连接前)、T5的编码器-解码器结构等。通过这次动手搭建你应该对Transformer不再感到神秘。它就像一个精心设计的乐高套装每个模块都有明确的功能和接口。注意力机制固然是其灵魂但残差连接、层归一化、位置编码和前馈网络这些“基础设施”同样不可或缺它们共同保证了深度网络的稳定训练和强大表达能力。下一步你可以尝试用这个框架在一个真实的小数据集如IWSLT德语-英语翻译数据集上运行配置优化器、损失函数和训练循环观察一个真正的翻译模型是如何被训练出来的。或者尝试修改架构比如将ReLU激活函数换成GELU或者实现相对位置编码如T5的Relative Position Bias这将让你对模型的细节有更深的掌控力。