Decoder 拆解：掩码自注意力、交叉注意力与 KV Cache

很多人啃 Transformer 源码啃到 Encoder 还能撑住一进 Decoder 就卡壳了三个子层、两种注意力、一个掩码代码里还塞着 KV Cache 的分支判断读着读着就不知道 Q 是从哪来的了。我自己第一次手写 Transformer Decoder 的时候交叉注意力那块的张量形状对错了三次训练 loss 卡在 4.2 一动不动排查了一整晚才发现是把 Encoder 的输出和 Decoder 的输入顺序传反了。这篇就把 Transformer Decoder 结构从设计动机到代码实现彻底拆一遍它到底比 Encoder 多了什么、掩码为什么必须那样生成、训练和推理为什么走两条完全不同的路径、KV Cache 的显存该怎么估。不管你是刚看完论文想动手复现还是已经在调模型但总在解码阶段出问题下面这些内容应该都能直接用上。1. Decoder 在整个架构里的位置与设计初衷1.1 从一句话生成说起Decoder 到底在干什么先把 Encoder 和 Decoder 的分工说清楚不然后面全是糊的。Encoder 干的事是“理解”把一句源语言句子压成一组带上下文的向量每个位置都融合了全句信息。Decoder 干的事是“生成”它要一个词一个词地把目标句子吐出来而且每吐一个词都必须同时参考两样东西——已经吐出来的那些词自身历史以及 Encoder 给的那堆上下文向量源端信息。这就是 Decoder 和 Encoder 最本质的区别。Encoder 的每个位置在算注意力时能看到整句所有位置没有任何限制。Decoder 不行它在生成第 t 个词的时候物理上还没生成第 t1 个词所以必须人为地把“未来”遮住。这个“遮”的动作就是掩码Mask也是 Decoder 一切复杂度的源头。所以你在读 Decoder 结构图时脑子里要始终挂着一条时间线当前这一步能看什么、不能看什么。带掩码的自注意力负责“看历史”交叉注意力负责“看源端”前馈网络负责“做非线性变换”。三个子层串起来才是一次完整的解码。我习惯把 Decoder 类比成一个边听边写的翻译员耳朵听着原文交叉注意力的 K、V手在纸上写译文自注意力的 Q 来自已写部分但眼睛只能看已经写下的字不能偷看还没写的内容。这个类比虽然有简化但对理解 Q、K、V 的来源分配特别管用。1.2 为什么 Decoder 不能直接复用 Encoder 的结构经常有人问既然 Encoder 和 Decoder 都是 N 层堆叠为什么不干脆用同一套 Layer 代码多个开关切换一下理论上可以但实践中有两处根本性差异硬套会出问题。第一处是注意力的输入来源。Encoder 的自注意力里Q、K、V 全部来自上一层输出三者同源。Decoder 的自注意力也是三者同源但它多了一个交叉注意力模块这个模块里 Q 来自 Decoder 自己上一层的输出K 和 V 来自 Encoder 的最终输出。Q 和 KV 不同源这是 Encoder 里完全不存在的结构。你不能简单用一个self_attn(x, x, x)就糊过去得明确区分“谁查谁”。第二处是因果性约束。Encoder 可以全并行一次前向就把整句算完。Decoder 在训练时可以借助 Teacher Forcing 做到并行但推理时必须串行——每步只生成一个 token然后把它拼回输入再走下一步。这意味着 Decoder 的代码路径天然分为训练分支和推理分支两者共享权重但计算图完全不同。我在早期实现时贪图省事推理也走训练那条路每步重新算全序列结果生成 200 个 token 花了四十多秒改成增量解码后降到两秒以内。还有一处容易被忽略的差异Decoder 的输出层通常挂一个线性投影加 softmax映射到词表维度而 Encoder 的输出直接喂给下一层或下游任务头。这个投影层往往和输入 Embedding 做权重共享weight tying能省下vocab_size × d_model的参数词表五万、维度五百的时候就是两千五百万参数不是小数目。注意权重共享这件事在 Encoder-Decoder 架构里只共享“源端 Embedding 与 Decoder 输出投影”或者“目标端 Embedding 与输出投影”别搞混。共享错了模型也能训但收敛会明显变慢且很难通过 loss 曲线看出来。2. 三明治结构的逐层拆解2.1 带掩码的自注意力训练时并行、推理时串行先看第一个子层。输入是目标序列的 Embedding 加上位置编码形状[batch, tgt_len, d_model]。进来之后先做三路线性变换得到 Q、K、V然后计算缩放点积注意力最后过一个掩码。掩码的构造是这块的核心。我们要的是一个下三角矩阵位置 i 只能 attend 到位置 j ≤ i。实现上先造一个全 1 的矩阵取上三角不含对角线置为负无穷或者用torch.triu(torch.ones(len, len), diagonal1).bool()生成一个布尔掩码在 softmax 之前把对应位置填成-1e9或者float(-inf)。这里有个细节很多人踩过填-1e9和填-inf在 fp16 下行为不一样。-inf在经过 softmax 后是精确的 0但某些实现里-inf参与梯度计算会产生 NaN-1e9在 fp16 下会溢出成-inf反而更危险。稳妥做法是用torch.finfo(dtype).min它是当前 dtype 能表示的最小有限数softmax 之后接近 0 又不会产生 NaN。我在 fp16 混合精度训练时就是因为随手写了-1e9跑了三百步突然 loss 变 NaN查了两天才定位到掩码这行。还有形状对齐的问题。掩码要能广播到[batch, num_heads, tgt_len, tgt_len]。最简单是生成[tgt_len, tgt_len]然后 unsqueeze 成[1, 1, tgt_len, tgt_len]让 batch 和 head 维度自动广播。如果你做 padding 掩码还要额外乘上一个[batch, 1, 1, tgt_len]的有效位掩码把 padding 位置也遮掉。两个掩码是逻辑与的关系别写成加法。训练阶段的并行性来自哪里来自 Teacher Forcing。我们把完整的目标序列右移一位前面补 BOS一次性喂进去因为掩码保证了每个位置只能看到自己及之前的位置所以一次前向就能同时算出所有位置的预测不需要循环。这就是为什么训练能高效并行而推理不行。2.2 交叉注意力Q 从哪来、KV 从哪来第二个子层是 Decoder 独有的。它的 Q 来自上一子层的输出也就是 Decoder 自身的表示K 和 V 来自 Encoder 的最终输出memory形状[batch, src_len, d_model]。为什么这么设计直观理解Decoder 当前要生成一个词它需要拿“我现在想表达什么”Q去源句子里检索“哪部分信息对我有用”K 做匹配打分V 提供内容。这跟数据库查询是一个逻辑Q 是查询条件K 是索引V 是实际数据。交叉注意力不需要因果掩码因为源端是完整可见的没有“未来”的概念。但它需要 padding 掩码把源句子里补齐的无效位置遮掉否则模型会 attend 到一堆填充符上长句短句混在一个 batch 里时这个影响特别明显。我做过对比不加源端 padding 掩码同一批数据上 BLEU 掉了将近 2 个点而且注意力可视化里能看到大量权重落在句尾的 pad 上。参数量上交叉注意力有四个投影矩阵W_qd_model × d_model、W_k、W_v、W_o。注意 W_q 的输入维度是 Decoder 的d_modelW_k 和 W_v 的输入维度是 Encoder 的输出维度。如果 Encoder 和 Decoder 的隐藏维度不一样有些架构会这样设计这两个矩阵就是长方形而非正方形形状要单独处理。标准的 base 配置里两者都是 512所以是四个 512×512 的方阵。还有一个常被问到的点交叉注意力能不能多层共享同一份 Encoder 输出能而且几乎所有实现都是这么做的——Encoder 只跑一次输出被每一层 Decoder 的交叉注意力重复使用通常还会缓存 K、V 的投影结果避免重复计算。这也是 Encoder-Decoder 架构推理时的一个优化点后面讲 KV Cache 会展开。2.3 前馈网络与残差归一化的排布顺序第三个子层就是位置前馈网络结构上是两层线性加一个激活Linear(d_model, d_ff) → ReLU/GELU → Linear(d_ff, d_model)。原论文取d_ff 4 × d_model这个 4 倍不是随便定的它让 FFN 承担了整层大约一半的参数量起到“记忆容量”的作用。然后是残差连接和 LayerNorm 的排布这块的争论持续了好几年。原论文用的是 Post-LNx LayerNorm(x Sublayer(x))残差加完再归一化。这种排布在浅层网络里没问题但层数一多梯度在 LayerNorm 处容易出问题必须配 warmup 才能训起来而且对初始化极其敏感。后来大家普遍改成 Pre-LNx x Sublayer(LayerNorm(x))先归一化再进子层残差主干上是一条干净的直通路梯度可以直接回流到底层。我自己实测下来Pre-LN 在 12 层以上的配置里优势非常明显不做 warmup 也能收敛学习率可以设得更激进前 1000 步的 loss 下降曲线明显更陡。代价是最终效果可能略逊于调好参数的 Post-LN但在工程上省下的调参时间远超这点差距。现在主流实现基本都是 Pre-LN或者用 RMSNorm 替代 LayerNorm 做进一步的简化。最后还别忘了在 Decoder 堆叠完之后加一个最终的 LayerNorm。Pre-LN 结构下这一步是必需的因为残差主干的数值没有被归一化过直接送进输出投影会有尺度问题。用一张表把三个子层的差异摆清楚子层Q 来源K/V 来源需要因果掩码需要源端 padding 掩码掩码自注意力当前层输入当前层输入是否目标端 padding 另算交叉注意力上一子层输出Encoder 输出否是前馈网络无注意力无否否3. 从零手写一个 Decoder 层代码与参数3.1 掩码的生成与形状对齐先写掩码。这段代码看着简单但几乎每个人第一次写都会在维度上栽跟头。我把它拆成一个独立函数方便单测。import torch import torch.nn as nn def build_causal_mask(tgt_len, device, dtypetorch.float32): # 上三角不含对角线为 True表示需要被遮住的位置 mask torch.triu(torch.ones(tgt_len, tgt_len, devicedevice, dtypetorch.bool), diagonal1) # 用当前 dtype 的最小有限值填充避免 fp16 溢出成 -inf fill_value torch.finfo(dtype).min causal torch.zeros(tgt_len, tgt_len, devicedevice, dtypedtype) causal causal.masked_fill(mask, fill_value) return causal # [tgt_len, tgt_len]用的时候在注意力计算里 unsqueeze 到[1, 1, tgt_len, tgt_len]再和注意力分数相加。如果你同时还要处理目标端的 padding可以再乘一个布尔矩阵def combine_masks(causal, tgt_pad_mask): # causal: [1, 1, T, T] tgt_pad_mask: [B, 1, 1, T] pad (1.0 - tgt_pad_mask.to(causal.dtype)) * torch.finfo(causal.dtype).min return causal pad这里要注意加法组合两个-min值会导致溢出成-inf虽然 softmax 结果还是对的但在某些混合精度实现里会出现梯度 NaN。更保险的做法是把 padding 掩码直接masked_fill上去或者用torch.where判断。提示写完掩码一定要做个肉眼检查。取一个长度 5 的例子打印出来应该是下三角为 0、上三角为极小值。我见过有人把diagonal0参数写错结果对角线自己被遮了模型永远看不到当前位置的输入loss 卡在词表熵附近下不去。3.2 多头注意力的拆分与合并多头注意力的核心是把d_model拆成num_heads × d_k每个头独立算注意力再拼回来。实现上有两种写法一种是写循环逐头算清晰但慢一种是用 reshape 加 transpose 一次算完快但形状绕。生产代码一律用后者。class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): super().__init__() assert d_model % num_heads 0 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, query, key, value, maskNone): B, T, _ query.shape S key.shape[1] # 拆头: [B, T, H, d_k] - [B, H, T, d_k] q self.w_q(query).view(B, T, self.num_heads, self.d_k).transpose(1, 2) k self.w_k(key).view(B, S, self.num_heads, self.d_k).transpose(1, 2) v self.w_v(value).view(B, S, self.num_heads, self.d_k).transpose(1, 2) scores torch.matmul(q, k.transpose(-2, -1)) / (self.d_k ** 0.5) if mask is not None: scores scores mask attn torch.softmax(scores, dim-1) out torch.matmul(attn, v) # [B, H, T, d_k] out out.transpose(1, 2).contiguous().view(B, T, self.d_model) return self.w_o(out)几个要点。第一缩放因子是sqrt(d_k)而不是sqrt(d_model)因为点积是在单个头的维度上做的。dim 越大点积方差越大不缩放的话 softmax 会变得极其尖锐梯度几乎为零。第二transpose之后必须contiguous()才能view否则会报 stride 错误这个坑几乎人人都踩。第三交叉注意力里T ! Sreshape 时 query 用 T、key/value 用 S别统一写成一个变量。参数计算也顺手说一下。单个注意力模块四个d_model × d_model矩阵参数量4d²。当d_model512时是 1,048,576约 1M。三个子层加起来自注意力 4d²、交叉注意力 4d²、FFN 是d × 4d 4d × d 8d²合计 16d² ≈ 4.2M 参数每层。6 层 Decoder 就是约 25M加上 Embedding词表三万时约 15M和输出投影整体在 40M 上下和原论文 base 版 65M 的总量含 Encoder是对得上的。3.3 完整 DecoderLayer 与 Decoder 堆叠把三个子层拼起来注意 Pre-LN 的顺序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 nn.Sequential( nn.Linear(d_model, d_ff), nn.GELU(), nn.Dropout(dropout), nn.Linear(d_ff, d_model), ) 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, memory, causal_mask, src_pad_maskNone): # 子层一掩码自注意力 h self.norm1(x) x x self.dropout(self.self_attn(h, h, h, causal_mask)) # 子层二交叉注意力Q 来自 DecoderKV 来自 Encoder h self.norm2(x) x x self.dropout(self.cross_attn(h, memory, memory, src_pad_mask)) # 子层三前馈 h self.norm3(x) x x self.dropout(self.ffn(h)) return x堆叠就是循环 N 次最后接一个nn.LayerNorm和一个输出投影nn.Linear(d_model, vocab_size)。这里有个容易忽略的点memory在每一层都是同一份不需要重新计算但每层有各自的w_k、w_v投影所以不能把投影结果跨层复用除非你显式做权重共享那是另一回事。初始化也值得说一句。原论文用 Xavier 初始化实践中 PyTorch 的默认初始化在小模型上够用但层数上到 12 层以上建议对线性层显式做 Xavier uniform并且把输出投影的权重和输入 Embedding 绑定。绑定之后如果 Embedding 有 padding_idx记得那个位置的向量要单独置零并冻结梯度否则 padding 会污染共享权重。4. 训练与推理两条路径的差异4.1 Teacher Forcing 与移位输入的构造训练时 Decoder 的输入是目标序列右移一位。举个例子目标句子是BOS 我 爱 学 习 EOS那么 Decoder 输入是BOS 我 爱 学 习监督标签是我 爱 学 习 EOS。这样位置 0 根据BOS预测“我”位置 1 根据BOS 我预测“爱”以此类推。构造的时候最稳的做法是同一份张量做切片dec_input tgt[:, :-1]labels tgt[:, 1:]。别手工拼字符串再 tokenize那样容易在 BOS 和 EOS 上出错。label 里的 padding 位置要在算 loss 时用ignore_index排除掉CrossEntropyLoss 直接设ignore_indexpad_id就行。Teacher Forcing 的本质是“用真实历史作为输入”而不是用模型自己上一步的预测。好处是训练稳定、可以完全并行坏处是训练和推理存在分布差异——推理时模型看到的是自己生成的、可能有错的 token这叫 exposure bias。缓解手段有 scheduled sampling按概率混入模型自己的预测、label smoothing软化目标分布防止模型过度自信后者实现成本极低通常直接把label_smoothing0.1传给损失函数就行。我在几个小数据集上试过加了 label smoothing 之后生成结果的重复率明显下降。注意做 label smoothing 时如果同时用ignore_index要确认框架版本里这两者的交互是否正确。早期某些版本里 label smoothing 会作用到被忽略的位置上导致 loss 出现奇怪的偏移。4.2 KV Cache增量解码的核心推理阶段每生成一个 token 都要重新跑一遍 Decoder。如果每步都把整个已生成序列重新前向一次计算量是 O(n²) 的重复劳动。KV Cache 的作用就是把每层每步算过的 K、V 存下来下一步只算新 token 的 Q、K、V然后和缓存拼接。实现上的关键改动有两处。一是缓存张量的形状管理通常维护一个[B, H, max_len, d_k]的预分配缓冲区用位置索引写入避免每步torch.cat带来的内存碎片。二是位置编码的偏移新 token 的位置是cache_len不是 0。显存估算有个粗略公式cache_bytes 2 × num_layers × batch × num_heads × seq_len × d_k × bytes_per_elem拿一个具体配置算32 层、batch1、32 头、d_k128、上下文 4096、fp16 存储。代入得2 × 32 × 1 × 32 × 4096 × 128 × 2 2,147,483,648字节约 2GB。这只是单条请求如果并发 8 路就是 16GB显存压力立刻上来了。所以实际部署里 batch 大小往往是被 KV Cache 卡住的而不是被模型权重卡住的。这也是后来各种分页缓存、量化缓存方案出现的原因。我在实现第一版 KV Cache 时犯过一个典型错误把 K 和 V 的缓存在所有层之间共享了一份。表面上看省了显存实际上每层学到的投影完全不同共享之后输出直接乱掉生成的全是无意义的高频词。缓存的粒度必须是“每层独立一份”这个不能省。4.3 采样策略对 Decoder 输出的影响Decoder 最后输出的是词表上的 logits怎么从这里选 token直接决定了生成质量。贪心解码每步取 argmax确定性最强但容易陷入重复循环beam search 维护 k 条候选路径效果通常更好但计算量翻 k 倍而且容易生成过于保守的“安全”句子。温度采样的做法是把 logits 除以温度系数 T 再做 softmax。T 1 时分布变尖锐输出更确定T 1 时分布变平坦多样性提升但容易跑偏。top-k 是只保留概率最高的 k 个候选重新归一化top-pnucleus是保留累积概率达到 p 的最小集合。实践中 top-p 通常比 top-k 更好用因为它能自适应分布的形态——模型很确定时候选集自动收窄不确定时自动放宽。有个组合我用了很久温度 0.7 top-p 0.9 重复惩罚 1.1。这套参数在通用生成任务上比较稳但如果你做的是需要严格事实性的任务比如结构化抽取老老实实用贪心或者 beam search别上采样。采样带来的多样性在那些场景里是纯噪声。还有一点训练时的 dropout 和推理时的model.eval()必须配对。我见过有人推理忘了切 eval 模式dropout 一直在随机丢神经元输出每次都不一样还以为是自己采样参数没调好查了半天。5. 常见问题与排查技巧实录5.1 损失不降、输出重复、注意力全黑现象一loss 从第一步开始就不动稳定在ln(vocab_size)附近。这个数值说明模型输出接近均匀分布什么也没学到。最常见的原因是掩码写错了把对角线也遮住了导致每个位置连自己都看不到。第二个可能是学习率太大直接炸了检查前 10 步的梯度范数如果出现几千上万的数值先把学习率调小一个量级。第三个可能是标签和输入错位了一位模型在学一个恒等映射永远学不会。排查方法很直接取一个 batch手工 decode 输入和标签看是不是差一位。现象二训练 loss 正常下降但生成结果全是重复的短语。这是自回归模型的经典病。先排除是不是推理没切 eval。然后看是不是没加重复惩罚。如果都加了还是重复检查位置编码——如果位置编码没加到 Decoder 输入上或者加了但每一步都用同一个位置模型就分不清词的顺序只能靠高频词凑数。我遇到过一次是位置编码的缓冲区在增量解码时没更新偏移每步都从头开始编生成的前十个词还行后面就崩了。现象三注意力权重可视化出来全是一张黑图。通常是 softmax 之前的分数数值范围太大导致分布极度集中在一个位置。检查有没有做sqrt(d_k)缩放检查 LayerNorm 是不是漏了某个子层。还有一种情况是 fp16 下的数值溢出改成 bf16 或者加 loss scaling 试试。现象四batch 里长短句混合时效果明显变差。八成是 padding 掩码没做全。目标端的 padding 要遮在自注意力里源端的 padding 要遮在交叉注意力里两个地方都要管。单独跑短句没问题、一混合就崩基本就是这个原因。5.2 常见问题速查表症状可能原因排查动作loss 卡在 ln(V)掩码遮了对角线 / 标签错位 / 学习率过大打印掩码矩阵手工核对输入标签位移loss 变 NaN掩码填了-1e9在 fp16 下溢出 / 梯度爆炸改用torch.finfo(dtype).min加梯度裁剪生成结果重复未切 eval / 缺重复惩罚 / 位置编码失效检查model.eval()加 repetition penalty长句效果骤降源端或目标端 padding 掩码缺失检查两个掩码是否都传进注意力推理速度极慢未实现 KV Cache每步全量重算加增量缓存预分配缓冲区显存 OOMbatch 很小时KV Cache 随上下文线性增长按公式估算缓存大小降 batch 或量化缓存交叉注意力权重集中在 pad 上源端掩码未生效确认 mask 形状能广播到[B, H, T, S]5.3 几个只有踩过才知道的细节位置编码的加法位置。位置编码要加在 Embedding 之后、进入第一个 Decoder 层之前。但如果你用 KV Cache 做增量解码新 token 的位置编码必须用它的绝对位置而不是在长度为 1 的序列上的位置 0。这个错误极其隐蔽因为不报错、形状也对就是效果差。dropout 加在哪。原论文在每个子层的输出上、残差相加之前加 dropout注意力权重上也可以加一份。但残差主干上不要加加了会破坏恒等通路。我早期把 dropout 加错了位置训了三天发现效果比不加还差。输出投影的 bias。大部分实现会去掉输出投影的 bias因为和 Embedding 权重绑定之后bias 的作用很有限反而多一组参数。但如果你的词表里有大量低频词保留 bias 有时能帮助校准基础频率。梯度裁剪的阈值。Transformer 训练基本都要裁剪常用值是 1.0。但如果你用了 Pre-LN 且层数不深裁剪阈值可以放宽到 2.0 甚至更高过度裁剪会拖慢收敛。判断标准是看裁剪前后的梯度范数比例如果长期有超过 30% 的步被裁说明阈值设小了。推理时的 batch 排布。同一个 batch 里的序列长度不同KV Cache 要按最长序列预分配短序列会产生浪费。实践中可以按长度分桶把相近长度的请求放在一批能省下不少显存。这个优化在线上服务里几乎是必做的。我个人在实际操作中的体会是Decoder 这块的坑绝大多数不在“结构理解”上而在“形状对齐”和“训练推理路径不一致”这两件事上。结构图看一遍就懂但掩码广播错一个维度、缓存偏移差一位、eval 忘了切这些都不会报错只会让你在 loss 曲线上干瞪眼。所以我的习惯是每写完一个 Decoder 组件立刻拿一个长度为 3 的假数据跑一遍把每层的输入输出形状、掩码矩阵、注意力权重的和是否等于 1全部打印出来对一遍。前期多花十分钟后期能省一整晚。