★ FEATURED ARTICLE Transformer 这几年被聊得太多了但大部分讲解要么停在“Self-Attention 就是让每个词看其他词”这种一句话层面要么直接甩公式把人劝退。我自己带过几个刚入行的同学发现一个很普遍的现象单独问 Attention 怎么算能说个大概但一问“训练的时候到底在算什么、推理的时候又省掉了什么、为什么同一个模型训练和推理表现不一样”就卡住了。问题不在于公式没背熟而在于没有把输入、Attention、训练、推理这四件事串成一条线。这篇东西就是想把这条线捋直。我会从一条输入句子怎么一步步变成输出讲起把 Attention 的矩阵运算、训练时的并行与损失、推理时的自回归与 KV Cache 全部串起来中间穿插参数量估算、显存占用、常见报错排查这些实操层面的东西。适合已经能跑通 Transformer 代码、但对其内部机制还停留在“大概知道”阶段的读者也适合想从 RNN 思维切换到 Attention 思维的从业者。看完你应该能自己回答为什么训练能并行而推理不能、为什么 Attention 是 O(n²)、KV Cache 到底缓存了什么。1. 从一条输入到输出Transformer 的整体数据流1.1 输入是怎么变成向量的先别急着看 Attention先把输入这条线走通。假设我们输入一句话“我 爱 学习”模型第一步不是直接处理这三个字而是先做Tokenization把文本切成 token再通过词表映射成整数 ID。比如“我”对应 1024“爱”对应 2048“学习”对应 4096具体数字取决于词表这里只是示意。这一步的产物是一个长度为 3 的整数序列。接下来是Embedding 层把每个整数 ID 映射成一个 d_model 维的向量。d_model 是 Transformer 的核心超参之一常见取值 512、768、1024。假设 d_model512那么这一步之后我们得到一个 3×512 的矩阵。你可以把它理解成每个 token 变成了一个 512 维空间里的点语义相近的词在这个空间里距离更近。但光有词向量还不够因为 Attention 本身是位置无关的——它只看 token 之间的关系不看顺序。如果打乱输入顺序Attention 算出来的结果几乎一样除了位置编码。所以必须加上位置编码Positional Encoding。原始论文用的是正弦余弦函数现在很多模型改用可学习的位置嵌入或者 RoPE。加完之后输入矩阵还是 3×512但每个向量里已经混入了“我在第几个位置”的信息。注意位置编码不是可选项。我见过有人做文本分类时图省事去掉位置编码结果模型在需要顺序信息的任务上直接崩掉准确率掉十几个点。哪怕你的任务看起来“词袋就够”也建议保留。到这里输入部分就完成了整数序列 → Embedding → 加位置编码 → 得到一个 n×d_model 的矩阵n 是序列长度。这个矩阵就是后面所有层的输入。1.2 Encoder 和 Decoder 的分工Transformer 原始结构是 Encoder-Decoder但现在的用法分化很严重。Encoder-only如 BERT 类擅长理解任务Decoder-only如 GPT 类擅长生成任务Encoder-Decoder如 T5 类擅长序列到序列任务。不管哪种核心积木都是Multi-Head Attention Feed Forward Network区别在于 Attention 的掩码方式和连接方式。Encoder 的每一层做两件事先做 Self-Attention让每个位置看到所有位置再过 FFN对每个位置独立做非线性变换。注意 Encoder 里没有掩码所有位置互相可见。Decoder 的每一层做三件事Masked Self-Attention只能看当前位置及之前、Cross-Attention看 Encoder 的输出、FFN。Decoder-only 模型则去掉了 Cross-Attention只保留 Masked Self-Attention。这里有个容易混淆的点Decoder 的 Masked Self-Attention 和 Encoder 的 Self-Attention 计算方式完全一样只是多了一个上三角掩码把未来位置的注意力分数设成负无穷softmax 之后这些位置权重为 0。这个掩码是训练能并行、推理必须串行的根源后面会详细讲。1.3 残差连接和 LayerNorm 的位置每个子层Attention、FFN外面都包了一层残差连接 LayerNorm。原始论文是 Post-LN先残差再加 LN后来大家发现 Pre-LN先 LN 再进子层最后残差相加训练更稳定现在主流模型基本都用 Pre-LN。这个改动看起来小但对深层模型能不能训起来影响巨大。残差连接的作用是给梯度一条高速公路让深层网络不至于梯度消失。LayerNorm 则是对每个 token 的向量做归一化稳定训练。我实测下来Pre-LN 在 12 层以上的模型里学习率可以设得更大收敛也更稳Post-LN 则需要 warmup 很久否则容易发散。2. Attention 到底在算什么从公式到矩阵形状2.1 Q、K、V 的来历Attention 的核心是三个矩阵Query、Key、Value。它们都是从同一个输入矩阵 Xn×d_model通过三个不同的线性变换得到的Q X · W_Q形状 n×d_kK X · W_K形状 n×d_kV X · W_V形状 n×d_vW_Q、W_K、W_V 是可学习参数。d_k 和 d_v 通常等于 d_model / num_heads。比如 d_model512、num_heads8那么每个头的 d_kd_v64。为什么要有 Q、K、V 三个而不是一个你可以这样理解Q 是“我在找什么”K 是“我有什么可以被找到”V 是“我实际提供的内容”。用 Q 和 K 算相似度得到权重再用权重对 V 加权求和。这个设计让模型可以灵活地决定“关注谁”和“取什么信息”分开。2.2 缩放点积注意力的计算过程公式是Attention(Q, K, V) softmax(Q·Kᵀ / √d_k) · V拆开看每一步的形状变化以单头、n3、d_k64 为例Q·Kᵀ3×64 乘 64×3得到 3×3 的分数矩阵。第 i 行第 j 列表示第 i 个 token 对第 j 个 token 的注意力分数。除以 √d_k防止点积结果过大导致 softmax 梯度消失。d_k64 时 √d_k8相当于把分数缩小 8 倍。softmax对每一行做归一化使每行和为 1得到注意力权重。乘 V3×3 乘 3×64得到 3×64 的输出。为什么要除以 √d_k因为当 d_k 较大时点积的方差会随 d_k 线性增长softmax 会变得非常尖锐梯度趋近于 0。除以 √d_k 把方差拉回 1 附近保证梯度健康。这个细节很多人背了公式但没想过为什么面试被问到就露馅。2.3 多头注意力为什么要拆成多个头Multi-Head Attention 就是把上面的过程并行做 h 次每次用不同的 W_Q、W_K、W_V最后把 h 个头的输出拼接起来再过一次线性变换 W_O。拆多头的好处是不同的头可以关注不同的模式。比如有的头关注语法依赖有的头关注语义相似有的头关注位置邻近。如果只用单头这些模式会挤在同一个注意力分布里互相干扰。我做过一个对比实验同样参数量下8 头的模型比 1 头的在长句理解任务上高 3-5 个点。计算上多头并不是真的把矩阵拆开算而是把 d_model 维切成 h 份用一次大的矩阵乘法完成再 reshape。这样 GPU 利用率更高。2.4 复杂度为什么是 O(n²)注意力分数矩阵是 n×n所以计算量和显存都是 O(n²)。n512 时是 26 万n4096 时是 1600 万n32768 时是 10 亿。这就是长文本处理的瓶颈所在。实际显存占用估算假设 batch8、heads12、n4096、d_k64单层注意力分数矩阵占用 8×12×4096×4096×4 字节 ≈ 6.4 GB。这还只是一层12 层就是 77 GB直接爆显存。所以长上下文模型必须用 FlashAttention、稀疏注意力、滑动窗口等优化手段。提示如果你在训练时遇到 OOM先算一下注意力矩阵的显存往往它就是元凶。降低 batch size 或序列长度是最直接的办法但会牺牲吞吐需要权衡。3. 训练阶段并行、损失与反向传播3.1 训练为什么能并行这是 Transformer 相比 RNN 最大的优势。RNN 必须按时间步串行第 t 步依赖第 t-1 步的隐藏状态无法并行。Transformer 的 Self-Attention 在训练时因为有了 Mask每个位置的计算只依赖输入矩阵和掩码所有位置可以同时算。具体来说训练时我们把完整的目标序列右移一位一次性喂给 DecoderMask 保证第 i 个位置只能看到前 i-1 个位置。这样一次前向就能算出所有位置的输出和损失GPU 可以吃满。这也是为什么 Transformer 能在大规模数据上训练——并行度决定了吞吐。3.2 损失函数与标签对齐训练时用的是交叉熵损失对每个位置的输出做 softmax 得到词表上的概率分布然后和真实的下一个 token 算交叉熵。注意标签是输入序列右移一位输入 [BOS, 我, 爱, 学习]标签是 [我, 爱, 学习, EOS]。这里有个细节padding 位置要 mask 掉不参与损失计算。否则模型会学着去预测 padding浪费容量。实现上用一个 loss_mask把 padding 位置的损失置零最后求平均时除以有效 token 数而不是总长度。3.3 反向传播与梯度问题反向传播沿着计算图回传更新所有参数。Transformer 训练中常见的梯度问题有两个梯度消失和梯度爆炸。梯度消失靠残差连接和 Pre-LN 缓解梯度爆炸靠梯度裁剪gradient clipping通常设 max_norm1.0。学习率调度也很关键。原始论文用 warmup inverse sqrt decay前几千步线性升温之后按步数平方根倒数衰减。现在很多模型用 cosine decay 或带 warmup 的 cosine。我实测下来warmup 步数设总步数的 1%-5% 比较稳太少容易早期发散太多浪费训练时间。3.4 参数量与显存估算以 d_model512、layers6、heads8、词表 30000 为例粗略估算组件参数量Embedding30000×512 ≈ 15.4M每层 AttentionW_Q/W_K/W_V/W_O4×512×512 ≈ 1.05M每层 FFN两层中间 20482×512×2048 ≈ 2.1M每层合计≈ 3.15M6 层合计≈ 18.9M总计≈ 34M训练显存除了参数还要存梯度、优化器状态Adam 是参数的 2 倍、激活值。激活值往往是大头和 batch、序列长度成正比。经验公式训练显存 ≈ 参数量×4fp32 激活值。激活值可以用 gradient checkpointing 换时间省显存代价是前向多算一次。4. 推理阶段自回归与 KV Cache4.1 推理为什么必须串行训练能并行推理却不行原因在于推理时没有真实标签只能一个一个生成。生成第 t 个 token 时输入是前 t-1 个 token模型输出第 t 个 token 的概率分布采样得到 token 后拼到输入里再生成下一个。这个过程天然串行无法像训练那样一次算完。这就是所谓的自回归生成。每一步都要做一次完整前向序列越长步数越多总延迟越高。这也是为什么推理优化KV Cache、批处理、投机解码这么重要。4.2 KV Cache 缓存了什么如果每生成一个 token 都重新算所有位置的 K、V那计算量会爆炸。KV Cache 的思路是已经算过的位置的 K、V 不会变缓存起来新 token 只需要算自己的 Q、K、V然后用自己的 Q 去和缓存的 K 算注意力。具体来说生成第 t 个 token 时新 token 的 Q1×d_k缓存的 Kt×d_k缓存的 Vt×d_v注意力分数1×t输出1×d_v这样每步的计算量从 O(t²) 降到 O(t)总生成 n 个 token 的复杂度从 O(n³) 降到 O(n²)。KV Cache 的显存占用是 2×layers×heads×n×d_k×batch×精度字节数。n4096、layers32、heads32、d_k128、batch1、fp16 时约 2×32×32×4096×128×2 ≈ 2.1 GB。长上下文时这个数字很可观。注意KV Cache 是推理显存的主要占用之一。如果并发请求多batch 一大显存很快见底。生产环境常用 PagedAttention 之类的技术来管理 KV Cache 的显存碎片。4.3 采样策略贪心、束搜索与温度推理时怎么从概率分布里选 token直接影响生成质量贪心每步选概率最大的简单快但容易重复、缺乏多样性。束搜索保留 top-k 条候选路径最后选整体概率最高的。质量通常更好但计算量是 k 倍。温度采样把 logits 除以温度 T 再 softmax。T1 更确定T1 更随机。T0 等价于贪心。top-k / top-p只从概率最高的 k 个或累计概率 p 的集合里采样过滤掉低概率的垃圾 token。我一般对话任务用 top-p0.9 温度 0.7代码生成用温度 0.2 甚至贪心。这些参数没有万能值得按任务调。4.4 训练与推理的差异排查很多人遇到“训练时 loss 降得很好推理结果却很差”的情况。常见原因训练/推理不一致比如训练用了 dropout推理忘了 eval 模式dropout 还在生效。位置编码外推问题训练长度 512推理长度 1024位置编码没做过外推训练性能骤降。KV Cache 实现 bug缓存了错误的 K、V或者位置索引对不上。采样参数不当温度太高导致胡言乱语太低导致重复。排查时先固定随机种子、用贪心解码排除采样干扰再对比训练和推理的输入预处理是否一致最后检查位置编码和 Cache 逻辑。5. 常见问题与排查速查5.1 训练不收敛怎么办先看 loss 曲线。如果 loss 一开始就 NaN多半是学习率太大或数据里有异常值先把学习率降 10 倍试试。如果 loss 震荡不降检查 warmup 是否够、batch size 是否太小、LayerNorm 位置是否合适。如果 loss 降了又升可能是过拟合加 dropout 或 weight decay。我踩过的一个坑数据里混入了空序列导致 attention 全被 masksoftmax 输出 NaN整个 batch 的 loss 变 NaN。后来在数据管道里加了长度过滤才好。5.2 显存不够怎么优化按优先级排序降低 batch size 或序列长度最直接用混合精度fp16/bf16显存减半开 gradient checkpointing用时间换显存用梯度累积模拟大 batch换更省显存的注意力实现如 FlashAttention分布式训练数据并行、张量并行5.3 推理速度慢怎么加速上 KV Cache必做批处理多个请求用更快的注意力 kernel量化int8/int4牺牲一点精度换速度投机解码用小模型草稿、大模型验证5.4 常见报错速查表报错可能原因解决lossNaN学习率过大、数据异常降 lr、清洗数据CUDA OOMbatch/序列过长、KV Cache 过大降 batch、开 checkpointing输出重复温度过低、缺 repetition penalty调高温度、加惩罚推理比训练慢很多没上 KV Cache、串行无批处理加 Cache、批处理长文本效果差位置编码外推不足用 RoPE 外推或滑动窗口6. 把整条链路串起来的几个关键认知走到这里输入、Attention、训练、推理四块应该能连成一条线了。我最后再补几个容易被忽略但很关键的认知点。第一Attention 的本质是加权平均权重由 Q 和 K 的相似度决定。理解了这一点就能理解为什么它位置无关、为什么需要位置编码、为什么复杂度是 O(n²)。第二训练和推理用的是同一套参数但计算图完全不同。训练是 teacher forcing 并行 全序列损失推理是自回归 串行 单步输出。KV Cache 是连接两者的桥梁——它把训练时算过的 K、V 复用起来避免推理时重复计算。第三所有优化都围绕两个瓶颈O(n²) 的注意力和 O(n) 的串行步数。长上下文优化针对前者推理加速针对后者。想清楚你的瓶颈在哪再选优化手段不要盲目堆技术。第四参数量不等于能力数据质量和训练策略往往更关键。我见过参数量翻倍但效果没提升的模型也见过小模型靠更好的数据和调度打赢大模型。调参之前先把数据管道和评估做扎实。最后分享一个我自己的习惯每次改模型结构或训练配置先在一个小数据集上跑通全流程确认 loss 能降、推理能出合理结果再上大规模。这样能省下大量试错时间也能快速定位是结构问题还是数据问题。
阅读完成 · 觉得有帮助？
