2026/8/30 10:39:39

Transformer架构详解:从注意力机制到PyTorch实现

Transformer架构详解:从注意力机制到PyTorch实现 我先把一个判断放在最前面如果你现在只打算学一个深度学习架构Transformer 大概率是性价比最高的选择。这不是因为它“热门”而是因为它把序列建模的底层逻辑换掉了。从 NLP 到 CV从语音到推荐系统甚至最近几年大模型背后的 GPT 系列核心骨干都是 Transformer。与其一个模型一个模型去追不如把 Transformer 本身吃透。很多初学者学 Transformer 时会遇到三个问题第一注意力机制的公式能看懂但不知道为什么要有 Q、K、V 三个向量也不清楚它们和“相似度”有什么关系第二看完原版论文或者各种图解明白了架构图但一打开代码发现自己写不出一个可运行的前向传播第三知道 Transformer 在 NLP 里很火但不理解它凭什么能用在图像、时间序列这些非文本场景里。这篇文章会从上面三个问题展开先讲清楚注意力机制的设计动机再拆解标准 Transformer 的完整架构然后用 PyTorch 从零实现一个可运行的简化版本最后把视觉 Transformer、Swin Transformer、TCNTransformer 这类延伸方向串起来让你既能理解原理也能在自己的项目里动手改。1. 为什么最后是 Transformer要理解 Transformer 为什么重要先看它之前的主流方案长什么样。在 Transformer 出现之前序列建模主要靠两类架构。一类是 RNN包括 LSTM、GRU。它们的思路是“按顺序处理”当前时刻的隐藏状态依赖上一个时刻的输出所以天然适合文本、语音这类有时间顺序的数据。但问题也很明显串行计算导致训练速度慢而且长距离依赖会随着序列变长而衰减即使 LSTM 加了门控机制也只是缓解没有根治。另一类是纯 CNN。CNN 在处理图像时很擅长提取局部特征但如果用它处理序列往往需要堆很多层才能扩大感受野。层数越深训练难度越大长距离信息的传递效率也越低。Transformer 的做法是彻底抛弃“逐步递归”改成“并行计算 全局注意力”。它不再按时间顺序一步步读数据而是让序列中任意两个位置直接建立联系。这意味着什么第一个变化是速度。并行计算让训练时间大幅缩短这是它能支撑大规模预训练模型的前提。第二个变化是能力。自注意力机制让每个 token 都能“看到”序列里的所有 token长距离依赖不再是问题。第三个变化是通用性。只要能把数据表示成一组向量Transformer 就能处理文本、图像、时序数据都可以。所以“为什么最后是 Transformer”这个问题的答案不是因为它某个指标特别高而是它在速度、能力和通用性三个维度上同时改变了游戏规则。2. 核心概念注意力机制与 QKV 的直觉理解Transformer 的地基是注意力机制。网上关于注意力机制的图解已经很多了这里换一个角度从“你到底在算什么”出发来理解。假设你有一句话“今天天气很好我决定去公园散步。”如果要判断“今天”和哪个词关系最紧密你会下意识去关注“天气”和“散步”而不是“我”。注意力机制做的就是这个事情给序列里的每个词分配一个权重权重越高表示当前词在编码时应该越关注它。Q、K、V 这三个向量是这里的关键。可以把它们类比成一次检索过程Query查询是你在搜索框里输入的内容代表“我现在想找什么”。Key键是每个候选文档的标题代表“我能提供什么”。Value值是文档正文代表“如果匹配上了真正的内容是什么”。注意力权重的计算方式是用 Query 和所有 Key 做点积得到相似度分数再经过 Softmax 归一化成概率分布最后用这个概率对 Value 做加权求和。用公式表示就是Attention(Q, K, V) softmax(Q * K^T / sqrt(d_k)) * V除以根号 d_k 是为了防止点积结果过大导致 Softmax 进入饱和区梯度变得非常小。多头注意力做的事情更简单不只用一组 Q、K、V而是把向量拆成多组并行计算注意力最后拼接起来。这样做的意义是让模型能够同时关注不同位置、不同表示子空间的信息。比如在一个句子里一个头可能关注语法关系另一个头关注指代关系还有一个头关注语义相近的词。多头机制本质上是让模型拥有多种“看问题”的视角。到这里可以记住一个判断注意力机制最核心的思想不是“计算相似度”这个动作而是“让模型自己决定该关注哪里”的能力。3. Transformer 架构拆解从输入到输出标准 Transformer 采用 Encoder-Decoder 结构。原文中 Encoder 和 Decoder 各有 6 层但层数可以根据任务调整。下面按数据流向拆解一遍。3.1 输入处理Embedding 与位置编码Transformer 本身不感知顺序如果把“我喜欢你”和“你喜欢我”输入模型得到的向量一模一样。要让模型理解顺序必须主动加入位置信息。原始论文使用正弦位置编码PE(pos, 2i) sin(pos / 10000^(2i / d_model)) PE(pos, 2i 1) cos(pos / 10000^(2i / d_model))其中 pos 是位置下标i 是维度下标。这种编码方式的好处是不同位置的编码向量不同而且模型可以通过线性变换捕捉相对位置信息。现代实现中很多模型改用可学习位置编码也就是把位置编码当成一个可训练参数。两者各有优劣正弦位置编码可以处理任意长度序列可学习位置编码在固定长度内效果更稳定。3.2 Encoder 层自注意力 前馈 残差每个 Encoder 层由两个子层组成多头自注意力层前馈神经网络层Feed-Forward NetworkFFN每个子层外面都套了残差连接和 Layer Normalization。残差连接解决的是深层网络梯度消失问题让信息在跨层传递时更顺畅。Layer Normalization 则对每个样本的特征维度做归一化让训练更稳定。这里有一个容易混淆的点Batch Normalization 和 Layer Normalization 的区别。简单理解Batch Norm 对同一个特征维度跨样本做归一化Layer Norm 对同一个样本跨特征维度做归一化。Transformer 使用 Layer Norm因为序列长度经常变化而每个样本的特征维度通常是固定的Layer Norm 更稳定。FFN 通常是一个两层全连接网络FFN(x) ReLU(x * W1 b1) * W2 b2中间层的维度一般是 d_model 的 4 倍例如 d_model 是 512 时FFN 中间层是 2048。3.3 Decoder 层掩码自注意力 交叉注意力Decoder 的结构比 Encoder 复杂一些包含三个子层掩码多头自注意力层多头交叉注意力层前馈神经网络层掩码自注意力是 Decoder 和 Encoder 最大的区别。训练时我们不能让模型“看到”未来的词。比如预测“今天”后面的“天气”时模型不能提前知道“天气”这个词。实现方式是在计算注意力分数时对未来的位置加上一个极大的负数通常是 -inf这样 Softmax 之后这些位置的权重几乎为零。交叉注意力层的 Query 来自 Decoder 自己上一层的输出Key 和 Value 来自 Encoder 的最终输出。这一步的作用是让 Decoder 在生成每个词时都能从源序列中获取相关信息。3.4 输出层与 SoftmaxDecoder 最后一层的输出经过一个线性层映射到词表大小再用 Softmax 得到每个词的概率分布。训练时通过交叉熵损失和真实标签计算误差推理时通过自回归方式逐个生成 token。4. 手写简化版 TransformerPyTorch 代码实现只看架构图是不够的真正理解一个模型最快的方式是动手写代码。下面用 PyTorch 实现一个简化但完整的 Transformer。4.1 环境准备本文代码依赖以下环境Python 3.8 PyTorch 1.10 NumPy如果没有安装 PyTorch可以使用以下命令pip install torch numpy4.2 多头注意力实现先实现多头注意力模块。这里有一个容易出错的地方多头注意力的数据形状变化。输入是 [batch_size, seq_len, d_model]需要先拆成 [batch_size, num_heads, seq_len, head_dim]计算完注意力后再合并回 [batch_size, seq_len, d_model]。import torch import torch.nn as nn import math class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): super(MultiHeadAttention, self).__init__() assert d_model % num_heads 0 self.d_model d_model self.num_heads num_heads self.head_dim d_model // 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.fc_out nn.Linear(d_model, d_model) def forward(self, query, key, value, maskNone): batch_size query.shape[0] # 1. 线性变换并拆成多头 Q self.w_q(query).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) K self.w_k(key).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) V self.w_v(value).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) # 2. 计算注意力分数 scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim) # 3. 应用掩码如果存在 if mask is not None: scores scores.masked_fill(mask 0, -1e9) # 4. Softmax 归一化 attention torch.softmax(scores, dim-1) # 5. 加权求和 out torch.matmul(attention, V) # 6. 合并多头 out out.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) out self.fc_out(out) return out这段代码的关键点在第 2 步和第 4 步先算 Q 和 K 的点积得到相似度矩阵再用 Softmax 转成权重分布。mask 的作用是在计算时屏蔽不需要关注的位置。4.3 位置编码与前馈网络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, dtypetorch.float).unsqueeze(1) 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): return x self.pe[:, :x.size(1), :] class FeedForward(nn.Module): def __init__(self, d_model, d_ff2048, dropout0.1): super(FeedForward, self).__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(torch.relu(self.linear1(x))))位置编码里注册为 buffer 而不是 parameter是为了让它在模型保存和加载时跟着走但不会被优化器更新。4.4 Encoder 层class EncoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff2048, dropout0.1): super(EncoderLayer, self).__init__() self.self_attn MultiHeadAttention(d_model, num_heads) self.ffn FeedForward(d_model, d_ff, dropout) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) def forward(self, x, maskNone): # 自注意力 残差 LayerNorm attn_out self.self_attn(x, x, x, mask) x self.norm1(x self.dropout(attn_out)) # 前馈网络 残差 LayerNorm ffn_out self.ffn(x) x self.norm2(x self.dropout(ffn_out)) return x class Encoder(nn.Module): def __init__(self, d_model, num_heads, num_layers6, d_ff2048, dropout0.1): super(Encoder, self).__init__() self.layers nn.ModuleList([ EncoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers) ]) def forward(self, x, maskNone): for layer in self.layers: x layer(x, mask) return x有一个小细节值得注意残差连接和 LayerNorm 的顺序。原版论文是先加残差再 LayerNorm也就是 Post-LN。后来的实践中发现 Pre-LN先 LayerNorm 再进子层在深层网络上更稳定。上面代码按原版实现实际工程中可以根据需要调整。4.5 Decoder 层class DecoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff2048, dropout0.1): super(DecoderLayer, self).__init__() self.self_attn MultiHeadAttention(d_model, num_heads) self.cross_attn MultiHeadAttention(d_model, num_heads) self.ffn FeedForward(d_model, d_ff, dropout) 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): # 掩码自注意力 attn_out self.self_attn(x, x, x, tgt_mask) x self.norm1(x self.dropout(attn_out)) # 交叉注意力Q 来自 decoderK、V 来自 encoder attn_out self.cross_attn(x, enc_output, enc_output, src_mask) x self.norm2(x self.dropout(attn_out)) # 前馈网络 ffn_out self.ffn(x) x self.norm3(x self.dropout(ffn_out)) return x4.6 完整 Transformer 模型class Transformer(nn.Module): def __init__(self, src_vocab_size, tgt_vocab_size, d_model512, num_heads8, num_layers6, d_ff2048, max_len100, dropout0.1): super(Transformer, self).__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, max_len) self.encoder Encoder(d_model, num_heads, num_layers, d_ff, dropout) self.decoder Decoder(d_model, num_heads, num_layers, d_ff, dropout) self.fc_out nn.Linear(d_model, tgt_vocab_size) def forward(self, src, tgt, src_maskNone, tgt_maskNone): # 编码器输入 src_embedded self.positional_encoding(self.encoder_embedding(src)) enc_output self.encoder(src_embedded, src_mask) # 解码器输入 tgt_embedded self.positional_encoding(self.decoder_embedding(tgt)) dec_output self.decoder(tgt_embedded, enc_output, src_mask, tgt_mask) # 输出层 output self.fc_out(dec_output) return output需要说明的是上面的代码是教学用简化版缺失了 Embedding 缩放原版会在 Embedding 后乘以 sqrt(d_model)和 Dropout 的具体位置。实际训练中这些细节会影响效果。5. Transformer 如何扩展到视觉领域如果说 Transformer 在 NLP 的成功还能用“序列建模能力强”来解释那它在视觉领域的应用就更有意思了图像不是天然序列为什么要用 Transformer5.1 ViT把图像当成词序列ViTVision Transformer的核心思路非常直接把一张图像切成固定大小的 Patch比如 16x16 像素一个 Patch然后每个 Patch 拉平成一个向量加上位置编码送入标准 Transformer Encoder 训练。这里有一个关键操作Patch Embedding。class PatchEmbedding(nn.Module): def __init__(self, img_size224, patch_size16, in_channels3, embed_dim768): super(PatchEmbedding, self).__init__() num_patches (img_size // patch_size) ** 2 self.proj nn.Conv2d(in_channels, embed_dim, kernel_sizepatch_size, stridepatch_size) def forward(self, x): # x: [batch_size, in_channels, img_size, img_size] x self.proj(x) # [batch_size, embed_dim, num_patches_height, num_patches_width] x x.flatten(2) # [batch_size, embed_dim, num_patches] x x.transpose(1, 2) # [batch_size, num_patches, embed_dim] return x这里用 Conv2d 实现 Patch Embedding 是工程上的常用技巧。把卷积核大小和步长都设为 patch_size一次卷积就能完成切块和线性映射两个操作比手写切块再线性变换效率更高。ViT 在 ImageNet 这类大数据集上表现很好但在小数据集上效果不如 CNN。原因是 Transformer 缺少 CNN 的归纳偏置局部性和平移等变性需要更多数据才能学到视觉特征。这是一个非常实际的工程结论。5.2 Swin Transformer引入层级和窗口注意力ViT 的问题在于全局注意力计算量太大图像分辨率升高后开销急剧增加。Swin Transformer 的解决方案是“窗口注意力”只在局部窗口内计算自注意力窗口之间通过移动窗口操作实现信息交互。这样做有两个好处计算复杂度从 O(N²) 降为 O(N)N 是 Patch 数量。形成类似 CNN 的层级结构可以自然地作为骨干网络用在目标检测、分割等任务上。Swin Transformer 在工程上特别适合做“通用视觉骨干”这也是它被广泛采用的原因。5.3 视觉 Transformer 的适用场景判断如果你要处理的任务是图像分类、目标检测、语义分割并且数据量足够Transformer 类模型是值得尝试的方案。但如果训练数据很少或者推理延迟要求极高CNN 仍然是更稳妥的选择。在实际项目中一个常见的做法是用 ImageNet 预训练的 Swin Transformer 作为骨干网络然后在自己的下游任务上微调。这样既利用了 Transformer 的表征能力又避免了从零训练的成本。6. CNN、RNN 与 Transformer 的对比总结把三种架构放在一起横向对比更能看出 Transformer 的特点。维度RNN (LSTM)CNNTransformer计算方式串行递归卷积滑动窗口并行全局注意力长距离依赖弱有梯度衰减弱需堆层数扩大感受野强任意位置直接连接训练速度慢快快可并行位置信息天然存储在序列顺序中由卷积核相对位置隐式编码需要显式添加位置编码适合场景短序列、流式数据图像局部特征提取长序列、大规模数据、跨模态主要短板无法并行、长序列性能差感受野受限、全局建模弱计算复杂度高、小数据易过拟合这个表格可以当作选型参考但不要把它当成绝对标准。实际中经常把几种架构混着用比如 TCNTransformer 用于时间序列预测CNNTransformer 用于视觉任务都是取长补短的组合。7. Transformer 在时间序列预测中的应用Transformer 在时间序列预测领域也越来越常见。一个典型的做法是用 TCN时间卷积网络提取局部时序特征再用 Transformer 建模长距离依赖。这种组合解决了一个实际问题纯 TCN 擅长捕捉局部模式但感受野有限纯 Transformer 擅长全局依赖但对局部特征不敏感。两者结合可以同时处理短期模式和长期趋势。简化流程如下import torch import torch.nn as nn class TCNTransformer(nn.Module): def __init__(self, input_dim, d_model, num_heads, num_layers, kernel_size3): super(TCNTransformer, self).__init__() # TCN 层提取局部时序特征 self.tcn nn.Sequential( nn.Conv1d(input_dim, d_model, kernel_size, paddingkernel_size // 2), nn.ReLU(), nn.Conv1d(d_model, d_model, kernel_size, paddingkernel_size // 2), nn.ReLU(), ) # Transformer Encoder建模长距离依赖 encoder_layer nn.TransformerEncoderLayer(d_modeld_model, nheadnum_heads) self.transformer nn.TransformerEncoder(encoder_layer, num_layersnum_layers) def forward(self, x): # x: [batch_size, seq_len, input_dim] x_tcn self.tcn(x.transpose(1, 2)).transpose(1, 2) x_out self.transformer(x_tcn) return x_out在股票预测、气象预测、电力负荷预测等场景中这类模型往往比单一架构表现更好。但要注意时间序列预测是出了名的“结果容易波动”的任务模型结构只是因素之一数据预处理、特征工程、训练验证方式都会显著影响结果。8. Transformer 的改进方向与“涨点”技巧很多人在自己的任务上用 Transformer 效果不好第一反应是换模型。事实上不少改进空间藏在细节里。8.1 架构层面的常见改进相对位置编码相比绝对位置编码相对位置编码能更好地建模两个 token 之间的距离关系。Pre-LN 替代 Post-LN深层网络训练更稳定是当前工程实现的主流选择。FlashAttention通过 IO 感知的注意力计算优化降低显存占用提升训练速度。稀疏注意力只在部分 token 之间计算注意力降低长序列的计算复杂度。8.2 训练层面的实用技巧学习率预热Transformer 训练初期用较小的学习率避免模型震荡后面再逐步增大。梯度裁剪防止梯度爆炸尤其是深层 Transformer。标签平滑缓解模型过拟合提高泛化能力。Warmup Cosine Decay效果比固定学习率好很多。# 学习率预热示例 def get_lr_scheduler(optimizer, warmup_steps, total_steps): def lr_lambda(step): if step warmup_steps: return step / max(1.0, warmup_steps) progress (step - warmup_steps) / max(1.0, total_steps - warmup_steps) return 0.5 * (1.0 math.cos(math.pi * progress)) return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)8.3 数据层面的常见问题Transformer 是数据饥渴型模型。数据量不够、数据噪声大、标签不一致都会让 Transformer 的优势发挥不出来。如果你在小数据集上训练优先考虑使用预训练模型做微调而不是从头训练。9. 常见问题与排查方法Transformer 训练和部署中下面几个问题出现频率最高。问题现象可能原因排查方式解决方案训练 loss 不下降学习率过大或过小打印 loss 曲线检查学习率使用学习率预热调整初始学习率显存不够序列太长注意力矩阵 O(N²) 爆炸监控显存占用使用梯度累积、减少 batch size、改用稀疏注意力过拟合严重数据量不足模型太大观察训练集和验证集 loss 差异增加 dropout、数据增强、使用预训练模型推理速度慢自回归生成逐个 token分析瓶颈使用 KV Cache、批量推理、模型量化位置编码外推失效训练长度和推理长度不一致测试不同长度输入改用相对位置编码或 ALiBi9.1 一个容易被忽略的问题Mask 的实现很多初学者在实现 Decoder 的掩码自注意力时只对注意力分数矩阵的上三角做掩码却忘了把掩码放在正确的位置上。def generate_tgt_mask(seq_len): mask torch.triu(torch.ones(seq_len, seq_len) * float(-inf), diagonal1) return mask这段代码生成一个上三角矩阵对角线以下为 0表示保留对角线以上为 -inf表示屏蔽。应用时直接加到注意力分数矩阵上再经过 Softmax被屏蔽的位置权重就会变成 0。9.2 训练不稳定的原因Transformer 训练不稳定通常最先怀疑两个地方学习率设置和 LayerNorm 位置。Post-LN 在深层网络训练中容易出现梯度问题换成 Pre-LN 往往能解决。此外初始化方式也很重要尤其注意残差分支的初始化避免输出方差过大。10. 工程落地与最佳实践10.1 用现成库还是手写如果是学习原理手写一遍是不可替代的。如果是项目落地直接用 Hugging Face Transformers 或 PyTorch 官方实现不要在业务项目里从零造轮子。原因很简单成熟库已经处理了大量边界情况、优化和测试你和团队可以用更少的时间做出更稳定的效果。10.2 配置管理Transformer 模型参数多不同实验的配置差异很容易混乱。建议把模型超参、训练超参、数据路径统一放在配置文件中。model: d_model: 512 num_heads: 8 num_layers: 6 d_ff: 2048 dropout: 0.1 train: batch_size: 32 learning_rate: 0.0001 warmup_steps: 4000 total_steps: 100000 gradient_clip: 1.0使用 YAML 或 JSON 配置能显著提高实验的可复现性。10.3 显存优化策略混合精度训练用 FP16 减少显存占用和计算时间。梯度累积不增大 batch size 的情况下模拟更大的 batch。激活检查点以时间换空间用重计算减少激活值存储。序列长度裁剪如果序列长度分布不均匀可以考虑动态 padding减少无效计算。10.4 安全与合规提醒训练数据里可能包含敏感信息部署模型时要考虑数据脱敏、访问控制和推理日志的合规问题。微调开源模型时留意模型许可证和数据使用条款。这部分看起来琐碎但生产环境出问题往往就在这些地方。11. 总结与后续学习方向Transformer 的核心思想可以浓缩成一句话用自注意力机制替代递归和卷积让模型能够并行处理序列并在任意两个位置之间建立直接联系。这篇文章讲清楚了四个问题注意力机制里的 Q、K、V 到底在做什么以及为什么需要多头机制。标准 Transformer 的 Encoder-Decoder 结构中每个模块的作用和顺序。用 PyTorch 从零实现一个简化版 Transformer 的完整代码流程。Transformer 从 NLP 扩展到视觉和时间序列的路径以及改进和落地技巧。下一步值得深入的方向有三个第一读一遍原论文《Attention Is All You Need》重点看实验部分的参数设置和消融分析很多设计决策都有实验证据支撑。第二读一遍 Hugging Face Transformers 中 BERT 或 GPT 的源码对比自己手写的版本看工程实现做了哪些优化。第三找一个具体的下游任务比如文本分类、命名实体识别或者股票预测用预训练模型和微调流程完整跑一遍把从数据到部署的全链路串起来。建议把文章里的核心代码先复制下来跑通一个最小 demo再逐步往自己的任务上迁移。收藏这篇文章以后遇到 Transformer 相关的问题时可以回来快速翻一遍排查思路。