资讯动态

Transformer模型详解:从注意力公式到可训练脚本的工程实践

发布时间:2026/10/17 20:39:07 来源:尧图企业网站定制
简介这份PDF资料面向深度学习入门者、NLP方向学生及需要理解Transformer架构的算法工程师系统讲解2017年Google论文提出的Transformer模型。内容从Encoder-Decoder整体架构切入逐层拆解自注意力层、位置前馈网络与编码器-解码器注意力层并借助“The animal didnt cross the street because it was too tired”等实例说明Self-Attention如何让模型在处理每个词时整合全句信息同时覆盖词嵌入维度512、Query/Key/Value向量计算、缩放点积注意力与Softmax归一化等关键细节。资源包共1个PDF文件约2.01MB结构完整、图示丰富适合作为论文精读的辅助笔记或课程复习材料。目前已有313人学习读者可借此建立对Transformer并行计算优势、长文本处理能力及其对BERT、GPT系列影响的整体认知并了解长距离依赖建模与计算资源需求等现存挑战。1. Transformer 模型详解从注意力公式到能跑通的训练脚本很多团队第一次把 Transformer 模型详解翻出来看都是因为手里的 LSTM 或 CNN 在某类任务上撞了墙长距离依赖学不动、并行度上不去、显存吃满但效果不涨。Transformer 用一套自注意力机制把序列建模从「按时间步递归」改成「全序列一次算」训练能并行、长依赖能直连这是它真正值钱的地方。这篇不是论文复述而是按一线落地的顺序拆注意力到底怎么算、位置编码为什么不能省、一个最小可训练模型怎么写、训练时哪些参数一动就翻车。适合已经会 PyTorch、想自己从零搭一遍而不是只会调库的工程师也适合要判断「这个方向值不值得投入」的技术负责人。2. 自注意力到底在算什么从公式到张量形状Transformer 的核心不是「有很多层」而是每一层里那个缩放点积注意力。理解它的关键不是背公式而是把张量形状在脑子里跑一遍形状对了代码基本不会错。2.1 缩放点积注意力的张量推导给定输入序列表示 $X \in \mathbb{R}^{B \times L \times D}$其中 B 是 batchL 是序列长度D 是模型维度。用三组线性变换得到 Query、Key、Value$$Q XW_Q,\quad K XW_K,\quad V XW_V$$注意力输出为$$\text{Attention}(Q,K,V) \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$$这里 $d_k$ 是每个头的维度。除以 $\sqrt{d_k}$ 不是玄学是因为当 $d_k$ 变大时$QK^\top$ 的方差随维度线性增长softmax 会被推到饱和区梯度接近 0。这个缩放是让训练能收敛的必要操作去掉它深层模型基本训不动。形状上$QK^\top$ 得到 $B \times L \times L$ 的注意力分数矩阵softmax 在最后一维做归一化再乘 $V$ 得到 $B \times L \times d_k$。多头就是把 D 切成 h 份每份独立算最后拼接再投影。2.2 一个能直接跑的最小注意力实现import torch import torch.nn as nn import math 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 # 一次性投影出 QKV减少三次矩阵乘的调度开销 self.qkv_proj nn.Linear(d_model, 3 * d_model) self.out_proj nn.Linear(d_model, d_model) def forward(self, x, maskNone): B, L, D x.shape # qkv: [B, L, 3*D] - [B, L, 3, h, d_k] - [3, B, h, L, d_k] qkv self.qkv_proj(x).reshape(B, L, 3, self.num_heads, self.d_k) qkv qkv.permute(2, 0, 3, 1, 4) q, k, v qkv[0], qkv[1], qkv[2] # 缩放点积scores [B, h, L, L] scores torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: # mask 中 True 的位置被屏蔽填一个极小值 scores scores.masked_fill(mask, float(-inf)) attn torch.softmax(scores, dim-1) # 加权求和 - [B, h, L, d_k] - [B, L, D] out torch.matmul(attn, v) out out.transpose(1, 2).reshape(B, L, D) return self.out_proj(out)逻辑说明qkv_proj把三次线性变换合并成一次是常见工程优化数学上等价。permute把 head 维度提到前面是为了让matmul在最后两维上做批量矩阵乘。masked_fill用-inf而不是 0是因为 softmax 里 $e^{-\infty}0$填 0 反而会分走概率质量。参数说明d_model一般取 256、512、768必须是num_heads的整数倍num_heads常见 4、8、16头数越多单头维度越小表达越细但计算调度开销上升。d_k太小比如小于 16时注意力分布会过于尖锐实践中不建议。2.3 位置编码为什么去掉它模型就退化自注意力本身是置换等变的把输入序列打乱输出只是跟着打乱模型感知不到顺序。所以必须显式注入位置信息。常见两种做法方式实现优点代价正弦位置编码固定公式不可学习无需参数可外推到更长序列表达固定未必适配任务可学习位置嵌入nn.Embedding(max_len, d_model)灵活训练中自适应受 max_len 限制超长要插值正弦编码公式$$PE_{(pos,2i)} \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right),\quad PE_{(pos,2i1)} \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right)$$我一般先用可学习位置嵌入因为实现简单、效果稳只有当序列长度会动态变化、需要外推时才换正弦或旋转位置编码。这里有个血泪经验位置编码的初始化方差要和词嵌入匹配否则训练初期位置信号会被词嵌入淹没表现为 loss 前几百步几乎不动。3. 从零搭一个可训练的 Transformer 编码器光有注意力层不够要能训练还得把残差、归一化、前馈网络、mask 都接对。这一章给一个完整可跑的编码器并说明每个模块为什么这么放。3.1 编码器层的标准结构一个编码器层 多头注意力 残差 LayerNorm 前馈网络 残差 LayerNorm。顺序上有两种流派Post-LN原始论文归一化在残差之后和 Pre-LN归一化在子层之前。Pre-LN 训练更稳深层时几乎不用 warmup 也能收敛现在多数实现默认 Pre-LN。class FeedForward(nn.Module): def __init__(self, d_model, d_ff, dropout0.1): super().__init__() self.net nn.Sequential( nn.Linear(d_model, d_ff), nn.GELU(), # 比 ReLU 更平滑深层更稳 nn.Dropout(dropout), nn.Linear(d_ff, d_model), ) def forward(self, x): return self.net(x) class EncoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout0.1): super().__init__() self.norm1 nn.LayerNorm(d_model) self.attn MultiHeadAttention(d_model, num_heads) self.norm2 nn.LayerNorm(d_model) self.ff FeedForward(d_model, d_ff, dropout) self.dropout nn.Dropout(dropout) def forward(self, x, maskNone): # Pre-LN先归一化再进子层残差直连 x x self.dropout(self.attn(self.norm1(x), mask)) x x self.dropout(self.ff(self.norm2(x))) return x逻辑说明d_ff通常是d_model的 4 倍这是原始论文的经验值前馈层承担了大部分参数容量。GELU 替代 ReLU 是常见做法负半轴有非零梯度深层网络不容易出现神经元死亡。残差连接保证梯度能直通底层是深层 Transformer 能训起来的前提。参数说明dropout在中小数据集上取 0.1大数据集可以降到 0.00.1d_ff取 4 倍是默认显存紧张时可以降到 2 倍但表达力会下降需要相应增加层数补偿。3.2 完整编码器与 mask 构造class TransformerEncoder(nn.Module): def __init__(self, vocab_size, d_model256, num_heads8, num_layers4, d_ff1024, max_len512, dropout0.1): super().__init__() self.token_emb nn.Embedding(vocab_size, d_model) self.pos_emb nn.Embedding(max_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) def forward(self, input_ids, pad_maskNone): B, L input_ids.shape positions torch.arange(L, deviceinput_ids.device).unsqueeze(0) x self.token_emb(input_ids) self.pos_emb(positions) # pad_mask: [B, L]True 表示 padding 位置 # 扩展成 [B, 1, 1, L]让所有 query 都屏蔽这些 key attn_mask None if pad_mask is not None: attn_mask pad_mask[:, None, None, :] for layer in self.layers: x layer(x, attn_mask) return self.norm(x)逻辑说明词嵌入和位置嵌入相加而不是拼接是为了保持维度一致、减少参数。pad_mask的扩展维度很关键注意力分数形状是[B, h, L, L]mask 要广播到[B, 1, 1, L]才能屏蔽所有 query 对 padding key 的关注。这一步写错模型会把 padding 当成有效信息表现为验证集指标虚高但推理时崩掉。参数说明num_layers在中小任务上 46 层够用d_model256 起步max_len要覆盖实际最长序列超了会索引越界。vocab_size由分词器决定不要拍脑袋设。3.3 训练循环里必须盯的三个量搭好模型只是开始训练时我固定盯三个量loss 曲线、梯度范数、注意力熵。loss 不降先看学习率和 warmup梯度范数突然飙到几十多半是某个 batch 里有超长序列或异常样本注意力熵持续走低说明注意力塌缩到少数位置通常是 mask 写错或位置编码没生效。optimizer torch.optim.AdamW(model.parameters(), lr3e-4, weight_decay0.01) scheduler torch.optim.lr_scheduler.OneCycleLR( optimizer, max_lr3e-4, total_stepstotal_steps, pct_start0.1 ) for step, batch in enumerate(loader): input_ids, pad_mask, labels batch logits model(input_ids, pad_mask) loss criterion(logits.view(-1, vocab_size), labels.view(-1)) optimizer.zero_grad() loss.backward() # 梯度裁剪是深层 Transformer 的后悔药 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() scheduler.step()逻辑说明AdamW 把权重衰减和梯度更新解耦比 Adam 加 L2 更规范。OneCycleLR 前 10% 步数做 warmup是因为训练初期注意力分数方差大直接上大学习率容易发散。梯度裁剪max_norm1.0是标配能挡住偶发的梯度爆炸。参数说明lr3e-4 是中小模型常用起点大模型要降到 1e-4 甚至更低weight_decay0.01 是常规值任务数据量小可以加大到 0.05 抑制过拟合pct_start0.1 表示 warmup 占比层数越多可以适当加大。4. 训练 Transformer 最容易翻车的五个坑这一章全是踩过的坑按「现象 → 原因 → 解决」写能帮你省掉几天调参时间。4.1 loss 前几百步几乎不动现象训练启动后 loss 长时间在初始值附近震荡看不出下降趋势。原因warmup 太短或没有 warmup加上 Post-LN 结构初期梯度被 LayerNorm 缩放得极小。解决改用 Pre-LNwarmup 步数设为总步数的 5%10%学习率峰值先降到 1e-4 试。4.2 验证集指标虚高推理时输出乱码现象训练时验证 loss 很低但拿模型做推理输出重复或完全无关。原因mask 构造错误padding 位置参与了注意力模型在训练时「偷看」了 padding 之外的填充信息。解决检查 mask 的广播维度是否为[B, 1, 1, L]并写一个单元测试构造含 padding 的输入确认 padding 位置的注意力权重为 0。4.3 显存爆掉但 batch 已经很小现象batch 降到 1 还是 OOM。原因注意力分数矩阵是 $O(L^2)$序列长度翻倍显存翻四倍长序列是真正的显存杀手。解决先确认最大序列长度必要时用梯度累积模拟大 batch再考虑分块注意力或降低d_model。不要盲目加显卡先算清楚 $B \times h \times L^2$ 的量级。4.4 训练后期 loss 突然变成 NaN现象训练到中途 loss 突然 NaN之前一切正常。原因某个 batch 出现极端值或学习率在后期仍然偏高梯度爆炸。解决加梯度裁剪max_norm1.0检查数据里有没有空序列或超长序列把学习率调度改成带衰减的余弦或线性下降。4.5 模型对位置完全不敏感现象把输入序列顺序打乱输出几乎不变。原因位置编码初始化方差太小或被后续层「洗掉」也可能是位置嵌入和词嵌入相加后量级不匹配。解决检查位置嵌入的初始化确保其标准差和词嵌入接近可以临时把位置编码换成正弦固定编码做对照实验如果效果变好说明可学习位置嵌入没训起来。5. 进阶技巧把 Transformer 训得更稳更快的三个具体手段到这一步模型能跑了接下来是让它跑得好。我一般按「先稳后快」的顺序调稳不住的速度没有意义。第一个手段是权重初始化。Transformer 对初始化比 CNN 敏感得多nn.Linear默认的 Kaiming 初始化在深层里会让激活方差逐层漂移。常见做法是对残差分支的最后一层投影做缩放初始化把输出方差压小让残差主路占主导。具体就是在out_proj和前馈第二层上把权重乘一个 $1/\sqrt{2N}$ 的系数N 是层数。这个改动很小但能让 12 层以上的模型训练曲线明显更平滑。第二个手段是混合精度训练。用torch.cuda.amp把前向和反向放到 fp16显存能省 30%50%速度提升 1.5 倍左右。但要注意 softmax 和 LayerNorm 在 fp16 下容易溢出autocast会自动把这些算子留在 fp32不用手动改。配合 GradScaler 做损失缩放避免小梯度下溢。scaler torch.cuda.amp.GradScaler() for batch in loader: optimizer.zero_grad() with torch.cuda.amp.autocast(dtypetorch.float16): logits model(input_ids, pad_mask) loss criterion(logits.view(-1, vocab_size), labels.view(-1)) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) scaler.step(optimizer) scaler.update()逻辑说明scaler.scale把 loss 放大避免 fp16 下小梯度变成 0unscale_在裁剪前把梯度还原回真实尺度否则裁剪阈值会失真scaler.step会检查梯度里有没有 inf/nan有就跳过这一步相当于自动的后悔药。第三个手段是学习率预热加余弦退火。warmup 解决初期不稳定余弦退火让后期学习率平滑降到接近 0模型能收敛到更平坦的极小值。我一般设max_lr3e-4、min_lr1e-5、warmup 占总步数 10%剩下的按余弦下降。这套组合在中小规模任务上几乎不用改就能用。验证方法上我习惯在训练中途固定一个小的验证集每隔若干步算一次 loss 和准确率同时把注意力权重可视化出来看。注意力图如果出现明显的对角结构或任务相关的对齐模式说明模型学到了东西如果一片均匀多半是没训好。这个习惯帮我提前发现过好几次 mask 写错的问题比等训练跑完再看指标省时间。最后说个我自己的教训早期我总想一步到位把模型调大结果显存不够、训练不稳、调参周期拉长反而拖慢进度。后来固定流程是先用一个 4 层、d_model256的小模型把数据管线和训练循环跑通确认 loss 能正常下降、推理结果合理再逐步加层加宽。这个顺序看起来慢实际是最快的。希望帮到你。本文还有配套的精品资源点击获取

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑