Transformer模型详解:从注意力公式到能跑通的极简实现
简介这份《Transformer 模型详解》PDF 面向自然语言处理初学者、深度学习入门者及需要系统梳理注意力机制的算法工程师帮助读者从零理解 Transformer 的整体架构与核心原理。资源为单个 PDF 文件压缩包约 2.01MB内容以图文结合方式展开涵盖 Encoder-Decoder 结构、词嵌入维度设定、Self-Attention 计算流程以及 Query、Key、Value 向量的作用等关键知识点。文中通过机器翻译实例逐步演示自注意力如何让模型在处理代词 it 时关联到 animal 或 street并说明编码器与解码器各子层的输入输出流动过程。目前已有 313 人学习适合作为入门读物或复习提纲帮助读者建立对 Transformer 的直观认知并为后续理解 BERT、GPT 等衍生模型打下基础。1. Transformer 模型详解从注意力公式到能跑通的极简实现很多人第一次翻开 Transformer 的论文会被那张 encoder-decoder 结构图劝退觉得这东西离自己很远。但真到动手时你会发现核心不过是一个缩放点积注意力加上前馈网络剩下的都是工程堆叠。我在带新人时反复验证过一件事只要能把单头注意力用几十行代码写出来并跑通后面多头、位置编码、残差这些概念都会顺理成章。这篇笔记就按这个思路走先讲清楚每个模块为什么这么设计再给能直接复现的最小代码最后把训练和推理阶段最容易翻车的地方摊开讲。适合已经会 PyTorch 基础、想真正吃透 Transformer 而不是只会调库的从业者。全文不依赖任何特定框架版本代码在常见环境里都能跑。2. 把注意力公式拆成可运行的张量操作Transformer 的一切都建立在注意力机制上。论文里那个公式写得简洁但真正落到代码时形状变换和掩码处理才是血泪经验的集中地。这一章先把单头注意力写出来再扩展到多头每一步都对应一个可以打印形状的中间结果。2.1 缩放点积注意力的三个矩阵从哪来输入序列先经过三个线性层分别得到 Query、Key、Value。这里的关键是理解它们各自的角色Query 代表「我在找什么」Key 代表「我有什么可以被匹配」Value 代表「匹配上之后我实际取走的信息」。三个线性层的权重是独立学习的不是共享的。缩放因子是 1/sqrt(d_k)其中 d_k 是每个头的维度。这个缩放不是可有可无的玄学当 d_k 较大时点积结果的方差会随维度增长softmax 会进入饱和区梯度变得极小。除以 sqrt(d_k) 把方差拉回 1 附近训练才稳定。import torch import torch.nn as nn import math class SingleHeadAttention(nn.Module): def __init__(self, d_model, d_k): super().__init__() self.d_k d_k # 三个独立线性层分别生成 Q、K、V self.w_q nn.Linear(d_model, d_k) self.w_k nn.Linear(d_model, d_k) self.w_v nn.Linear(d_model, d_k) def forward(self, x, maskNone): # x: (batch, seq_len, d_model) Q self.w_q(x) # (batch, seq_len, d_k) K self.w_k(x) V self.w_v(x) # 缩放点积Q 乘 K 的转置 scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) # scores: (batch, seq_len, seq_len) if mask is not None: # 把需要屏蔽的位置填一个极大负数softmax 后趋近 0 scores scores.masked_fill(mask 0, -1e9) attn torch.softmax(scores, dim-1) out torch.matmul(attn, V) # (batch, seq_len, d_k) return out, attn这段代码里最容易被忽略的是 mask 的形状。mask 需要能广播到 (batch, seq_len, seq_len)常见做法是生成一个下三角矩阵用于解码器的自回归掩码。如果 mask 维度不对PyTorch 不会报错但结果全错这是排查时第一个要检查的地方。参数方面d_model 是输入维度d_k 是投影后的维度。单头场景下 d_k 通常等于 d_model多头场景下 d_k d_model / num_heads。如果 d_model 不能被 num_heads 整除后面拼接时会直接报形状错误。2.2 多头注意力为什么不是简单的并行多头注意力的直觉是让模型在不同子空间里关注不同的模式。比如一个头可能关注语法依赖另一个头关注位置邻近关系。实现上就是把 d_model 切成 num_heads 份每份独立做注意力最后拼接再过一个线性层。class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): super().__init__() assert d_model % num_heads 0, d_model 必须能被 num_heads 整除 self.d_model d_model self.num_heads num_heads self.d_k 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.w_o nn.Linear(d_model, d_model) # 输出投影 def forward(self, x, maskNone): batch, seq_len, _ x.shape # 投影后拆成多头(batch, seq_len, num_heads, d_k) - (batch, num_heads, seq_len, d_k) Q self.w_q(x).view(batch, seq_len, self.num_heads, self.d_k).transpose(1, 2) K self.w_k(x).view(batch, seq_len, self.num_heads, self.d_k).transpose(1, 2) V self.w_v(x).view(batch, seq_len, self.num_heads, self.d_k).transpose(1, 2) 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) out torch.matmul(attn, V) # (batch, num_heads, seq_len, d_k) # 拼回(batch, seq_len, d_model) out out.transpose(1, 2).contiguous().view(batch, seq_len, self.d_model) return self.w_o(out)这里有个实操细节transpose 之后张量在内存里不连续直接 view 会报错必须先调用 contiguous()。这个坑我在早期实现时踩过不止一次报错信息是「view size is not compatible with input tensor‘s size and stride」看到这个基本就是忘了 contiguous。多头注意力的参数量是 4 * d_model^2Q、K、V、O 四个投影矩阵和头数无关。头数增加只是让每个头维度变小总参数量不变。这也是为什么调头数时不需要担心模型容量爆炸。2.3 位置编码正弦方案和可学习方案怎么选自注意力本身对序列顺序无感知打乱输入顺序结果不变。位置编码就是往输入里注入位置信息。常见两种做法正弦固定编码和可学习编码。正弦编码用不同频率的 sin 和 cos 函数生成公式里 pos 是位置i 是维度索引。它的好处是可以外推到训练时没见过的长度因为函数形式固定。可学习编码就是 nn.Embedding(max_len, d_model)简单直接但超过 max_len 就没法处理。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) self.register_buffer(pe, pe.unsqueeze(0)) # (1, max_len, d_model) def forward(self, x): # x: (batch, seq_len, d_model) return x self.pe[:, :x.size(1)]register_buffer 的作用是把 pe 注册成模型状态的一部分会跟着模型保存和加载但不会被优化器更新。如果直接用普通属性保存模型时容易丢加载后位置编码变成随机值推理结果完全不对。选型上如果你的任务序列长度固定且不长可学习编码够用。如果要做长度外推或者序列长度变化大正弦编码更稳。现在很多新架构用旋转位置编码但那是另一个话题先把基础两种吃透。3. 搭一个能跑通的编码器层和解码器层理解了注意力之后编码器和解码器就是把注意力、前馈网络、残差连接、层归一化按特定顺序组装起来。这一章给出完整的编码器层和解码器层实现并说明每个组件的位置为什么不能随便换。3.1 前馈网络和残差连接的位置关系每个编码器层包含两个子层多头自注意力和前馈网络。每个子层外面都包着残差连接和层归一化。原始论文的顺序是「子层 - 残差 - 层归一化」后来很多实现改成「层归一化 - 子层 - 残差」也就是 Pre-LN。Pre-LN 训练更稳定不需要 warmup 也能收敛现在更常用。前馈网络就是两层线性变换加激活函数中间维度通常是 d_model 的 4 倍。激活函数原始论文用 ReLU现在常见 GELU 或 SwiGLU。class FeedForward(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.act nn.GELU() def forward(self, x): return self.linear2(self.dropout(self.act(self.linear1(x)))) 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.ffn FeedForward(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, x, maskNone): # Pre-LN先归一化再进子层 attn_out self.self_attn(self.norm1(x), mask) x x self.dropout1(attn_out) # 残差连接 ffn_out self.ffn(self.norm2(x)) x x self.dropout2(ffn_out) return x残差连接的作用是给梯度一条高速公路让深层网络也能训练。如果没有残差堆到 6 层以上梯度就会严重衰减。dropout 加在残差分支上而不是主分支上这是标准做法别搞反。d_ff 一般设为 4 * d_model。这个比例不是硬性规定但偏离太多会影响表达能力和计算量的平衡。如果显存紧张可以降到 2 倍但效果会打折扣。3.2 解码器的交叉注意力掩码怎么设解码器层比编码器多一个交叉注意力子层Query 来自解码器自身Key 和 Value 来自编码器输出。同时解码器的自注意力需要因果掩码防止看到未来位置。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.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.dropout1 nn.Dropout(dropout) self.dropout2 nn.Dropout(dropout) self.dropout3 nn.Dropout(dropout) def forward(self, x, enc_out, src_maskNone, tgt_maskNone): # 自注意力带因果掩码 attn_out self.self_attn(self.norm1(x), tgt_mask) x x self.dropout1(attn_out) # 交叉注意力Q 来自解码器K/V 来自编码器 cross_out self.cross_attn(self.norm2(x), src_mask) x x self.dropout2(cross_out) ffn_out self.ffn(self.norm3(x)) x x self.dropout3(ffn_out) return x def generate_causal_mask(seq_len): # 下三角矩阵1 表示可见0 表示屏蔽 mask torch.tril(torch.ones(seq_len, seq_len)).unsqueeze(0).unsqueeze(0) return mask # (1, 1, seq_len, seq_len)因果掩码必须是下三角对角线为 1。如果写成上三角模型就只能看到未来训练 loss 会降不下去。交叉注意力的 mask 通常用源序列的 padding mask形状是 (batch, 1, 1, src_len)广播到所有头。这里有个容易混淆的点自注意力的 mask 形状是 (batch, 1, tgt_len, tgt_len)交叉注意力是 (batch, 1, 1, src_len)。两者不能混用混用后要么报形状错误要么静默产生错误结果。3.3 把编码器和解码器拼成完整模型有了层实现完整模型就是 embedding 加位置编码再堆 N 层编码器和 N 层解码器最后接一个输出线性层。class Transformer(nn.Module): def __init__(self, src_vocab, tgt_vocab, d_model512, num_heads8, num_layers6, d_ff2048, dropout0.1, max_len5000): super().__init__() self.src_embed nn.Embedding(src_vocab, d_model) self.tgt_embed nn.Embedding(tgt_vocab, d_model) self.pos_enc PositionalEncoding(d_model, max_len) self.encoder_layers nn.ModuleList([ EncoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers) ]) self.decoder_layers nn.ModuleList([ DecoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers) ]) self.fc_out nn.Linear(d_model, tgt_vocab) self.dropout nn.Dropout(dropout) self.d_model d_model def forward(self, src, tgt, src_maskNone, tgt_maskNone): # embedding 后乘 sqrt(d_model)让 embedding 和位置编码量级匹配 src_emb self.dropout(self.pos_enc(self.src_embed(src) * math.sqrt(self.d_model))) tgt_emb self.dropout(self.pos_enc(self.tgt_embed(tgt) * math.sqrt(self.d_model))) enc_out src_emb for layer in self.encoder_layers: enc_out layer(enc_out, src_mask) dec_out tgt_emb for layer in self.decoder_layers: dec_out layer(dec_out, enc_out, src_mask, tgt_mask) return self.fc_out(dec_out) # (batch, tgt_len, tgt_vocab)embedding 乘以 sqrt(d_model) 这一步经常被省略但省略后 embedding 的数值范围远小于位置编码位置信息会主导输入训练初期 loss 震荡明显。加上之后两者量级接近收敛更平滑。num_layers 原始论文用 6现在大模型动辄几十层。层数增加带来的收益会递减而且训练成本线性增长。如果只是做小规模任务2 到 4 层往往就够没必要盲目堆深。4. 训练 Transformer 时最容易翻车的五个地方代码能跑通不代表能训出结果。这一章记录我在实际项目里反复遇到的五个问题每个都按现象、原因、解决来写都是真金白银换来的教训。4.1 loss 不下降先查学习率和 warmup现象模型跑了几百步loss 几乎不动或者先降后剧烈震荡。原因Transformer 对学习率非常敏感。原始论文用了 warmup 策略前几千步学习率从 0 线性增加到峰值之后按步数平方根衰减。如果直接用固定学习率初期梯度不稳定容易把参数推到坏区域。解决实现 warmup 调度。峰值学习率一般设 1e-4 到 5e-4warmup 步数 4000 左右。小数据集可以缩短到几百步。def lr_lambda(step, d_model512, warmup4000): if step 0: step 1 return (d_model ** -0.5) * min(step ** -0.5, step * warmup ** -1.5) scheduler torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)这个公式来自原始论文d_model 越大峰值学习率越小。如果换了 d_model记得同步调整。4.2 推理时输出重复或乱码检查 mask 和位置编码现象训练 loss 正常但推理生成的序列重复同一个词或者完全不通顺。原因最常见的是推理时没加因果掩码或者位置编码在加载模型后丢失。另一个可能是训练和推理的序列长度不一致位置编码切片出错。解决推理时每一步都要重新生成因果掩码长度等于当前已生成序列长度。位置编码用 register_buffer 保存加载后打印一下确认数值正常。如果做自回归生成记得每步只取最后一个位置的输出。4.3 显存爆炸从 batch size 和序列长度入手现象训练几步后 OOM或者 batch 稍微调大就崩。原因注意力的显存占用是 O(seq_len^2)序列长度翻倍显存翻四倍。另外中间激活值如果不及时释放也会累积。解决优先降序列长度其次降 batch size。可以用梯度累积模拟大 batch。如果还不行考虑梯度检查点用计算换显存。# 梯度累积示例 accum_steps 4 for i, batch in enumerate(loader): loss model(batch) / accum_steps loss.backward() if (i 1) % accum_steps 0: optimizer.step() optimizer.zero_grad()梯度累积时 loss 要除以累积步数否则等效学习率会放大。4.4 过拟合dropout 和权重衰减怎么配现象训练 loss 持续下降验证 loss 先降后升差距越来越大。原因模型容量相对数据量过大或者正则化不足。解决dropout 设 0.1 到 0.3权重衰减设 0.01 左右。如果数据量很小考虑减小 d_model 和层数。早停也是有效手段验证 loss 连续几轮不降就停。4.5 训练和推理结果不一致排查 dropout 和 norm现象同一批数据train 模式和 eval 模式输出差异巨大。原因忘了调 model.eval()dropout 还在随机丢弃。或者 LayerNorm 在 batch 维度上的行为不符合预期。解决推理前必须 model.eval()并用 torch.no_grad() 包住。如果用了自定义 norm确认它在 eval 模式下行为正确。这个坑很隐蔽因为 train 模式下 loss 正常只有对比推理结果时才会发现。5. 用一个小任务验证你的实现是否正确代码写完了怎么确认它真的对我的习惯是找一个极小的任务让模型过拟合如果连过拟合都做不到说明实现有问题。这一章给出一个完整的验证流程包括数据构造、训练循环和结果检查。5.1 构造一个序列反转任务序列反转是个经典验证任务输入 [1,2,3,4,5]输出 [5,4,3,2,1]。这个任务需要模型学会位置对应关系能跑通说明注意力机制工作正常。import torch from torch.utils.data import DataLoader, TensorDataset def make_reverse_data(num_samples, seq_len, vocab_size): src torch.randint(1, vocab_size, (num_samples, seq_len)) tgt torch.flip(src, dims[1]) return TensorDataset(src, tgt) vocab_size 20 seq_len 10 dataset make_reverse_data(2000, seq_len, vocab_size) loader DataLoader(dataset, batch_size64, shuffleTrue) model Transformer(src_vocabvocab_size, tgt_vocabvocab_size, d_model64, num_heads4, num_layers2, d_ff256) optimizer torch.optim.Adam(model.parameters(), lr1e-3, betas(0.9, 0.98), eps1e-9) criterion nn.CrossEntropyLoss(ignore_index0)d_model 设 64、层数设 2 是为了快速验证真实任务再放大。Adam 的 betas 和 eps 用 Transformer 论文推荐值比默认值更稳。5.2 训练循环和关键监控指标训练循环里除了 loss还要监控准确率。反转任务用 token 级准确率就能判断模型是否学会。for epoch in range(20): model.train() total_loss, total_correct, total_tokens 0, 0, 0 for src, tgt in loader: tgt_input tgt[:, :-1] # 输入去掉最后一个 token tgt_output tgt[:, 1:] # 目标去掉第一个 token tgt_mask generate_causal_mask(tgt_input.size(1)) logits model(src, tgt_input, tgt_masktgt_mask) loss criterion(logits.reshape(-1, vocab_size), tgt_output.reshape(-1)) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() total_loss loss.item() preds logits.argmax(dim-1) total_correct (preds tgt_output).sum().item() total_tokens tgt_output.numel() acc total_correct / total_tokens print(fepoch {epoch}, loss {total_loss/len(loader):.4f}, acc {acc:.4f})梯度裁剪设 1.0 是防止梯度爆炸的常规操作。如果 loss 突然变成 nan先检查梯度裁剪有没有生效。5.3 怎么判断结果是真学会了而不是碰巧训练 20 轮后反转任务的准确率应该接近 1.0。如果只有 0.3 左右说明模型没学到位置对应关系。这时候可以打印几个样本的输入输出对比看错误模式。model.eval() with torch.no_grad(): src, tgt next(iter(loader)) tgt_input tgt[:, :-1] tgt_mask generate_causal_mask(tgt_input.size(1)) logits model(src, tgt_input, tgt_masktgt_mask) preds logits.argmax(dim-1) for i in range(3): print(输入:, src[i].tolist()) print(期望:, tgt[i, 1:].tolist()) print(预测:, preds[i].tolist()) print(---)如果预测结果和期望完全无关检查 mask 方向是否反了。如果预测结果总是同一个 token检查输出层和 loss 的维度是否匹配。如果预测结果部分正确但顺序错乱检查位置编码有没有正确加上。这个验证流程我每次重写 Transformer 都会跑一遍花不了十分钟但能省下后面几小时的排查时间。希望帮到你。本文还有配套的精品资源点击获取